%0 = "tf.TensorListSetItem"(%arg0, %arg1, %arg2) : (tensor<!tf_type.variant<tensor<*xi32>>>, tensor<i32>, tensor<*xi32>) -> tensor<!tf_type.variant<tensor<*xi32>>>

def LegalizeTensorListSetItem : Pat<(TF_TensorListSetItemOp $input, $index, $item, $unused_resize_attr),
    (TFL_CustomOp (Size3InputRange $input, $index, $item),
    (CreateStringAttr<"\"TensorListSetItem\"">), (EmptyCustomOptions))>;

def LegalizeTensorListFromTensor : Pat<(TF_TensorListFromTensorOp $input, $element_shape),
    (TFL_CustomOp (Size2InputRange $input, $element_shape),

}

// -----

// CHECK-LABEL: tensorlistSetItem
func.func @tensorlistSetItem(%arg0: tensor<3x10xf32>, %arg1: tensor<1xi32>, %arg2: tensor<i32>, %arg3: tensor<10xf32>) -> tensor<3x10xf32> {
  %0 = "tf.TensorListFromTensor"(%arg0, %arg1) : (tensor<3x10xf32>, tensor<1xi32>) -> tensor<!tf_type.variant<tensor<10xf32>>>

// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.Mul
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.Mul
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf.TensorListSetItem
// CHECK-NEXT: tf.TensorListSetItem
// CHECK-NEXT: tf_mlrt.tf_promise
// CHECK-NEXT: tf_mlrt.tf_promise
// CHECK-NEXT: return

// CHECK-LABEL: tf.MultiListWhileRegion_cond

      %element = "tf.Const"() <{value = dense<0.0> : tensor<64x1xbf16>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<64x1xbf16>
      // Results in a bad mismatch of shapes.
      %updated_list = "tf.TensorListSetItem"(%list, %index, %element) : (tensor<!tf_type.variant<tensor<64x1xbf16>>>, tensor<i32>, tensor<64x1xbf16>) -> tensor<!tf_type.variant<tensor<64x1xbf16>>>

      return %updated_list : tensor<!tf_type.variant<tensor<64x1xbf16>>>

// NOLINTNEXTLINE
opt<bool> enable_dynamic_update_slice(
    "enable-dynamic-update-slice",
    llvm::cl::desc("Whether to enable dynamic update slice op to convert "
                   "TensorListSetItem op."),
    llvm::cl::init(false));

// NOLINTNEXTLINE
opt<bool> import_hlo("import-hlo",
                     llvm::cl::desc("Whether the input file is hlo file."),
                     llvm::cl::init(false));

        // CHECK: tf.TensorListSetItem{{.*}}: (tensor<!tf_type.variant<tensor<2x2xf32>>>, tensor<i32>, tensor<2x2xf32>) -> tensor<!tf_type.variant<tensor<2x2xf32>>>
        %6 = "tf.TensorListSetItem"(%3, %4, %5) {device = ""} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<2x2xf32>)-> tensor<*x!tf_type.variant>

      %outputs_8, %control_9 = tf_executor.island wraps "tf.Identity"(%outputs_6) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
      %outputs_10, %control_11 = tf_executor.island wraps "tf.TensorListSetItem"(%arg3, %arg2, %outputs_4) {device = "", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant>, tensor<i32>, tensor<*xf32>) -> tensor<*x!tf_type.variant>

        "tensorlist op lowering (default true)">,
      Option<"enable_dynamic_update_slice_", "enable-dynamic-update-slice",
             "bool", "false", "When specified to true, lower TensorListSetItem with "
             "DynamicUpdateSlice op (default false).">,
  ];
}

def ModifyIONodesPass : Pass<"tfl-modify-io-nodes", "mlir::func::FuncOp"> {
  let summary = "Modify the type of the model io nodes";

  // CHECK-NEXT: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[BROADCAST]], %[[UPDATE_SLICE]], %[[SET_INDEX]]) : (tensor<10xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<10xf32>
  %set = "tf.TensorListSetItem"(%tl, %arg0, %elem) : (tensor<!tf_type.variant<tensor<f32>>>, tensor<i32>, tensor<f32>) -> tensor<!tf_type.variant<tensor<f32>>>
  // CHECK-NEXT: %[[SIZE_SHAPE2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>