| 26 | |
| 27 | template <typename T, typename IntT, typename Context> |
| 28 | void SumCooGradCPUKernel(const Context& dev_ctx, |
| 29 | const SparseCooTensor& x, |
| 30 | const SparseCooTensor& dout, |
| 31 | const IntArray& axis, |
| 32 | bool keep_dim, |
| 33 | SparseCooTensor* dx) { |
| 34 | EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx); |
| 35 | unsigned int n_dim = axis.size(); |
| 36 | |
| 37 | const DenseTensor& x_indices = x.indices(); |
| 38 | const DenseTensor& dout_indices = dout.indices(); |
| 39 | const DenseTensor& dout_values = dout.values(); |
| 40 | const auto* dout_indices_data = dout_indices.data<int64_t>(); |
| 41 | const auto* dout_values_data = dout_values.data<T>(); |
| 42 | |
| 43 | DenseTensor* dx_indices = dx->mutable_indices(); |
| 44 | DenseTensor* dx_values = dx->mutable_values(); |
| 45 | *dx_indices = x_indices; |
| 46 | |
| 47 | const auto* dx_indices_data = dx_indices->data<int64_t>(); |
| 48 | auto* dx_values_data = dx_values->data<T>(); |
| 49 | |
| 50 | funcs::SetConstant<Context, T> set_constant; |
| 51 | if (n_dim == 0) { |
| 52 | T value = dout_values.data<T>()[0]; |
| 53 | set_constant(dev_ctx, dx_values, value); |
| 54 | if (dx_values->dtype() != dx->dtype()) { |
| 55 | *dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype()); |
| 56 | } |
| 57 | return; |
| 58 | } |
| 59 | |
| 60 | auto dim = axis[0] < 0 ? x.dims().size() + axis[0] : axis[0]; |
| 61 | auto sparse_dim = x.sparse_dim(); |
| 62 | if (dim >= sparse_dim) { |
| 63 | dim = dim - sparse_dim + 1; |
| 64 | phi::ReduceSumGradKernel<T, Context>( |
| 65 | dev_ctx, x.values(), dout.values(), {dim}, keep_dim, false, dx_values); |
| 66 | if (dx_values->dtype() != dx->dtype()) { |
| 67 | *dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype()); |
| 68 | } |
| 69 | return; |
| 70 | } |
| 71 | // Ensure the sparse_dim is not less than 1. |
| 72 | if (sparse_dim == 1) { |
| 73 | keep_dim = true; |
| 74 | } |
| 75 | |
| 76 | int64_t dense_dim = 1; |
| 77 | for (auto i = 1; i < x.values().dims().size(); ++i) { |
| 78 | dense_dim *= x.values().dims()[i]; |
| 79 | } |
| 80 | |
| 81 | std::map<std::vector<IntT>, int64_t> indices_map; |
| 82 | for (auto j = 0; j < dout_indices.dims()[1]; ++j) { |
| 83 | std::vector<IntT> pos; |
| 84 | pos.reserve(dout_indices.dims()[0]); |
| 85 | for (int i = 0; i < dout_indices.dims()[0]; ++i) { |
nothing calls this directly
no test coverage detected