// 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>

    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r) {
      return ArraysAreCastCompatible(l, r);
    }
  }];
}

def TF_SpaceToDepthOp : TF_Op<"SpaceToDepth", [Pure]> {
  let summary = "SpaceToDepth for tensors of type T.";

  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`