| 46 | |
| 47 | template <typename T, typename Context> |
| 48 | void DistGradKernel(const Context& dev_ctx, |
| 49 | const DenseTensor& x, |
| 50 | const DenseTensor& y, |
| 51 | const DenseTensor& out, |
| 52 | const DenseTensor& out_grad, |
| 53 | float p, |
| 54 | DenseTensor* x_grad, |
| 55 | DenseTensor* y_grad) { |
| 56 | if ((!x_grad) && (!y_grad)) { |
| 57 | return; |
| 58 | } |
| 59 | |
| 60 | if ((x_grad && x_grad->numel() == 0) || (y_grad && y_grad->numel() == 0)) { |
| 61 | if (x_grad) { |
| 62 | dev_ctx.template Alloc<T>(x_grad); |
| 63 | if (x_grad->numel() != 0) { |
| 64 | Full<T, Context>(dev_ctx, x_grad->dims(), 0, x_grad); |
| 65 | } |
| 66 | } |
| 67 | if (y_grad) { |
| 68 | dev_ctx.template Alloc<T>(y_grad); |
| 69 | if (y_grad->numel() != 0) { |
| 70 | Full<T, Context>(dev_ctx, y_grad->dims(), 0, y_grad); |
| 71 | } |
| 72 | } |
| 73 | return; |
| 74 | } |
| 75 | |
| 76 | auto t = Subtract<T, Context>(dev_ctx, x, y); |
| 77 | DenseTensor x_grad_tmp; |
| 78 | x_grad_tmp.Resize(t.dims()); |
| 79 | DenseTensor y_grad_tmp; |
| 80 | y_grad_tmp.Resize(t.dims()); |
| 81 | PNormGradKernel<T, Context>( |
| 82 | dev_ctx, t, out, out_grad, p, -1, 1e-12, false, true, &x_grad_tmp); |
| 83 | |
| 84 | if (x_grad) { |
| 85 | // do reduce, the implementation of cpu SumKernel has bug, it changes |
| 86 | // the dims of output internally, so we Resize x/y_grad twice. |
| 87 | auto res_x = GetReduceDims(x_grad_tmp.dims(), x.dims()); |
| 88 | if (!std::get<0>(res_x).empty()) { |
| 89 | x_grad->Resize(std::get<1>(res_x)); |
| 90 | SumKernel<T, Context>( |
| 91 | dev_ctx, x_grad_tmp, std::get<0>(res_x), x.dtype(), false, x_grad); |
| 92 | x_grad->Resize(x.dims()); |
| 93 | } else { |
| 94 | x_grad->ShareBufferWith(x_grad_tmp); |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | if (y_grad) { |
| 99 | ScaleKernel<T, Context>(dev_ctx, x_grad_tmp, -1.0, 0.0, false, &y_grad_tmp); |
| 100 | auto res_y = GetReduceDims(y_grad_tmp.dims(), y.dims()); |
| 101 | if (!std::get<0>(res_y).empty()) { |
| 102 | y_grad->Resize(std::get<1>(res_y)); |
| 103 | SumKernel<T, Context>( |
| 104 | dev_ctx, y_grad_tmp, std::get<0>(res_y), y.dtype(), false, y_grad); |
| 105 | y_grad->Resize(y.dims()); |
nothing calls this directly
no test coverage detected