// graph mode). Calling this on a context tells it to start a step.
TF_CAPI_EXPORT extern void TFE_ContextStartStep(TFE_Context* ctx);

// Ends a step. When there is no active step (that is, every started step has
// been ended) step containers will be cleared. Note: it is not safe to call
// TFE_ContextEndStep while ops that rely on the step container may be running.

  // Tensor associated with `tensor_id` is deleted.
  void DeleteGradient(int64_t tensor_id);

  // Runs forward autodiff. Should be called whenever a new operation is
  // available and the accumulator is active.
  //
  // Like GradientTape::RecordOperation, this method takes the operation type
  // `op_type` (e.g. "Add"), the operation's inputs (`input_tensors`,
  // `input_tensor_id`, and `input_dtypes`; the latter two are somewhat

//       populated by core. The plugin is responsible for setting
//       `struct_size` as well, along with all other fields.
//     * Refer to "TensorFlow Versioning Strategy" section at
//       https://github.com/tensorflow/community/pull/257/files.
//     * Note that the API is still under active development and doesn't have
//       versioning guarantees yet.
//   * `void* ext` is a free-form field that can be populated by

// TF_ExecuteOperation will, if in eager mode, execute, if in graph mode, maybe
// capture some inputs and then add a node in the graph. The output tensors are
// returned through the provided TF_OutputList.
// Any active tape will observe the effects of this execution.
void TF_ExecuteOperation(TF_AbstractOp* op, int num_inputs,
                         TF_AbstractTensor* const* inputs, TF_OutputList* o,
                         TF_Status* s);

  // Maps from name of an operation to the Node* in 'graph'.
  std::unordered_map<tensorflow::string, tensorflow::Node*> name_map
      TF_GUARDED_BY(mu);

  // The keys of this map are all the active sessions using this graph. Each
  // value records whether the graph has been mutated since the corresponding
  // session has been run (this is detected in RecordMutation function). If the

#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_status.h"
#include "xla/pjrt/c/pjrt_c_api.h"

// --------------------------------------------------------------------------
// C API for device. The API is under active development and eventually
// should allow registering a plugin device with TensorFlow.

#ifdef __cplusplus
extern "C" {
#endif

// TF_Device is a C wrapper to the C++ TF Device class. This is to be passed

  struct CacheStats {
    int64_t kernel_cache_size;
    int64_t device_cache_size;
    std::map<std::string, int64_t> func_kernel_cache_entries;
    int64_t local_rendezvous_cache_active_size;
  };
  virtual CacheStats GetCacheStats() = 0;

  //===--------------------------------------------------------------------===//
  // Distributed runtime related functions.