%8 = stablehlo.convert %7 : (tensor<1x2xi32>) -> tensor<1x2xf32>
    %9 = stablehlo.convert %2 : (tensor<2x3xi8>) -> tensor<2x3xf32>
    %10 = stablehlo.dot_general %8, %9, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
    %11 = stablehlo.convert %3 : (tensor<1x3xi32>) -> tensor<1x3xf32>

  %11 = "tf.Reshape"(%10, %0) {device = ""} : (tensor<4x1x1x2x1xi32>, tensor<2xi32>) -> tensor<4x2xi32>
  %12 = "tf.Transpose"(%11, %6) {device = ""} : (tensor<4x2xi32>, tensor<2xi32>) -> tensor<2x4xi32>
  %13 = "tf.Reshape"(%arg0, %1) {device = ""} : (tensor<1x1x2x1xf32>, tensor<1xi32>) -> tensor<2xf32>
  %14 = "tf.ScatterNd"(%12, %13, %3) {device = ""} : (tensor<2x4xi32>, tensor<2xf32>, tensor<4xi32>) -> tensor<1x2x4x1xf32>

}

func.func @addN(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>, %arg2: tensor<2x3xi32>) -> tensor<2x3xi32> {
  %0 = "tf.AddN"(%arg0, %arg1, %arg2) : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
  func.return %0 : tensor<2x3xi32>

// CHECK-LABEL: addN
// CHECK:  "tfl.add_n"(%arg0, %arg1, %arg2) : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
// CHECK:  return
}

  %2 = "tf.AddN"(%arg0, %1, %1) : (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  %3 = "tf.AddN"(%1, %arg0, %1) : (tensor<2xi32>, tensor<2xi32> , tensor<2xi32>) -> tensor<2xi32>
  %4 = "tf.AddN"(%1, %1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  %5 = "tf.AddN"(%arg0, %1, %0) : (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %2, %3, %4, %5: tensor<2xi32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>
}

  %0 = "tf_device.cluster_func"(%arg0, %arg1) {func = @_func, use_spmd_for_xla_partitioning = true, use_tpu = true, num_cores_per_replica = 2 : i64} : (tensor<2x4xf32>, tensor<4x2xf32>) -> tensor<2x2xf32>
  %1:2 = "tf.TPUPartitionedOutputV2"(%0) {device = "", partition_dims = [2, 1]} : (tensor<2x2xf32>) -> (tensor<1x2xf32>, tensor<1x2xf32>)

  %1 = "tfl.reshape"(%0, %cst3) : (tensor<4x2xf32>, tensor<2xi32>) -> tensor<1x8xf32>
  %2 = "tfl.mul"(%0, %cst2) {fused_activation_function = "RELU6"} : (tensor<4x2xf32>, tensor<2xf32>) -> tensor<4x2xf32>

  func.return %1, %2 : tensor<1x8xf32>, tensor<4x2xf32>

  // CHECK-NEXT:     [[SCALAR:%.*]] = "tfl.pseudo_const"() <{value = dense<10> : tensor<i32>}> : () -> tensor<i32>
  // CHECK-NEXT:     [[ADD:%.*]] = tfl.add([[SCALAR]], [[SUB]]) {fused_activation_function = "NONE"} : (tensor<i32>, tensor<3x2xi32>) -> tensor<3x2xi32>
  // CHECK-NEXT:     return [[ADD]] : tensor<3x2xi32>

  %0 = "tfl.pseudo_const" () {value = dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>} : () -> tensor<3x2xi32> loc("Const")

func.func @main(%arg0: tensor<3x2xf32>) -> tensor<3x2xi32> {
  // CHECK-LABEL: @main
  // CHECK: "tfl.bucketize"(%arg0) <{boundaries = [0.000000e+00 : f32, 1.000000e+01 : f32, 1.000000e+02 : f32]}> : (tensor<3x2xf32>) -> tensor<3x2xi32>
  %0 = "tfl.bucketize"(%arg0) {boundaries = [0.0 : f32, 10.0 : f32, 100.0 : f32]} : (tensor<3x2xf32>) -> tensor<3x2xi32>
  func.return %0 : tensor<3x2xi32>

  // CHECK: %[[RES_INT_2:.*]] = mhlo.convert %[[RES_INT_1]] : tensor<3x2xi32>
  %result_int = "tf.Cast"(%result) {Truncate = false} : (tensor<3x2x!tf_type.qint32>) -> tensor<3x2xi32>
  // CHECK: return %[[RES_INT_2]] : tensor<3x2xi32>
  func.return %result_int : tensor<3x2xi32>

    %12 = "tfl.logistic"(%11) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %13 = tfl.mul %arg4, %12 {fused_activation_function = "NONE"} : tensor<4x2xf32>
    %14 = "tfl.relu"(%10#1) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %15 = "tfl.logistic"(%10#0) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %16 = tfl.mul %15, %14 {fused_activation_function = "NONE"} : tensor<4x2xf32>