// is used as an input to an op (so we know when we're done computing gradients
// for that Tensor). We also count, for each tape entry, how many of its output
// Tensors need gradients to be computed (Tensors which are not used do not need
// any gradients to be computed).
//
// Finally, we start a backprop stack with a set of tape entries for which we
// have all gradients available. This set usually is a subset of the set of

// Struct to be filled in to define a custom device tensor handle. Fields are
// required except where indicated.
typedef struct TFE_CustomDeviceTensorHandleMethods {
  int version = TFE_CUSTOM_DEVICE_VERSION;

  // Computes the rank of the tensor handle.
  //
  // Shapes are specified via callbacks because retrieving the shape of a tensor
  // is a blocking operation for async eager; custom devices should avoid

        dtype_(dtype) {}

  const ParallelDevice& device_;
  std::vector<TensorHandlePtr> tensors_;
  // Parallel tensors are immutable but compute their shape lazily unless it is
  // provided on construction. The optional has a value if the lazy computation
  // has been completed or the shape was provided on construction.
  mutable absl::optional<std::vector<int64_t>> shape_;
  const TF_DataType dtype_;
};

// will block till the operation that produces `h` has completed.
TF_CAPI_EXPORT extern const char* TFE_TensorHandleDeviceName(
    TFE_TensorHandle* h, TF_Status* status);

// Returns the name of the device in whose memory `h` resides.
//
// This function will block till the operation that produces `h` has completed.
TF_CAPI_EXPORT extern const char* TFE_TensorHandleBackingDeviceName(