attr {
    key: "element_dtype"
    value {
      type: DT_FLOAT
    }
  }
  experimental_debug_info {
  }
}
node {
  name: "TensorListStack"
  op: "TensorListStack"
  input: "TensorListSetItem"
  input: "TensorListStack/element_shape"
  attr {
    key: "element_dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "num_elements"
    value {

  // CHECK: "tf.Const"() <{value = dense<> : tensor<0x3xi32>}> : () -> tensor<0x3xi32>
  // CHECK-NOT: tf.TensorListStack
  %1 = "tf.TensorListStack"(%0, %arg0) : (tensor<!tf_type.variant<tensor<3xi32>>>, tensor<1xi32>) -> tensor<0x3xi32>
  func.return %1 : tensor<0x3xi32>
}

// -----

// CHECK-LABEL: func @tensorlistGetItem

    %elem_shape_2 = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32>
    %tls = "tf.TensorListStack"(%tl_set_item, %elem_shape_2) {num_elements = 300 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<300x?xf32>

    %elem_2 = "tf.Const"() {value = dense<10.0> : tensor<9xf32>} : () -> tensor<9xf32>

def LegalizeTensorListStack : Pat<(TF_TensorListStackOp $input, $shape, $unused_num_elements),
    (TFL_CustomOp (Size2InputRange $input, $shape),
    (CreateStringAttr<"\"TensorListStack\"">), (EmptyCustomOptions))>;

def LegalizeTensorListSetItem : Pat<(TF_TensorListSetItemOp $input, $index, $item, $unused_resize_attr),
    (TFL_CustomOp (Size3InputRange $input, $index, $item),

// CHECK-LABEL: listStack
func.func @listStack(%arg0: tensor<!tf_type.variant<tensor<*xi32>>>, %arg1: tensor<i32>) -> tensor<*xi32> {
  %0 = "tf.TensorListStack"(%arg0, %arg1) : (tensor<!tf_type.variant<tensor<*xi32>>>, tensor<i32>) -> tensor<*xi32>
  // CHECK: %0 = "tfl.custom"(%arg0, %arg1) <{custom_code = "TensorListStack", custom_option = #tfl<const_bytes : "0x">}> : (tensor<!tf_type.variant<tensor<*xi32>>>, tensor<i32>) -> tensor<*xi32>
  func.return %0 : tensor<*xi32>

  %1 = "tf.TensorListSetItem"(%0, %arg2, %arg3) : (tensor<!tf_type.variant<tensor<10xf32>>>, tensor<i32>, tensor<10xf32>) -> tensor<!tf_type.variant<tensor<10xf32>>>
  %2 = "tf.TensorListStack"(%1, %arg1) : (tensor<!tf_type.variant<tensor<10xf32>>>, tensor<1xi32>) -> tensor<3x10xf32>
  func.return %2 : tensor<3x10xf32>

  // CHECK-NEXT: TensorListStack
  %2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<2x8xf32>

  %addn2 = "tf.AddN"(%addn, %zeros-like) : (tensor<!tf_type.variant<tensor<f32>>>, tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>>
  %stack = "tf.TensorListStack"(%addn2, %elem_shape) : (tensor<!tf_type.variant<tensor<f32>>>, tensor<0xi32>) -> tensor<10xf32>
  // CHECK-NEXT: %[[LEN:.*]] = "tf.Const"() <{value = dense<10> : tensor<i32>}> : () -> tensor<i32>

      %outputs_55, %control_56 = tf_executor.island wraps "tf.TensorListStack"(%outputs_53#3, %outputs_27) {_tpu_replicate = "cluster_sample_sequence", device = "", num_elements = 1024 : i64} : (tensor<!tf_type.variant>, tensor<1xi32>) -> tensor<*xf32>

        %8 = "tf.StopGradient"(%6) : (tensor<*x!tf_type.variant>) -> tensor<*x!tf_type.variant>
        // CHECK: tf.TensorListStack{{.*}}: (tensor<!tf_type.variant<tensor<2x2xf32>>>, tensor<i32>) -> tensor<?x2x2xf32>
        %9 = "tf.TensorListStack"(%8, %7) {device = "", num_elements = -1 : i64} : (tensor<*x!tf_type.variant>, tensor<i32>) -> tensor<*xf32>
        tf_executor.yield
      }