omniloader.collate

Collate functions for batching unified samples.

Two collators are provided:

• unified_collate() — the default: every sample already shares the same tensor shapes (the unified schema, with sequences padded to a fixed time_dim), so it simply stacks tensors and gathers string metadata into lists.

• DynamicCollator — pads each sequence key only to the batch’s longest sample instead of a fixed time_dim. Use it with OmniLoader(pad_features=False) to avoid wasting memory/compute on padding when sequence lengths vary a lot across the batch.

Masks vs. cost: a <name>_mask guarantees correctness (padded steps never corrupt the model output or the loss), but a padded step still costs compute and memory (self-attention is O(T²)). Dynamic padding targets that cost — so it is only useful in the native-length regime. Under the default fixed time_dim (pad_features=True) every batch is already one length and DynamicCollator is a no-op.

Classes

DynamicCollator

Collate with per-batch dynamic padding of sequence keys.

Functions

unified_collate(→ dict[str, Any])

Collate a list of unified samples into a batched dict.

Module Contents

omniloader.collate.unified_collate(batch: collections.abc.Sequence[dict[str, Any]]) → dict[str, Any]

Collate a list of unified samples into a batched dict.

Parameters:

batch – Samples produced by OmniLoader.

Returns:

A dict where each tensor key holds a stacked tensor of shape (B, ...) and each non-tensor key holds a list of length B.

Raises:

ValueError – If the batch is empty.

class omniloader.collate.DynamicCollator(schema: omniloader.schema.spec.UnifiedSchema, keys: collections.abc.Sequence[str] | None = None)

Collate with per-batch dynamic padding of sequence keys.

Sequence values (those whose spec sets time_dim) are padded to the longest sample in the batch — (B, L_batch, F) instead of a fixed (B, time_dim, F) — and their <name>_mask is rebuilt accordingly. Vectors, scalars and metadata are batched exactly as unified_collate() does. Pair with OmniLoader(pad_features=False) so samples reach the collator at their native length; under the default fixed time_dim every sample already shares one length, so this collator has nothing to shrink and behaves like unified_collate().

Parameters:

• schema – The unified schema, used to identify sequence keys and their placeholder fill values.

• keys – Optional subset of sequence keys to pad dynamically. When None, every sequence key in the schema is padded dynamically.

__call__(batch: collections.abc.Sequence[dict[str, Any]]) → dict[str, Any]

Collate batch, dynamically padding the configured sequence keys.

Parameters:

batch – Samples to collate.

Returns:

The batched dict.

Raises:

ValueError – If the batch is empty.