cugraph_pyg.data.feature_store.WholeFeatureStore#

class cugraph_pyg.data.feature_store.WholeFeatureStore(memory_type='distributed', location='cpu')[source]#

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.

Methods

get_all_tensor_attrs()

Returns all registered tensor attributes.

__init__(memory_type='distributed', location='cpu')[source]#
Parameters:
memory_type: str (optional, default=’distributed’)

The memory type of this store.

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

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

Methods

__init__([memory_type, location])

Parameters:

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.