mapping_networks.distributed.benchmark_strategies

mapping_networks.distributed.benchmark_strategies(spec: ParameterSpec, latent_dim: int = 64, device: str = 'cpu', n_warmup: int = 2, n_repeat: int = 5) → dict[str, dict[str, Any]]

Compare SLVT, layerwise, and grouped strategies on a given spec.

Builds one generator per strategy, runs n_warmup warm-up calls, then measures peak memory and wall-clock generation time over n_repeat calls. The reported peak is the maximum across all measured calls.

Parameters:

• spec – Target ParameterSpec to benchmark against.

• latent_dim – Latent dimension to use for all strategies.

• device – "cpu" or "cuda".

• n_warmup – Number of warm-up calls before measurement (default 2).

• n_repeat – Number of timed calls (default 5).

Returns:

{strategy_name: {"peak_bytes": float, "generate_time_ms": float}}

On error (e.g. SLVT projection too large) the dict for that strategy additionally contains an "error" key with the exception message.

Example:

results = benchmark_strategies(spec, latent_dim=64)
# {"slvt": {"peak_bytes": ..., "generate_time_ms": ...}, ...}