torch_geometric.data.FeatureStore

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

Bases: ABC

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.

  • *args – Arguments passed to TensorAttr.

  • **kwargs – Keyword arguments passed to TensorAttr.

Raises:

ValueError – If the input TensorAttr is not fully specified.

get_tensor(*args, convert_type: bool = False, **kwargs) → Union[Tensor, ndarray][source]

Synchronously obtains a tensor from the FeatureStore.

Parameters:

  • *args – Arguments passed to TensorAttr.

  • convert_type (bool, optional) – Whether to convert the type of the output tensor to the type of the attribute index. (default: False)

  • **kwargs – Keyword arguments passed to TensorAttr.

Raises:

ValueError – If the input TensorAttr is not fully specified.

multi_get_tensor(attrs: List[TensorAttr], convert_type: bool = False) → 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.

  • convert_type (bool, optional) – Whether to convert the type of the output tensor to the type of the attribute index. (default: False)

Raises:

ValueError – If any input TensorAttr is not fully specified.

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

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

Parameters:
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.

  • *args – Arguments passed to TensorAttr.

  • **kwargs – Keyword arguments passed to TensorAttr.

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]

Returns all registered tensor attributes.

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

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