|
|
|
@ -41,77 +41,14 @@ class ElementwiseDivKernel : public framework::OpKernel<T> {
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
template <typename T>
|
|
|
|
|
struct ElementwiseDivGradFunctor {
|
|
|
|
|
template <typename Device, typename X, typename Y, typename Z, typename dX,
|
|
|
|
|
typename dY, typename dZ>
|
|
|
|
|
void operator()(Device d, X x, Y y, Z z, dX dx, dY dy, dZ dz) {
|
|
|
|
|
auto y_e = framework::EigenVector<T>::Flatten(*y);
|
|
|
|
|
auto z_e = framework::EigenVector<T>::Flatten(*z);
|
|
|
|
|
auto dz_e = framework::EigenVector<T>::Flatten(*dz);
|
|
|
|
|
|
|
|
|
|
if (dx) {
|
|
|
|
|
auto dx_e = framework::EigenVector<T>::Flatten(*dx);
|
|
|
|
|
dx_e.device(d) = dz_e / y_e;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if (dy) {
|
|
|
|
|
auto dy_e = framework::EigenVector<T>::Flatten(*dy);
|
|
|
|
|
dy_e.device(d) = -1.0 * dz_e * z_e / y_e;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
template <typename T>
|
|
|
|
|
struct ElementwiseDivBroadCastGradFunctor {
|
|
|
|
|
template <typename Device, typename X, typename Y, typename Z, typename dX,
|
|
|
|
|
typename dY, typename dZ, typename Pre, typename N>
|
|
|
|
|
void operator()(Device d, X x, Y y, Z z, dX dx, dY dy, dZ dz, Pre pre, N n) {
|
|
|
|
|
auto x_e = framework::EigenVector<T>::Flatten(*x);
|
|
|
|
|
auto y_e = framework::EigenVector<T>::Flatten(*y);
|
|
|
|
|
auto dz_e = framework::EigenVector<T>::Flatten(*dz);
|
|
|
|
|
|
|
|
|
|
auto y_e_bcast = y_e.reshape(Eigen::DSizes<int, 2>(1, n))
|
|
|
|
|
.broadcast(Eigen::DSizes<int, 2>(pre, 1))
|
|
|
|
|
.reshape(Eigen::DSizes<int, 1>(x_e.size()));
|
|
|
|
|
|
|
|
|
|
if (dx) {
|
|
|
|
|
auto dx_e = framework::EigenVector<T>::Flatten(*dx);
|
|
|
|
|
dx_e.device(d) = dz_e / y_e_bcast;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if (dy) {
|
|
|
|
|
auto dy_e = framework::EigenVector<T>::Flatten(*dy);
|
|
|
|
|
dy_e.device(d) = (-1.0 * (x_e * dz_e) / (y_e_bcast * y_e_bcast))
|
|
|
|
|
.reshape(Eigen::DSizes<int, 2>(pre, n))
|
|
|
|
|
.sum(Eigen::array<int, 1>{{0}});
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
struct DivGradDX {
|
|
|
|
|
HOSTDEVICE T operator()(T x, T y, T out, T dout) const { return dout / y; }
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
template <typename T>
|
|
|
|
|
struct ElementwiseDivBroadCast2GradFunctor {
|
|
|
|
|
template <typename Device, typename X, typename Y, typename Z, typename dX,
|
|
|
|
|
typename dY, typename dZ, typename Pre, typename N, typename Post>
|
|
|
|
|
void operator()(Device d, X x, Y y, Z z, dX dx, dY dy, dZ dz, Pre pre, N n,
|
|
|
|
|
Post post) {
|
|
|
|
|
auto x_e = framework::EigenVector<T>::Flatten(*x);
|
|
|
|
|
auto y_e = framework::EigenVector<T>::Flatten(*y);
|
|
|
|
|
auto dz_e = framework::EigenVector<T>::Flatten(*dz);
|
|
|
|
|
|
|
|
|
|
auto y_e_bcast = y_e.reshape(Eigen::DSizes<int, 3>(1, n, 1))
|
|
|
|
|
.broadcast(Eigen::DSizes<int, 3>(pre, 1, post))
|
|
|
|
|
.reshape(Eigen::DSizes<int, 1>(x_e.size()));
|
|
|
|
|
if (dx) {
|
|
|
|
|
auto dx_e = framework::EigenVector<T>::Flatten(*dx);
|
|
|
|
|
dx_e.device(d) = dz_e / y_e_bcast;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if (dy) {
|
|
|
|
|
auto dy_e = framework::EigenVector<T>::Flatten(*dy);
|
|
|
|
|
dy_e.device(d) = (-1.0 * (x_e * dz_e) / (y_e_bcast * y_e_bcast))
|
|
|
|
|
.reshape(Eigen::DSizes<int, 3>(pre, n, post))
|
|
|
|
|
.sum(Eigen::array<int, 2>{{0, 2}});
|
|
|
|
|
}
|
|
|
|
|
struct DivGradDY {
|
|
|
|
|
HOSTDEVICE T operator()(T x, T y, T out, T dout) const {
|
|
|
|
|
return -dout * x / (y * y);
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
@ -128,10 +65,8 @@ class ElementwiseDivGradKernel : public framework::OpKernel<T> {
|
|
|
|
|
auto* dx = ctx.Output<Tensor>(framework::GradVarName("X"));
|
|
|
|
|
auto* dy = ctx.Output<Tensor>(framework::GradVarName("Y"));
|
|
|
|
|
int axis = ctx.Attr<int>("axis");
|
|
|
|
|
ElementwiseGradCompute<DeviceContext, T, ElementwiseDivGradFunctor<T>,
|
|
|
|
|
ElementwiseDivBroadCastGradFunctor<T>,
|
|
|
|
|
ElementwiseDivBroadCast2GradFunctor<T>>(
|
|
|
|
|
ctx, x, y, out, dout, axis, dx, dy);
|
|
|
|
|
ElemwiseGradCompute<DeviceContext, T, DivGradDX<T>, DivGradDY<T>>(
|
|
|
|
|
ctx, *x, *y, *out, *dout, axis, dx, dy, DivGradDX<T>(), DivGradDY<T>());
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|