%cst_min = stablehlo.constant dense<0.0> : tensor<f32>
    %cst_max = stablehlo.constant dense<6.0> : tensor<f32>
    %7 = "stablehlo.clamp"(%cst_min, %6, %cst_max) {device = ""} : (tensor<f32>, tensor<1x2x2x2xf32>, tensor<f32>) -> tensor<1x2x2x2xf32>

  %1 = "tfl.mul"(%0, %arg2) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>

    // conversion, then dequantized during inference.
    // Activation: f32, Weight: qi8, Bias: f32
    WEIGHT_ONLY = 1;

    // Apply default dynamic range quantization. Quantized tensor value's
    // ranges are determined during graph runtime.
    // Activation: f32, Weight: qi8, Bias: f32
    POST_TRAINING_QUANTIZATION_DYNAMIC_RANGE = 2;

  %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

def FuseTransposeFCLhsToBatchMatmul : Pat<
  (TFL_BatchMatMulOp
    (TFL_FullyConnectedOp:$fc_output
      (TFL_TransposeOp TensorOf<[F32]>:$fc_lhs,
                       (Arith_ConstantOp:$perm_value $p0)),
      TensorOf<[F32]>:$fc_rhs,
      $bias, $TFL_AF_None, $TFL_FCWO_Default,
      $keep_num_dims, $asymmetric_quantize_inputs_fc
    ),
    $bmm_rhs, ConstBoolAttrTrue, $transpose_rhs,

    func.return %tls, %tls_2 : tensor<300x?xf32>, tensor<300x?xf32>
  }
}

// CHECK-LABEL: HloModule main
// CHECK:       ENTRY %main.{{[0-9]+}} () -> (f32[300,8], f32[300,9]) {
// CHECK:       %tuple.{{[0-9]+}} = (f32[300,8]{1,0}, f32[300,9]{1,0})

/// Examples:
/// 1. realValue is a primitive value attribute:
/// (realValue: FloatAttr, quantizedElementType: UniformQuantizedType[i8:f32])
///   -> (IntegerAttr, outConvertedType: i8)
/// 2. realValue is an elements attribute:
/// (realValue: DenseElementsAttr[tensor<2x2xf32>],
///  quantizedElementType: UniformQuantizedType[i8:f32])
///   -> (DenseElementsAttr[tensor<2x2xi8>], outConvertedType: tensor<2x2xi8>)
Attribute quantizeAttr(Attribute realValue,

  // CHECK: device = "gpu"
  %1 = "tf.MatMul"(%arg0, %0) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
  // CHECK: device = "cpu"
  %2 = "tf.Relu"(%1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "cpu"} : (tensor<3x3xf32>) -> tensor<3x3xf32>

  %y, %batch_mean, %batch_var, %reserve_1, %reserve_2, %reserve_3
    = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg1, %arg1, %arg1)
       {
         data_format = "NHWC",
         epsilon = 1.001 : f32,
         exponential_avg_factor = 1.0 : f32,
         is_training = true
       }
        : (tensor<1x28x28x64xf32>, tensor<64xf32>, tensor<64xf32>,
           tensor<64xf32>, tensor<64xf32>)

    });
  }
};

// This helper function makes legality check easier. Both convert ops in the
// patterns below are considered legal:
//  - `BitcastConvertOp` (i32 -> f32) + `ConvertOp` (f32 -> bf16)
//  - `ConvertOp` (bf16 -> f32) -> `BitcastConvertOp` (f32 -> i32)
template <typename ConvertOp, typename OtherConvertOp>
bool IsConvertOpLegal(ConvertOp convert_op, BFloat16TypeConverter& converter) {