%cst = arith.constant dense<[1, 2]> : tensor<2xi32>
  %0 = "tfl.mean"(%arg0, %cst) {keep_dims = false} : (tensor<1x48x48x512xf32>, tensor<2xi32>) -> tensor<1x512xf32>
  func.return %0 : tensor<1x512xf32>
}

// CHECK:       func @mean_4d_no_keepdim([[VAL_0:%.*]]: tensor<1x48x48x512xf32>) -> tensor<1x512xf32> {
// CHECK:           [[VAL_1:%.*]] = "tfl.pseudo_const"(){{.*}}dense<[1, 512]> : tensor<2xi32>

  %2 = "tfl.add"(%arg0, %arg3) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  %3 = "tfl.pack"(%1, %2) {axis = 0 : i32, per_device_costs = {CPU = 2.0 : f32, GPU = -1.0 : f32}, values_count = 2 : i32, tac.device = "CPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>

func.func @main(%arg0 : tensor<!tf_type.variant<tensor<2xi32>>>, %arg1: tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<2xi32>>> {
  func.return %arg0 : tensor<!tf_type.variant<tensor<2xi32>>>
}

// CHECK:          func.func @main(%[[ARG0:.*]]: tensor<!tf_type.variant<tensor<2xi32>>>, %[[ARG1:.*]]: tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<2xi32>>>

  // CHECK: %{{.*}} = "tfl.reshape"(%{{.*}}, %[[cst]]) : (tensor<4xf32>, tensor<2xi32>) -> tensor<2x2xf32>
  %cst = arith.constant dense<[2, 2]> : tensor<2xi32>
  %0 = "tfl.pseudo_const" () {value = dense<1.0> : tensor<4xf32>} : () -> tensor<4xf32> loc("Const")
  %1 = "tfl.reshape" (%0, %cst) : (tensor<4xf32>, tensor<2xi32>) -> tensor<2x2xf32> loc("reshape")

      } : (tensor<!tf_type.resource<tensor<2xi32>>>) -> tensor<*xi32>
    return %1 : tensor<*xi32>
  }

  func.func private @called_by_batch_func_assign(%arg: tensor<*x!tf_type.resource>) -> tensor<?xi32> {
    %0 = "tf.Const"() {value = dense<4> : tensor<2xi32>} : () -> tensor<2xi32>
    "tf.AssignVariableOp"(%arg, %0) : (tensor<*x!tf_type.resource>, tensor<2xi32>) -> ()

      %0  = "tf.Const"() {value = dense<[56, 904]> : tensor<2xi32>} : () -> tensor<2xi32>
      // CHECK: %{{.*}} = stablehlo.custom_call @tf.ResizeBilinear(%arg0, %{{.*}}) {align_corners = false, device = "", half_pixel_centers = true} : (tensor<1x32x32x128xf32>, tensor<2xi32>) -> tensor<1x64x64x128xf32>

  } : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
  %1 = "tf.UniformDequantize"(%0, %scales, %zps) {
    quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
  } : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
  func.return %1 : tensor<1xf32>
}

// -----

  %cst_0 = "tf.Const"() {value = dense<0.0235294122> : tensor<f32>} : () -> tensor<f32>
  %cst_1 = "tf.Const"() {value = dense<0.00708661414> : tensor<1xf32>} : () -> tensor<1xf32>
  %cst_2 = "tf.Const"() {value = dense<1.799000e+03> : tensor<1xf32>} : () -> tensor<1xf32>

// -----

func.func @foldDivI32() -> tensor<2xi32> {
  %0 = stablehlo.constant dense<[9, 3]> : tensor<2xi32>
  %1 = stablehlo.constant dense<[4, 6]> : tensor<2xi32>
  %2 = stablehlo.divide %0, %1 : tensor<2xi32>
  return %2 : tensor<2xi32>
}

// CHECK-LABEL: foldDivI32
// CHECK: stablehlo.constant dense<[2, 0]> : tensor<2xi32>

// -----

func.func @main(%arg0: tensor<32x4x4x128xf32>, %arg1: tensor<1x32x42x128xf32>, %arg2: tensor<4xi32>) -> tensor<1x64x84x32xf32> {
  %0 = "tfl.custom"(%arg0, %arg1, %arg2) {custom_code = "Convolution2DTransposeBias", custom_option = #tfl<const_bytes : "0x010000000200000002000000">} : (tensor<32x4x4x128xf32>, tensor<1x32x42x128xf32>, tensor<4xi32>) -> tensor<1x64x84x32xf32>
  func.return %0 : tensor<1x64x84x32xf32>
}
// CHECK-LABEL: main