}

// -----

// CHECK-LABEL: func @testMul
func.func @testMul(%arg0: tensor<2xui16>) -> (tensor<2xui16>) {
  %0 = "tf.Mul"(%arg0, %arg0) {T = "tfdtype$DT_UINT16", device = "/device:CPU:0", name = "Mul"} : (tensor<2xui16>, tensor<2xui16>) -> tensor<2xui16>
  func.return %0 : tensor<2xui16>
}

// -----

// Test error message for incompatible element types.

// CHECK: name: "const2"
// CHECK-NEXT: op: "Const"
// CHECK: dtype: DT_HALF
// CHECK: half_val: 15360
// CHECK-NEXT: half_val: 16384
    %2:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = bf16, value = dense<[4.900000e+01, 8.200000e+02]> : tensor<2xbf16>} : () -> tensor<bf16> loc("const3")
// CHECK: name: "const3"

    %3:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = bf16, value = dense<0.000000e+00> : tensor<bf16>} : () -> tensor<bf16> loc("const4")
    tf_executor.fetch %0#0, %1#0 : tensor<1x2xf16>, tensor<2xf16>
  }
  func.return %graph#0, %graph#1 : tensor<1x2xf16>, tensor<2xf16>
}

// CHECK: node {
// CHECK-NEXT: name: "const1"

  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

# CHECK-DAG: tf.IteratorGetNextSync{{.+}}-> (tensor<*xi16>, tensor<3x24xi32>, tensor<?x4x32xi64>)
# CHECK-DAG: tf.MultiDeviceIteratorGetNextFromShard{{.+}}-> (tensor<5x40xf16>, tensor<*xcomplex<f32>>, tensor<6x48x?xcomplex<f64>>)
# CHECK-DAG: tf.InfeedDequeueTuple{{.+}}-> (tensor<?x?x?xui16>, tensor<*xui32>, tensor<7x56xui64>)
# CHECK-DAG: tf.InfeedDequeue{{.+}}-> tensor<?x8x?xi8>

// RUN: tf-opt -tfl-legalize-random %s | FileCheck %s


// CHECK-LABEL:   func @tfl_wrapped_jax_random_normal(
// CHECK-SAME:                                        %[[RNG:.*]]: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
// CHECK:           %[[VAL_0:.*]] = stablehlo.constant dense<[3, 4]> : tensor<2xi32>

)

var (
	i8 int8 = 0x12
	i16 int16 = 0x1234
	i32 int32 = 0x12345678
	i64 int64 = 0x123456789abcdef0
	ui8 uint8 = 0x12
	ui16 uint16 = 0x1234
	ui32 uint32 = 0x12345678
	ui64 uint64 = 0x123456789abcdef0

	ni8 = ^i8
	ni16 = ^i16
	ni32 = ^i32
	ni64 = ^i64
	nui8 = ^ui8
	nui16 = ^ui16
	nui32 = ^ui32
	nui64 = ^ui64
)

var nfail = 0

func main() {
	if nfail > 0 {
		fmt.Printf("BUG\n")
	}
}

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.func @cast_bf16_add_v2_to_fp32(%arg0: tensor<2xbf16>, %arg1: tensor<2xbf16>) -> tensor<2xf32> {
  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<2xbf16>, tensor<2xbf16>) -> tensor<2xbf16>
  %1 = "tf.Cast"(%0) {Truncate = false} : (tensor<2xbf16>) -> tensor<2xf32>
  return %1 : tensor<2xf32>
}
// The signature of the function is not changed.
// CHECK: func @cast_bf16_add_v2_to_fp32(%[[ARG_0:.*]]: tensor<2xbf16>, %[[ARG_1:.*]]: tensor<2xbf16>) -> tensor<2xf32>