cugraph_pyg.data.feature_store.TensorDictFeatureStore

class cugraph_pyg.data.feature_store.TensorDictFeatureStore

A basic implementation of the PyG FeatureStore interface that stores feature data in a single TensorDict. This type of feature store is not distributed, so each node will have to load the entire graph’s features into memory.

Methods

get_all_tensor_attrs()

Returns all registered tensor attributes.

__init__()

Constructs an empty TensorDictFeatureStore.

Methods

__init__()	Constructs an empty TensorDictFeatureStore.
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.