SR-Forge

Collation

GeneralCollation

A general-purpose collator for batches of Entry objects.

Collation stacks to add a batch dimension:

• Tensor (same shape) — torch.stack(batch, dim=0)
• Tensor (different shapes) — kept as a list (not stackable)
• list / tuple — collected into a list ([list1, list2])
• dict — recursively collate each key
• None (all) — collected into a list ([None, None, ...])
• None (mixed with other types) — collected into a list
• str — collected into a list (["s1", "s2"])
• int / float — torch.tensor([v1, v2])
• bool — torch.tensor([v1, v2], dtype=torch.bool)

__call__(batch: List[Entry]) -> Entry

Collate a batch of entries.

Ensures all entries in the batch are of the same type and then calls the collate class method on that type.

Parameters:

Name	Type	Description	Default
batch	List[Entry]	A list of Entry objects to collate.	required

Returns:

Name	Type	Description
Entry	Entry	A single Entry object containing the collated batch.

collate(batch: List[Any], key: str = None) -> Any

Recursively collate a list of items by stacking.

Adds a batch dimension to each field. Tensors are stacked (torch.stack), lists/tuples are collected into a list, strings are collected, and scalars become 1-D tensors.

Parameters:

Name	Type	Description	Default
batch	List[Any]	The list of items to collate.	required
key	str	The key corresponding to the batch, used for error messages. Defaults to None.	None

Returns:

Name	Type	Description
Any	Any	The collated batch.

Raises:

Type	Description
TypeError	If the batch contains mixed types for the same key, or an unsupported type is encountered.