/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators/elementwise_sub_op.h (2025B)
#ifndef CAFFE2_OPERATORS_ELEMENTWISE_SUB_OP_H_
#define CAFFE2_OPERATORS_ELEMENTWISE_SUB_OP_H_
#include
#include
#include
#include "caffe2/operators/elementwise_ops.h"
#include "caffe2/operators/elementwise_ops_utils.h"
#include "caffe2/utils/math.h"
namespace caffe2 {
template
struct SubFunctor {
explicit SubFunctor(bool allow_broadcast_fastpath=false)
: allow_broadcast_fastpath_(allow_broadcast_fastpath) {}
template
bool Forward(
const std::vector& A_dims,
const std::vector& B_dims,
const TIn* A,
const TIn* B,
TOut* C,
Context* context) const {
math::Sub(
A_dims.size(),
A_dims.data(),
B_dims.size(),
B_dims.data(),
A,
B,
C,
context);
return true;
}
template
bool Backward(
const std::vector& A_dims,
const std::vector& B_dims,
const TGrad* dC,
const TIn* /* A */,
const TIn* /* B */,
const TOut* /* C */,
TGrad* dA,
TGrad* dB,
Context* context) const {
const std::vector C_dims =
elementwise_ops_utils::ComputeBinaryBroadcastForwardDims(
A_dims, B_dims);
std::vector A_back_dims;
std::vector B_back_dims;
elementwise_ops_utils::ComputeBinaryBroadcastBackwardDims(
A_dims, B_dims, &A_back_dims, &B_back_dims);
math::ReduceSum(
C_dims.size(),
C_dims.data(),
A_back_dims.data(),
TGrad(1),
dC,
dA,
context,
allow_broadcast_fastpath_);
math::ReduceSum(
C_dims.size(),
C_dims.data(),
B_back_dims.data(),
TGrad(-1),
dC,
dB,
context,
allow_broadcast_fastpath_);
return true;
}
const bool allow_broadcast_fastpath_;
};
} // namespace caffe2
#endif // CAFFE2_OPERATORS_ELEMENTWISE_SUB_OP_H_