}

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
  // CHECK: "tfl.space_to_batch_nd"(%arg0, %arg1, %arg2) : (tensor<1x4x4x3xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<*xf32>
}

  %cst = arith.constant dense<[1, -1, 4]> : tensor<3xi32>
  // CHECK: "tfl.reshape"(%arg0, %cst)
  %0 = "tfl.reshape"(%arg0, %cst) : (tensor<3x4xi32>, tensor<3xi32>) -> tensor<1x3x4xi32>
  func.return %0 : tensor<1x3x4xi32>
}

// -----

func.func @invalid_two_dynamic_dims_on_reshape(%arg0: tensor<3x4xi32>) -> tensor<1x3x4xi32> {

  %0 = stablehlo.constant() {value = dense<3> : tensor<2x2x1x3xi8>} : () -> tensor<2x2x1x3x!quant.uniform<i8:f32:3, {3.000000e+00, 3.000000e+00, 3.000000e+00}>>
  %1 = stablehlo.constant() {value = dense<5> : tensor<1x1x1x3xi32>} : () -> tensor<1x1x1x3x!quant.uniform<i32:f32:3, {2.000000e+00, 2.000000e+00, 2.000000e+00}>>

module {
// expected-warning @+1 {{Invalid number of results from MaxUnpooling2D}}
func.func private @max_unpooling_2d_invalid_results(%arg0: tensor<1x1x2x1xf32>, %arg1: tensor<1x1x2x1xi32>) -> (tensor<1x2x4x1xf32>, tensor<1x2x4x1xi32>) attributes {tf._implements = #tf_type.func<@"addons:MaxUnpooling2D", {padding = "SAME", pool_size = [2, 2], strides = [2, 2]}>}

// expected-warning @+1 {{Invalid number of arguments to MaxUnpooling2D}}