cugraph_pyg.data.feature_store.TensorDictFeatureStore#

class cugraph_pyg.data.feature_store.TensorDictFeatureStore[source]#

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__()[source]#

Methods

__init__()

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.