// 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>}>

  let description = [{
    If we pass same TensorList to a while op as multiple arguments or just use
    the same TensorList at multiple places and assign different
    TensorListSetItem to elements of TensorList, the shape inference is then
    unable to identify the Shape of these args and thus the input TensorList
    shape is unidentifiable.

  // |item| is expanded to have the same dimension as input_handle and
  // |index| is expanded to [index, 0, 0, ...] as the indices to input_handle.
  // On a high level, it's doing something like:
  // def : Pat<(TensorListSetItem($input_handle, $index, $item)),
  //           (XlaDynamicUpdateSlice($input_handle, ExpandDims($item, 0),
  //              Concat(ExpandDims($index, 0), [0, 0, 0, ...])))>

_Replicate TensorList init ops for correct shape assignments in shape inference_

If we pass same TensorList to a while op as multiple arguments or just use
the same TensorList at multiple places and assign different
TensorListSetItem to elements of TensorList, the shape inference is then
unable to identify the Shape of these args and thus the input TensorList
shape is unidentifiable.

      "TensorListFromTensor",
      "TensorListGather",
      "TensorListGetItem",
      "TensorListLength",
      "TensorListPopBack",
      "TensorListPushBack",
      "TensorListReserve",
      "TensorListSetItem",
      "TensorListSplit",
      "TensorListStack",
      "TensorScatterAdd",
      "TensorScatterMax",
      "TensorScatterMin",
      "TensorScatterSub",
      "TensorScatterUpdate",

    }
  };

  while (!worklist.empty()) {
    tensorlist = worklist.top();
    worklist.pop();

    // TensorLists are semantically immutable. For example, TensorListSetItem
    // takes a TensorList as input and produces a TensorList as output. So to
    // traverse modifications to TensorList and verify that all elements written

  let results = (outs
    TF_VariantTensor:$output_handle
  );

  TF_DerivedOperandTypeAttr element_dtype = TF_DerivedOperandTypeAttr<1>;
}

def TF_TensorListSetItemOp : TF_Op<"TensorListSetItem", [Pure]> {
  let summary = "";

  let arguments = (ins
    TF_VariantTensor:$input_handle,
    TF_Int32Tensor:$index,
    TF_Tensor:$item,