mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
9f9244aabe
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/62369 This diff is a big no-op that just sets up scaffolding for passing the "allow_broadcast_fastpath" from caffe2 operator protos created in Python down to C++. To facilitate this, we create helper template wrappers that pass a flag for "allow_broadcast_fastpath" down to elementwise functors. This flag will determine whether to try and take the broadcast fastpath, which we will add in subsequent diffs. Test Plan: sandcastle + let github CI run Differential Revision: D28154475 fbshipit-source-id: 15750a0bcd2994fbc6a61fb5653d8cae6b0177dd
115 lines
2.6 KiB
C++
115 lines
2.6 KiB
C++
#ifndef CAFFE2_UTILS_MATH_REDUCE_H_
|
|
#define CAFFE2_UTILS_MATH_REDUCE_H_
|
|
|
|
#include "caffe2/core/common.h"
|
|
#include "caffe2/core/types.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
class Tensor;
|
|
|
|
namespace math {
|
|
|
|
template <typename T, class Context>
|
|
TORCH_API void
|
|
ReduceMin(const int N, const T* X, T* y, Tensor* scratch_ptr, Context* context);
|
|
|
|
template <typename T, class Context>
|
|
TORCH_API void
|
|
ReduceMax(const int N, const T* X, T* y, Tensor* scratch_ptr, Context* context);
|
|
|
|
// In all of the reduce functions, X_dims and Y_dims should have ndim elements.
|
|
// Each dimension of Y_dims must match the corresponding dimension of X_dims or
|
|
// must be equal to 1. The dimensions equal to 1 indicate the dimensions of X to
|
|
// be reduced.
|
|
|
|
// Y = alpha * ReduceMin(X)
|
|
template <typename T, class Context>
|
|
TORCH_API void ReduceMin(
|
|
const int ndim,
|
|
const int* X_dims,
|
|
const int* Y_dims,
|
|
const T alpha,
|
|
const T* X,
|
|
T* Y,
|
|
Context* context,
|
|
bool allow_broadcast_fastpath=false);
|
|
|
|
// Y = alpha * ReduceMax(X)
|
|
template <typename T, class Context>
|
|
TORCH_API void ReduceMax(
|
|
const int ndim,
|
|
const int* X_dims,
|
|
const int* Y_dims,
|
|
const T alpha,
|
|
const T* X,
|
|
T* Y,
|
|
Context* context,
|
|
bool allow_broadcast_fastpath=false);
|
|
|
|
// Y = alpha * ReduceSum(X)
|
|
template <typename T, class Context>
|
|
TORCH_API void ReduceSum(
|
|
const int ndim,
|
|
const int* X_dims,
|
|
const int* Y_dims,
|
|
const T alpha,
|
|
const T* X,
|
|
T* Y,
|
|
Context* context,
|
|
bool allow_broadcast_fastpath=false);
|
|
|
|
// Y = alpha * ReduceMean(X)
|
|
template <typename T, class Context>
|
|
TORCH_API void ReduceMean(
|
|
const int ndim,
|
|
const int* X_dims,
|
|
const int* Y_dims,
|
|
const T alpha,
|
|
const T* X,
|
|
T* Y,
|
|
Context* context,
|
|
bool allow_broadcast_fastpath=false);
|
|
|
|
// Y = alpha * ReduceL1(X)
|
|
template <typename T, class Context>
|
|
TORCH_API void ReduceL1(
|
|
const int ndim,
|
|
const int* X_dims,
|
|
const int* Y_dims,
|
|
const T alpha,
|
|
const T* X,
|
|
T* Y,
|
|
Context* context,
|
|
bool allow_broadcast_fastpath=false);
|
|
|
|
// Y = alpha * ReduceL2(X)
|
|
template <typename T, class Context>
|
|
TORCH_API void ReduceL2(
|
|
const int ndim,
|
|
const int* X_dims,
|
|
const int* Y_dims,
|
|
const T alpha,
|
|
const T* X,
|
|
T* Y,
|
|
Context* context,
|
|
bool allow_broadcast_fastpath=false);
|
|
|
|
// Computes mean and variance over axes.
|
|
template <typename T, class Context>
|
|
TORCH_API void Moments(
|
|
const int ndims,
|
|
const int* X_dims,
|
|
const int* Y_dims,
|
|
const T* X,
|
|
T* mean,
|
|
T* var,
|
|
Context* context,
|
|
bool allow_broadcast_fastpath=false);
|
|
|
|
} // namespace math
|
|
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_UTILS_MATH_REDUCE_H_
|