Function BoundSamples

tensorflow/compiler/tf2xla/kernels/resampler_ops.cc:195–231

Bounds samples to 0 if the warp image indices are out of the (-1, image_size) bound. The resulting dimension is given by 'result_dims'.

// bound.
// The resulting dimension is given by 'result_dims'.
XlaOp BoundSamples(XlaOpKernelContext* ctx, XlaOp warp,
                   xla::PrimitiveType warp_type, TensorShape warp_shape,
                   std::vector<int64> result_dims,
                   std::vector<int64> broadcasted_dims, int64 last_warp_dim,
                   xla::Shape data_shape, XlaOp sample) {
  auto is_gt_minus_one =
      xla::Gt(warp,
              xla::ConvertElementType(
                  xla::ConstantR1<float>(ctx->builder(), {-1, -1}), warp_type),
              /*broadcast_dimensions=*/{warp_shape.dims() - 1});
  auto is_lt_image_size = xla::Lt(
      warp,
      xla::ConvertElementType(
          xla::ConstantR1<float>(
              ctx->builder(),
              {/*width=*/static_cast<float>(data_shape.dimensions(2)),
               /*height=*/static_cast<float>(data_shape.dimensions(1))}),
          warp_type),
      /*broadcast_dimensions=*/{warp_shape.dims() - 1});

  auto is_in_bound_padded_x_y = xla::And(is_gt_minus_one, is_lt_image_size);
  // Reduce along last dimension. The resulting dimension is:
  // [batch, dim_0, ...dim_n].
  auto is_in_bound = xla::Reduce(
      is_in_bound_padded_x_y, xla::ConstantR0<bool>(ctx->builder(), true),
      xla::CreateScalarAndComputation(xla::PrimitiveType::PRED, ctx->builder()),
      {last_warp_dim});

  // Broadcast 'is_in_bound' to the same dimension as 'result_dims'.
  auto broadcasted_is_in_bound =
      xla::BroadcastInDim(is_in_bound, result_dims, broadcasted_dims);

  // Set out of bound samples to zero.
  auto zeros =
      xla::Broadcast(xla::Zero(ctx->builder(), warp_type), result_dims);
  return xla::Select(broadcasted_is_in_bound, sample, zeros);
}

// Build computation the backprop into input 'data'.
// Where input: