SameOperandsAndResultsScale,
    PredOpTrait<"input and output must have same element type",
      TFL_TCresVTEtIsSameAsOp<0, 0>>,
    TFL_OperandHasRankAtMost<0, 4>
  ]> {
  let summary = "SpaceToDepth operator";

  let description = [{
    Rearranges blocks of spatial data, into depth. More specifically,
    this op outputs a copy of the input tensor where values from the `height`

    ```mlir
    module {
      func @while_body {
        %input = "tf.IteratorGetNext"(...) {device = "/CPU:0"} -> tensor<2x224x224x3xf32>
        %space_to_depth = "tf.SpaceToDepth"(%input) {block_size = 2, ...}: (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32>
        %device_launch = "tf_device.cluster_func"(%space_to_depth,...) {func = @_func,...)
        return ...
      }

// CHECK-LABEL: argmax64
// CHECK:  "tfl.arg_max"(%arg0, %arg1) : (tensor<3xi32>, tensor<i32>) -> tensor<i64>
}

func.func @space_to_depth(%arg0: tensor<1x2x2x1xf32>) -> tensor<?xf32> {
  %0 = "tf.SpaceToDepth"(%arg0) {block_size = 2: i64,  data_format = "NHWC"}: (tensor<1x2x2x1xf32>) -> tensor<?xf32>
  func.return %0 : tensor<?xf32>

  // CHECK-LABEL: space_to_depth
  // CHECK: %[[ARG:.*]]: tensor<1x2x2x1xf32>