// CHECK: %0 = "quantfork.qcast"(%[[CONST]]) : (tensor<2x1024xf32>) -> tensor<2x1024x!quant.uniform<i8<-127:127>:f32, 3.9370078740157481E-9>>
// CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x1024x!quant.uniform<i8<-127:127>:f32, 3.9370078740157481E-9>>) -> tensor<2x1024xf32>

// CHECK-SAME: effective_hidden_scale_intermediate = tensor<*x!quant.calibrated<f32<-5.000000e-01:5.000000e-01>>>
// CHECK-SAME: input_to_cell_intermediate = tensor<*x!quant.calibrated<f32<-4.000000e+00:4.000000e+00>>>
// CHECK-SAME: input_to_forget_intermediate = tensor<*x!quant.calibrated<f32<-1.600000e+01:1.600000e+01>>>
// CHECK-SAME: input_to_input_intermediate = tensor<*x!quant.calibrated<f32<-3.200000e+01:3.200000e+01>>>

// CHECK-NEXT: %[[Y:[a-z0-9]*]], {{.*}}_FusedBatchNormEx
// CHECK-NEXT: return %[[Y]]
  %y:6 = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg2, %arg3, %arg4) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>)

// CHECK: %[[q0:.*]] = "quantfork.qcast"(%[[cst]]) {volatile}
// CHECK-SAME: tensor<2x!quant.uniform<i32:f32:0, {0.044169864606680966,0.042867627733627671}>>
// CHECK: %[[dq0:.*]] = "quantfork.dcast"(%[[q0]])

// CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[cst_1]]) {volatile}
// CHECK-SAME: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.075176584439014829,0.072960192762960605}>>
// CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]])

  %1 = "quantfork.dcast"(%0) : (tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0011764706057660721:-43>>) -> tensor<1x3x4x3xf32>
  %q_w = "quantfork.qcast"(%cst) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8:f32, 0.0125:-24>>
  %dq_w = "quantfork.dcast"(%q_w) : (tensor<2x3x3x2x!quant.uniform<i8:f32, 0.0125:-24>>) -> tensor<2x3x3x2xf32>

  %0 = corert.get_op_handler %arg0 "cpu"
  // CHECK: %[[RESULT:.*]] = corert.executeop(%[[ARG_0:.*]]) "tf.MatMul"(%[[ARG_1:.*]], %[[ARG_1]]) {T = f32, transpose_a = false, transpose_b = false} : 1
  %1 = corert.executeop(%0) "tf.MatMul"(%arg1, %arg1) {T = f32, device = "cpu", transpose_a = false, transpose_b = false} : 1
  tfrt.return %arg0, %1 : !tfrt.chain, !corert.tensorhandle

  %2 = corert.create_dense_tensor.i32 {shape = [0], value = []}
  // CHECK-NEXT: %[[RESULT_3:.*]] = corert.executeop(%[[RESULT_0]]) "tf.AddV2"(%[[RESULT_2]], %[[RESULT_2]])
  %3 = corert.executeop(%0) "tf.AddV2"(%2, %2) {T = f32, device = "/device:CPU:0"} : 1
  // CHECK-NEXT: %[[RESULT_4:.*]] = tfrt.get_device %[[ARG_0]] {device_name = "/device:GPU:0"}
  // CHECK-NEXT: %[[RESULT_5:.*]] = corert.get_dst_tensor_type %[[RESULT_3]], %[[RESULT_4]]

// CHECK: %0 = "quantfork.qcast"(%[[CONST]]) : (tensor<2x1024xf32>) -> tensor<2x1024x!quant.uniform<i8<-127:127>:f32, 3.9370078740157481E-9>>
// CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x1024x!quant.uniform<i8<-127:127>:f32, 3.9370078740157481E-9>>) -> tensor<2x1024xf32>

  %cpu = corert.get_op_handler %ch "cpu"
  %0 = corert.executeop(%cpu) "tf.Relu"(%arg0) { T = f32 } : 1
  %arg1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %arg1_th {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
  %1 = tfrt_fallback_async.executeop key(0) cost(100) device("/CPU:0") "tf.Relu"(%arg1) { T = f32 } : 1
  tfrt.return

    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "quant-cast-bf16-ops-to-f32";
  }
  StringRef getDescription() const final {
    return "Cast BF16 operations to F32.";
  }

  void runOnOperation() override;
};

class CastBf16OpsToF32 : public RewritePattern {
 public:
  explicit CastBf16OpsToF32(MLIRContext* context)