diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index c51e9f4d026..7db3539fcef 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -820,6 +820,10 @@ followed by cropping along the `height` and `width` dimensions.
TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
TF_DerivedOperandTypeAttr Tidx = TF_DerivedOperandTypeAttr<1>;
+
+ let verifier = [{
+ return Verify(*this);
+ }];
}
def TF_BatchToSpaceNDOp : TF_Op<"BatchToSpaceND", [NoSideEffect]> {
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
index f4e5dc05eb0..8410929a19f 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
@@ -695,6 +695,149 @@ void BatchMatMulV2Op::getCanonicalizationPatterns(
results.insert<BatchMatMulV2ToMatMul>(context);
}
+//===----------------------------------------------------------------------===//
+// BatchToSpaceOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult Verify(BatchToSpaceOp op) {
+ // Op already has a constraint that block_size >= 2.
+ int64_t block_size = op.block_size().getSExtValue();
+
+ llvm::SmallVector<int64_t, 4> input_shape(4, ShapedType::kDynamicSize);
+ auto input_type = op.input().getType().cast<TensorType>();
+ if (input_type.hasRank()) {
+ if (input_type.getRank() != 4)
+ return op.emitOpError()
+ << "requires input to be a 4D tensor, but got " << input_type;
+
+ int64_t input_batch = input_type.getDimSize(0);
+ if (input_batch != ShapedType::kDynamicSize &&
+ input_batch % (block_size * block_size) != 0) {
+ return op.emitOpError()
+ << "requires input batch (dimension 0) to be evenly divisible "
+ "by (block_size * block_size), but got input batch "
+ << input_batch << " and block_size " << block_size;
+ }
+
+ input_shape.assign(input_type.getShape().begin(),
+ input_type.getShape().end());
+ }
+
+ auto crops_type = op.crops().getType().cast<TensorType>();
+ if (crops_type.hasRank()) {
+ if (crops_type.getRank() != 2)
+ return op.emitOpError()
+ << "requires crops to be a 2D tensor, but got " << crops_type;
+
+ auto dim_of_size = [&](int64_t dim, int64_t size) {
+ if (crops_type.isDynamicDim(dim)) return true;
+ return crops_type.getDimSize(dim) == size;
+ };
+ if (!dim_of_size(0, 2) || !dim_of_size(1, 2))
+ return op.emitOpError()
+ << "requires crops to be a tensor<2x2>, but got " << crops_type;
+ }
+
+ DenseIntElementsAttr crops_attr;
+ // Crops are defined as [[crop_top, crop_bottom], [crop_left, crop_right]],
+ // and flattened as [crop_top, crop_bottom, crop_left, crop_right]
+ llvm::SmallVector<int64_t, 4> crops_values;
+ if (matchPattern(op.crops(), m_Constant(&crops_attr))) {
+ assert(crops_attr.getNumElements() == 4 &&
+ "tf.BatchToSpace crops must have 4 elements");
+
+ auto crops_range = crops_attr.getIntValues();
+ for (const auto &crops_value : crops_range) {
+ int64_t crops_value_int = crops_value.getSExtValue();
+ if (crops_value_int < 0)
+ return op.emitOpError()
+ << "requires all crop values to be nonnegative, but got "
+ << crops_attr;
+
+ crops_values.push_back(crops_value_int);
+ }
+ }
+
+ auto output_type = op.output().getType().cast<TensorType>();
+ if (output_type.hasRank()) {
+ if (output_type.getRank() != 4)
+ return op.emitOpError()
+ << "requires output to be a 4D tensor, but got " << output_type;
+
+ auto static_dims = [](int64_t dim_a, int64_t dim_b) {
+ return dim_a != ShapedType::kDynamicSize &&
+ dim_b != ShapedType::kDynamicSize;
+ };
+
+ auto output_shape = output_type.getShape();
+
+ // output batch = input batch / (block_size * block_size).
+ int64_t input_batch = input_shape[0];
+ int64_t output_batch = output_shape[0];
+ if (static_dims(input_batch, output_batch) &&
+ (output_batch * block_size * block_size) != input_batch)
+ return op.emitOpError()
+ << "requires output batch (dimension 0) to be equal to input "
+ "batch (dimension 0) / (block_size * block_size), but got "
+ "output batch "
+ << output_batch << ", input batch " << input_batch
+ << ", and block_size " << block_size;
+
+ auto check_spatial_dim = [&](int64_t spatial_dim_index,
+ llvm::StringRef dim_name,
+ llvm::StringRef crop_a_name,
+ llvm::StringRef crop_b_name) -> LogicalResult {
+ int64_t input_dim = input_shape[spatial_dim_index];
+ int64_t output_dim = output_shape[spatial_dim_index];
+ if (!static_dims(input_dim, output_dim)) return success();
+
+ int64_t input_dim_pad = input_dim * block_size;
+ // If crops are unknown, the maximum output spatial dim size is input
+ // spatial dim size * block_size, as crops can be minimum 0.
+ if (crops_values.empty() && output_dim > input_dim * block_size)
+ return op.emitOpError()
+ << "requires output " << dim_name << " (dimension "
+ << spatial_dim_index << ") to be less than or equal to input "
+ << dim_name << " (dimension " << spatial_dim_index
+ << ") * block_size, but got output " << dim_name << " "
+ << output_dim << ", input " << dim_name << " " << input_dim
+ << ", and block_size " << block_size;
+
+ if (!crops_values.empty()) {
+ // output spatial dim = input spatial dim * block_size - crops.
+ int64_t crop_a = crops_values[2 * (spatial_dim_index - 1)];
+ int64_t crop_b = crops_values[2 * (spatial_dim_index - 1) + 1];
+ if (output_dim != input_dim_pad - crop_a - crop_b)
+ return op.emitOpError()
+ << "requires output " << dim_name << " (dimension "
+ << spatial_dim_index << ") to be equal to input " << dim_name
+ << " (dimension " << spatial_dim_index << ") * block_size - "
+ << crop_a_name << " - " << crop_b_name << ", but got output "
+ << dim_name << " " << output_dim << ", input " << dim_name
+ << " " << input_dim << ", " << crop_a_name << " " << crop_a
+ << ", " << crop_b_name << " " << crop_b << ", and block_size "
+ << block_size;
+ }
+
+ return success();
+ };
+
+ if (failed(check_spatial_dim(1, "height", "crop_top", "crop_bottom")) ||
+ failed(check_spatial_dim(2, "width", "crop_left", "crop_right")))
+ return failure();
+
+ int64_t input_depth = input_shape[3];
+ int64_t output_depth = output_shape[3];
+ if (static_dims(input_depth, output_depth) && output_depth != input_depth)
+ return op.emitOpError()
+ << "requires output depth (dimension 3) to be equal to input "
+ "depth (dimension 3), but got output depth "
+ << output_depth << " and input depth " << input_depth;
+ }
+
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// BiasAddOp
//===----------------------------------------------------------------------===//
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir
index 058585b41d7..4ba2e83300b 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir
@@ -2872,4 +2872,141 @@ func @testSendTPUEmbeddingGradients(%x: tensor<512x256xf32>) {
return
}
+// -----
+//===--------------------------------------------------------------------===//
+// tf.BatchToSpace
+//===--------------------------------------------------------------------===//
+
+func @testBatchToSpaceDynamic(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) {
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
+ return
+}
+
+func @testBatchToSpaceRankedInput(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<*xi32>) {
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<?x?x?x?xf32>, tensor<*xi32>) -> tensor<*xf32>
+ return
+}
+
+func @testBatchToSpaceRankedCrops(%arg0: tensor<*xf32>, %arg1: tensor<?x?xi32>) {
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<?x?xi32>) -> tensor<*xf32>
+ return
+}
+
+func @testBatchToSpaceRankedOutput(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) {
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<?x?x?x?xf32>
+ return
+}
+
+func @testBatchToSpaceStatic(%arg0: tensor<36x8x8x8xf32>) {
+ %crops = "tf.Const"() {value = dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+ %0 = "tf.BatchToSpace"(%arg0, %crops) {block_size = 3 : i64} : (tensor<36x8x8x8xf32>, tensor<2x2xi32>) -> tensor<4x21x17x8xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidInputRank(%arg0: tensor<8xf32>, %arg1: tensor<*xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires input to be a 4D tensor, but got 'tensor<8xf32>'}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<8xf32>, tensor<*xi32>) -> tensor<*xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidInputBatch(%arg0: tensor<2x4x6x8xf32>, %arg1: tensor<*xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires input batch (dimension 0) to be evenly divisible by (block_size * block_size), but got input batch 2 and block_size 2}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<2x4x6x8xf32>, tensor<*xi32>) -> tensor<*xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidCropsRank(%arg0: tensor<*xf32>, %arg1: tensor<?x?x?xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires crops to be a 2D tensor, but got 'tensor<?x?x?xi32>'}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<?x?x?xi32>) -> tensor<*xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidCropsFirstDim(%arg0: tensor<*xf32>, %arg1: tensor<3x?xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires crops to be a tensor<2x2>, but got 'tensor<3x?xi32>'}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<3x?xi32>) -> tensor<*xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidCropsSecondDim(%arg0: tensor<*xf32>, %arg1: tensor<?x3xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires crops to be a tensor<2x2>, but got 'tensor<?x3xi32>'}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<?x3xi32>) -> tensor<*xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceBadCropValues(%arg0: tensor<*xf32>) {
+ %crops = "tf.Const"() {value = dense<[[-1, -2], [-3, -4]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+ // expected-error @+1 {{'tf.BatchToSpace' op requires all crop values to be nonnegative, but got dense<[[-1, -2], [-3, -4]]> : tensor<2x2xi32>}}
+ %0 = "tf.BatchToSpace"(%arg0, %crops) {block_size = 2 : i64} : (tensor<*xf32>, tensor<2x2xi32>) -> tensor<*xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputRank(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires output to be a 4D tensor, but got 'tensor<8xf32>'}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<8xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputBatch(%arg0: tensor<16x8x8x3xf32>, %arg1: tensor<*xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires output batch (dimension 0) to be equal to input batch (dimension 0) / (block_size * block_size), but got output batch 8, input batch 16, and block_size 2}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<*xi32>) -> tensor<8x8x8x3xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputHeight(%arg0: tensor<16x8x8x3xf32>, %arg1: tensor<*xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires output height (dimension 1) to be less than or equal to input height (dimension 1) * block_size, but got output height 17, input height 8, and block_size 2}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<*xi32>) -> tensor<4x17x8x3xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputHeightCrops(%arg0: tensor<16x8x8x3xf32>) {
+ %crops = "tf.Const"() {value = dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+ // expected-error @+1 {{'tf.BatchToSpace' op requires output height (dimension 1) to be equal to input height (dimension 1) * block_size - crop_top - crop_bottom, but got output height 8, input height 8, crop_top 1, crop_bottom 2, and block_size 2}}
+ %0 = "tf.BatchToSpace"(%arg0, %crops) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<2x2xi32>) -> tensor<4x8x9x3xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputWidth(%arg0: tensor<16x4x4x3xf32>, %arg1: tensor<*xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires output width (dimension 2) to be less than or equal to input width (dimension 2) * block_size, but got output width 9, input width 4, and block_size 2}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x4x4x3xf32>, tensor<*xi32>) -> tensor<4x4x9x3xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputWidthCrops(%arg0: tensor<16x8x8x3xf32>) {
+ %crops = "tf.Const"() {value = dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+ // expected-error @+1 {{'tf.BatchToSpace' op requires output width (dimension 2) to be equal to input width (dimension 2) * block_size - crop_left - crop_right, but got output width 8, input width 8, crop_left 3, crop_right 4, and block_size 2}}
+ %0 = "tf.BatchToSpace"(%arg0, %crops) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<2x2xi32>) -> tensor<4x13x8x3xf32>
+ return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputDepth(%arg0: tensor<16x8x8x3xf32>, %arg1: tensor<*xi32>) {
+ // expected-error @+1 {{'tf.BatchToSpace' op requires output depth (dimension 3) to be equal to input depth (dimension 3), but got output depth 8 and input depth 3}}
+ %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<*xi32>) -> tensor<4x8x8x8xf32>
+ return
+}