// CHECK: %[[lstm:.*]] = "tfl.unidirectional_sequence_lstm"(%[[input_0]], %[[input_1]], %[[input_2]], %[[input_3]], %[[input_4]], %[[input_5]], %[[input_6]], %[[input_7]], %[[input_8]],
// CHECK-SAME: %[[input_9]], %[[input_10]], %[[input_11]], %[[input_12]], %[[input_13]], %[[input_14]], %[[input_15]], %[[input_16]], %[[input_17]], %[[input_18]], %[[input_19]]

// CHECK-SAME: %[[input_0]],
// CHECK-SAME: %[[input_1]], %[[input_2]], %[[input_3]], %[[input_4]],
// CHECK-SAME: %[[input_5]], %[[input_6]], %[[input_7]], %[[input_8]],
// CHECK-SAME: %[[input_9]], %[[input_9]], %[[input_9]],
// CHECK-SAME: %[[input_10]], %[[input_11]], %[[input_12]], %[[input_13]],
// CHECK-SAME: %[[input_9]], %[[input_9]],
// CHECK-SAME: %[[input_14]], %[[input_15]],

  std::vector<TFE_TensorHandle*> compute_fn_inputs;
  TFE_TensorHandle* input_a = TestScalarTensorHandle(ctx, 2.0f);
  TFE_TensorHandle* input_b = TestScalarTensorHandle(ctx, 1.0f);
  compute_fn_inputs.push_back(input_a);
  compute_fn_inputs.push_back(input_b);

  TFE_Op* compute_fn_op = TF_ConcreteFunctionMakeCallOp(

  // TFL lstm only supports time-majored inputs, so if it's not time-majored,
  // we will transpose the inputs and outputs.
  auto time_major_attr = func_op->getAttrOfType<BoolAttr>("tf.time_major");
  if (time_major_attr == nullptr) return failure();

  bool time_majored = time_major_attr.getValue();
  auto input_type = mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
  if (!input_type) {

void TF_AddInput(TF_OperationDescription* desc, TF_Output input) {
  desc->node_builder.Input(&input.oper->node, input.index);
}

void TF_AddInputList(TF_OperationDescription* desc, const TF_Output* inputs,
                     int num_inputs) {
  std::vector<NodeBuilder::NodeOut> input_list;
  input_list.reserve(num_inputs);
  for (int i = 0; i < num_inputs; ++i) {
    input_list.emplace_back(&inputs[i].oper->node, inputs[i].index);
  }

}

// CHECK-LABEL: func @tensor_array_while_test
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain
func.func @tensor_array_while_test(%indices: tensor<?xi32>, %input_0: tensor<?x?x?xf32>, %input_1: tensor<?x?x?xf32>) -> (tensor<?x?x512xf32>, tensor<?x?x512xf32>) {
  %index = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> (tensor<i32>)

      : inputs_(inputs.begin(), inputs.end()) {}

  InputList(const absl::Span<const Input>& inputs)  // NOLINT(runtime/explicit)
      : inputs_(inputs.begin(), inputs.end()) {}

  InputList(
      const std::initializer_list<Output>& out) {  // NOLINT(runtime/explicit)
    for (auto const& x : out) {
      inputs_.push_back(x);
    }
  }

                    (MHLO_CompositeOp:$old_val
                    (variadic $a_input), 
                    ConstantStrAttr<StrAttr, "aten.avg_pool2d.default">, $attrs, $_, $_),
                    (TFL_TransposeOp 
                        (TFL_AveragePool2DOp 
                            /*input*/ (TFL_TransposeOp $a_input,
                                (Arith_ConstantOp

    return errors::InvalidArgument(
        "Must specify a gradient input for each output.");
  }
  std::vector<bool> reachable_nodes = GetReachableNodes();
  for (const Output& input : inputs_) {
    if (!reachable_nodes[input.node()->id()]) {
      return errors::InvalidArgument(
          "Cannot compute the partial derivative for node '",
          input.node()->name(),
          "' as it's unreachable from the output node(s).");

  Init(2);
  // Create nested loop:
  //  while (input1 < 6) {
  //    inner_input1 = input1
  //    while (inner_input1 < 3) {
  //      input2 += 1
  //      inner_input1 += 2
  //    }
  //    input1 += input2
  //  }
  //
  // Expected execution with initial values input1 = input2 = 0:
  //
  // outer inner               inner_
  // step# step# input1 input2 input1
  // ------------------------------------