return %2 : tensor<1x300x10xf32>
  }

  func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> {
    %cst = "tf.Const"() {value = dense<0.0787401571> : tensor<f32>} : () -> tensor<f32>
    %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32>
    %1 = "tf.Mul"(%0, %cst) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
    return %1 : tensor<*xf32>
  }

// CHECK: return %[[ORIGINAL_CAST]] : tensor<?x?x?x?xbf16>

// CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32>
// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0787401571> : tensor<f32>
}

// -----

module {
  func.func @quantize_sharded_weights_with_xladot_with_identity(%arg0: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {