Feature

class dgl.graphbolt.Feature[source]

Bases: object

A wrapper of feature data for access.

count()[source]

Get the count of the feature.

Returns:

The count of the feature.

Return type:

int

metadata()[source]

Get the metadata of the feature.

Returns:

The metadata of the feature.

Return type:

Dict

read(ids: Tensor | None = None)[source]

Read from the feature.

Parameters:

ids (torch.Tensor, optional) – The index of the feature. If specified, only the specified indices of the feature are read. If None, the entire feature is returned.

Returns:

The read feature.

Return type:

torch.Tensor

read_async(ids: Tensor)[source]

Read the feature by index asynchronously.

Parameters:

ids (torch.Tensor) – The index of the feature. Only the specified indices of the feature are read.

Returns:

  • A generator object. – The returned generator object returns a future on read_async_num_stages(ids.device)`th invocation. The return result can be accessed by calling `.wait(). on the returned future object. It is undefined behavior to call .wait() more than once.

  • Example Usage

  • ——–

  • >>> import dgl.graphbolt as gb

  • >>> feature = gb.Feature(…)

  • >>> ids = torch.tensor([0, 2])

  • >>> for stage, future in enumerate(feature.read_async(ids))

  • … pass

  • >>> assert stage + 1 == feature.read_async_num_stages(ids.device)

  • >>> result = future.wait() # result contains the read values.

read_async_num_stages(ids_device: device)[source]

The number of stages of the read_async operation. See read_async function for directions on its use. This function is required to return the number of yield operations when read_async is used with a tensor residing on ids_device.

Parameters:

ids_device (torch.device) – The device of the ids parameter passed into read_async.

Returns:

The number of stages of the read_async operation.

Return type:

int

size()[source]

Get the size of the feature.

Returns:

The size of the feature.

Return type:

torch.Size

update(value: Tensor, ids: Tensor | None = None)[source]

Update the feature.

Parameters:

  • value (torch.Tensor) – The updated value of the feature.

  • ids (torch.Tensor, optional) – The indices of the feature to update. If specified, only the specified indices of the feature will be updated. For the feature, the ids[i] row is updated to value[i]. So the indices and value must have the same length. If None, the entire feature will be updated.