// RUN: tf-opt -tfl-legalize-tensorlist -canonicalize -split-input-file %s | FileCheck %s

// CHECK-LABEL: listReserveScalarShapeI32
func.func @listReserveScalarShapeI32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi32>>> {
  %0 = "tf.TensorListReserve"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi32>>>

// Returns a `llvm::DenseMap` which maps from the index of tensorlist in the
// result, to the index of the same tensorlist in the arguments. For `If` op's
// branch functions, the results and arguments are not usually matched 1-1. This
// will let us konw which tensorlist result maps to which tensorlist in the
// arguments. Once we know this info it will help us decide the types of the
// result tensorlist based on the operand's of the `If` op.

      });
  return has_variant_input || has_variant_output;
}

// There are 2 standard tf ops which are not TensorList ops that may take as
// input a tensorlist. These are tf.AddN and tf.ZeroesLike. Since the runtime
// implementation of a tensorlist are not compatible between tf and tflite
// we cannot use tflite tensorlist kernels until these cases are handled.
bool IsNonTensorListVariantOp(Operation* op) {

  if (a.TypeId() == tensorflow::TypeIndex::Make<::tensorflow::TensorList>()) {
    TF_RETURN_IF_ERROR(ValidateVariantType<::tensorflow::TensorList>(a));
    *out = ::tensorflow::TensorList();

    return ::tensorflow::TensorListBinaryAdd(
        cc_ctx, *a.get<::tensorflow::TensorList>(),
        *b.get<::tensorflow::TensorList>(),
        out->get<::tensorflow::TensorList>(), cc_binary_add);

             "legalized by this pass, then the IR won't be changed so that "
             "tensorlist ops can pass through (default false).">,
      Option<
        "default_to_single_batch_", "default-to-single-batch", "bool", "false",
        "When specified to true, if the tensorlist ops has unspecified batch "
        "size, this pass will assume that the batch size is one to proceed "

  OpBuilder builder(list);
  llvm::SmallVector<int64_t, 8> element_shape;
  // Infer TensorList element shape from the return type first, and then from
  // the const element shape operand. We first check the return type because
  // shape inference might have successfully inferred the element shape from
  // write operations on the TensorList.
  if (failed(GetElementShapeFromResultType(list.getType(), &element_shape))) {

  // If `emit_builtin_tflite_ops` is true, TF Lite legalization passes will be
  // added, which produces TF Lite ops.
  bool emit_builtin_tflite_ops = true;
  // If `lower_tensor_list_ops` is true, tensorlist ops will be lowered to basic
  // TF ops before legalization to TF Lite dialect.
  bool lower_tensor_list_ops = false;
  // The allowlist of functions that would be preserved after trimming.

    }
  };

  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

    // meaningless value -1 for non resource outputs.
    out_desc.input_index =
        it != output_to_input_alias.end() ? it->getSecond() : -1;
    // MLIR-based TF-Compiler bridge doesn't support tensorlist output yet.
    // TODO(ycao): Support tensorlist-type output.
    out_desc.is_tensor_list = false;
  }

  // XLA computation always uses Tuple shape.
  *xla_output_shape = xla::ShapeUtil::MakeTupleShape(shapes);

  let summary =
    "Replicate TensorList init ops for correct shape assignments in shape inference";

  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