| 22 | namespace { |
| 23 | |
| 24 | void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp& input, |
| 25 | DataType input_dtype, const TensorShape& input_tensor_shape, |
| 26 | absl::Span<const int64> block_shape, |
| 27 | const xla::Literal& paddings) { |
| 28 | const int input_rank = input_tensor_shape.dims(); |
| 29 | const absl::InlinedVector<int64, 4> input_shape = |
| 30 | input_tensor_shape.dim_sizes(); |
| 31 | const int block_rank = block_shape.size(); |
| 32 | |
| 33 | OP_REQUIRES( |
| 34 | ctx, input_rank >= 1 + block_rank, |
| 35 | errors::InvalidArgument("input rank should be >= ", 1 + block_rank, |
| 36 | " instead of ", input_rank)); |
| 37 | absl::Span<const int64> remainder_shape(input_shape); |
| 38 | remainder_shape.remove_prefix(1 + block_rank); |
| 39 | |
| 40 | OP_REQUIRES( |
| 41 | ctx, |
| 42 | paddings.shape().rank() == 2 && |
| 43 | block_rank == xla::ShapeUtil::GetDimension(paddings.shape(), 0) && |
| 44 | 2 == xla::ShapeUtil::GetDimension(paddings.shape(), 1), |
| 45 | errors::InvalidArgument("paddings should have shape [", block_rank, |
| 46 | ", 2] instead of ", |
| 47 | xla::ShapeUtil::HumanString(paddings.shape()))); |
| 48 | |
| 49 | xla::XlaBuilder* b = ctx->builder(); |
| 50 | |
| 51 | // 1. Zero-pad the start and end of dimensions `[1, ..., M]` of the |
| 52 | // input according to `paddings` to produce `padded` of shape `padded_shape`. |
| 53 | xla::PaddingConfig padding_config; |
| 54 | std::vector<int64> padded_shape(input_shape.begin(), input_shape.end()); |
| 55 | int64 block_num_elems = 1LL; |
| 56 | padding_config.add_dimensions(); // Don't pad the batch dimension. |
| 57 | for (int i = 0; i < block_rank; ++i) { |
| 58 | auto* dim = padding_config.add_dimensions(); |
| 59 | int64 pad_start = paddings.Get<int64>({i, 0}); |
| 60 | int64 pad_end = paddings.Get<int64>({i, 1}); |
| 61 | OP_REQUIRES(ctx, pad_start >= 0 && pad_end >= 0, |
| 62 | errors::InvalidArgument("Paddings must be non-negative")); |
| 63 | dim->set_edge_padding_low(pad_start); |
| 64 | dim->set_edge_padding_high(pad_end); |
| 65 | padded_shape[1 + i] += pad_start + pad_end; |
| 66 | block_num_elems *= block_shape[i]; |
| 67 | } |
| 68 | // Don't pad the remainder dimensions. |
| 69 | for (int i = 0; i < remainder_shape.size(); ++i) { |
| 70 | padding_config.add_dimensions(); |
| 71 | } |
| 72 | OP_REQUIRES(ctx, block_num_elems > 0, |
| 73 | errors::InvalidArgument( |
| 74 | "The product of the block dimensions must be positive")); |
| 75 | |
| 76 | xla::XlaOp padded = |
| 77 | xla::Pad(input, XlaHelpers::Zero(b, input_dtype), padding_config); |
| 78 | |
| 79 | // 2. Reshape `padded` to `reshaped_padded` of shape: |
| 80 | // |
| 81 | // [batch] + |