}

  Status AddInput(AbstractTensorHandle* input) override {
    GraphTensor* t = dyn_cast<GraphTensor>(input);
    if (!t) {
      return tensorflow::errors::InvalidArgument(
          "Unable to cast input to GraphTensor");
    }
    TF_AddInput(op_.get(), t->output_);
    return absl::OkStatus();
  }
  Status AddInputList(absl::Span<AbstractTensorHandle* const> inputs) override {

// 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]],

is used instead.

Operands are replicated inputs and packed inputs.

replicated_inputs: each group of `n` inputs corresponds to an input for a single
individual replica and is mapped to a single region argument. Inside one group
the operands are matching in order the `devices` attribute. Each replicated
input must have compatible shapes and types.
packed_inputs: each input corresponds to an input broadcasted across all

  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) {

    Node* input =
        ops::SourceOp("FakeNullary", builder.opts().WithName("Input"));
    Node* clustered_producer =
        ops::BinaryOp("FakeBinary", input, input,
                      builder.opts().WithName("ClusteredProducer"));
    // The first input is hostmem and the second input is devicemem.
    Node* consumer_in_different_cluster =
        ops::BinaryOp("FakeBinary", input, clustered_producer,

}

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

class MatMulGradientFunction : public GradientFunction {
 public:
  explicit MatMulGradientFunction(vector<AbstractTensorHandle*> f_inputs,
                                  AttrBuilder f_attrs)
      : forward_inputs_(f_inputs), forward_attrs_(f_attrs) {
    for (auto input : forward_inputs_) {
      if (input) {
        input->Ref();
      }
    }
  }

  Status Compute(AbstractContext* ctx,

    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).");