Scatter 'updates' tensor to 'grad_data' based on 'indices'. Returns the resulting tensor of dimension: [batch, dim_0, ...dim_n, 2, 2, data_channels]. This function can also be seen as the inverse of 'Gather2by2Neighbors'.
| 165 | // resulting tensor of dimension: [batch, dim_0, ...dim_n, 2, 2, data_channels]. |
| 166 | // This function can also be seen as the inverse of 'Gather2by2Neighbors'. |
| 167 | XlaOp ScatterToGradData(XlaOpKernelContext* ctx, XlaOp grad_data, XlaOp indices, |
| 168 | XlaOp updates, int64 warp_dims, |
| 169 | xla::PrimitiveType xla_type) { |
| 170 | xla::ScatterDimensionNumbers scatter_dim_numbers; |
| 171 | const int64 neighbor_data_dimensions = warp_dims + 2; |
| 172 | // Since the Scatter output dimensions are [batch, dim_0, ... dim_n, 2, 2, |
| 173 | // data_channels], the update window dimensions is the last 3 dimensions. |
| 174 | scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 3); |
| 175 | scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 2); |
| 176 | scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 1); |
| 177 | scatter_dim_numbers.set_index_vector_dim(warp_dims - 1); |
| 178 | |
| 179 | scatter_dim_numbers.add_inserted_window_dims(0); |
| 180 | scatter_dim_numbers.add_scatter_dims_to_operand_dims(0); |
| 181 | // Since input is of dimension [batch, height(y), width(x), channel], and warp |
| 182 | // is of dimension [batch, x, y], the ordering of x, y here needs to be |
| 183 | // swapped when scattering. |
| 184 | scatter_dim_numbers.add_scatter_dims_to_operand_dims(2); |
| 185 | scatter_dim_numbers.add_scatter_dims_to_operand_dims(1); |
| 186 | |
| 187 | return xla::Scatter(grad_data, indices, updates, |
| 188 | xla::CreateScalarAddComputation(xla_type, ctx->builder()), |
| 189 | scatter_dim_numbers); |
| 190 | } |
| 191 | |
| 192 | // Bounds samples to 0 if the warp image indices are out of the (-1, image_size) |
| 193 | // bound. |
no test coverage detected