TF_Status* s);

// Creates a new TF_AbstractFunction from the current tracing states in the
// context. The provided `ctx` is consumed by this API call and deleted.
// The returned TF_AbstractFunction must be deleted by the client,
// TODO(aminim): clarify the contract on the state of the context after this
// call.
TF_AbstractFunction* TF_FinalizeFunction(TF_ExecutionContext* ctx,

#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/platform/refcount.h"
namespace tensorflow {
namespace tracing {
namespace graph {
using tensorflow::AbstractFunction;
// Thin wrapper around a FunctionDef.
class GraphFunction : public AbstractFunction {
 public:
  explicit GraphFunction(FunctionDef fdef);
  ~GraphFunction() override;

  // GraphFunction maybe stay alive for the duration of the returned

TF_ExecutionContext* TF_CreateFunction(const char* fn_name, TF_Status* s) {
  return wrap(CreateTracingExecutionContext(fn_name, s));
}

TF_AbstractFunction* TF_FinalizeFunction(TF_ExecutionContext* ctx,
                                         TF_OutputList* outputs, TF_Status* s) {
  AbstractFunction* func;
  TracingContext* tracing_ctx = dyn_cast<TracingContext>(unwrap(ctx));
  if (!tracing_ctx) {
    tsl::Set_TF_Status_from_Status(

 public:
  // Returns which subclass is this instance of.
  AbstractOpAttrsKind getKind() const { return kind_; }
  virtual ~AbstractOpAttrs() = default;

  // Returns the AbstractFunction as a FunctionDef.
  virtual void GetNameAttrList(
      tensorflow::NameAttrList* name_and_attrs) const = 0;

  virtual bool GetInt(absl::string_view, int64_t* result) const = 0;

// function.
class AbstractFunction : public core::RefCounted {
 protected:
  enum AbstractFunctionKind { kGraph, kMlir };
  explicit AbstractFunction(AbstractFunctionKind kind) : kind_(kind) {}

 public:
  // Returns which subclass is this instance of.
  AbstractFunctionKind getKind() const { return kind_; }

  // Returns the AbstractFunction as a FunctionDef.

}  // namespace tracing

DEFINE_CONVERSION_FUNCTIONS(AbstractContext, TF_ExecutionContext)
DEFINE_CONVERSION_FUNCTIONS(AbstractTensorHandle, TF_AbstractTensor)
DEFINE_CONVERSION_FUNCTIONS(AbstractFunction, TF_AbstractFunction)
DEFINE_CONVERSION_FUNCTIONS(AbstractOperation, TF_AbstractOp)
DEFINE_CONVERSION_FUNCTIONS(OutputList, TF_OutputList)
}  // namespace tensorflow

  ASSERT_EQ(h->FullType().type_id(), TFT_UNSET);
  ASSERT_EQ(unwrap(inputs[0])->FullType().type_id(), TFT_UNSET);

  // Clean up operation and inputs.
  TF_DeleteAbstractOp(add_op);

  TF_AbstractFunction* func =
      TF_FinalizeFunction(graph_ctx, add_outputs, status.get());
  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
  // Note: TF_OutputList does not own the underlying AbstractTensors, those

                absl::Span<AbstractTensorHandle* const> inputs,
                absl::Span<AbstractTensorHandle*> outputs, bool use_function) {
  if (use_function) {
    const char* fn_name = "test_fn";
    core::RefCountPtr<AbstractFunction> scoped_func;
    // Returning null tensors from a tf.function is not supported, so we keep
    // track of indices in the model's outputs are nullptr in this set.

#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/platform/status.h"

namespace tensorflow {
namespace tracing {
namespace graph {
GraphFunction::GraphFunction(FunctionDef fdef)
    : AbstractFunction(kGraph),
      func_record_(new FunctionRecord(std::move(fdef), {}, true)) {}
GraphFunction::~GraphFunction() {}
Status GraphFunction::GetFunctionDef(const FunctionDef **fdef) {
  *fdef = &(func_record_->fdef());

  // parent_ctx_->Unref();
}
TapeOperation* TapeContext::CreateOperation() {
  return new TapeOperation(parent_ctx_->CreateOperation(), tape_, registry_);
}
Status TapeContext::RegisterFunction(AbstractFunction* f) {
  return parent_ctx_->RegisterFunction(f);
}
Status TapeContext::RemoveFunction(const string& func) {
  return parent_ctx_->RemoveFunction(func);
}

}  // namespace gradients