Class TensorValue

tensorflow/core/framework/op_kernel.h:538–569

Holds a tensor or tensor reference. For tensor references, we need a mutex to prevent concurrent access to the tensor.

// Holds a tensor or tensor reference. For tensor references, we need
// a mutex to prevent concurrent access to the tensor.
struct TensorValue {
  TensorValue() : mutex_if_ref(nullptr), tensor(nullptr) {}
  explicit TensorValue(Tensor* t) : mutex_if_ref(nullptr), tensor(t) {}
  TensorValue(mutex* mu, Tensor* t) : mutex_if_ref(mu), tensor(t) {}
  Tensor* operator->() const { return tensor; }
  bool is_ref() const { return mutex_if_ref != nullptr; }

  // Return the dtype of the Tensor. For references, return the underlying type.
  DataType dtype() const {
    if (is_ref()) {
      return MakeRefType(tensor->dtype());
    } else {
      return tensor->dtype();
    }
  }

  // Return the dtype of the Tensor. For references, return the underlying type.
  // This variation on the dtype() acquires the lock for references.
  //
  // TODO(b/133843385): Disallow dtype modifications
  DataType dtype_safe() const {
    if (is_ref()) {
      tf_shared_lock ml(*mutex_if_ref);
      return MakeRefType(tensor->dtype());
    } else {
      return tensor->dtype();
    }
  }

  mutex* mutex_if_ref;  // nullptr if not a ref, != nullptr if a ref
  Tensor* tensor;
};

// Used to store partitioned graphs from function-calling ops.
struct GraphCollector {