}

// -----

module {
  func.func @propagate_xladotv2_bf16(%arg0: tensor<1x2x2x2xbf16>) -> tensor<1x2x2x1024xbf16> {
    %cst = "tf.Const"() {value = dense<127> : tensor<2x1024xi8>} : () -> tensor<2x1024xi8>
    %0 = "tf.Identity"(%cst) : (tensor<2x1024xi8>) -> tensor<2x1024xi8>

}

// -----

module {
  func.func @quantize_xladotv2_bf16(%arg0: tensor<1x2x2x2xbf16>) -> (tensor<1x2x2x1024xbf16>) {
    %cst_0 = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x1024xbf16>} : () -> tensor<2x1024xbf16>
    %0 = "tf.XlaDotV2"(%arg0, %cst_0) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor<1x2x2x2xbf16>, tensor<2x1024xbf16>) -> tensor<1x2x2x1024xbf16>
    // Check dequantize performed in bf16.

  %2 = "tf.BiasAdd"(%1, %cst) {data_format = "NHWC", device = ""} : (tensor<1x3x2x2xbf16>, tensor<2xbf16>) -> tensor<1x3x2x2xbf16>
  %3 = "tf.Cast"(%2) {Truncate = false} : (tensor<1x3x2x2xbf16>) -> tensor<1x3x2x2xf32>
  %4 = "tf.IdentityN"(%3) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>

  %1 = "tf.Conv2D"(%0, %cst) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 2, 1]}
       : (tensor<1x3x4x3xbf16>, tensor<2x3x3x2xbf16>) -> tensor<1x3x2x2xbf16> loc("Model/conv2d")
  %2 = "tf.Cast"(%1) {Truncate = false} : (tensor<1x3x2x2xbf16>) -> tensor<1x3x2x2xf32>
  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
  return %3 : tensor<1x3x2x2xf32>
// CHECK: tf.Conv2D

  %6 = "tf.TPUReplicatedOutput"(%5) {device = ""} : (tensor<1x3x2x2xbf16>) -> tensor<1x3x2x2xbf16>
  %7 = "tf.Cast"(%6) {Truncate = false} : (tensor<1x3x2x2xbf16>) -> tensor<1x3x2x2xf32>
  func.return %7 : tensor<1x3x2x2xf32>
}

// CHECK: func @tpu_conv(%[[ARG0:.*]]: tensor<1x3x4x3xf32>)