// CHECK:   %0 = "tfl.while"(%arg0) ({
// CHECK:   ^bb0(%arg1: tensor<*xf32>):
// CHECK:     %[[RES0:.*]] = func.call @cond(%arg1) : (tensor<*xf32>) -> tensor<*xf32>
// CHECK:     "tfl.yield"(%[[RES0]]) : (tensor<*xf32>) -> ()
// CHECK:   },  {
// CHECK:   ^bb0(%arg1: tensor<*xf32>):
// CHECK:     %[[RES1:.*]] = func.call @body(%arg1) : (tensor<*xf32>) -> tensor<*xf32>
// CHECK:     "tfl.yield"(%[[RES1]]) : (tensor<*xf32>) -> ()

  func.return %0: tensor<?x?xui32>
}

// -----

// CHECK-LABEL: func @floordiv_int
func.func @floordiv_int(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> tensor<?xi32> {
  // CHECK-DAG: [[DIV:%.+]] = chlo.broadcast_divide %arg0, %arg1 : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>

// CHECK:     }) : (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>)

func.func @main(%arg0: tensor<i32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  // While %arg0 is greater than zero, element wise add %arg1 with itself.
  %0:2 = "tfl.while"(%arg0, %arg1) ({
  ^bb0(%arg2: tensor<*xi32>, %arg3: tensor<*xf32>):
    %1 = func.call @cond(%arg2, %arg3) : (tensor<*xi32>, tensor<*xf32>) -> tensor<i1>

    return %1 : tensor<1x3xf32>
  }

  func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
    return %0 : tensor<1x3xf32>
  }
}

}

// -----

// CHECK-LABEL:func @partitioned_2d
// CHECK-SAME: ([[ARG0:%.*]]: tensor<10x3xf32>, [[ARG1:%.*]]: tensor<10x3xf32>)
func.func @partitioned_2d(%arg0: tensor<10x3xf32>, %arg1: tensor<10x3xf32>) -> tensor<20x3xf32> {
  // CHECK: [[PI_0:%.*]] = "tf.TPUPartitionedInputV2"([[ARG0]], [[ARG1]])
  // CHECK-SAME: _XlaSharding = "123"
  // CHECK-SAME: partition_dims = [2, 1]

  %arg0: tensor<1x28x28x64xf32>,
  %arg1: tensor<64xf32>
) -> tensor<1x28x28x64xf32> {

  // CHECK: %[[ARG_PERM:.*]] = "tf.Const"()
  // CHECK-SAME: <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}>
  // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]])

  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: (%[[ARG_TRANSPOSE]], %arg1, %arg1, %arg1, %arg1)
  // CHECK-SAME: data_format = "NCHW"

func.func private @testIf1Else(tensor<*xf32>) -> tensor<*xf32>

// CHECK-LABEL: func @testIf1Result(%arg0: tensor<i1>, %arg1: tensor<*xf32>)
func.func @testIf1Result(%arg0: tensor<i1>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0 = "tf.If"(%arg0, %arg1) {
    then_branch = @testIf1Then, else_branch = @testIf1Else, is_stateless = false,
    _attr0 = 10, _attr1 = true, attr2 = "hello"

     // CHECK-NEXT:   tf_device.return %[[READ_OUT]]
     %0 = "tf.ReadVariableOp"(%arg1) : (tensor<*x!tf_type.resource<tensor<4xf32>>>) -> tensor<4xf32>
     %1 = "tf.A"() : () -> (tensor<2x!tf_type.string>)
     "tf_device.launch"() ({

	f    func(uint64, uint64) uint64
	arg1 uint64
	arg2 uint64
	res  uint64
	name string
}

var extTests = [...]extTest{
	{f: func(a, b uint64) uint64 { op1 := int32(a); op2 := int32(b); return uint64(uint32(op1 / op2)) }, arg1: 0x1, arg2: 0xfffffffeffffffff, res: 0xffffffff, name: "div"},
	{f: func(a, b uint64) uint64 { op1 := int32(a); op2 := int32(b); return uint64(uint32(op1 * op2)) }, arg1: 0x1, arg2: 0x100000001, res: 0x1, name: "mul"},

  func.return %3 : tensor<f32>
}

func.func private @add(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32>  {
  %0 = "tf.Add"(%arg0, %arg1): (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

func.func private @sub(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32>  {
  %0 = "tf.Sub"(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
}