func.return %arg0: tensor<1x!tf_type.quint8>
}

// -----

// CHECK-LABEL: func @id_quint16(%arg0: tensor<1xui16>) -> tensor<1xui16> {
func.func @id_quint16(%arg0: tensor<1x!tf_type.quint16>) -> tensor<1x!tf_type.quint16> {
  // CHECK-NEXT: return %arg0 : tensor<1xui16>
  func.return %arg0: tensor<1x!tf_type.quint16>
}

// -----

// CHECK-LABEL: func @uniform_quantize

func.func @rng_bit_generator(%arg0: tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>) {
  %output_state, %output = "vhlo.rng_bit_generator_v1"(%arg0) <{rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>}> : (tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>)
  func.return %output_state, %output : tensor<2xui64>, tensor<10x12xui32>
}

// CHECK:func.func private @rng_bit_generator(%arg0: tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>) {

  // CHECK:  %cst = arith.constant dense<0> : tensor<1xi32>
  // CHECK:  %4 = "tfl.reduce_min"(%arg0, %cst) <{keep_dims = false}> : (tensor<2xi16>, tensor<1xi32>) -> tensor<i16>
  // CHECK:  %5 = "tfl.arg_min"(%arg0, %cst) : (tensor<2xi16>, tensor<1xi32>) -> tensor<i32>
  // CHECK:  return %4, %5 : tensor<i16>, tensor<i32>
}


// CHECK-LABEL: func @convert_argmin_constant

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

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

func.func @i32() -> tensor<4xi32> {
  // CHECK-LABEL: @i32

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

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

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

func.func @i32() -> tensor<4xi32> {
  // CHECK-LABEL: @i32

  %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> {
  %cst = arith.constant dense<[false, true]> : tensor<2xi1>

    <I16Neg>  -6703  </I16Neg>
    <I32>  266703  </I32>
    <I32Neg>  -266703  </I32Neg>
    <I64>  266703  </I64>
    <I64Neg>  -266703  </I64Neg>
    <UI>   266703   </UI>
    <UI8>  112  </UI8>
    <UI16>  6703  </UI16>
    <UI32>  266703  </UI32>
    <UI64>  266703  </UI64>
    <F32>  266.703  </F32>
    <F32Neg>  -266.703  </F32Neg>
    <F64>  266.703  </F64>
    <F64Neg>  -266.703  </F64Neg>
</WhitespaceValuesParent>

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