Loaders

class PooledBatch(*args: Any, **kwargs: Any)

Bases: Batch

A custom Batch class for handling graph data with pooled-graph data attributes in tgp.

This class extends Batch to support a batch of graphs along with precomputed pooled representations. It stores additional information needed to reconstruct individual graph data objects and manage multiple levels of pooled data.

classmethod from_data_list(data_list: List[BaseData], follow_batch: List[str] | None = None, exclude_keys: List[str] | None = None) → PooledBatch

Constructs a PooledBatch from a list of graph Data objects.

This method collates node and edge attributes, as well as any sub-Data object storing pooled data, from each graph in data_list into a single PooledBatch. It handles attribute increments and batch assignments, and stores metadata required to separate individual graphs later.

Parameters:

• data_list (list) – A list of Data or HeteroData objects to batch.

• follow_batch (Optional[List[str]]) – Keys for which to create additional batch assignment vectors (e.g., node-level attributes to track).

• exclude_keys (Optional[List[str]]) – Attributes to exclude from collation.

Returns:

A batched object containing all graphs from

data_list, with _slice_dict and _inc_dict attributes set for reconstruction.

Return type:

PooledBatch

get_example(idx: int) → BaseData

Retrieves an individual graph data object from the batch.

This method separates the batched data at the specified index idx, using the stored _slice_dict and _inc_dict to reconstruct the original BaseData object.

Parameters:

idx (int) – Index of the graph to extract from the batch.

Returns:

The reconstructed

Data or HeteroData object at position idx.

Return type:

BaseData

Raises:

RuntimeError – If the batch was not created via from_data_list(), making reconstruction impossible.

class PoolCollater(dataset: Dataset | Sequence[BaseData] | DatasetAdapter, follow_batch: List[str] | None = None, exclude_keys: List[str] | None = None)

Bases: Collater

A custom collate function for pooling dataloaders.

This class extends the PyG collater to produce PooledBatch objects when collating a list of Data or HeteroData objects. It invokes from_data_list() to perform the batching with optional follow_batch and exclude_keys arguments.

class PoolDataLoader(dataset: Dataset | Sequence[BaseData] | DatasetAdapter, batch_size: int = 1, shuffle: bool = False, follow_batch: List[str] | None = None, exclude_keys: List[str] | None = None, **kwargs)

Bases: DataLoader

A DataLoader for pooled graph datasets, returning PooledBatch objects.

This class extends DataLoader to integrate with PoolCollater, automatically batching pooled graph data. It accepts follow_batch and exclude_keys parameters to propagate to the collate function.

Parameters:

• dataset (object) – The source dataset from which to load graph data.

• batch_size (int) – Number of graphs per batch. (default: 1)

• shuffle (bool) – Whether to shuffle the data each epoch. (default: False)

• follow_batch (list, optional) – Keys for which to create assignment vectors in the batch. (default: None)

• exclude_keys (list, optional) – Attributes to exclude from collation. (default: None)

• **kwargs – Additional keyword arguments forwarded to torch.utils.data.DataLoader.