%1 = "tfl.add"(%cst_2, %cst_1) {fused_activation_function = "NONE"} : (tensor<2x2x2xf32>, tensor<  2x2xf32>) -> tensor<2x2x2xf32>
  %2 = "tfl.add"(%cst_0, %cst_2) {fused_activation_function = "NONE"} : (tensor<    2xf32>, tensor<2x2x2xf32>) -> tensor<2x2x2xf32>

  func.return %0, %1, %2 : tensor<2x2xf32>, tensor<2x2x2xf32>, tensor<2x2x2xf32>

  constexpr absl::string_view kProgram = R"mlir(
func.func @main(
    %arg0: tensor<10x10xf32>, %scale: tensor<10xf32>, %zp: tensor<10xi32>
  ) -> tensor<10x10xi8> {
  %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) {
    quantization_axis = 1 : i64,
    quantization_min_val = -128 : i64,
    quantization_max_val = 127 : i64
  } : (tensor<10x10xf32>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10x!tf_type.qint8>

  %1 = "tf.Identity"(%0) : (tensor<2x2xf32>) -> tensor<2x2xf32>
  return %1 : tensor<2x2xf32>
}

// -----
// The following op sharding is used in the following test case:
// Proto debug string:
//   type: MAXIMAL
//   tile_assignment_dimensions: 1
//   tile_assignment_devices: 0
// Serialized string:
//   "\08\01\1A\01\01\22\01\00"
// Proto debug string:
//   type: MAXIMAL

    %cst = arith.constant dense<[2, 100]> : tensor<2xi64>
    %0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<200xf32>
    %1 = "tfl.reshape"(%0, %cst) : (tensor<200xf32>, tensor<2xi64>) -> tensor<2x100xf32>
    func.return %1 : tensor<2x100xf32>
  }

// CHECK: [[UNPACK:%.*]]:3 = "tf.Unpack"(%arg0) <{axis = 0 : i64}> : (tensor<3x2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>, tensor<2x2xf32>)
// CHECK: [[SCALAR_ZERO:%.*]] = arith.constant dense<0> : tensor<i32>
// CHECK: [[CONCAT:%.*]] = "tf.Concat"([[SCALAR_ZERO]], [[UNPACK]]#0, [[UNPACK]]#1, [[UNPACK]]#2) : (tensor<i32>, tensor<2x2xf32>, tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<?x2xf32>
// CHECK: [[LENGTHS:%.*]] = arith.constant dense<0> : tensor<0xi64>

          tensor<f32>, tensor<i32>) -> tensor<2x!tf_type.qint32>
    func.return %1 : tensor<2x!tf_type.qint32>
}

// -----

// CHECK-LABEL: func @while_region_qint
func.func @while_region_qint(%arg0: tensor<2x2xf32>) -> (tensor<2x?xf32>, tensor<?x2xf32>) {
  %scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
  %zps2 = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>

// CHECK-NOT: quantfork.qcast
// CHECK-NOT: quantfork.dcast

// -----

// Test that q/dq pair is not inserted when constant has multiple users.

func.func @no_qdq_for_multiple_users(%arg0: tensor<2x2xf32>) -> tensor<2x3xf32> attributes {tf._original_func_name = "main_0"} {
  %cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
  %0 = "tf.XlaCallModule"(%arg0, %cst) {

%arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %arg18, %arg19, %arg20, %arg21, %arg22, %arg23) <{fused_activation_function = "NONE", time_major = false}> : (tensor<600x?x20xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<40xf32>, tensor<?x?xf32>,...