torch_geometric.data.FeatureStore

class FeatureStore(tensor_attr_cls: Optional[Any] = None)[source]

Bases: object

An abstract base class to access features from a remote feature store.

Parameters

tensor_attr_cls (TensorAttr, optional) – A user-defined TensorAttr class to customize the required attributes and their ordering to unique identify tensor values. (default: None)

put_tensor(tensor: Union[Tensor, ndarray], *args, **kwargs) bool[source]

Synchronously adds a tensor to the FeatureStore. Returns whether insertion was successful.

Parameters
  • tensor (torch.Tensor or np.ndarray) – The feature tensor to be added.

  • **kwargs (TensorAttr) – Any relevant tensor attributes that correspond to the feature tensor. See the TensorAttr documentation for required and optional attributes.

Raises

ValueError – If the input TensorAttr is not fully specified.

get_tensor(*args, **kwargs) Union[Tensor, ndarray][source]

Synchronously obtains a tensor from the FeatureStore.

Parameters

**kwargs (TensorAttr) – Any relevant tensor attributes that correspond to the feature tensor. See the TensorAttr documentation for required and optional attributes.

Raises
multi_get_tensor(attrs: List[TensorAttr]) List[Union[Tensor, ndarray]][source]

Synchronously obtains a list of tensors from the FeatureStore for each tensor associated with the attributes in attrs.

Note

The default implementation simply iterates over all calls to get_tensor(). Implementor classes that can provide additional, more performant functionality are recommended to to override this method.

Parameters

attrs (List[TensorAttr]) – A list of input TensorAttr objects that identify the tensors to obtain.

Raises
remove_tensor(*args, **kwargs) bool[source]

Removes a tensor from the FeatureStore. Returns whether deletion was successful.

Parameters

**kwargs (TensorAttr) – Any relevant tensor attributes that correspond to the feature tensor. See the TensorAttr documentation for required and optional attributes.

Raises

ValueError – If the input TensorAttr is not fully specified.

update_tensor(tensor: Union[Tensor, ndarray], *args, **kwargs) bool[source]

Updates a tensor in the FeatureStore with a new value. Returns whether the update was succesful.

Note

Implementor classes can choose to define more efficient update methods; the default performs a removal and insertion.

Parameters
  • tensor (torch.Tensor or np.ndarray) – The feature tensor to be updated.

  • **kwargs (TensorAttr) – Any relevant tensor attributes that correspond to the feature tensor. See the TensorAttr documentation for required and optional attributes.

get_tensor_size(*args, **kwargs) Optional[Tuple[int, ...]][source]

Obtains the size of a tensor given its TensorAttr, or None if the tensor does not exist.

abstract get_all_tensor_attrs() List[TensorAttr][source]

Obtains all tensor attributes stored in this FeatureStore.

view(*args, **kwargs) AttrView[source]

Returns a view of the FeatureStore given a not yet fully-specified TensorAttr.