func.func @main() -> (tensor<5x11xf32>) {
      %arg0 = "tf.Const"() {value = dense<-3.0> : tensor<5x7xf32>} : () -> tensor<5x7xf32>
      %arg1 = "tf.Const"() {value = dense<-3.0> : tensor<11x7xf32>} : () -> tensor<11x7xf32>

      %1 = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = true} : (tensor<5x7xf32>, tensor<11x7xf32>) -> tensor<5x11xf32>

      func.return %1 : tensor<5x11xf32>
    }
  })";

    %12 = stablehlo.subtract %10, %11 : tensor<1x3xf32>  // q1 * q2 - z1 * q2
    %13 = stablehlo.multiply %12, %4 : tensor<1x3xf32>  // s1 * s2
    %14 = call @uniform_quantize_1(%13, %5, %6) : (tensor<1x3xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x3xi8>
    %15 = call @uniform_dequantize_0(%14, %5, %6) : (tensor<1x3xi8>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x3xf32>
    return %15 : tensor<1x3xf32>
  }
// CHECK-NOT: stablehlo.uniform_quantize

  // CHECK: return %[[RESULT]]
  func.return %0: tensor<10x2xf32>
}

// -----

func.func @select_first(%arg0: tensor<10x2xf32>, %arg1: tensor<10x2xf32>) -> tensor<10x2xf32> {
  func.return %arg0: tensor<10x2xf32>
}

// CHECK-LABEL: testMultiInputLegacyCallOp
func.func @testMultiInputLegacyCallOp(%arg0: tensor<10x2xf32>, %arg1: tensor<10x2xf32>) -> tensor<10x2xf32> {

  %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
// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<10x2xf32>}> : () -> tensor<10x2xf32>
// CHECK: %[[batch_matmul:.*]] = "tf.BatchMatMulV2"(%arg0, %[[cst]])

  "tfl.custom"(%0, %2) {custom_code = "FlexAssignVariableOp", custom_option = #tfl<const_bytes : "0x1041737369676E5661726961626C654F70003B121041737369676E5661726961626C654F701A001A002A0B0A056474797065120230012A140A0E76616C69646174655F736861706512022800320000024F3E1414042801">} : (tensor<!tf_type.resource<tensor<1x10xf32>>>, tensor<1x10xf32>) -> ()

// CHECK:           }) : (tensor<1x?xf32>, tensor<1x0xf32>, tensor<1x0xf32>, tensor<1x0xf32>, tensor<1x0xf32>, tensor<1x3xf32>, tensor<1x3xf32>, tensor<1x3xf32>, tensor<1x3xf32>, none, none, none, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<3x1xf32>, tensor<3xf32>, tensor<1x3xf32>, tensor<1x1xf32>, none, none, none, none) -> tensor<1x3xf32>
// CHECK:           [[VAL_50:%.*]] = tensor.cast [[VAL_51:%.*]] : tensor<1x3xf32> to tensor<1x?xf32>

// CHECK:  "tfl.cos"(%arg0) : (tensor<f32>) -> tensor<f32>
}

func.func @elu(%arg0: tensor<11x16xf32>) -> tensor<11x16xf32> {
  %0 = "tf.Elu"(%arg0) : (tensor<11x16xf32>) -> tensor<11x16xf32>
  func.return %0 : tensor<11x16xf32>

// CHECK-LABEL:elu
// CHECK:  "tfl.elu"(%arg0) : (tensor<11x16xf32>) -> tensor<11x16xf32>
}

func.func @expandDims(%arg0: tensor<2x2xf32>, %arg1: tensor<i32>) -> tensor<1x2x2xf32> {

  // time constant.
  %0 = "tf.ConcatV2"(%arg0, %arg0, %arg0, %arg0, %axis) : (tensor<16x1xf32>, tensor<16x1xf32>, tensor<16x1xf32>, tensor<16x1xf32>, tensor<i32>) -> tensor<16x4xf32>

  // Just to introduce an extra use for %cst.
  %1 = "tf.AddV2"(%axis, %axis) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>

  return %0 : tensor<16x4xf32>
}
// Check that the constant is cloned with same value.

      %1 = "quantfork.qcast"(%0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
      %2 = "quantfork.dcast"(%1) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
      %3 = stablehlo.reshape %2 : (tensor<1x3xf32>) -> tensor<3x1xf32>

      return %0 : tensor<1x1x4xf32>
    }
  )mlir";

  const OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kXlaCallModuleOpWithQuantizationMethodAttr);
  ASSERT_TRUE(module_op);