torch_brain.batching

Utilities for collating samples of variable length into a batch.

The core entry point is collate(), a drop-in extension of PyTorch’s default_collate that handles tensors of differing lengths. Instead of configuring it directly, you wrap objects with a collation recipe:

• pad() (and its variants) to pad a dimension to the batch maximum.

• chain() to concatenate along the first dimension.

• Optionally, track_mask() or track_batch() to emit a padding mask or per-element batch index.

When used with PyTorch’s default DataLoader, collate() applies the wrapped recipes when batching. Below is an example of how these are used in practice.

import torch
from torch_brain.datasets import Dataset
from torch_brain.batching import collate, chain, pad, track_mask

class MyDataset(Dataset):
    ...

    def __getitem__(self, index):
        # logic to gather data
        ...

        return {
            "spike_times": chain(spike_times),
            "spike_seqlen": len(spike_times),
            "query_times": pad(query_times),
            "query_mask": track_mask(query_times),
        }

ds = MyDataset(...)

loader = torch.utils.data.DataLoader(
    dataset=ds,
    sampler=...,
    collate_fn=collate,  # <- specify collate from torch_brain.batching
    num_workers=4,
    ...
)

Collate Function

collate

Extension of PyTorch's default_collate function to enable more advanced collation of samples of variable lengths.

Collation Recipes

Wrap objects with these functions to collate them in different ways.

chain	Wrap an object to specify that it (or any of its members) should be stacked along the first dimension when batching.
pad	Wrap an object to specify that it (or any of its members) should be padded to the maximum length in the batch.
pad8	Wrap an object to specify that it (or any of its members) should be padded to the maximum length in the batch.
pad2d	Wrap an object to specify that it (or any of its members) should be padded to the maximum length in the batch along its first two dimensions.
pad2d8	Wrap an object to specify that it (or any of its members) should be padded to the maximum length in the batch along its first two dimensions.

Mask & Index Tracking

Use these alongside a collation wrapper to track padding masks or batch indices for a particular array or tensor.

track_batch	Wrap an array or tensor to track the batch_index.
track_mask	Wrap an array or tensor to specify that its padding mask should be tracked.
track_mask8	Wrap an array or tensor to specify that its padding mask should be tracked.
track_mask2d	Wrap an array or tensor to specify that its padding mask should be tracked.
track_mask2d8	Wrap an array or tensor to specify that its padding mask should be tracked.