Status MlirAbstractOp::Execute(absl::Span<AbstractTensorHandle*> retvals,
                               int* num_retvals) {
  OperationState* state;
  TF_RETURN_IF_ERROR(Create(operands_, &state));
  Operation* op = function_context_->CreateOperationFromState(*state);
  *num_retvals = op->getNumResults();
  for (int i = 0; i < *num_retvals; i++)
    retvals[i] = new MlirTensor(op->getResult(i));
  return absl::OkStatus();
}

                           ForwardOperation*);

// Make the call to `Tape::RecordOperation`.
Status Execute(AbstractOperation*, AbstractContext*,
               absl::Span<AbstractTensorHandle*> retvals, int* num_retvals,
               ForwardOperation*, Tape*, const GradientRegistry&);

}  // namespace internal
}  // namespace gradients
}  // namespace tensorflow

                             int* num_retvals) override {
    std::vector<TFE_TensorHandle*> outputs(*num_retvals);
    TF_Status status;
    device_.execute(tensorflow::wrap(op), num_retvals, outputs.data(), &status,
                    info_);
    if (status.status.ok()) {
      for (int i = 0; i < *num_retvals; ++i) {
        retvals[i] = tensorflow::unwrap(outputs[i]);

Status Execute(AbstractOperation* op_, AbstractContext* ctx,
               absl::Span<AbstractTensorHandle*> retvals, int* num_retvals,
               ForwardOperation* forward_op_, Tape* tape,
               const GradientRegistry& registry) {
  TF_RETURN_IF_ERROR(op_->Execute(retvals, num_retvals));
  for (int i = 0; i < *num_retvals; i++) {
    // TODO(srbs): Manage refcount of ForwardOperation's inputs/outputs.

    TF_RETURN_IF_ERROR(function_op->SetDeviceName(resource.device.c_str()));

    AbstractTensorHandle* resource_handle = nullptr;
    int num_retvals = 1;
    TF_RETURN_IF_ERROR(function_op->Execute(
        absl::MakeSpan(&resource_handle, num_retvals), &num_retvals));
    AbstractTensorHandlePtr owned_resource_handle(resource_handle);
    if (!tensorflow::isa<ImmediateExecutionTensorHandle>(

  virtual Status AddInputList(
      absl::Span<AbstractTensorHandle* const> inputs) = 0;
  virtual Status Execute(absl::Span<AbstractTensorHandle*> retvals,
                         int* num_retvals) = 0;

  virtual Status SetAttrString(const char* attr_name, const char* data,
                               size_t length) = 0;
  virtual Status SetAttrInt(const char* attr_name, int64_t value) = 0;

  ASSERT_TRUE(tensorflow::errors::IsInvalidArgument(v.status())) << v.status();
  ASSERT_TRUE(absl::StrContains(v.status().message(),
                                "Expecting 2 outputs, but *num_retvals is 1"));
}

TEST_P(FunctionTest, IncorrectDtypeInOutputSignatureFails) {
  // Construct a scalar.
  impl::TaggedValueTensor x = CreateScalarTensor<float, TF_FLOAT>(2.0f);
  FunctionDef fdef = AddFunc();