%1 = "tfl.add"(%cst_2, %cst_1) {fused_activation_function = "NONE"} : (tensor<2x2x2xf32>, tensor<  2x2xf32>) -> tensor<2x2x2xf32>
  %2 = "tfl.add"(%cst_0, %cst_2) {fused_activation_function = "NONE"} : (tensor<    2xf32>, tensor<2x2x2xf32>) -> tensor<2x2x2xf32>

  func.return %0, %1, %2 : tensor<2x2xf32>, tensor<2x2x2xf32>, tensor<2x2x2xf32>

  func.return %0 : tensor<2x2xf32>

// CHECK-LABEL:squeezeDefault
// CHECK:  "tfl.squeeze"(%arg0) <{squeeze_dims = []}> : (tensor<1x2x2xf32>) -> tensor<2x2xf32>
}

func.func @squeezeSingleAxis(%arg0: tensor<2x1x2xf32>) -> tensor<2x2xf32> {
  %0 = "tf.Squeeze"(%arg0) {squeeze_dims = [1]} : (tensor<2x1x2xf32>) -> tensor<2x2xf32>
  func.return %0 : tensor<2x2xf32>

// CHECK-LABEL:squeezeSingleAxis

    %14 = stablehlo.multiply %12, %13 : tensor<1x4x3xf32>  // s1 * s2
    %15 = call @uniform_quantize_1(%14, %5, %6) : (tensor<1x4x3xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x3xi8>
    %16 = call @uniform_dequantize_0(%15, %5, %6) : (tensor<1x4x3xi8>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x3xf32>
    return %16 : tensor<1x4x3xf32>
  }

// CHECK-LABEL: QuantizeUnidirectionalLstmFullPerTensor
func.func @QuantizeUnidirectionalLstmFullPerTensor(%arg0: tensor<1x2x3xf32>) -> (tensor<1x2x3xf32>) {
  %input = "quantfork.stats"(%arg0) {layerStats = dense<[0.0, 1.0]> : tensor<2xf32>} : (tensor<1x2x3xf32>) -> tensor<1x2x3xf32>
  %1 = "tfl.pseudo_const"() {value = dense<[[0.1]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>

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

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

    func.func @main() -> (tensor<1x4x4xf32>) {
      %%arg0 = "tf.Const"() {value = dense<-3.0> : tensor<1x4x2xf32>} : () -> tensor<1x4x2xf32>
      %%arg1 = "tf.Const"() {value = dense<-3.0> : tensor<1x2x4xf32>} : () -> tensor<1x2x4xf32>

      %%1 = "tf.%s"(%%arg0, %%arg1) {T = f32, adj_x = false, adj_y = false, grad_x = false, grad_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<1x2x4xf32>) -> tensor<1x4x4xf32>

  func.return %17 : tensor<1x2x3xf32>

  // CHECK: %[[NONE:.*]] = "tfl.no_value"() <{value}> : () -> none
  // CHECK: %[[DQ_1:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x1xf16>) -> tensor<1x1xf32>
  // CHECK: %[[DQ_2:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x1xf16>) -> tensor<1x1xf32>
  // CHECK: %[[DQ_3:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x1xf16>) -> tensor<1x1xf32>

  %50 = "tfl.read_variable"(%5) : (tensor<!tf_type.resource>) -> tensor<1x2x3xf32>
  %51 = "quantfork.stats"(%50) {layerStats = dense<[0.0, 1.0]> : tensor<2xf32>} : (tensor<1x2x3xf32>) -> tensor<1x2x3xf32>
  %52 = "tfl.concatenation"(%51, %0) {axis = 1 : i32, fused_activation_function = "NONE"} : (tensor<1x2x3xf32>, tensor<1x2x3xf32>) -> tensor<1x4x3xf32>

  module {
    func.func @main(%arg0: tensor<2x2x2xf32>, %arg1: tensor<2x2x2xf32>) -> tensor<2x2x2xf32> attributes {_from_xla_call_module} {
      %0 = stablehlo.dot_general %arg0, %arg1,
        batching_dims = [0] x [0],
        contracting_dims = [2] x [1],
        precision = [DEFAULT, DEFAULT]
      : (tensor<2x2x2xf32>, tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
        return %0 : tensor<2x2x2xf32>
    }
  }
)mlir";

func.func @prepareAdd(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
  %cst = arith.constant dense<[[0.0, 1.0], [2.0, 255.0]]> : tensor<2x2xf32>
  %add = "tfl.add"(%arg0, %cst) {fused_activation_function="NONE"} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
  func.return %add : tensor<2x2xf32>

// CHECK: %[[cst:.*]] = arith.constant dense<[{{\[}}0.000000e+00, 1.000000e+00], [2.000000e+00, 2.550000e+02]]>