cugraph_pyg.data.feature_store.FeatureStore

class cugraph_pyg.data.feature_store.FeatureStore(memory_type=None, location='cpu')

A basic implementation of the PyG FeatureStore interface that stores feature data in WholeGraph WholeMemory. This type of feature store is distributed, and avoids data replication across workers.

Data should be sliced before being passed into this feature store. That means each worker should have its own partition and put_tensor should be called for each worker’s local partition. When calling get_tensor, multi_get_tensor, etc., the entire tensor can be accessed regardless of what worker’s partition the desired slice of the tensor is on.

Methods

get_all_tensor_attrs()

Returns all registered tensor attributes.

__init__(memory_type=None, location='cpu')

Constructs an empty FeatureStore.

Parameters:

memory_type: str (optional, default=None)

Has no effect. Retained for compatibility purposes.

location: str(optional, default=’cpu’)

The location (‘cpu’ or ‘cuda’) where data is stored.

Methods

__init__([memory_type, location])	Constructs an empty FeatureStore.
get_all_tensor_attrs()	Returns all registered tensor attributes.
get_tensor(*args[, convert_type])	Synchronously obtains a tensor from the FeatureStore.
get_tensor_size(*args, **kwargs)	Obtains the size of a tensor given its TensorAttr, or None if the tensor does not exist.
multi_get_tensor(attrs[, convert_type])	Synchronously obtains a list of tensors from the FeatureStore for each tensor associated with the attributes in attrs.
put_tensor(tensor, *args, **kwargs)	Synchronously adds a tensor to the FeatureStore.
remove_tensor(*args, **kwargs)	Removes a tensor from the FeatureStore.
update_tensor(tensor, *args, **kwargs)	Updates a tensor in the FeatureStore with a new value.
view(*args, **kwargs)	Returns a view of the FeatureStore given a not yet fully-specified TensorAttr.