%7 = "tf.AddV2"(%2, %6) {device = ""} : (tensor<1x1024xbf16>, tensor<1x1024xbf16>) -> tensor<1x1024xbf16>
    %8 = "tf.Identity"(%4) {device = ""} : (tensor<i32>) -> tensor<i32>
    %9 = "tf.Identity"(%arg1) {device = ""} : (tensor<i32>) -> tensor<i32>
    func.return %8, %9, %1, %7, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<1x1024xbf16>, tensor<1024x1024xbf16>
  }

  %1 = "tf.BatchMatMulV2"(%0, %cst) {adj_x = false, adj_y = false, device = ""} : (tensor<1x1x10xbf16>, tensor<10x2xbf16>) -> tensor<1x1x2xbf16>
  %2 = "tf.Cast"(%1) {Truncate = false} : (tensor<1x1x2xbf16>) -> tensor<1x1x2xf32>
  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x1x2xf32>) -> tensor<1x1x2xf32>
  return %3 : tensor<1x1x2xf32>
}

// CHECK: func @cast_bf16_batch_matmul_v2_to_fp32

    %1 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x1024xi8>) -> tensor<2x1024xbf16>
    %2 = "tf.XlaDotV2"(%arg0, %1) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor<1x2x2x2xbf16>, tensor<2x1024xbf16>) -> tensor<1x2x2x1024xbf16>
    %3 = "tf.Identity"(%2) : (tensor<1x2x2x1024xbf16>) -> tensor<1x2x2x1024xbf16>
    return %3 : tensor<1x2x2x1024xbf16>

  // Move binary op batched RHS before reshape:
  // binary(reshape(lhs), rhs) => reshape(binary(lhs, flatten(rhs)))
  // Pattern targetted here is as follows-
  // [input, lhr, rhs] == [<1x1024x128>, <1x1024x8x16>, <1x1x8x16xf32>]
  // This is valid only when the-
  // 1.last dimension of lhs is equal to the number of elements in constant rhs.
  // 2.Reduded shape of rhs, here <8x16> is equal to last dimensions of lhs.