%cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = ""} : () -> tensor<i32>
  %0 = "tf.MultiDeviceIteratorFromStringHandle"(%arg0) <{output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]}> {device = ""} : (tensor<!tf_type.string>) -> tensor<!tf_type.resource>

  %cst = "tf.Const"() {device = "", value = dense<[2.167590e-01]> : tensor<1xf32>} : () -> tensor<1xf32>
  %cst_1 = "tf.Const"() {device = "", value = dense<1.000000e-03> : tensor<f32>} : () -> tensor<f32>
  %0 = "tf.Identity"(%cst) {device = ""} : (tensor<1xf32>) -> tensor<1xf32>

// the quantization parameters as a TypeAttr and "quant.dcast" op used to
// convert the output type to the next op. Here are the transformations:
//
// input   min cst       max cst              input
//  \       |             |                     |
//   \  (tf.Identity) (tf.Identity)   =>   quant.qcast
//    \     |             |                     |

    // CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<3> : tensor<i32>}> : () -> tensor<i32>
    // CHECK-NEXT: %[[UQ:.*]] = "tf.UniformQuantize"(%arg0, %cst, %cst_0) <{quantization_axis = -1 : i64, quantization_max_val = 127 : i64, quantization_min_val = -128 : i64}> : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>

// CHECK-LABEL: FuseTransposeFCRhsToBatchMatmul
func.func @FuseTransposeFCRhsToBatchMatmul(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128xf32>, %arg2: none) -> tensor<16x128xf32> {
  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tfl.transpose"(%arg1, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>
  // CHECK: "tfl.batch_matmul"(%arg0, %arg1)

// CHECK: %[[conv:.*]] = "tfl.conv_2d"(%arg0, %arg1, %arg2)
// CHECK-SAME: -> tensor<1x112x112x32x!quant.uniform<u8:f32, 0.0078431372549019607:128>>
// CHECK: %[[cst:.*]] = "tfl.pseudo_qconst"()
// CHECK: %[[add:.*]] = tfl.add(%[[conv]], %[[cst]])
// CHECK-SAME: -> tensor<1x112x112x32x!quant.uniform<u8:f32, 1.000000e+00>>
// CHECK: return %[[add]]
}

// CHECK-LABEL: test_conv_2d_activation_and_bias

  %0 = "tf.EmptyTensorList"(%arg0, %arg1) : (tensor<?xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi64>>>
  // CHECK: %cst = arith.constant dense<0> : tensor<i32>
  // CHECK: %0 = "tfl.custom"(%arg0, %cst) <{custom_code = "TensorListReserve", custom_option = #tfl<const_bytes : "0x04">}> : (tensor<?xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi64>>>

     }) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
     "tf_device.launch"() ({
       %cst = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
       "tf.OpA"(%0#1) {mini_batch_splits = ""} : (tensor<3x!tf_type.string>) -> ()
       tf_device.return

  // CHECK-SAME:      <[-1.000000e+00, 1.000000e+00]> : tensor<2xf32>}>
  // CHECK-SAME:      : (tensor<40x37xf32>) -> tensor<40x37xf32>
  // CHECK-NEXT:    %[[cst:.*]] = "tfl.pseudo_const"() <{value = dense<
  // CHECK-SAME:      1.000000e+00> : tensor<40xf32>}> : () -> tensor<40xf32>
  // CHECK-NEXT:    %[[fc:.*]]:2 = "tfl.fully_connected"(%[[stat]], %arg1,

      /*infer_tensor_range=*/true, /*legacy_float_scale=*/false,
      /*is_qdq_conversion=*/false);

  quantization_driver.Initialize();

  int64_t num_constant_op = 0;
  main_fn.walk([&](arith::ConstantOp cst) { ++num_constant_op; });
  EXPECT_EQ(num_constant_op, 4);
}

TEST_F(ApplyQuantizationParamsPropagationTest,
       PropagateParamsCreatesQuantState) {
  const OwningOpRef<ModuleOp> module_op_ref =