// CHECK-DAG:   [[ITEM_RANK:%.*]] = "tf.Rank"([[ITEM]]) : (tensor<10xf32>) -> tensor<i32>
// CHECK-DAG:   [[ZERO:%cst.*]] = arith.constant dense<0> : tensor<i32>
// CHECK-DAG:   [[VECTOR_RANK:%.*]] = "tf.ExpandDims"([[ITEM_RANK]], [[ZERO]]) : (tensor<i32>, tensor<i32>) -> tensor<1xi32>

func.func @island_no_wrap_print_multiple_ops(%arg0: tensor<i32>, %arg1: tensor<f32>) {
  tf_executor.graph {
    // CHECK: tf_executor.island
    // CHECK-NOT: wraps
    %0:3 = tf_executor.island {
      %1:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>) loc("identity@some_function")
      %2:2 = "tf.IdentityN"(%1#0, %1#1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>) loc("identity@some_function")

  %one = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
  %elem_shape = "tf.Const"() {value = dense<[-1, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
  %size = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
  // CHECK: TensorListReserve
  // CHECK: TensorListReserve
  %tl = "tf.TensorListReserve"(%elem_shape, %size) : (tensor<2xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<?x1xf32>>>

  // CHECK: "tfl.average_pool_2d"(%{{.*}}) <{filter_height = 3 : i32, filter_width = 6 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 3 : i32, stride_w = 1 : i32}> : (tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32>
  %0 = "tfl.average_pool_2d"(%arg0) {filter_height = 3 : i32, filter_width = 6 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 3 : i32, stride_w = 1 : i32} : (tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32> loc("avgpool")

  %1 = "tf.Add"(%z, %0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %cond, %x, %y, %1 : tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>
}

// CHECK-LABEL: func @while_test_if
// CHECK-SAME: -> !tf_mlrt.tensor
func.func @while_test_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>) -> (tensor<i32>) {

  tf_executor.graph {
    %c, %control0 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %a, %control1 = tf_executor.island wraps "tf.A"() {_class = ["loc:@class"]} : () -> (tensor<2xf32>)
    %s:2, %control2 = tf_executor.island wraps "tf.Split"(%c, %a) {num_split = 2 : i32} : (tensor<i32>, tensor<2xf32>) -> (tensor<1xf32>, tensor<1xf32>)
    tf_executor.fetch
  }
  func.return
}

// -----

!tf_res_i32 = tensor<*x!tf_type.resource<tensor<i32>>>
!tf_res_f32 = tensor<*x!tf_type.resource<tensor<f32>>>

// CHECK-LABEL: func @hoist_tail_assign
// CHECK-SAME:  ([[ARG0:%.*]]: tensor<*x!tf_type.resource<tensor<f32>>>)
func.func @hoist_tail_assign(%arg0: !tf_res_f32) {
  // CHECK: [[REPLICATE:%.*]]:4 = tf_device.replicate {n = 2 : i32}
  %replicate:2 = tf_device.replicate {n = 2 : i32} {
    // CHECK: [[OP_A:%.*]] = "tf.OpA"

  %input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
  %bias_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
  %bias_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
  %output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
  %output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1399 : i32}} {

// CHECK-LABEL: stateless_op
func.func @stateless_op() -> tensor<i32> {
  // CHECK: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
  %0 = "tf.StatelessRandomGetAlg"() {device = ""} : () -> tensor<i32>
  return %0 : tensor<i32>
}

^bb0(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>):
  // CHECK: tac.cost = 0x4D5C0000
  %0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32, fused_activation_function = "RELU6", tac.device = "CPU"} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32>
  func.return %0 : tensor<256x32x32x16xf32>