lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Extract giant train/new-run's CLI override mapping into a table-driven function (issues.md Issues 3 & 4)

Browse Source 
train()'s ~140-line hand-written flag->config translation (three different
ad hoc "more specific flag wins" patterns) and new_run()'s near-verbatim
copy are replaced by a shared FlagSpec/FLAG_SPECS table and
overrides_from_flags() in config.py, reused by both commands. This makes
the override/precedence logic directly unit-testable without CliRunner,
closing coverage gaps that had zero tests (e.g. --emb-dim/--conditioning
dual-axis fan-out, three of four WGAN knob legs, --stage2-generator
overriding --mode, router's stage1-only asymmetry).

No CLI flags, help text, or precedence semantics changed --
`giant train --help`/`giant new-run --help` are byte-identical before and
after, and all previously-passing CliRunner tests still pass unmodified.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-08-12 15:04:49 +02:00
parent 01acbfed61
commit 2bfb1ab056
4 changed files with 381 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/cli.py
+48 -175
View File
@@ -653,9 +653,8 @@ def train(
)
raise typer.Exit(1)
cli_train = {
k: v
for k, v in {
cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
flag_values: dict[str, object] = {
"epochs": epochs,
"batch_size": batch_size_value,
"lr": lr,
@@ -672,122 +671,39 @@ def train(
"wandb_project": wandb_project,
"wandb_run_name": wandb_run_name,
"wandb_log_every": wandb_log_every,
}.items()
if v is not None
}
# --hidden-dim/--n-blocks/--dropout are stage-1-only shorthands kept for
# backward compatibility (they predate stage2_model having its own
# flags); --stage1-*/--stage2-* below are the explicit, discoverable
# per-stage flags, and take precedence when both are given.
cli_stage1_model: dict[str, object] = {
k: v
for k, v in {
"hidden_dim": hidden_dim,
"n_res_blocks": n_blocks,
"n_blocks": n_blocks,
"dropout": dropout,
}.items()
if v is not None
}
cli_stage1_model.update(
{
k: v
for k, v in {
"hidden_dim": stage1_hidden_dim,
"n_res_blocks": stage1_n_res_blocks,
"dropout": stage1_dropout,
}.items()
if v is not None
}
)
cli_stage2_model: dict[str, object] = {
k: v
for k, v in {
"hidden_dim": stage2_hidden_dim,
"n_res_blocks": stage2_n_res_blocks,
"dropout": stage2_dropout,
"decoder": stage2_decoder.value if stage2_decoder is not None else None,
"k_max": stage2_k_max,
"context_dim": stage2_context_dim,
"stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
}.items()
if v is not None
}
cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
if cli_router:
cli_stage1_model["router"] = cli_router
# --emb-dim/--conditioning set both conditioning axes (v0.2 had one
# shared value for particle+material).
cli_conditioning: dict[str, dict] = {}
if emb_dim is not None:
cli_conditioning["particle"] = {"emb_dim": emb_dim}
cli_conditioning["material"] = {"emb_dim": emb_dim}
if conditioning is not None:
cli_conditioning.setdefault("particle", {})["type"] = conditioning.value
cli_conditioning.setdefault("material", {})["type"] = conditioning.value
overrides: dict[str, dict] = {}
if cli_train:
overrides["train"] = cli_train
if cli_stage1_model:
overrides["stage1_model"] = cli_stage1_model
if cli_stage2_model:
overrides["stage2_model"] = cli_stage2_model
if cli_conditioning:
overrides["conditioning"] = cli_conditioning
# --mode/--n-critic/--gp-weight/--critic-lr/--noise-dim apply to BOTH
# stages by default (v0.2 had one global mode/wgan config shared by both
# — see giant.config.migrate_config's train.mode /
# train.{n_critic,gp_weight,critic_lr} precedent); the --stage1-*/
# --stage2-* variants below override a single stage independently (decision
# 7), which is what actually enables e.g. `--stage1-generator flow
# --stage2-generator wgan`.
if mode is not None:
overrides.setdefault("stage1_model", {})["generator"] = mode.value
overrides.setdefault("stage2_model", {})["generator"] = mode.value
if stage1_generator is not None:
overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
if stage2_generator is not None:
overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
shared_wgan_overrides = {
k: v
for k, v in {
"stage1_hidden_dim": stage1_hidden_dim,
"stage1_n_res_blocks": stage1_n_res_blocks,
"stage1_dropout": stage1_dropout,
"stage2_hidden_dim": stage2_hidden_dim,
"stage2_n_res_blocks": stage2_n_res_blocks,
"stage2_dropout": stage2_dropout,
"stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None,
"stage2_k_max": stage2_k_max,
"stage2_context_dim": stage2_context_dim,
"stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
"mode": mode.value if mode is not None else None,
"stage1_generator": stage1_generator.value if stage1_generator is not None else None,
"stage2_generator": stage2_generator.value if stage2_generator is not None else None,
"conditioning": conditioning.value if conditioning is not None else None,
"emb_dim": emb_dim,
"router_config": cli_router or None,
"n_critic": n_critic,
"gp_weight": gp_weight,
"noise_dim": noise_dim,
"critic_lr": critic_lr,
}.items()
if v is not None
"stage1_n_critic": stage1_n_critic,
"stage1_gp_weight": stage1_gp_weight,
"stage1_noise_dim": stage1_noise_dim,
"stage1_critic_lr": stage1_critic_lr,
"stage2_n_critic": stage2_n_critic,
"stage2_gp_weight": stage2_gp_weight,
"stage2_noise_dim": stage2_noise_dim,
"stage2_critic_lr": stage2_critic_lr,
}
stage1_wgan_overrides = {
k: v
for k, v in {
"n_critic": stage1_n_critic,
"gp_weight": stage1_gp_weight,
"noise_dim": stage1_noise_dim,
"critic_lr": stage1_critic_lr,
}.items()
if v is not None
}
stage2_wgan_overrides = {
k: v
for k, v in {
"n_critic": stage2_n_critic,
"gp_weight": stage2_gp_weight,
"noise_dim": stage2_noise_dim,
"critic_lr": stage2_critic_lr,
}.items()
if v is not None
}
for stage_name, stage_specific in (
("stage1_model", stage1_wgan_overrides),
("stage2_model", stage2_wgan_overrides),
):
stage_wgan = {**shared_wgan_overrides, **stage_specific}
if stage_wgan:
overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(stage_wgan)
overrides = gconfig.overrides_from_flags(flag_values)
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
gconfig.validate_config(cfg)
@@ -912,75 +828,32 @@ def new_run(
(with the full dataset-derived meta section), so this scaffold's meta
section is just a placeholder recording what was asked for and when.
"""
cli_train = {
k: v
for k, v in {
cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
flag_values: dict[str, object] = {
"epochs": epochs,
"batch_size": batch_size,
"lr": lr,
}.items()
if v is not None
}
cli_stage1_model: dict[str, object] = {
k: v
for k, v in {
"hidden_dim": hidden_dim,
"n_res_blocks": n_blocks,
"n_blocks": n_blocks,
"dropout": dropout,
}.items()
if v is not None
"stage1_hidden_dim": stage1_hidden_dim,
"stage1_n_res_blocks": stage1_n_res_blocks,
"stage1_dropout": stage1_dropout,
"stage2_hidden_dim": stage2_hidden_dim,
"stage2_n_res_blocks": stage2_n_res_blocks,
"stage2_dropout": stage2_dropout,
"stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None,
"stage2_k_max": stage2_k_max,
"stage2_context_dim": stage2_context_dim,
"stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
"mode": mode.value if mode is not None else None,
"stage1_generator": stage1_generator.value if stage1_generator is not None else None,
"stage2_generator": stage2_generator.value if stage2_generator is not None else None,
"conditioning": conditioning.value if conditioning is not None else None,
"emb_dim": emb_dim,
"router_config": cli_router or None,
}
cli_stage1_model.update(
{
k: v
for k, v in {
"hidden_dim": stage1_hidden_dim,
"n_res_blocks": stage1_n_res_blocks,
"dropout": stage1_dropout,
}.items()
if v is not None
}
)
cli_stage2_model: dict[str, object] = {
k: v
for k, v in {
"hidden_dim": stage2_hidden_dim,
"n_res_blocks": stage2_n_res_blocks,
"dropout": stage2_dropout,
"decoder": stage2_decoder.value if stage2_decoder is not None else None,
"k_max": stage2_k_max,
"context_dim": stage2_context_dim,
"stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
}.items()
if v is not None
}
cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
if cli_router:
cli_stage1_model["router"] = cli_router
cli_conditioning: dict[str, dict] = {}
if emb_dim is not None:
cli_conditioning["particle"] = {"emb_dim": emb_dim}
cli_conditioning["material"] = {"emb_dim": emb_dim}
if conditioning is not None:
cli_conditioning.setdefault("particle", {})["type"] = conditioning.value
cli_conditioning.setdefault("material", {})["type"] = conditioning.value
overrides: dict[str, dict] = {}
if cli_train:
overrides["train"] = cli_train
if cli_stage1_model:
overrides["stage1_model"] = cli_stage1_model
if cli_stage2_model:
overrides["stage2_model"] = cli_stage2_model
if cli_conditioning:
overrides["conditioning"] = cli_conditioning
if mode is not None:
overrides.setdefault("stage1_model", {})["generator"] = mode.value
overrides.setdefault("stage2_model", {})["generator"] = mode.value
if stage1_generator is not None:
overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
if stage2_generator is not None:
overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
overrides = gconfig.overrides_from_flags(flag_values)
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
gconfig.validate_config(cfg)
giant/config.py
+108
View File
@@ -893,6 +893,114 @@ def _deep_merge(base: dict, override: dict) -> dict:
return result
@dataclass(frozen=True)
class FlagSpec:
"""One CLI flag's mapping into the config-overrides tree.
`paths` lists every dotted config path this flag writes (>1 means fan-out
to multiple stages/axes, e.g. `--mode` -> both stages' `generator`).
`precedence` controls write order when two flags target the same path:
specs are applied in ascending precedence, so a higher-precedence (more
specific) flag overwrites a lower-precedence (shared/shorthand) one —
this is the "build a shared dict, then let a more specific dict win"
pattern `giant train`/`giant new-run` need (e.g. `--hidden-dim` vs
`--stage1-hidden-dim`, or `--n-critic` vs `--stage1-n-critic`),
generalized to one mechanism instead of three different ad hoc ones.
"""
name: str
paths: tuple[str, ...]
precedence: int = 0
# Flag -> config-path table shared by `giant train`/`giant new-run`
# (giant/cli.py) so both commands resolve CLI overrides identically. See
# issues.md Issue 3: this replaces ~140 lines of hand-written, imperative
# dict-building in cli.py with one declarative table plus
# `overrides_from_flags` below.
FLAG_SPECS: tuple[FlagSpec, ...] = (
# train block -- flat pass-through, unique paths, precedence irrelevant.
FlagSpec("epochs", ("train.epochs",)),
FlagSpec("batch_size", ("train.batch_size",)),
FlagSpec("lr", ("train.lr",)),
FlagSpec("weight_decay", ("train.weight_decay",)),
FlagSpec("ema_decay", ("train.ema_decay",)),
FlagSpec("warmup_epochs", ("train.warmup_epochs",)),
FlagSpec("val_fraction", ("train.val_fraction",)),
FlagSpec("num_workers", ("train.num_workers",)),
FlagSpec("seed", ("train.seed",)),
FlagSpec("validate_every", ("train.validate_every",)),
FlagSpec("validate_steps", ("train.validate_steps",)),
FlagSpec("max_val_batches", ("train.max_val_batches",)),
FlagSpec("wandb", ("train.wandb",)),
FlagSpec("wandb_project", ("train.wandb_project",)),
FlagSpec("wandb_run_name", ("train.wandb_run_name",)),
FlagSpec("wandb_log_every", ("train.wandb_log_every",)),
# --hidden-dim/--n-blocks/--dropout are stage-1-only backward-compat
# shorthands (they predate stage2_model having its own flags);
# --stage1-* wins when both are given.
FlagSpec("hidden_dim", ("stage1_model.hidden_dim",), precedence=0),
FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1),
FlagSpec("n_blocks", ("stage1_model.n_res_blocks",), precedence=0),
FlagSpec("stage1_n_res_blocks", ("stage1_model.n_res_blocks",), precedence=1),
FlagSpec("dropout", ("stage1_model.dropout",), precedence=0),
FlagSpec("stage1_dropout", ("stage1_model.dropout",), precedence=1),
# stage2-only knobs.
FlagSpec("stage2_hidden_dim", ("stage2_model.hidden_dim",)),
FlagSpec("stage2_n_res_blocks", ("stage2_model.n_res_blocks",)),
FlagSpec("stage2_dropout", ("stage2_model.dropout",)),
FlagSpec("stage2_decoder", ("stage2_model.decoder",)),
FlagSpec("stage2_k_max", ("stage2_model.k_max",)),
FlagSpec("stage2_context_dim", ("stage2_model.context_dim",)),
FlagSpec("stage2_stage1_context", ("stage2_model.stage1_context",)),
# --mode applies to both stages by default (v0.2 had one shared
# mode/wgan config); --stage{1,2}-generator override a single stage.
FlagSpec("mode", ("stage1_model.generator", "stage2_model.generator"), precedence=0),
FlagSpec("stage1_generator", ("stage1_model.generator",), precedence=1),
FlagSpec("stage2_generator", ("stage2_model.generator",), precedence=1),
# --emb-dim/--conditioning set both conditioning axes (v0.2 had one
# shared value for particle+material).
FlagSpec("conditioning", ("conditioning.particle.type", "conditioning.material.type")),
FlagSpec("emb_dim", ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")),
# Pre-aggregated router override dict (built by `_router_cli_overrides`
# in cli.py from --router/--router-type/--n-experts/--router-axis).
# Router overrides only ever land on stage1_model -- this asymmetry is
# deliberate (see cli.py) and must not be "fixed" into a fan-out here.
FlagSpec("router_config", ("stage1_model.router",)),
# WGAN: shared knobs apply to both stages by default (v0.2 had one
# shared wgan config); --stage{1,2}-* override a single stage.
FlagSpec("n_critic", ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic"), precedence=0),
FlagSpec("stage1_n_critic", ("stage1_model.wgan.n_critic",), precedence=1),
FlagSpec("stage2_n_critic", ("stage2_model.wgan.n_critic",), precedence=1),
FlagSpec("gp_weight", ("stage1_model.wgan.gp_weight", "stage2_model.wgan.gp_weight"), precedence=0),
FlagSpec("stage1_gp_weight", ("stage1_model.wgan.gp_weight",), precedence=1),
FlagSpec("stage2_gp_weight", ("stage2_model.wgan.gp_weight",), precedence=1),
FlagSpec("noise_dim", ("stage1_model.wgan.noise_dim", "stage2_model.wgan.noise_dim"), precedence=0),
FlagSpec("stage1_noise_dim", ("stage1_model.wgan.noise_dim",), precedence=1),
FlagSpec("stage2_noise_dim", ("stage2_model.wgan.noise_dim",), precedence=1),
FlagSpec("critic_lr", ("stage1_model.wgan.critic_lr", "stage2_model.wgan.critic_lr"), precedence=0),
FlagSpec("stage1_critic_lr", ("stage1_model.wgan.critic_lr",), precedence=1),
FlagSpec("stage2_critic_lr", ("stage2_model.wgan.critic_lr",), precedence=1),
)
def overrides_from_flags(values: dict[str, object]) -> dict:
"""Build the nested, section-keyed config-overrides dict
`merge_cli_overrides` expects, from `{flag_name: value}`.
Flags absent from `values`, or mapped to `None` (= not given on the
CLI), are skipped. See `FlagSpec`/`FLAG_SPECS` above for the precedence
rule applied when two flags target the same path.
"""
overrides: dict = {}
for spec in sorted(FLAG_SPECS, key=lambda s: s.precedence):
if spec.name not in values or values[spec.name] is None:
continue
for path in spec.paths:
_set_path(overrides, path, values[spec.name])
return overrides
# v0.2 [train] keys that pass through to v0.3 [train] unchanged (same name,
# same meaning) when present in the loaded file — everything model-shaped
# moved to the stage/conditioning blocks instead (see the rest of
tests/test_cli_train_overrides.py
+80
View File
@@ -94,3 +94,83 @@ def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5
def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
result = runner.invoke(
cli.app,
["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
)
assert result.exit_code == 1
assert "--batch-size must be an integer or 'auto'" in result.output
def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_path):
captured: dict = {}
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
captured["out_dir"] = out_dir
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
resume_dir = tmp_path / "resumed_run"
resume_dir.mkdir()
(resume_dir / "last.pt").touch()
explicit_out = tmp_path / "explicit_run"
result = runner.invoke(
cli.app,
["train", "dummy.parquet", "--out", str(explicit_out), "--resume", str(resume_dir / "last.pt")],
)
assert result.exit_code == 0, result.output
assert captured["out_dir"] == explicit_out
def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
captured: dict = {}
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
captured["out_dir"] = out_dir
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
resume_dir = tmp_path / "resumed_run"
resume_dir.mkdir()
(resume_dir / "last.pt").touch()
result = runner.invoke(cli.app, ["train", "dummy.parquet", "--resume", str(resume_dir / "last.pt")])
assert result.exit_code == 0, result.output
assert captured["out_dir"] == resume_dir
def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypatch, tmp_path):
captured: dict = {}
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
captured["out_dir"] = out_dir
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
monkeypatch.chdir(tmp_path)
result = runner.invoke(cli.app, ["train", "dummy.parquet"])
assert result.exit_code == 0, result.output
assert captured["out_dir"] == Path("checkpoints") / cli.gconfig.default_out_dir_name(cli.gconfig.DEFAULT_CONFIG)
def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
captured: dict = {}
def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
captured["batch_size"] = cfg["train"]["batch_size"]
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
result = runner.invoke(
cli.app,
["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "auto"],
)
assert result.exit_code == 0, result.output
assert captured["batch_size"] == 123
assert "batch_size: 123 (auto-estimated from free GPU memory)" in result.output
tests/test_config.py
+119
View File
@@ -889,6 +889,125 @@ def test_merge_cli_overrides_real_config_fixtures_pass_key_validation(fixture_na
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / fixture_name, {}) # must not raise
# ---------------------------------------------------------------------------
# overrides_from_flags (issues.md Issue 3): the flag -> config-path table
# shared by `giant train`/`giant new-run`. Each test below pins one
# precedence rule directly, without CliRunner — see also
# tests/test_cli_train_overrides.py for the thin end-to-end smoke coverage.
# ---------------------------------------------------------------------------
def test_overrides_from_flags_empty_values_yield_empty_overrides():
assert gconfig.overrides_from_flags({}) == {}
assert gconfig.overrides_from_flags({"epochs": None, "hidden_dim": None}) == {}
def test_overrides_from_flags_train_block_passthrough():
overrides = gconfig.overrides_from_flags({"epochs": 5, "lr": 1e-3, "hidden_dim": None})
assert overrides == {"train": {"epochs": 5, "lr": 1e-3}}
@pytest.mark.parametrize(
("shorthand", "explicit", "path_key"),
[
("hidden_dim", "stage1_hidden_dim", "hidden_dim"),
("n_blocks", "stage1_n_res_blocks", "n_res_blocks"),
("dropout", "stage1_dropout", "dropout"),
],
)
def test_overrides_from_flags_stage1_explicit_overrides_shorthand(shorthand, explicit, path_key):
overrides = gconfig.overrides_from_flags({shorthand: 1, explicit: 2})
assert overrides["stage1_model"][path_key] == 2
@pytest.mark.parametrize(
("shorthand", "path_key"),
[("hidden_dim", "hidden_dim"), ("n_blocks", "n_res_blocks"), ("dropout", "dropout")],
)
def test_overrides_from_flags_stage1_shorthand_alone(shorthand, path_key):
overrides = gconfig.overrides_from_flags({shorthand: 7})
assert overrides["stage1_model"][path_key] == 7
def test_overrides_from_flags_stage2_only_knobs():
overrides = gconfig.overrides_from_flags(
{
"stage2_hidden_dim": 32,
"stage2_n_res_blocks": 4,
"stage2_dropout": 0.1,
"stage2_decoder": "one_shot",
"stage2_k_max": 8,
"stage2_context_dim": 16,
"stage2_stage1_context": "sampled",
}
)
assert overrides["stage2_model"] == {
"hidden_dim": 32,
"n_res_blocks": 4,
"dropout": 0.1,
"decoder": "one_shot",
"k_max": 8,
"context_dim": 16,
"stage1_context": "sampled",
}
assert "stage1_model" not in overrides
def test_overrides_from_flags_mode_fans_to_both_stages():
overrides = gconfig.overrides_from_flags({"mode": "wgan"})
assert overrides["stage1_model"]["generator"] == "wgan"
assert overrides["stage2_model"]["generator"] == "wgan"
def test_overrides_from_flags_stage1_generator_overrides_mode_for_stage1_only():
overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage1_generator": "flow"})
assert overrides["stage1_model"]["generator"] == "flow"
assert overrides["stage2_model"]["generator"] == "wgan"
def test_overrides_from_flags_stage2_generator_overrides_mode_for_stage2_only():
overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage2_generator": "flow"})
assert overrides["stage1_model"]["generator"] == "wgan"
assert overrides["stage2_model"]["generator"] == "flow"
def test_overrides_from_flags_emb_dim_sets_both_conditioning_axes():
overrides = gconfig.overrides_from_flags({"emb_dim": 24})
assert overrides["conditioning"]["particle"]["emb_dim"] == 24
assert overrides["conditioning"]["material"]["emb_dim"] == 24
def test_overrides_from_flags_conditioning_sets_both_axes_type():
overrides = gconfig.overrides_from_flags({"conditioning": "onehot"})
assert overrides["conditioning"]["particle"]["type"] == "onehot"
assert overrides["conditioning"]["material"]["type"] == "onehot"
def test_overrides_from_flags_router_config_only_touches_stage1():
overrides = gconfig.overrides_from_flags({"router_config": {"enabled": True, "type": "energy"}})
assert overrides["stage1_model"]["router"] == {"enabled": True, "type": "energy"}
assert "stage2_model" not in overrides
@pytest.mark.parametrize(
("shared", "stage1_specific", "stage2_specific", "path_key"),
[
("n_critic", "stage1_n_critic", "stage2_n_critic", "n_critic"),
("gp_weight", "stage1_gp_weight", "stage2_gp_weight", "gp_weight"),
("noise_dim", "stage1_noise_dim", "stage2_noise_dim", "noise_dim"),
("critic_lr", "stage1_critic_lr", "stage2_critic_lr", "critic_lr"),
],
)
def test_overrides_from_flags_wgan_knobs_split_per_stage(shared, stage1_specific, stage2_specific, path_key):
overrides = gconfig.overrides_from_flags({shared: 5.0, stage1_specific: 3.0})
assert overrides["stage1_model"]["wgan"][path_key] == 3.0
assert overrides["stage2_model"]["wgan"][path_key] == 5.0
overrides = gconfig.overrides_from_flags({shared: 5.0, stage2_specific: 2.5})
assert overrides["stage1_model"]["wgan"][path_key] == 5.0
assert overrides["stage2_model"]["wgan"][path_key] == 2.5
# ---------------------------------------------------------------------------
# checkpoint config-mismatch warnings (unchanged surface, still exercised)
# ---------------------------------------------------------------------------