Function CalculateGradData

tensorflow/compiler/tf2xla/kernels/resampler_ops.cc:248–329

Build computation the backprop into input 'data'. Where input: grad_output is of dimension [batch, dim_0, ...dim_n, channel] ratio is of dimension [batch, dim_0, ...dim_n, 2] gather_indices is of dimension [batch, dim_0, ...dim_n, 3] data_shape is of dimension [batch, x(width), y(height), channel] Output: scatter-add to each 2x2 grad_data neighbor: grad_data[fx, fy, chan] += output_grad * dx * dy

// where (dx, dy) is (1 - ratio). If (dx, dy) is out of bound, then the their
// contribution is 0 to 'grad_data'.
XlaOp CalculateGradData(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio,
                        XlaOp gather_indices, XlaOp warp,
                        xla::PrimitiveType warp_type, TensorShape warp_shape,
                        int64 last_warp_dim, int64 data_channels,
                        xla::Shape data_shape) {
  // Weights tensor has dimension [batch, dim_0, ... dim_n, 4].
  auto weights = BilinearWeights(ctx, ratio, warp_shape, warp_type);

  auto warp_dims = warp_shape.dim_sizes();
  std::vector<int64> warp_dims_without_last_dims(warp_dims.begin(),
                                                 warp_dims.end() - 1);

  std::vector<int64> reshaped_weights_dims = warp_dims_without_last_dims;
  // Reshape the last dimension of size 4 to two dimensions [2, 2].
  reshaped_weights_dims.push_back(2);
  reshaped_weights_dims.push_back(2);
  std::vector<int64> reshape_dims(warp_shape.dims());
  std::iota(reshape_dims.begin(), reshape_dims.end(), 0);
  // The dimension is [batch, dim_0,..., dim_n, 2, 2].
  auto reshaped_weights = xla::Reshape(weights, /*dimensions=*/reshape_dims,
                                       /*new_sizes=*/reshaped_weights_dims);

  std::vector<int64> weights_with_channels_dims = reshaped_weights_dims;
  weights_with_channels_dims.push_back(data_channels);
  std::vector<int64> reshaped_weights_indices(reshaped_weights_dims.size());
  std::iota(reshaped_weights_indices.begin(), reshaped_weights_indices.end(),
            0);

  // Set out of bound weights to 0.
  // The dimension of the reshaped_weight: [batch, dim_0, ...dim_n, 2, 2].
  std::vector<int64> reshaped_result_dims(warp_dims.begin(),
                                          warp_dims.end() - 1);
  reshaped_result_dims.push_back(2);
  reshaped_result_dims.push_back(2);
  std::vector<int64> broadcasted_dims(warp_dims.size() - 1);
  std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0);
  reshaped_weights = BoundSamples(ctx, warp, warp_type, warp_shape,
                                  reshaped_result_dims, broadcasted_dims,
                                  last_warp_dim, data_shape, reshaped_weights);

  // The dimension is [batch, dim_0, ..., dim_n, 2, 2, data_channel].
  auto broadcast_reshaped_weights = xla::BroadcastInDim(
      reshaped_weights, weights_with_channels_dims, reshaped_weights_indices);

  std::vector<int64> grad_output_indices(warp_dims_without_last_dims.size());
  std::iota(grad_output_indices.begin(), grad_output_indices.end(), 0);
  grad_output_indices.push_back(weights_with_channels_dims.size() - 1);
  XlaOp broadcast_grad_output = xla::BroadcastInDim(
      grad_output, weights_with_channels_dims, grad_output_indices);

  auto grad_output_multiply_weights =
      broadcast_grad_output * broadcast_reshaped_weights;

  auto grad_data = xla::ConstantLiteral(
      ctx->builder(), xla::Literal::CreateFromShape(data_shape));

  // Pad grad data then slice it back.
  //