// CHECK-SAME: %[[ARG0:.*]]: tensor<3x4x5xf32>
  // CHECK-SAME: %[[ARG1:.*]]: tensor<3x5x2xf32>
  // CHECK-SAME: %[[ARG2:.*]]: tensor<2x3x2xi64>
  func.func private @composite_and_gather(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x2xf32>, %arg2: tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32> {
    // CHECK: %[[Q1:.*]] = "quantfork.qcast"(%[[ARG0]]) {volatile} : (tensor<3x4x5xf32>) -> tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>

  } : (tensor<2x7xi64>) -> (tensor<2x7xi32>)
  %1 = "tf.Identity"(%0) : (tensor<2x7xi32>) -> tensor<2x7xi32>
  func.return %1 : tensor<2x7xi32>
}

// CHECK-LABEL: func @func
// CHECK-SAME: %arg0: tensor<2x7xi64> {mhlo.sharding = "\08\03\1A\03\01\01\02\22\02\00\01B\01\00"
func.func @func(%arg0: tensor<2x7xi64>) -> (tensor<2x7xi32>) {

    %11 = "tf.Pad"(%arg0, %10) : (tensor<2x224x224x3xf32>, tensor<4x2xi32>) -> tensor<2x230x230x3xf32>
    %12 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2x1xf32>) -> tensor<2x1xi64>
    %13 = "tf.Reshape"(%12, %9) : (tensor<2x1xi64>, tensor<1xi32>) -> tensor<2xi64>
    %14 = "tf.Squeeze"(%arg1) {squeeze_dims = [-1]} : (tensor<2x1xf32>) -> tensor<2xf32>
    // CHECK: "tf.Conv2D"
    // CHECK-SAME: strides = [1, 1, 1, 1]

// -----

// Test that q/dq pair is not inserted between constant and XlaCallModule op
// whose entry function name does not include conv nor dot_general.

func.func @no_qdq_except_conv_and_dot_general(%arg0: tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32> attributes {tf._original_func_name = "main_0"} {
  %cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<3x4x2xf32>} : () -> tensor<3x4x2xf32>
  %0 = "tf.XlaCallModule"(%cst, %arg0) {

  func.func @main() -> () {
    %cst0 = "tf.Const"(){ value = dense<0> : tensor<3x5xi1>} : () -> tensor<3x5xi1>
    %0 = "tf.Where"(%cst0) : (tensor<3x5xi1>) -> tensor<?x2xi64>
    func.return
  }
})";

absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
    const char* mlir_module_str,
    ConfigProto::Experimental::MlirBridgeRollout rollout_state) {

      stablehlo.return %3 : tensor<f32>
  }) {
    window_dimensions = array<i64: 1, 1, 2, 2>,
    window_strides = array<i64: 1, 1, 2, 2>,
    padding = dense<[[0, 0], [0, 0], [1, 1], [1, 1]]> : tensor<4x2xi64>
  } : (tensor<1x4x16x16xf32>, tensor<f32>) -> tensor<1x4x9x9xf32>
  return %2 : tensor<1x4x9x9xf32>
}
// CHECK-SAME: %[[ARG:.+]]: tensor<1x16x16x4xf32>
// CHECK-DAG: stablehlo.constant