| 22 | namespace { |
| 23 | |
| 24 | void BatchToSpace(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& crops) { |
| 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 | crops.shape().rank() == 2 && |
| 43 | block_rank == xla::ShapeUtil::GetDimension(crops.shape(), 0) && |
| 44 | 2 == xla::ShapeUtil::GetDimension(crops.shape(), 1), |
| 45 | errors::InvalidArgument("crops should have shape [", block_rank, |
| 46 | ", 2] instead of ", |
| 47 | xla::ShapeUtil::HumanString(crops.shape()))); |
| 48 | |
| 49 | const int64 batch_size = input_shape[0]; |
| 50 | |
| 51 | // Compute the product of the block_shape values. |
| 52 | int64 block_num_elems = 1; |
| 53 | for (int i = 0; i < block_rank; ++i) { |
| 54 | block_num_elems *= block_shape[i]; |
| 55 | } |
| 56 | OP_REQUIRES(ctx, block_num_elems > 0, |
| 57 | errors::InvalidArgument( |
| 58 | "The product of the block dimensions must be positive")); |
| 59 | |
| 60 | // 1. Reshape `input` to `reshaped` of shape: |
| 61 | // [block_shape[0], ..., block_shape[M-1], |
| 62 | // batch / prod(block_shape), |
| 63 | // input_shape[1], ..., input_shape[N-1]] |
| 64 | |
| 65 | OP_REQUIRES( |
| 66 | ctx, batch_size % block_num_elems == 0, |
| 67 | errors::InvalidArgument("Input batch dimension (", batch_size, |
| 68 | ") is not divisible by product of block sizes (", |
| 69 | block_num_elems, ")")); |
| 70 | std::vector<int64> reshaped_shape(input_rank + block_rank); |
| 71 | std::copy(block_shape.begin(), block_shape.end(), reshaped_shape.begin()); |
| 72 | reshaped_shape[block_rank] = batch_size / block_num_elems; |
| 73 | std::copy(input_shape.begin() + 1, input_shape.end(), |
| 74 | reshaped_shape.begin() + block_rank + 1); |
| 75 | xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape); |
| 76 | |
| 77 | // 2. Permute dimensions of `reshaped` to produce `permuted` of shape |
| 78 | // [batch / prod(block_shape), |
| 79 | // |
| 80 | // input_shape[1], block_shape[0], |
| 81 | // ..., |