mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-18 21:21:17 +00:00
Introduce Rowwise/Colwise Sum to math util (#656)
This commit is contained in:
parent
600dc9ecc5
commit
25f45cb2db
2 changed files with 32 additions and 0 deletions
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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>( \
|
||||
|
|
|
|||
Loading…
Reference in a new issue