ItemSet

class dgl.graphbolt.ItemSet(items: int | Tensor | Tuple[Tensor], names: str | Tuple[str] | None = None)

Bases: object

A wrapper of a tensor or tuple of tensors.

Parameters:

• items (Union[int, torch.Tensor, Tuple[torch.Tensor]]) –

  The tensors to be wrapped. - If it is a single scalar (an integer or a tensor that holds a single

  value), the item would be considered as a range_tensor created by torch.arange.

  • If it is a multi-dimensional tensor, the indexing will be performed along the first dimension.

  • If it is a tuple, each item in the tuple must be a tensor.

• names (Union[str, Tuple[str]], optional) – The names of the items. If it is a tuple, each name must corresponds to an item in the items parameter. The naming is arbitrary, but in general practice, the names should be chosen from [‘labels’, ‘seeds’, ‘indexes’] to align with the attributes of class dgl.graphbolt.MiniBatch.

Examples

>>> import torch
>>> from dgl import graphbolt as gb

1. Integer: number of nodes.

>>> num = 10
>>> item_set = gb.ItemSet(num, names="seeds")
>>> list(item_set)
[tensor(0), tensor(1), tensor(2), tensor(3), tensor(4), tensor(5),
 tensor(6), tensor(7), tensor(8), tensor(9)]
>>> item_set[:]
tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> item_set.names
('seeds',)

2. Torch scalar: number of nodes. Customizable dtype compared to Integer.

>>> num = torch.tensor(10, dtype=torch.int32)
>>> item_set = gb.ItemSet(num, names="seeds")
>>> list(item_set)
[tensor(0, dtype=torch.int32), tensor(1, dtype=torch.int32),
 tensor(2, dtype=torch.int32), tensor(3, dtype=torch.int32),
 tensor(4, dtype=torch.int32), tensor(5, dtype=torch.int32),
 tensor(6, dtype=torch.int32), tensor(7, dtype=torch.int32),
 tensor(8, dtype=torch.int32), tensor(9, dtype=torch.int32)]
>>> item_set[:]
tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int32)
>>> item_set.names
('seeds',)

3. Single tensor: seed nodes.

>>> node_ids = torch.arange(0, 5)
>>> item_set = gb.ItemSet(node_ids, names="seeds")
>>> list(item_set)
[tensor(0), tensor(1), tensor(2), tensor(3), tensor(4)]
>>> item_set[:]
tensor([0, 1, 2, 3, 4])
>>> item_set.names
('seeds',)

4. Tuple of tensors with same shape: seed nodes and labels.

>>> node_ids = torch.arange(0, 5)
>>> labels = torch.arange(5, 10)
>>> item_set = gb.ItemSet(
...     (node_ids, labels), names=("seeds", "labels"))
>>> list(item_set)
[(tensor(0), tensor(5)), (tensor(1), tensor(6)), (tensor(2), tensor(7)),
 (tensor(3), tensor(8)), (tensor(4), tensor(9))]
>>> item_set[:]
(tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9]))
>>> item_set.names
('seeds', 'labels')

5. Tuple of tensors with different shape: seeds and labels.

>>> seeds = torch.arange(0, 10).reshape(-1, 2)
>>> labels = torch.tensor([1, 1, 0, 0, 0])
>>> item_set = gb.ItemSet(
...     (seeds, labels), names=("seeds", "lables"))
>>> list(item_set)
[(tensor([0, 1]), tensor([1])),
 (tensor([2, 3]), tensor([1])),
 (tensor([4, 5]), tensor([0])),
 (tensor([6, 7]), tensor([0])),
 (tensor([8, 9]), tensor([0]))]
>>> item_set[:]
(tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
 tensor([1, 1, 0, 0, 0]))
>>> item_set.names
('seeds', 'labels')

6. Tuple of tensors with different shape: hyperlink and labels.

>>> seeds = torch.arange(0, 10).reshape(-1, 5)
>>> labels = torch.tensor([1, 0])
>>> item_set = gb.ItemSet(
...     (seeds, labels), names=("seeds", "lables"))
>>> list(item_set)
[(tensor([0, 1, 2, 3, 4]), tensor([1])),
 (tensor([5, 6, 7, 8, 9]), tensor([0]))]
>>> item_set[:]
(tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
 tensor([1, 0]))
>>> item_set.names
('seeds', 'labels')

property names: Tuple[str]

Return the names of the items.

property num_items: int

Return the number of the items.