// -----

// CHECK-LABEL: @if
func.func @if() {
  // CHECK-NOT: Optional
  %0 = builtin.unrealized_conversion_cast to tensor<5xi1>
  %1 = "tf.OptionalNone"() : () -> tensor<!tf_type.variant<tensor<f32>>>
  %2 = "tf.If"(%0, %1) <{else_branch = @false, is_stateless = false, then_branch = @true}>
      : (tensor<5xi1>, tensor<!tf_type.variant<tensor<f32>>>) -> (tensor<!tf_type.variant<tensor<f32>>>)
  return
}

            %cst = arith.constant dense<3.0> : tensor<32xf32>
            %cmp = "tf.Less"(%read, %cst) : (tensor<32xf32>, tensor<32xf32>) -> tensor<32xi1>
            %dims = arith.constant dense<0> : tensor<1xi32>
            %reduce = "tf.All"(%cmp, %dims) {keep_dims = false} : (tensor<32xi1>, tensor<1xi32>) -> tensor<i1>
            "tf.Yield"(%reduce) : (tensor<i1>) -> ()
        }, {

  %0 = "tfl.pseudo_sparse_qconst"() {compressed_data = dense<[1, 2, 3, 4, -1, -2, -3, -4]> : tensor<8xi8>, qtype = tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>>, s_param = #tfl.sparsity_parameter<traversal_order = [0, 1, 2, 3], block_map...

  %0 = "tfl.pseudo_sparse_qconst"() {compressed_data = dense<[1, 2, 3, 4, -1, -2, -3, -4]> : tensor<8xi8>, qtype = tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>>, s_param = #tfl.sparsity_parameter<traversal_order = [0, 1, 2, 3], block_map...

// CHECK-LABEL: @quantize_constant
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x3xf32>
func.func @quantize_constant(%arg0: tensor<1x3xf32>) -> tensor<1x2xf32> {
  // CHECK-DAG: %[[QCST:.*]] = stablehlo.constant() <{value = dense<-78> : tensor<3x2xi8>}> : () -> tensor<3x2x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>
  // CHECK-DAG: %[[Q1:.*]] = stablehlo.uniform_quantize %[[ARG0]]
  // CHECK-NOT: "quantfork.qcast"
  // CHECK: %[[DOT:.*]] = stablehlo.dot %[[Q1]], %[[QCST]]

// CHECK-LABEL: test_gather
func.func @test_gather(%arg0 : tensor<1x2x20xf32>, %arg1 : tensor<3x5xi32>) -> (tensor<1x3x5x20xf32>){
    // expected-remark@+1 {{Cost: 1}}
    %0 = "tf.Const"() { value = dense<[1]> : tensor<1xi32> } : () -> tensor<1xi32>
    // expected-remark@+1 {{Cost: 300}}
    %1 = "tf.GatherV2"(%arg0, %arg1, %0) : (tensor<1x2x20xf32>, tensor<3x5xi32>, tensor<1xi32>) -> tensor<1x3x5x20xf32>
    // expected-remark@+1 {{Cost: 40}}

  %0 = "tf.BroadcastTo"(%input, %shape) : (tensor<3xi16>, tensor<2xi32>) -> tensor<3x3xi16>
  func.return %0: tensor<3x3xi16>

// CHECK-LABEL: broadcast_to_i16_low_dim
// CHECK:    %0 = "tf.BroadcastTo"(%arg0, %arg1) : (tensor<3xi16>, tensor<2xi32>) -> tensor<3x3xi16>
// CHECK:    return %0 : tensor<3x3xi16>
}

  %zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

  // CHECK: %[[input:.*]] = "tf.ConcatV2"(%arg0, %arg1, %[[VAL:.*]]) : (tensor<3x3xi8>, tensor<3x3xi8>, tensor<i64>) -> tensor<6x3xi8>
  // CHECK: %[[input_qint:.*]] = "tf.Cast"(%[[input]]) <{Truncate = false}> : (tensor<6x3xi8>) -> tensor<6x3x!tf_type.qint8>
  // CHECK: %[[output:.*]] = "tf.UniformDequantize"(%[[input_qint]]

func.func @select(
    %arg0: tensor<1x3xi1>,
    %arg1: tensor<1x3x!quant.uniform<i8:f32, 2.000000e+00:-1>>,
    %arg2: tensor<1x3x!quant.uniform<i8:f32, 2.000000e+00:-1>>
  ) -> tensor<1x3x!quant.uniform<i8:f32, 2.000000e+00:-1>> {
  %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (
    tensor<1x3xi1>,
    tensor<1x3x!quant.uniform<i8:f32, 2.000000e+00:-1>>,

    %8 = "tf.Sub"(%7, %cst_4) {device = ""} : (tensor<2x3xi32>, tensor<i32>) -> tensor<2x3xi32>
    %9 = "tf.Identity"(%cst_2) {device = ""} : (tensor<3x4xi8>) -> tensor<3x4xi8>
    %10 = "tf.Cast"(%9) {Truncate = false, device = ""} : (tensor<3x4xi8>) -> tensor<3x4xi32>
    %11 = "tf.Einsum"(%8, %10) {device = "", equation = "ab,bc->ac"} : (tensor<2x3xi32>, tensor<3x4xi32>) -> tensor<2x4xi32>