| 22 | |
| 23 | template <typename T, typename Context> |
| 24 | void SoftmaxGradKernel(const Context& dev_ctx, |
| 25 | const DenseTensor& out, |
| 26 | const DenseTensor& out_grad, |
| 27 | int axis, |
| 28 | DenseTensor* x_grad) { |
| 29 | using XPUType = typename XPUTypeTrait<T>::Type; |
| 30 | const int rank = x_grad->dims().size(); |
| 31 | const int calc_axis = funcs::CanonicalAxis(axis, rank); |
| 32 | |
| 33 | // allocate memory on device. |
| 34 | dev_ctx.template Alloc<T>(x_grad); |
| 35 | if (x_grad->numel() == 0) { |
| 36 | return; |
| 37 | } |
| 38 | |
| 39 | // For 0D Tensor |
| 40 | if (rank == 0) { |
| 41 | funcs::set_constant(dev_ctx, x_grad, static_cast<T>(0.0)); |
| 42 | return; |
| 43 | } |
| 44 | |
| 45 | std::vector<int64_t> x_dims; |
| 46 | for (int i = 0; i < rank; i++) { |
| 47 | x_dims.push_back(x_grad->dims()[i]); |
| 48 | } |
| 49 | |
| 50 | int r = xpu::softmax_grad<XPUType>( |
| 51 | dev_ctx.x_context(), |
| 52 | reinterpret_cast<const XPUType*>(out.data<T>()), |
| 53 | reinterpret_cast<const XPUType*>(out_grad.data<T>()), |
| 54 | reinterpret_cast<XPUType*>(x_grad->data<T>()), |
| 55 | x_dims, |
| 56 | static_cast<int64_t>(calc_axis)); |
| 57 | PADDLE_ENFORCE_XDNN_SUCCESS(r, "softmax_grad"); |
| 58 | } |
| 59 | |
| 60 | } // namespace phi |
| 61 |
nothing calls this directly
no test coverage detected