Introduce Rowwise/Colwise Sum to math util (#656)

This commit is contained in:
Sherlock 2019-03-22 10:12:58 -07:00 committed by GitHub
parent 600dc9ecc5
commit 25f45cb2db
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 32 additions and 0 deletions

View file

@ -150,6 +150,18 @@ void AddStripedBatch(
int batch,
Provider* provider);
// Compute the row-wise sum of a N*D matrix X, and write it to a N
// dimensional vector y.
template <typename T, class Provider>
void RowwiseSum(int N, int D, const T* x, T* y,
Provider* provider);
// Compute the column-wise sum of a N*D matrix X, and write it to a D
// dimensional vector y.
template <typename T, class Provider>
void ColwiseSum(int N, int D, const T* x, T* y,
Provider* provider);
// Compute the row-wise max of a N*D matrix X, and write it to a N
// dimensional vector y.
template <typename T, class Provider>

View file

@ -780,6 +780,26 @@ SPECIALIZED_REDUCEMAX(int64_t)
#undef SPECIALIZED_REDUCEMAX
#define SPECIALIZED_ROWWISESUM(T) \
template <> \
void RowwiseSum<T, CPUMathUtil>( \
const int N, const int D, const T* x, T* y, CPUMathUtil*) { \
EigenVectorMap<T>(y, N) = \
ConstEigenMatrixMap<T>(x, D, N).colwise().sum(); \
}
SPECIALIZED_ROWWISESUM(float)
#undef SPECIALIZED_ROWWISESUM
#define SPECIALIZED_COLWISESUM(T) \
template <> \
void ColwiseSum<T, CPUMathUtil>( \
const int N, const int D, const T* x, T* y, CPUMathUtil*) { \
EigenVectorMap<T>(y, D) = \
ConstEigenMatrixMap<T>(x, D, N).rowwise().sum(); \
}
SPECIALIZED_COLWISESUM(float)
#undef SPECIALIZED_COLWISESUM
#define SPECIALIZED_ROWWISEMAX(T) \
template <> \
void RowwiseMax<T, CPUMathUtil>( \