Class PersistentTensor

tensorflow/core/framework/op_kernel.h:267–287

Wraps a tensor that is held by an Op across calls to Compute(). For memory safety when using asynchronous devices like GPUs, the system must be notified when a Tensor is used inside an Op execution. The wrapper ensures that all uses of the Tensor are tracked, because in order to retrieve the Tensor the caller must use AccessTensor which notifies the context.

// order to retrieve the Tensor the caller must use AccessTensor which
// notifies the context.
class PersistentTensor {
 public:
  PersistentTensor() {}
  explicit PersistentTensor(const Tensor& tensor) : tensor_(tensor) {}

  // Caller does not own the returned Tensor*.
  Tensor* AccessTensor(OpKernelConstruction* context);
  // Caller does not own the returned Tensor*.
  Tensor* AccessTensor(OpKernelContext* context);

  // The check for initialization does not need to access the
  // underlying tensor buffer.
  bool IsInitialized() const { return tensor_.IsInitialized(); }

  int64 NumElements() const { return tensor_.NumElements(); }

  int64 AllocatedBytes() const { return tensor_.AllocatedBytes(); }

 private:
  Tensor tensor_;
};

class OpKernelConstruction {
 public: