pytorch/caffe2/operators/elementwise_div_op.h
Xiaomeng Yang 9243b64bff
[Caffe2] Update elementwise ops to support numpy style boradcast (#8070)
* Update elementwise ops to support numpy style boradcast

Update elementwise ops to support numpy style boradcast

* Fix sqrt_op

* Fix compare ops

* Fix gradient test

* Fix optimizer legacy broadcast

* Fix legacy broadcast for elementwise ops

* Skip flaky test

* Fix eigen simple binary op

* Fix attention test

* Fix rnn test

* Fix LSTM test

* Fix tan grad

* Fix schema check
2018-06-05 15:49:16 -07:00

49 lines
1.0 KiB
C++

#ifndef CAFFE2_OPERATORS_ELEMENTWISE_DIV_OP_H_
#define CAFFE2_OPERATORS_ELEMENTWISE_DIV_OP_H_
#include <vector>
#include "caffe2/operators/elementwise_ops.h"
#include "caffe2/utils/math.h"
namespace caffe2 {
template <class Context>
struct DivFunctor {
template <typename TIn, typename TOut>
bool Forward(
const std::vector<int>& A_dims,
const std::vector<int>& B_dims,
const TIn* A,
const TIn* B,
TOut* C,
Context* context) const {
math::Div(
A_dims.size(),
A_dims.data(),
B_dims.size(),
B_dims.data(),
A,
B,
C,
context);
return true;
}
template <typename TGrad, typename TIn, typename TOut>
bool Backward(
const std::vector<int>& A_dims,
const std::vector<int>& B_dims,
const TGrad* dC_data,
const TIn* A_data,
const TIn* B_data,
const TOut* C_data,
TGrad* dA_data,
TGrad* dB_data,
Context* context) const;
};
} // namespace caffe2
#endif // CAFFE2_OPERATORS_ELEMENTWISE_DIV_OP_H_