// CHECK-LABEL:   func.func @convert_argmax_bool
func.func @convert_argmax_bool(%arg0: tensor<2xi1>) -> tensor<i32> {
  %0 = "mhlo.iota"() <{iota_dimension = 0 : i64}> : () -> tensor<2xi32>
  %1 = mhlo.constant dense<false> : tensor<i1>
  %2 = mhlo.constant dense<0> : tensor<i32>
  %3:2 = mhlo.reduce(%arg0 init: %1), (%0 init: %2) across dimensions = [0] : (tensor<2xi1>, tensor<2xi32>, tensor<i1>, tensor<i32>) -> (tensor<i1>, tensor<i32>)

}

// -----

// CHECK-LABEL: func @convolution_add_add_i8
func.func @convolution_add_add_i8(
    %lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
    %zp_offset: tensor<?x2x2x1xi8>, %bias: tensor<1xi8>
  ) -> tensor<?x2x2x1xi8> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution
  // CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]

// operate on quantized values.
//
// Examples from storage to quantized type:
//   i8 -> !quant<"uniform[i8:f32]{1.0}">
//   tensor<4xi8> -> tensor<4x!quant<"uniform[i8:f32]{1.0}">>
//   vector<4xi8> -> vector<4x!quant<"uniform[i8:f32]{1.0}">>
def Quantization_StorageCastOp : Quantization_Op<"scast", [Pure]> {
  let arguments = (ins quant_RealOrStorageValueType:$arg);

  // CHECK: %[[AND:.*]] = "tf.LogicalAnd"(%[[EQ]], %[[IS_ODD]]) : (tensor<2xi1>, tensor<2xi1>) -> tensor<2xi1>
  // CHECK: %[[OR:.*]] = "tf.LogicalOr"(%[[GT]], %[[AND]]) : (tensor<2xi1>, tensor<2xi1>) -> tensor<2xi1>
  // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[ROUND_VAL]], %[[ONE]]) : (tensor<2xf32>, tensor<f32>) -> tensor<2xf32>

// operate on quantized values.
//
// Examples from storage to quantized type:
//   i8 -> !quant<"uniform[i8:f32]{1.0}">
//   tensor<4xi8> -> tensor<4x!quant<"uniform[i8:f32]{1.0}">>
//   vector<4xi8> -> vector<4x!quant<"uniform[i8:f32]{1.0}">>
def quantfork_StorageCastOp : quantfork_Op<"scast", [Pure]> {
  let arguments = (ins quant_RealOrStorageValueType:$arg);

    %11 = stablehlo.reshape %2 : (tensor<1x1x1x1xi8>) -> tensor<1xi8>
    %12 = stablehlo.broadcast_in_dim %11, dims = [0] : (tensor<1xi8>) -> tensor<1x3x3x4xi8>
    %13 = stablehlo.convert %12 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
    %14 = stablehlo.convert %3 : (tensor<3x3x4x4xi8>) -> tensor<3x3x4x4xf32>

          )
      )
      module_str = self._extract_first_xla_call_module_op(
          self._output_saved_model_path
      )
      self.assertTrue(
          re.search('stablehlo.dot_general.*xi8>.*xi8>.*xi32>', module_str)
      )

    # Run quantization the first time, calibration is expected to be skipped.
    output_saved_model_path_2 = self.create_tempdir('output2').full_path

  if (!quantDenseAttr) {
    return nullptr;
  }

  // Cast from an expressed-type-based type to storage-type-based type,
  // preserving the sparse shape (i.e. tensor<4xf32> -> tensor<4xi8>).
  ShapedType newSparseType = mlir::dyn_cast_or_null<ShapedType>(
      quantizedElementType.castExpressedToStorageType(
          realSparseAttr.getType()));
  if (!newSparseType) {
    return nullptr;
  }

func.func @testLogicalAndDoesntFoldWithConstantFalseBroadcast(%arg0: tensor<2xi1>) -> (tensor<2xi1>) {
  %cst = arith.constant dense<false> : tensor<i1>

  // CHECK: %[[LOGICAL_AND:.*]] = "tf.LogicalAnd"
  %0 = "tf.LogicalAnd"(%cst, %arg0) : (tensor<i1>, tensor<2xi1>) -> tensor<2xi1>

  // CHECK: return %[[LOGICAL_AND]]
  func.return %0: tensor<2xi1>
}

// -----

// GlobalIterId should not be constant folded

  }

  // CHECK-LABEL: func @greater
  func.func @greater(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> {
    // CHECK-NEXT:  mhlo.compare GT, %arg0, %arg1
    %0 = "tf.Greater"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
    func.return %0: tensor<2xi1>
  }

  // CHECK-LABEL: batchmatmulv2