// 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>

      tf_executor.fetch %outputs#1, %outputs#0 : tensor<2x2x1x2xi32>, tensor<2x2x1x2xf32>
    }
    func.return %0#0, %0#1 : tensor<2x2x1x2xi32>, tensor<2x2x1x2xf32>
  }
}

// CHECK:        name: "serving_default_x"
// CHECK-NEXT:   op: "_Arg"

func.func @batchmatmulv2_basic(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
// CHECK-LABEL:   func @batchmatmulv2_basic
// CHECK-SAME:        ([[LHS:%.*]]: tensor<1x4x2xf32>, [[RHS:%.*]]: tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
// CHECK:           [[LHSSHAPE:%.*]] = shape.shape_of [[LHS]] : tensor<1x4x2xf32>
// CHECK:           [[RHSSHAPE:%.*]] = shape.shape_of [[RHS]] : tensor<3x2x4xf32>

  %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: return %[[DEQUANTIZED]] : tensor<*xf32>
}

// -----

module {
  func.func @not_quantize_matmul_without_const(%arg0: tensor<1x2x2x2xf32>, %arg1: tensor<2x1024xf32>) -> (tensor<*xf32>) {
    %arg0_identity = "tf.Identity"(%arg0) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
    %arg1_identity = "tf.Identity"(%arg1) {device = ""} : (tensor<2x1024xf32>) -> tensor<2x1024xf32>

  %2 = "tf.Relu6"(%1) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
  return %2 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: @composite_conv2d_with_bias_and_relu6_fn_1
// CHECK-NOT: "tf.CalibrationStatisticsSaver"

// -----

// Check the IfOp is set to stateful.

// CHECK-LABEL: func @testCase(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource<tensor<1x2x3xf32>>>)
func.func @testCase(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource<tensor<1x2x3xf32>>>) -> tensor<1x2x3xf32> {
  %0 = "tf.Case"(%arg0, %arg1) {branches = [@branch_0, @branch_1], is_stateless = false} : (tensor<i32>, tensor<!tf_type.resource<tensor<1x2x3xf32>>>) -> tensor<1x2x3xf32>

  // CHECK: [[Result0:%.*]] = "tf.CaseRegion"

    %7 = call @uniform_quantize_0(%arg0, %0, %1) : (tensor<1x4x2xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x2xi8>
    %8 = stablehlo.convert %7 : (tensor<1x4x2xi8>) -> tensor<1x4x2xf32>
    %9 = stablehlo.convert %2 : (tensor<2x3xi8>) -> tensor<2x3xf32>
    %10 = stablehlo.dot_general %8, %9, contracting_dims = [2] x [0] : (tensor<1x4x2xf32>, tensor<2x3xf32>) -> tensor<1x4x3xf32>

// fallback lowering is preferred for static shaped operands when available.

// CHECK-LABEL: batchmatmulv2
func.func @batchmatmulv2(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
  // NO_FALLBACK: mhlo.dynamic_broadcast_in_dim
  // NO_FALLBACK: mhlo.dot_general

  // SUPPORTED_FALLBACK_DEVICE: mhlo.reduce
  // SUPPORTED_FALLBACK_DEVICE: mhlo.dot_general

    indices_are_sorted = false
  } : (tensor<3x4x2xf32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32>
    %8 = "quantfork.qcast"(%7) {volatile} : (tensor<2x3x2x2xf32>) -> tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
    %9 = "quantfork.dcast"(%8) : (tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2x2xf32>
    return %9 : tensor<2x3x2x2xf32>
  }

  // CHECK: quantized_dot_general_fn_1