// CHECK: "tfl.cast"
  // CHECK: "tfl.batch_to_space_nd"
}

func.func @space_to_batch_nd(%arg0: tensor<1x4x4x3xf32>, %arg1: tensor<2xi32>, %arg2: tensor<2x2xi32>) -> tensor<*xf32> {
  %0 = "tf.SpaceToBatchND"(%arg0, %arg1, %arg2) : (tensor<1x4x4x3xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
  // CHECK-LABEL: space_to_batch_nd

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

  let description = [{
    This operation reshapes space dimensions into the "batch" dimension 0
  }];

  let arguments = (ins