func.func @cast_i1_to_i8() -> tensor<2xi8> {
  %cst = arith.constant dense<[false, true]> : tensor<2xi1>
  %0 = "tfl.cast"(%cst) : (tensor<2xi1>) -> tensor<2xi8>
  func.return %0 : tensor<2xi8>

// CHECK: %[[CST:.*]] = arith.constant dense<[0, 1]> : tensor<2xi8>
// CHECK:  return %[[CST]]
}

// CHECK-LABEL: @cast_i1_to_ui8
func.func @cast_i1_to_ui8() -> tensor<2xui8> {

  // CHECK: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[QUANTIZE]] : (tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<2xi8>
  // CHECK: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<2xi8>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>
  // CHECK: %[[DEQUANTIZE:.*]] = mhlo.uniform_dequantize %[[CONVERT_2]] : (tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<2xf32>

  func.return %0 : tensor<4xf64>
}

func.func @i8() -> tensor<4xi8> {
  // CHECK-LABEL: @i8
  // CHECK: value = dense<[1, 2, 3, 4]> : tensor<4xi8>
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 4]> : tensor<4xi8> } : () -> tensor<4xi8>
  func.return %0 : tensor<4xi8>
}

func.func @i16() -> tensor<4xi16> {
  // CHECK-LABEL: @i16

  func.return %0 : tensor<4xf64>
}

func.func @i8() -> tensor<4xi8> {
  // CHECK-LABEL: @i8
  // CHECK: value = dense<[1, 2, 3, 4]> : tensor<4xi8>
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 4]> : tensor<4xi8> } : () -> tensor<4xi8>
  func.return %0 : tensor<4xi8>
}

func.func @i16() -> tensor<4xi16> {
  // CHECK-LABEL: @i16

// CHECK-SAME:              %[[VAL_0:.*]]: tensor<2xi1>,
// CHECK-SAME:              %[[VAL_1:.*]]: tensor<2xi1>) -> tensor<2xi1> {
// CHECK:           %[[VAL_2:.*]] = "tf.LogicalAnd"(%[[VAL_0]], %[[VAL_1]]) : (tensor<2xi1>, tensor<2xi1>) -> tensor<2xi1>
// CHECK:           return %[[VAL_2]] : tensor<2xi1>
// CHECK:         }
func.func @and(%arg0: tensor<2xi1>, %arg1: tensor<2xi1>) -> tensor<2xi1> {
  %0 = mhlo.and %arg0, %arg1 : tensor<2xi1>

  func.return %0 : tensor<2x4xui8>
}

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

// CHECK-LABEL: func @and_unranked
func.func @and_unranked(%arg0: tensor<*xi1>, %arg1: tensor<*xi1>) -> tensor<*xi1> {

// CHECK: %[[BITCAST_CONVERT_0:.*]] = stablehlo.bitcast_convert %[[UQ]] : (tensor<1x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<1xi8>
// CHECK: %[[BITCAST_CONVERT_1:.*]] = stablehlo.bitcast_convert %[[BITCAST_CONVERT_0]] : (tensor<1xi8>) -> tensor<1x!quant.uniform<i8:f32, 1.000000e+00:3>>

// CHECK: %[[maxpool_i8:.*]] = "tf.MaxPool"(%[[sc1]])
// CHECK-SAME: (tensor<*xi8>) -> tensor<*xi8>
// CHECK: %[[reshape_i8:.*]] = "tf.Reshape"(%[[maxpool_i8]]
// CHECK-SAME: (tensor<*xi8>, tensor<2xi32>) -> tensor<*xi8>
// CHECK: %[[sc2:.*]] = "quantfork.scast"(%[[reshape_i8]])
// CHECK: %[[dq:.*]] = "quantfork.dcast"(%[[sc2]]) : (tensor<*x!quant.uniform<i8:f32, 5.000000e-02:-10>>)

  func.return %11 : tensor<*xf32>
}

// CHECK: %[[maxpool_i8:.*]] = "tf.MaxPool"
// CHECK-SAME: (tensor<*xi8>) -> tensor<*xi8>
// CHECK: %[[reshape_i8:.*]] = "tf.Reshape"(%[[maxpool_i8]]
// CHECK-SAME: (tensor<*xi8>, tensor<2xi32>) -> tensor<*xi8>
// CHECK: %[[scast:.*]] = "quantfork.scast"(%[[reshape_i8]]
// CHECK: %[[matmul:.*]] = "tf.PartitionedCall"(%[[scast]]
// CHECK-SAME: f = @composite_matmul_fn_1

  // Argument has rank 2, but its dimensions are unknown.
  constexpr absl::string_view kArgumentWithUnknownDims = R"mlir(
    func.func @unknown_dims_arg(%arg: tensor<?x?xi8>) -> tensor<?x?xi8> {
      return %arg : tensor<?x?xi8>
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kArgumentWithUnknownDims);
  ASSERT_TRUE(module_op);