omniloader.transforms.stats

Compute and persist feature normalization statistics over a dataset.

compute_feature_stats() estimates per-feature (mean, std) (streaming), while compute_stats() returns the richer set (mean/std/min/max/median/iqr) needed by every normalizer variant. compute_dataset_stats() computes that richer set separately per source dataset (grouped by the dataset metadata key) for dataset-conditional normalization. save_stats()/load_stats() round-trip either the flat ({key: {...}}) or the nested ({dataset: {key: {...}}}) form through JSON. All are modality-agnostic: they accumulate over the trailing feature axis of both vectors and sequences, counting only valid (unmasked) positions.

Functions

compute_feature_stats(→ dict[str, tuple[torch.Tensor, ...)	Estimate per-feature (mean, std) for each key over valid positions.
compute_stats(→ dict[str, dict[str, torch.Tensor]])	Compute the full statistic set per key, feeding every normalizer variant.
compute_dataset_stats(→ dict[str, dict[str, dict[str, ...)	Compute the full statistic set per key, grouped by source dataset.
save_stats(→ None)	Save a stats mapping to JSON (tensors are stored as nested lists).
load_stats(→ dict[str, Any])	Load a stats mapping saved by save_stats(), restoring float tensors.

Module Contents

omniloader.transforms.stats.compute_feature_stats(samples: collections.abc.Iterable[dict[str, torch.Tensor]], keys: collections.abc.Sequence[str], eps: float = 1e-06) → dict[str, tuple[torch.Tensor, torch.Tensor]]

Estimate per-feature (mean, std) for each key over valid positions.

Parameters:

• samples – An iterable of unified sample dicts (e.g. an OmniLoader). Each yields values and their <name>_mask companions.

• keys – The feature keys to compute statistics for.

• eps – Lower bound applied to the standard deviation to avoid zeros.

Returns:

Mapping of key to (mean, std) tensors of shape (F,) for features with a trailing feature axis, or () for scalar sequences.

Raises:

ValueError – If a key never has any valid position across the samples.

omniloader.transforms.stats.compute_stats(samples: collections.abc.Iterable[dict[str, torch.Tensor]], keys: collections.abc.Sequence[str], eps: float = 1e-06) → dict[str, dict[str, torch.Tensor]]

Compute the full statistic set per key, feeding every normalizer variant.

Unlike compute_feature_stats(), this collects all valid positions (to compute median/IQR), so it uses more memory; run it once and persist with save_stats().

Parameters:

• samples – An iterable of unified sample dicts.

• keys – The feature keys to compute statistics for.

• eps – Lower bound for std and iqr to avoid zeros.

Returns:

Mapping of key to {"mean", "std", "min", "max", "median", "iqr"} tensors of shape (F,) (or () for scalar sequences).

Raises:

ValueError – If a key never has any valid position across the samples.

omniloader.transforms.stats.compute_dataset_stats(samples: collections.abc.Iterable[dict[str, Any]], keys: collections.abc.Sequence[str], dataset_key: str = 'dataset', eps: float = 1e-06) → dict[str, dict[str, dict[str, torch.Tensor]]]

Compute the full statistic set per key, grouped by source dataset.

Feeds PerDatasetNormalize: each source dataset (identified by the dataset metadata key every unified sample carries) gets its own statistics, so datasets can be standardized by their own mean/std rather than a pooled estimate. As with compute_stats(), only valid (unmasked) positions contribute. Keys a dataset never provides are simply absent from that dataset’s stats.

Parameters:

• samples – An iterable of unified sample dicts (e.g. an OmniLoader). Each must carry dataset_key.

• keys – The feature keys to compute statistics for.

• dataset_key – Metadata key holding each sample’s source-dataset name.

• eps – Lower bound for std and iqr to avoid zeros.

Returns:

Mapping of dataset name to a stats mapping ({key: {"mean", "std", ...}}), matching compute_stats()’s per-key shape.

Raises:

ValueError – If a sample lacks dataset_key (required to group by dataset).

omniloader.transforms.stats.save_stats(stats: collections.abc.Mapping[str, Any], path: str | pathlib.Path) → None

Save a stats mapping to JSON (tensors are stored as nested lists).

Handles both the flat form from compute_stats() ({key: {field: tensor}}) and the nested per-dataset form from compute_dataset_stats() ({dataset: {key: {field: tensor}}}).

Parameters:

• stats – A stats mapping.

• path – Destination .json file.

omniloader.transforms.stats.load_stats(path: str | pathlib.Path) → dict[str, Any]

Load a stats mapping saved by save_stats(), restoring float tensors.

Restores either the flat or the nested per-dataset form (mirroring save_stats()).

Parameters:

path – Path to a JSON file written by save_stats().

Returns:

The stats mapping with leaf lists restored to float tensors.