SR-Forge

Core Data Structures

Entry

Bases: _DynamicStorage

The core data container in SR-Forge.

Entry is a dictionary-like container with attribute-style access (entry.lr is equivalent to entry['lr']). It holds all data for a sample as it flows through the pipeline — tensors, metadata, and intermediate results.

Stack-based convention: Dataset.__getitem__ returns fields in their natural form — tensors are [C, H, W], names are bare strings, scalars are Python types. Collation (via DataLoader) stacks samples to add a batch dimension: tensors become [B, C, H, W], names become lists, scalars become 1-D tensors.

Batched flag: entry.is_batched indicates whether the entry has a batch dimension. Set to True by :meth:collate and :meth:_slice_batch, defaults to False for entries from datasets or manual construction. Integer indexing (entry[i]) removes the flag; slicing (entry[0:1]) preserves it.

Batch operations (only meaningful when is_batched is True):

• entry.batch_size — number of samples in the batch.
• entry[i] — extract sample i, removing batch dim (PyTorch semantics).
• entry[start:stop] — extract a sub-batch (batch dim preserved).
• entry.unbatch() — split into a list of unbatched Entries.

Collation: Entry.collate([e1, e2, ...]) stacks entries to add a batch dimension. Sets _is_batched = True and stores _batch_size for O(1) lookup.

Round-trip invariant: Entry.collate([e1, e2])[i] recovers the original unbatched entry ei.

is_batched: bool property

Whether this Entry carries a batch dimension.

Set to True by :meth:collate and :meth:_slice_batch. Defaults to False for entries from datasets or manual construction. Integer indexing (entry[i]) clears the flag; slicing (entry[0:1]) preserves it.

batch_size: int property

Return the batch size.

Uses the stored _batch_size (set by :meth:collate and :meth:_slice_batch) when available, falling back to inference from the first tensor or list field.

Returns:

Name	Type	Description
int	int	The batch size.

Raises:

Type	Description
ValueError	If batch size cannot be inferred (no tensors or lists).

__init__(name: Union[str, list[str]] = None, *args, **kwargs)

Initializes the Entry.

Parameters:

Name	Type	Description	Default
name	str	An identifier for the entry, such as an image filename. Defaults to None.	None
*args	Any	Variable length argument list, passed to the dict constructor.	()
**kwargs	Any	Arbitrary keyword arguments, which become items in the entry.	{}

__getitem__(item: Any) -> Any

Dispatch on index type: int → sample, slice → sub-batch, str → field.

Integer indexing extracts a single sample, removing the batch dimension (PyTorch semantics: tensor[i] drops dim 0). Slice indexing extracts a sub-batch with batch dim preserved. String keys access fields as usual (dict behavior).

__len__()

Return batch size if batched.

Raises:

Type	Description
TypeError	If the Entry is unbatched.

Use entry.is_batched to check before calling len(), or wrap in a try/except to handle both cases in one call.

unbatch() -> List[Entry]

Split this batched Entry into a list of unbatched Entries.

Each element is obtained via integer indexing, so the batch dimension is removed (PyTorch semantics).

Returns:

Type	Description
List[Entry]	List[Entry]: One unbatched Entry per sample.

collate(batch: List[Entry]) -> Entry classmethod

Collate a list of Entries into a single batched Entry.

Stacks each field to add a batch dimension: tensors are stacked via torch.stack, lists are collected into a list-of-lists, strings are collected into a list, scalars become 1-D tensors. All entries must have the same keys.

The returned Entry has _is_batched = True and _batch_size = len(batch) for O(1) :attr:batch_size access.

Parameters:

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

Returns:

Name	Type	Description
Entry	Entry	A single batched Entry.

Raises:

Type	Description
ValueError	If entries have different keys.

merge_fields(fields: Mapping[str, Any], *, source: Optional[str] = None, allow_overwrite: bool = False) -> Entry

Merge fields into this entry in-place.

Parameters:

Name	Type	Description	Default
fields	Mapping[str, Any]	Mapping of field names to values. None is a no-op.	required
source	Optional[str]	Caller name for error messages (e.g. "model").	None
allow_overwrite	bool	If False (default), raises :class:KeyError when fields contains keys that already exist on the entry.	False

Returns:

Type	Description
Entry	This entry (modified in-place).

Raises:

Type	Description
TypeError	If fields is not a mapping (and not None).
KeyError	If a field conflict is detected and allow_overwrite is False.

GraphEntry

Stub for GraphEntry when torch-geometric is not installed.

This class exists so that isinstance(x, GraphEntry) checks work (always returning False) without requiring torch-geometric. Attempting to instantiate raises ImportError.

send(data: Any, device: torch.device) -> Any

Recursively moves data to a specified PyTorch device.

This function traverses nested data structures (lists, tuples, dicts, sets) and moves any found torch.Tensor to the target device. It also supports objects with a custom .to() method, like torch_geometric.data.Data or the Entry class itself.

Parameters:

Name	Type	Description	Default
data	Any	The data to move. Can be a tensor or a nested structure containing tensors.	required
device	device	The target device.	required

Returns:

Name	Type	Description
Any	Any	A copy of the data structure with all tensors moved to the specified device.

to_numpy(data: Any) -> Any

Recursively converts tensors to numpy arrays.

This function traverses nested data structures (lists, tuples, dicts, sets) and converts any found torch.Tensor to a numpy array via .detach().cpu().numpy().

Parameters:

Name	Type	Description	Default
data	Any	The data to convert. Can be a tensor or a nested structure containing tensors.	required

Returns:

Name	Type	Description
Any	Any	A copy of the data structure with all tensors converted to numpy arrays. Non-tensor values are returned as-is.