| 109 | |
| 110 | template <typename T, typename Context> |
| 111 | void SumCsrGradKernel(const Context& dev_ctx, |
| 112 | const SparseCsrTensor& x, |
| 113 | const SparseCsrTensor& dout, |
| 114 | const IntArray& axis, |
| 115 | bool keep_dim UNUSED, |
| 116 | SparseCsrTensor* dx) { |
| 117 | EmptyLikeCsrKernel<T, Context>(dev_ctx, x, dx); |
| 118 | unsigned int n_dim = axis.size(); |
| 119 | |
| 120 | const DenseTensor& x_crows = x.crows(); |
| 121 | const DenseTensor& x_cols = x.cols(); |
| 122 | const DenseTensor& dout_values = dout.values(); |
| 123 | const auto* x_crows_data = x_crows.data<int64_t>(); |
| 124 | |
| 125 | DenseTensor* dx_crows = dx->mutable_crows(); |
| 126 | DenseTensor* dx_cols = dx->mutable_cols(); |
| 127 | DenseTensor* dx_values = dx->mutable_values(); |
| 128 | |
| 129 | *dx_crows = x_crows; |
| 130 | *dx_cols = x_cols; |
| 131 | |
| 132 | funcs::SetConstant<Context, T> set_constant; |
| 133 | if (n_dim == 0) { |
| 134 | T value = dout_values.data<T>()[0]; |
| 135 | set_constant(dev_ctx, dx_values, value); |
| 136 | if (dx_values->dtype() != dx->dtype()) { |
| 137 | *dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype()); |
| 138 | } |
| 139 | return; |
| 140 | } |
| 141 | PADDLE_ENFORCE_EQ(axis[0], |
| 142 | -1, |
| 143 | common::errors::Unimplemented( |
| 144 | "`axis` of SumCsrKernel only support None or -1 now." |
| 145 | "More number will be supported in the future.")); |
| 146 | |
| 147 | if (x.dims().size() == 2) { |
| 148 | int value_index = 0; |
| 149 | for (int k = 0; k < x.dims()[0]; ++k) { |
| 150 | if (x_crows_data[k] == x_crows_data[k + 1]) { |
| 151 | continue; |
| 152 | } |
| 153 | T value = dout_values.data<T>()[value_index]; |
| 154 | set_constant(dev_ctx, dx_values, value); |
| 155 | value_index += 1; |
| 156 | } |
| 157 | } else { |
| 158 | int dout_value_index = 0; |
| 159 | int dx_value_index = 0; |
| 160 | for (auto batch = 0; batch < x.dims()[0]; ++batch) { |
| 161 | for (auto k = batch * (x.dims()[1] + 1); |
| 162 | k < batch * (x.dims()[1] + 1) + x.dims()[1]; |
| 163 | ++k) { |
| 164 | if (x_crows_data[k] == x_crows_data[k + 1]) { |
| 165 | continue; |
| 166 | } |
| 167 | T value = dout_values.data<T>()[dout_value_index]; |
| 168 | for (auto i = x_crows_data[k]; i < x_crows_data[k + 1]; ++i) { |
nothing calls this directly
no test coverage detected