From 1a3c90757148ececd53f6e31b3a6d7c222de356c Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 14 Aug 2026 14:37:58 +0200 Subject: [PATCH 1/3] Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base (gitea #39) Stage1Model, Stage2OneShot and Stage2Autoregressive each independently implemented ~90 near-identical lines of __init__ scaffolding: build-or-share cond_enc, particle_type_cfg normalisation, objective -> time_emb -> merged_cond_dim -> build_trunk, and the n_sec_head/type_head classifier heads (plus their identical RuntimeError guards). Now unblocked by #33 (trunk registry), #34 (block-conditioning registry) and #36 (build_mlp_head), which settled what belongs in the shared base. Adds StageModel(nn.Module) owning all of that: __init__ builds/shares cond_enc and normalises particle_type_cfg; _build_trunk_and_heads, called by each subclass after it sets up its own conditioning-assembly modules (cond_enc alone for Stage1Model, a context-fusion path for the two Stage2 classes), builds the objective/time embedding/trunk and the n_sec_head/type_head guarded by the shared _require_n_sec_head/ _require_type_head (Stage1Model overrides the n_sec guard since its message points at stage 2, not stage 1). Public __init__ signatures, attribute names, and forward/predict_* behaviour are unchanged. Verified with a pre/post state_dict-key-set diff against the pre-refactor classes (bit-identical) before writing this commit, plus new parametrized tests pinning each class's state_dict key set and the generator -> time_emb contract the base now owns. tests/test_migration_ v02_v03.py's existing bit-identical old-vs-new forward comparison and the rest of tests/test_network.py's per-class coverage pass unchanged. Co-Authored-By: Claude Opus 5 --- giant/model/models.py | 377 +++++++++++++++++++++++++---------------- giant/model/network.py | 2 + tests/test_network.py | 162 ++++++++++++++++++ 3 files changed, 393 insertions(+), 148 deletions(-) diff --git a/giant/model/models.py b/giant/model/models.py index ed22a99..d1f79c6 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -67,18 +67,136 @@ def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, em return k_max * CONT_SLOT_DIM -class Stage1Model(nn.Module): - """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs - move it to stage 2, except for a migrated v0.2 checkpoint - (`n_sec_head_k_max` given), where it stays attached here - since that's where its weights live and what conditioning it was trained - against (see `_migrate_legacy_model_config`). +class StageModel(nn.Module): + """Base owning the scaffolding common to `Stage1Model`, `Stage2OneShot`, + `Stage2Autoregressive` (gitea #39): build-or-share `cond_enc`, + `particle_type_cfg` normalisation, and — via `_build_trunk_and_heads`, + called by each subclass's `__init__` once its own conditioning-assembly + modules exist — the objective/time-embedding/trunk construction and the + `n_sec_head`/`type_head` classifier heads. A subclass supplies only its + own conditioning assembly (`Stage1Model` uses `cond_enc` directly; + `Stage2OneShot`/`Stage2Autoregressive` add a context-fusion path) and its + trunk's output width. `cond_enc`, if given, is used in place of building a fresh `ConditionEncoder` — `conditioning.share_stages = true`: `build_models` constructs one shared instance and passes it to both stages, halving the conditioning parameter count and forcing a common representation.""" + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + particle_cfg: dict, + material_cfg: dict, + cond_out_dim: int, + generator: str, + noise_dim: int, + k_max: int | None = None, + particle_type_cfg: dict | None = None, + cond_enc: ConditionEncoder | None = None, + ) -> None: + super().__init__() + self.generator_kind = generator + self.noise_dim = noise_dim + self.k_max = k_max + self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) + self.type_dim = stage2_type_dim( + self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"]) + ) + self.cond_enc = ( + cond_enc + if cond_enc is not None + else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + ) + + def _build_trunk_and_heads( + self, + *, + trunk_out_dim: int, + hidden_dim: int, + n_res_blocks: int, + cond_out_dim: int, + time_dim: int, + router: Router | None, + trunk_type: str, + block_conditioning: str, + dropout: float, + n_sec_head_k_max: int | None, + n_sec_head_cfg: dict | None, + type_head_out_dim: int | None, + type_head_cfg: dict | None, + ) -> None: + """Builds `self.time_emb`, `self.trunk`, `self.n_sec_head`, + `self.type_head`. Called by a subclass's `__init__` after it has set + up its own conditioning-assembly modules — `merged_cond_dim` below + must match the width that assembly (`_cond_embed`/`_base_cond`/ + `_token_cond`, or plain `cond_enc` for `Stage1Model`) actually + produces. + + `n_sec_head` is built iff `n_sec_head_k_max is not None` (output + width `n_sec_head_k_max + 1`) — `Stage1Model` passes this only for a + migrated v0.2 checkpoint, `Stage2OneShot`/`Stage2Autoregressive` pass + it whenever `build_n_sec_head=True`. `type_head` is built iff + `type_head_out_dim is not None` (the caller — only the two Stage2 + classes — passes `None` exactly when `particle_type_cfg["target"] == + "physical"`) *and* the objective doesn't fold the type slice into its + own trunk output (checked here, since `objective` is already needed + for the trunk itself). + """ + objective = build_objective(self.generator_kind) + has_time = objective.needs_time + self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None + merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim + in_dim = objective.trunk_in_dim(trunk_out_dim, self.noise_dim) + self.trunk = build_trunk( + router, + trunk_type, + in_dim, + trunk_out_dim, + hidden_dim, + n_res_blocks, + merged_cond_dim, + dropout, + block_conditioning, + ) + self.n_sec_head = None + if n_sec_head_k_max is not None: + head_cfg = HeadConfig.from_dict(n_sec_head_cfg) + hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) + self.n_sec_head = build_mlp_head(cond_out_dim, n_sec_head_k_max + 1, hidden, head_cfg.depth) + self.type_head = None + if type_head_out_dim is not None and not objective.folds_type_slice: + head_cfg = HeadConfig.from_dict(type_head_cfg) + hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) + self.type_head = build_mlp_head(cond_out_dim, type_head_out_dim, hidden, head_cfg.depth) + + def _require_n_sec_head(self) -> None: + if self.n_sec_head is None: + raise RuntimeError( + f"this {type(self).__name__} has no n_sec_head — it belongs to " + "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call " + "stage1.predict_n_sec(cond_cont, cond_cat) instead" + ) + + def _require_type_head(self) -> None: + if self.type_head is None: + raise RuntimeError( + f"this {type(self).__name__} has no type_head — either " + "particle_type.target='physical' (the type slice is part of " + "forward()'s own output) or generator='wgan' (the WGAN " + "trainer reads the type slice out of forward()'s output " + "directly instead)" + ) + + +class Stage1Model(StageModel): + """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs + move it to stage 2, except for a migrated v0.2 checkpoint + (`n_sec_head_k_max` given), where it stays attached here + since that's where its weights live and what conditioning it was trained + against (see `_migrate_legacy_model_config`).""" + def __init__( self, pdg_vocab: int, @@ -100,27 +218,31 @@ class Stage1Model(nn.Module): cond_enc: ConditionEncoder | None = None, n_sec_head_cfg: dict | None = None, ) -> None: - super().__init__() - self.generator_kind = generator - self.noise_dim = noise_dim - self.cond_enc = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + super().__init__( + pdg_vocab, + mat_vocab, + particle_cfg, + material_cfg, + cond_out_dim=cond_out_dim, + generator=generator, + noise_dim=noise_dim, + cond_enc=cond_enc, ) - objective = build_objective(generator) - has_time = objective.needs_time - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - in_dim = objective.trunk_in_dim(x_dim, noise_dim) - self.trunk = build_trunk( - router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout, block_conditioning + self._build_trunk_and_heads( + trunk_out_dim=x_dim, + hidden_dim=hidden_dim, + n_res_blocks=n_res_blocks, + cond_out_dim=cond_out_dim, + time_dim=time_dim, + router=router, + trunk_type=trunk_type, + block_conditioning=block_conditioning, + dropout=dropout, + n_sec_head_k_max=n_sec_head_k_max, + n_sec_head_cfg=n_sec_head_cfg, + type_head_out_dim=None, + type_head_cfg=None, ) - self.n_sec_head = None - if n_sec_head_k_max is not None: - head_cfg = HeadConfig.from_dict(n_sec_head_cfg) - hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) - self.n_sec_head = build_mlp_head(cond_out_dim, n_sec_head_k_max + 1, hidden, head_cfg.depth) def forward( self, @@ -133,21 +255,28 @@ class Stage1Model(nn.Module): cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb return self.trunk(x_t, cond, cond_cont, cond_cat) - def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only - valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0 - configs predict n_sec from Stage2OneShot instead.""" + def _require_n_sec_head(self) -> None: + """Overrides `StageModel`'s guard — a `Stage1Model` with no + `n_sec_head` points the caller to stage 2 (n_sec's default owner), + not to `stage1` as the base's message would.""" if self.n_sec_head is None: raise RuntimeError( "this Stage1Model has no n_sec_head — n_sec now lives on " "stage 2 by default; this method only exists " "for a migrated v0.2 checkpoint (n_sec.owner='stage1')" ) + + def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only + valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0 + configs predict n_sec from Stage2OneShot instead.""" + self._require_n_sec_head() + assert self.n_sec_head is not None c_emb = self.cond_enc(cond_cont, cond_cat) return self.n_sec_head(c_emb) -class Stage2OneShot(nn.Module): +class Stage2OneShot(StageModel): """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour, reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`, step 4/5, not implemented yet). @@ -169,9 +298,6 @@ class Stage2OneShot(nn.Module): Under a folding objective (wgan) the type slice stays folded into `sec_dim` (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation. - - `cond_enc`, if given, is used in place of building a fresh - `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). """ def __init__( @@ -200,54 +326,40 @@ class Stage2OneShot(nn.Module): n_sec_head_cfg: dict | None = None, type_head_cfg: dict | None = None, ) -> None: - super().__init__() - self.generator_kind = generator - self.noise_dim = noise_dim - self.k_max = k_max - self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) - self.type_dim = stage2_type_dim( - self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"]) - ) - self.cond_enc = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + super().__init__( + pdg_vocab, + mat_vocab, + particle_cfg, + material_cfg, + cond_out_dim=cond_out_dim, + generator=generator, + noise_dim=noise_dim, + k_max=k_max, + particle_type_cfg=particle_type_cfg, + cond_enc=cond_enc, ) self.context_adapter = ContextAdapter(x_dim, context_dim) self.fuse = nn.Sequential( nn.Linear(cond_out_dim + context_dim, cond_out_dim), nn.SiLU(), ) - objective = build_objective(generator) - has_time = objective.needs_time - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - in_dim = objective.trunk_in_dim(sec_dim, noise_dim) - self.trunk = build_trunk( - router, - trunk_type, - in_dim, - sec_dim, - hidden_dim, - n_res_blocks, - merged_cond_dim, - dropout, - block_conditioning, - ) - self.n_sec_head = None - if build_n_sec_head: - head_cfg = HeadConfig.from_dict(n_sec_head_cfg) - hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) - self.n_sec_head = build_mlp_head(cond_out_dim, k_max + 1, hidden, head_cfg.depth) - self.type_head = None target = self.particle_type_cfg.get("target", "physical") - if target != "physical" and not objective.folds_type_slice: - emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"]) - head_cfg = HeadConfig.from_dict(type_head_cfg) - hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) - self.type_head = build_mlp_head(cond_out_dim, k_max * emb_dim, hidden, head_cfg.depth) - self._type_k_max = k_max - self._type_emb_dim = emb_dim + type_head_out_dim = None if target == "physical" else k_max * self.type_dim + self._build_trunk_and_heads( + trunk_out_dim=sec_dim, + hidden_dim=hidden_dim, + n_res_blocks=n_res_blocks, + cond_out_dim=cond_out_dim, + time_dim=time_dim, + router=router, + trunk_type=trunk_type, + block_conditioning=block_conditioning, + dropout=dropout, + n_sec_head_k_max=k_max if build_n_sec_head else None, + n_sec_head_cfg=n_sec_head_cfg, + type_head_out_dim=type_head_out_dim, + type_head_cfg=type_head_cfg, + ) def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: base = self.cond_enc(cond_cont, cond_cat) @@ -272,12 +384,8 @@ class Stage2OneShot(nn.Module): cond_cat: torch.Tensor, stage1_out: torch.Tensor, ) -> torch.Tensor: - if self.n_sec_head is None: - raise RuntimeError( - "this Stage2OneShot has no n_sec_head — it belongs to a " - "migrated v0.2 checkpoint (n_sec.owner='stage1'); call " - "stage1.predict_n_sec(cond_cont, cond_cat) instead" - ) + self._require_n_sec_head() + assert self.n_sec_head is not None c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) return self.n_sec_head(c_emb) @@ -291,19 +399,13 @@ class Stage2OneShot(nn.Module): vectors (`target="embedding"`) — only under `generator in ("flow", "ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s own output instead (see class docstring).""" - if self.type_head is None: - raise RuntimeError( - "this Stage2OneShot has no type_head — either " - "particle_type.target='physical' (the type slice is part of " - "forward()'s own output) or generator='wgan' (the WGAN " - "trainer reads the type slice out of forward()'s output " - "directly instead)" - ) + self._require_type_head() + assert self.type_head is not None c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) - return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim) + return self.type_head(c_emb).view(-1, self.k_max, self.type_dim) -class Stage2Autoregressive(nn.Module): +class Stage2Autoregressive(StageModel): """Emits secondaries one at a time in descending-energy order, instead of `Stage2OneShot`'s simultaneous k_max-slot prediction. `history` selects `MarkovHistory` or @@ -324,9 +426,6 @@ class Stage2Autoregressive(nn.Module): on token position; `_token_cond` additionally fuses in the history encoding and two running scalars (remaining energy-budget fraction, normalized slot index), and feeds `forward`/`predict_type`/the trunk. - - `cond_enc`, if given, is used in place of building a fresh - `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). """ def __init__( @@ -357,20 +456,19 @@ class Stage2Autoregressive(nn.Module): n_sec_head_cfg: dict | None = None, type_head_cfg: dict | None = None, ) -> None: - super().__init__() - self.history_kind = history - self.generator_kind = generator - self.noise_dim = noise_dim - self.k_max = k_max - self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) - emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"]) - self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim) - - self.cond_enc = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + super().__init__( + pdg_vocab, + mat_vocab, + particle_cfg, + material_cfg, + cond_out_dim=cond_out_dim, + generator=generator, + noise_dim=noise_dim, + k_max=k_max, + particle_type_cfg=particle_type_cfg, + cond_enc=cond_enc, ) + self.history_kind = history self.context_adapter = ContextAdapter(x_dim, context_dim) self.base_fuse = nn.Sequential( nn.Linear(cond_out_dim + context_dim, cond_out_dim), @@ -391,35 +489,28 @@ class Stage2Autoregressive(nn.Module): nn.SiLU(), ) - objective = build_objective(generator) - has_time = objective.needs_time - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim) - in_dim = objective.trunk_in_dim(token_dim, noise_dim) - self.trunk = build_trunk( - router, - trunk_type, - in_dim, - token_dim, - hidden_dim, - n_res_blocks, - merged_cond_dim, - dropout, - block_conditioning, - ) - - self.n_sec_head = None - if build_n_sec_head: - head_cfg = HeadConfig.from_dict(n_sec_head_cfg) - hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) - self.n_sec_head = build_mlp_head(cond_out_dim, k_max + 1, hidden, head_cfg.depth) - self.type_head = None + # `self.type_dim` (set by StageModel.__init__) doubles as the raw + # `emb_dim` `stage2_trunk_sec_dim` wants: for a non-"physical" target + # `stage2_type_dim` already resolved `type_dim` to exactly that value; + # for "physical" the emb_dim argument goes unused anyway. + token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, self.type_dim) target = self.particle_type_cfg.get("target", "physical") - if target != "physical" and not objective.folds_type_slice: - head_cfg = HeadConfig.from_dict(type_head_cfg) - hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) - self.type_head = build_mlp_head(cond_out_dim, self.type_dim, hidden, head_cfg.depth) + type_head_out_dim = None if target == "physical" else self.type_dim + self._build_trunk_and_heads( + trunk_out_dim=token_dim, + hidden_dim=hidden_dim, + n_res_blocks=n_res_blocks, + cond_out_dim=cond_out_dim, + time_dim=time_dim, + router=router, + trunk_type=trunk_type, + block_conditioning=block_conditioning, + dropout=dropout, + n_sec_head_k_max=k_max if build_n_sec_head else None, + n_sec_head_cfg=n_sec_head_cfg, + type_head_out_dim=type_head_out_dim, + type_head_cfg=type_head_cfg, + ) def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: base = self.cond_enc(cond_cont, cond_cat) @@ -511,12 +602,8 @@ class Stage2Autoregressive(nn.Module): return out.view(B, K, -1) def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: - if self.n_sec_head is None: - raise RuntimeError( - "this Stage2Autoregressive has no n_sec_head — it belongs to " - "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call " - "stage1.predict_n_sec(cond_cont, cond_cat) instead" - ) + self._require_n_sec_head() + assert self.n_sec_head is not None return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out)) def predict_type( @@ -530,14 +617,8 @@ class Stage2Autoregressive(nn.Module): slot_idx: torch.Tensor, hist: torch.Tensor | None = None, ) -> torch.Tensor: - if self.type_head is None: - raise RuntimeError( - "this Stage2Autoregressive has no type_head — either " - "particle_type.target='physical' (the type slice is part of " - "forward()'s own output) or generator='wgan' (the WGAN " - "trainer reads the type slice out of forward()'s output " - "directly instead)" - ) + self._require_type_head() + assert self.type_head is not None c_emb = self._token_cond( cond_cont, cond_cat, diff --git a/giant/model/network.py b/giant/model/network.py index 547b39b..ffd037a 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -36,6 +36,7 @@ from giant.model.models import ( Stage1Model, Stage2Autoregressive, Stage2OneShot, + StageModel, resolve_type_n_classes, stage2_trunk_sec_dim, stage2_type_dim, @@ -102,6 +103,7 @@ __all__ = [ "Stage1Model", "Stage2Autoregressive", "Stage2OneShot", + "StageModel", "TRUNK_REGISTRY", "Trunk", "WganObjective", diff --git a/tests/test_network.py b/tests/test_network.py index 4a99ad7..73662eb 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -14,9 +14,11 @@ from giant.model.network import ( Stage1Model, Stage2Autoregressive, Stage2OneShot, + StageModel, build_critics, build_history, build_models, + build_objective, stage2_trunk_sec_dim, stage2_type_dim, ) @@ -987,3 +989,163 @@ def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_genera assert overridden is not None assert overridden.input_proj.out_features == 16 assert len(overridden.blocks) == 3 + + +# ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive +# scaffolding — construction order, and therefore fresh-init RNG draw order and +# state_dict key set, must stay byte-for-byte what it was before the base class +# existed. ------------------------------------------------------------------ + +_STAGE_HIDDEN_DIM = 32 +_STAGE_N_BLOCKS = 2 +_STAGE_COND_OUT_DIM = 16 + + +def _resblock_keys(prefix: str) -> set[str]: + return { + f"{prefix}.norm.weight", + f"{prefix}.norm.bias", + f"{prefix}.linear1.weight", + f"{prefix}.linear1.bias", + f"{prefix}.cond_proj.weight", + f"{prefix}.linear2.weight", + f"{prefix}.linear2.bias", + } + + +def _trunk_keys(prefix: str = "trunk") -> set[str]: + keys = { + f"{prefix}.input_proj.weight", + f"{prefix}.input_proj.bias", + f"{prefix}.out_proj.weight", + f"{prefix}.out_proj.bias", + } + for i in range(_STAGE_N_BLOCKS): + keys |= _resblock_keys(f"{prefix}.blocks.{i}") + return keys + + +def _cond_enc_keys() -> set[str]: + return { + "cond_enc.mlp.0.weight", + "cond_enc.mlp.0.bias", + "cond_enc.mlp.2.weight", + "cond_enc.mlp.2.bias", + "cond_enc.particle_mlp.0.weight", + "cond_enc.particle_mlp.0.bias", + "cond_enc.material_mlp.0.weight", + "cond_enc.material_mlp.0.bias", + } + + +def _fuse_keys(name: str) -> set[str]: + return {f"{name}.0.weight", f"{name}.0.bias"} + + +def _head_keys(name: str) -> set[str]: + return {f"{name}.0.weight", f"{name}.0.bias", f"{name}.2.weight", f"{name}.2.bias"} + + +def _expected_stage_keys(*, has_time: bool, extra: set[str]) -> set[str]: + keys = _cond_enc_keys() | _trunk_keys() | extra + if has_time: + keys.add("time_emb.freqs") + return keys + + +@pytest.mark.parametrize("generator", ["flow", "wgan"]) +def test_stage1_model_state_dict_keys_unchanged_by_stagemodel_refactor(generator): + model = Stage1Model( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=_STAGE_HIDDEN_DIM, + n_res_blocks=_STAGE_N_BLOCKS, + cond_out_dim=_STAGE_COND_OUT_DIM, + generator=generator, + time_dim=8, + noise_dim=8, + n_sec_head_k_max=15, + ) + expected = _expected_stage_keys( + has_time=build_objective(generator).needs_time, + extra=_head_keys("n_sec_head"), + ) + assert set(model.state_dict().keys()) == expected + + +@pytest.mark.parametrize("generator", ["flow", "wgan"]) +def test_stage2_oneshot_state_dict_keys_unchanged_by_stagemodel_refactor(generator): + model = Stage2OneShot( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=_STAGE_HIDDEN_DIM, + n_res_blocks=_STAGE_N_BLOCKS, + cond_out_dim=_STAGE_COND_OUT_DIM, + generator=generator, + time_dim=8, + noise_dim=8, + k_max=15, + ) + extra = _head_keys("n_sec_head") | {"context_adapter.proj.weight", "context_adapter.proj.bias"} | _fuse_keys("fuse") + expected = _expected_stage_keys(has_time=build_objective(generator).needs_time, extra=extra) + assert set(model.state_dict().keys()) == expected + + +@pytest.mark.parametrize("generator", ["flow", "wgan"]) +def test_stage2_autoregressive_state_dict_keys_unchanged_by_stagemodel_refactor(generator): + model = Stage2Autoregressive( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=_STAGE_HIDDEN_DIM, + n_res_blocks=_STAGE_N_BLOCKS, + cond_out_dim=_STAGE_COND_OUT_DIM, + generator=generator, + time_dim=8, + noise_dim=8, + k_max=15, + ) + extra = ( + _head_keys("n_sec_head") + | {"context_adapter.proj.weight", "context_adapter.proj.bias"} + | _fuse_keys("base_fuse") + | _fuse_keys("token_fuse") + | {"history_encoder.start", "history_encoder.mlp.0.weight", "history_encoder.mlp.0.bias"} + ) + expected = _expected_stage_keys(has_time=build_objective(generator).needs_time, extra=extra) + assert set(model.state_dict().keys()) == expected + + +@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive]) +def test_stage_classes_are_stagemodel_subclasses(cls): + assert issubclass(cls, StageModel) + + +@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive]) +@pytest.mark.parametrize("generator", ["flow", "ddpm", "wgan"]) +def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator): + kwargs = dict( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=_STAGE_HIDDEN_DIM, + n_res_blocks=_STAGE_N_BLOCKS, + cond_out_dim=_STAGE_COND_OUT_DIM, + generator=generator, + time_dim=8, + noise_dim=8, + ) + if cls is Stage1Model: + kwargs["n_sec_head_k_max"] = 15 + else: + kwargs["k_max"] = 15 + model = cls(**kwargs) + assert model.generator_kind == generator + assert model.noise_dim == 8 + assert (model.time_emb is not None) == build_objective(generator).needs_time From c4b12b5e7a52cf180ca4109277969c49b6c84c5c Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 14 Aug 2026 15:03:55 +0200 Subject: [PATCH 2/3] Pass ConditioningAxisConfig/ParticleTypeConfig themselves instead of raw dicts (gitea #38) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit build_models/build_critics parsed model_config into frozen dataclasses (ConditioningConfig, Stage2ModelConfig, ...) but then threw the parsed sub-objects away and passed the original raw dicts (conditioning["particle"], s2_spec.particle_type.to_dict()) down into ConditionEncoder/StageModel/etc, which re-read them with their own hardcoded .get(key, default) fallbacks — each an independent copy of a fact the dataclass already stated once. Worst instance: giant/training/trainers.py:236 converted an already-parsed ParticleTypeConfig back into a dict for no reason. Threads ConditioningAxisConfig (particle_cfg/material_cfg) and ParticleTypeConfig (particle_type_cfg) as the actual dataclass instances through every signature that used to type them dict: ConditionEncoder, StageModel/CriticModel, resolve_type_n_classes/stage2_type_dim/ stage2_trunk_sec_dim, giant/model/builders.py, giant/sample.py, giant/training/stage2_inputs.py, giant/training/trainers.py (StageSpec/ StageTrainer), giant/pipeline.py, giant/rollout.py, giant/validate.py — so ty now catches a misspelled field instead of it silently falling back. No config-schema change: config.toml/checkpoint model_config keep the same nested-dict shape; only what happens after the existing X.from_dict(...) parse changes. User-confirmed scope decision: both axes (particle_cfg/material_cfg and particle_type_cfg), not just the more heavily-duplicated particle_type_cfg axis, and not stopping at the two most literal parse-then-discard round trips — matching the issue's own proposal. Preserved-default decision: StageModel's particle_type_cfg=None sentinel (hit only by direct/test construction — build_models always passes an explicit particle_type) still resolves to ParticleTypeConfig(target= "physical"), not ParticleTypeConfig()'s own target="onehot" config-file default — switching it would have silently grown an unused, gradient-less type_head on every test that constructs Stage2OneShot/Stage2Autoregressive without particle_type_cfg=, breaking their "every param has a grad" checks. New tests in tests/test_network.py: ConditionEncoder/StageModel store the exact ConditioningAxisConfig/ParticleTypeConfig instance passed in (identity, not just equality) — no internal dict round-trip — and build_models's output carries real dataclass instances end to end, not the plain dicts it produced before this fix. Co-Authored-By: Claude Opus 5 --- giant/model/builders.py | 24 +++--- giant/model/encoders.py | 37 +++++---- giant/model/models.py | 64 ++++++++------- giant/pipeline.py | 6 +- giant/rollout.py | 4 +- giant/sample.py | 4 +- giant/training/stage2_inputs.py | 19 +++-- giant/training/trainers.py | 14 ++-- giant/validate.py | 4 +- tests/test_flow.py | 5 +- tests/test_network.py | 140 +++++++++++++++++++++----------- tests/test_objectives.py | 3 +- tests/test_phase2.py | 7 +- tests/test_rollout.py | 23 +++--- tests/test_router.py | 5 +- tests/test_sample.py | 13 +-- tests/test_train.py | 7 +- tests/test_validate.py | 18 ++-- tests/test_wgan.py | 5 +- 19 files changed, 233 insertions(+), 169 deletions(-) diff --git a/giant/model/builders.py b/giant/model/builders.py index 21a3191..de11da7 100644 --- a/giant/model/builders.py +++ b/giant/model/builders.py @@ -45,11 +45,10 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config) pdg_vocab = cfg["pdg_vocab"] mat_vocab = cfg["mat_vocab"] - conditioning = cfg["conditioning"] - particle_cfg = conditioning["particle"] - material_cfg = conditioning["material"] - particle_conditioning = particle_cfg["type"] - conditioning_cfg = ConditioningConfig.from_dict(conditioning) + conditioning_cfg = ConditioningConfig.from_dict(cfg["conditioning"]) + particle_cfg = conditioning_cfg.particle + material_cfg = conditioning_cfg.material + particle_conditioning = particle_cfg.type s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"]) s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"]) cond_out_dim = conditioning_cfg.out_dim @@ -108,7 +107,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: time_dim = getattr(s2_spec, generator).time_dim if objective.needs_time else 64 n_sec_owner = s2_spec.n_sec.owner k_max = s2_spec.k_max - particle_type_cfg = s2_spec.particle_type.to_dict() + particle_type_cfg = s2_spec.particle_type if decoder == "autoregressive": ar_cfg = s2_spec.autoregressive @@ -140,7 +139,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: ) else: sec_dim = stage2_trunk_sec_dim( - particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]) + particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg.emb_dim) ) result["stage2"] = Stage2OneShot( pdg_vocab=pdg_vocab, @@ -178,10 +177,9 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config) pdg_vocab = cfg["pdg_vocab"] mat_vocab = cfg["mat_vocab"] - conditioning = cfg["conditioning"] - particle_cfg = conditioning["particle"] - material_cfg = conditioning["material"] - conditioning_cfg = ConditioningConfig.from_dict(conditioning) + conditioning_cfg = ConditioningConfig.from_dict(cfg["conditioning"]) + particle_cfg = conditioning_cfg.particle + material_cfg = conditioning_cfg.material cond_out_dim = conditioning_cfg.out_dim s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"]) s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"]) @@ -204,12 +202,12 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: if s2_spec.active and build_objective(s2_spec.generator).is_adversarial: k_max = s2_spec.k_max - particle_type_cfg = s2_spec.particle_type.to_dict() + particle_type_cfg = s2_spec.particle_type in_dim = stage2_trunk_sec_dim( particle_type_cfg, s2_spec.generator, k_max, - resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]), + resolve_type_n_classes(particle_type_cfg, particle_cfg.emb_dim), ) result["stage2"] = CriticModel( pdg_vocab=pdg_vocab, diff --git a/giant/model/encoders.py b/giant/model/encoders.py index 5aeee3f..c82ab75 100644 --- a/giant/model/encoders.py +++ b/giant/model/encoders.py @@ -6,6 +6,7 @@ import torch.nn as nn import torch.nn.functional as F from giant.cond_layout import CondLayout +from giant.config import ConditioningAxisConfig from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM from giant.model.layers import _make_axis_mlp @@ -14,9 +15,9 @@ class ConditionEncoder(nn.Module): """Fuses continuous conditioning with particle/material identity. The particle and material axes are configured independently - (`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`) - and may mix freely, e.g. material "physical" with particle "embedding". - Three modes per axis: + (`particle_cfg`/`material_cfg`, each a `ConditioningAxisConfig`) and may + mix freely, e.g. material "physical" with particle "embedding". Three + modes per axis: - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s dense training-vocab index. Memorizes the training menu. - "physical": an `n_layers`-deep MLP over the axis's raw physical @@ -37,30 +38,30 @@ class ConditionEncoder(nn.Module): self, pdg_vocab: int, mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, + particle_cfg: ConditioningAxisConfig, + material_cfg: ConditioningAxisConfig, cont_dim: int = COND_DIM, out_dim: int = 128, ) -> None: super().__init__() - self.particle_cfg = dict(particle_cfg) - self.material_cfg = dict(material_cfg) + self.particle_cfg = particle_cfg + self.material_cfg = material_cfg # Also validates both axis types — an unknown one raises here. - self.layout = CondLayout.from_types(particle_cfg["type"], material_cfg["type"]) + self.layout = CondLayout.from_types(particle_cfg.type, material_cfg.type) - p_type = particle_cfg["type"] - p_emb_dim = particle_cfg["emb_dim"] + p_type = particle_cfg.type + p_emb_dim = particle_cfg.emb_dim if p_type == "embedding": self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim) elif p_type == "physical": - self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1)) + self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.n_layers) - m_type = material_cfg["type"] - m_emb_dim = material_cfg["emb_dim"] + m_type = material_cfg.type + m_emb_dim = material_cfg.emb_dim if m_type == "embedding": self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim) elif m_type == "physical": - self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1)) + self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.n_layers) in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim self.mlp = nn.Sequential( @@ -70,7 +71,7 @@ class ConditionEncoder(nn.Module): ) def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor): - p_type = self.particle_cfg["type"] + p_type = self.particle_cfg.type if p_type == "embedding": return self.pdg_emb(cond_cat[:, self.layout.PDG_COL]) if p_type == "physical": @@ -78,11 +79,11 @@ class ConditionEncoder(nn.Module): assert self.layout.particle_topn_col is not None return F.one_hot( cond_cat[:, self.layout.particle_topn_col], - num_classes=self.particle_cfg["emb_dim"], + num_classes=self.particle_cfg.emb_dim, ).float() def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor): - m_type = self.material_cfg["type"] + m_type = self.material_cfg.type if m_type == "embedding": return self.mat_emb(cond_cat[:, self.layout.MAT_COL]) if m_type == "physical": @@ -90,7 +91,7 @@ class ConditionEncoder(nn.Module): assert self.layout.material_topn_col is not None return F.one_hot( cond_cat[:, self.layout.material_topn_col], - num_classes=self.material_cfg["emb_dim"], + num_classes=self.material_cfg.emb_dim, ).float() def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: diff --git a/giant/model/models.py b/giant/model/models.py index d1f79c6..4f80fdf 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn -from giant.config import HeadConfig +from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM from giant.model.encoders import ConditionEncoder from giant.model.history import HistoryEncoder, build_history @@ -18,7 +18,7 @@ from giant.model.trunks import build_trunk # --------------------------------------------------------------------------- -def resolve_type_n_classes(particle_type_cfg: dict, particle_emb_dim: int) -> int: +def resolve_type_n_classes(particle_type_cfg: ParticleTypeConfig, particle_emb_dim: int) -> int: """Effective width fed to `stage2_type_dim`/`stage2_trunk_sec_dim` in place of a bare `conditioning.particle.emb_dim` read. Under `target = "onehot"` this is `stage2_model.particle_type.n_classes` (0 = @@ -29,22 +29,21 @@ def resolve_type_n_classes(particle_type_cfg: dict, particle_emb_dim: int) -> in apply — the width stays `conditioning.particle.emb_dim`, the embedding table's own dimensionality (`validate_config` requires `conditioning.particle.type = "embedding"` here).""" - if particle_type_cfg.get("target", "physical") == "onehot": - return particle_type_cfg.get("n_classes", 0) or particle_emb_dim + if particle_type_cfg.target == "onehot": + return particle_type_cfg.n_classes or particle_emb_dim return particle_emb_dim -def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int: +def stage2_type_dim(particle_type_cfg: ParticleTypeConfig, emb_dim: int) -> int: """Width of a single secondary slot's type slice — `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are this many classes/dims wide — callers resolve `emb_dim` via `resolve_type_n_classes` first).""" - target = particle_type_cfg.get("target", "physical") - return PARTICLE_PHYS_DIM if target == "physical" else emb_dim + return PARTICLE_PHYS_DIM if particle_type_cfg.target == "physical" else emb_dim -def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int) -> int: +def stage2_trunk_sec_dim(particle_type_cfg: ParticleTypeConfig, generator: str, k_max: int, emb_dim: int) -> int: """`Stage2OneShot`'s trunk output width. `target = "physical"` is untouched from v0.2/today: @@ -59,8 +58,7 @@ def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, em isn't part of this vector at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk only covers `k_max * CONT_SLOT_DIM`. """ - target = particle_type_cfg.get("target", "physical") - if target == "physical": + if particle_type_cfg.target == "physical": return k_max * SEC_SLOT_DIM if build_objective(generator).folds_type_slice: return k_max * (CONT_SLOT_DIM + emb_dim) @@ -87,22 +85,30 @@ class StageModel(nn.Module): self, pdg_vocab: int, mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, + particle_cfg: ConditioningAxisConfig, + material_cfg: ConditioningAxisConfig, cond_out_dim: int, generator: str, noise_dim: int, k_max: int | None = None, - particle_type_cfg: dict | None = None, + particle_type_cfg: ParticleTypeConfig | None = None, cond_enc: ConditionEncoder | None = None, ) -> None: super().__init__() self.generator_kind = generator self.noise_dim = noise_dim self.k_max = k_max - self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) + # `ParticleTypeConfig()`'s own dataclass default is target="onehot" + # (the config.toml default when [stage2_model.particle_type] is + # omitted) — a different question from "nobody passed anything to + # this constructor", which direct/test construction relies on + # defaulting to "physical" (build_models/build_critics always pass + # particle_type_cfg explicitly, so this sentinel is never hit there). + self.particle_type_cfg = ( + particle_type_cfg if particle_type_cfg is not None else ParticleTypeConfig(target="physical") + ) self.type_dim = stage2_type_dim( - self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"]) + self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg.emb_dim) ) self.cond_enc = ( cond_enc @@ -139,7 +145,7 @@ class StageModel(nn.Module): migrated v0.2 checkpoint, `Stage2OneShot`/`Stage2Autoregressive` pass it whenever `build_n_sec_head=True`. `type_head` is built iff `type_head_out_dim is not None` (the caller — only the two Stage2 - classes — passes `None` exactly when `particle_type_cfg["target"] == + classes — passes `None` exactly when `particle_type_cfg.target == "physical"`) *and* the objective doesn't fold the type slice into its own trunk output (checked here, since `objective` is already needed for the trunk itself). @@ -201,8 +207,8 @@ class Stage1Model(StageModel): self, pdg_vocab: int, mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, + particle_cfg: ConditioningAxisConfig, + material_cfg: ConditioningAxisConfig, hidden_dim: int = 256, n_res_blocks: int = 6, cond_out_dim: int = 128, @@ -285,7 +291,7 @@ class Stage2OneShot(StageModel): (a migrated v0.2 checkpoint, whose n_sec_head instead attaches to Stage1Model — see `_migrate_legacy_model_config`). - `particle_type_cfg["target"]` (default `"physical"`) selects the + `particle_type_cfg.target` (default `"physical"`) selects the secondary-type mechanism: `"physical"` keeps the type slice folded into the trunk's own flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by @@ -304,8 +310,8 @@ class Stage2OneShot(StageModel): self, pdg_vocab: int, mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, + particle_cfg: ConditioningAxisConfig, + material_cfg: ConditioningAxisConfig, hidden_dim: int = 256, n_res_blocks: int = 6, cond_out_dim: int = 128, @@ -321,7 +327,7 @@ class Stage2OneShot(StageModel): trunk_type: str = "resmlp", block_conditioning: str = "add", build_n_sec_head: bool = True, - particle_type_cfg: dict | None = None, + particle_type_cfg: ParticleTypeConfig | None = None, cond_enc: ConditionEncoder | None = None, n_sec_head_cfg: dict | None = None, type_head_cfg: dict | None = None, @@ -343,7 +349,7 @@ class Stage2OneShot(StageModel): nn.Linear(cond_out_dim + context_dim, cond_out_dim), nn.SiLU(), ) - target = self.particle_type_cfg.get("target", "physical") + target = self.particle_type_cfg.target type_head_out_dim = None if target == "physical" else k_max * self.type_dim self._build_trunk_and_heads( trunk_out_dim=sec_dim, @@ -432,8 +438,8 @@ class Stage2Autoregressive(StageModel): self, pdg_vocab: int, mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, + particle_cfg: ConditioningAxisConfig, + material_cfg: ConditioningAxisConfig, hidden_dim: int = 256, n_res_blocks: int = 6, cond_out_dim: int = 128, @@ -448,7 +454,7 @@ class Stage2Autoregressive(StageModel): trunk_type: str = "resmlp", block_conditioning: str = "add", build_n_sec_head: bool = True, - particle_type_cfg: dict | None = None, + particle_type_cfg: ParticleTypeConfig | None = None, history: str = "markov", attn_n_heads: int = 4, attn_n_layers: int = 2, @@ -494,7 +500,7 @@ class Stage2Autoregressive(StageModel): # `stage2_type_dim` already resolved `type_dim` to exactly that value; # for "physical" the emb_dim argument goes unused anyway. token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, self.type_dim) - target = self.particle_type_cfg.get("target", "physical") + target = self.particle_type_cfg.target type_head_out_dim = None if target == "physical" else self.type_dim self._build_trunk_and_heads( trunk_out_dim=token_dim, @@ -643,8 +649,8 @@ class CriticModel(nn.Module): self, pdg_vocab: int, mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, + particle_cfg: ConditioningAxisConfig, + material_cfg: ConditioningAxisConfig, in_dim: int, hidden_dim: int = 256, n_res_blocks: int = 6, diff --git a/giant/pipeline.py b/giant/pipeline.py index 3b567a7..ebfee47 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -188,8 +188,8 @@ def run_setup_stage( # independent of both. particle_cfg = cfg["conditioning"]["particle"] material_cfg = cfg["conditioning"]["material"] - particle_type_cfg_dict = cfg["stage2_model"].get("particle_type") or {} - particle_type_target = config.ParticleTypeConfig.from_dict(particle_type_cfg_dict).target + particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")) + particle_type_target = particle_type_cfg.target def _pdg_topn(n_classes: int) -> TopNMap: cache_key = setup_cache.topn_key("pdg", n_classes) @@ -210,7 +210,7 @@ def run_setup_stage( sec_type_topn_map: TopNMap | None = None if particle_type_target == "onehot": - sec_type_n_classes = resolve_type_n_classes(particle_type_cfg_dict, particle_cfg["emb_dim"]) + sec_type_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]) sec_type_topn_map = _pdg_topn(sec_type_n_classes) mat_topn_map: TopNMap | None = None diff --git a/giant/rollout.py b/giant/rollout.py index b1098d5..03ed04d 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -141,7 +141,7 @@ def decode_secondary_identity( Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, sec_type_l1_dist) — the last is `None` except under `"embedding"`. """ - target = sec_decoder.particle_type_cfg.get("target", "physical") + target = sec_decoder.particle_type_cfg.target if target == "physical": sec_full = torch.cat([sec_cont, sec_type], dim=-1).cpu().numpy() @@ -491,7 +491,7 @@ def rollout( "conditioning.particle.type='onehot' rollout needs pdg_topn_map " "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']" ) - if sec_decoder.particle_type_cfg.get("target") == "onehot" and sec_type_topn_map is None: + if sec_decoder.particle_type_cfg.target == "onehot" and sec_type_topn_map is None: raise RuntimeError( "stage2_model.particle_type.target='onehot' rollout needs sec_type_topn_map " "(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']" diff --git a/giant/sample.py b/giant/sample.py index fede756..374049b 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -131,7 +131,7 @@ def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int: def _type_folded(sec_decoder: torch.nn.Module) -> bool: - target = sec_decoder.particle_type_cfg.get("target", "physical") + target = sec_decoder.particle_type_cfg.target return target == "physical" or build_objective(sec_decoder.generator_kind).folds_type_slice @@ -282,7 +282,7 @@ def sample_secondaries_ar( k_max = sec_decoder.k_max type_dim = sec_decoder.type_dim objective = build_objective(sec_decoder.generator_kind) - target = sec_decoder.particle_type_cfg.get("target", "physical") + target = sec_decoder.particle_type_cfg.target type_folded = _type_folded(sec_decoder) token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py index c09011e..8bd3730 100644 --- a/giant/training/stage2_inputs.py +++ b/giant/training/stage2_inputs.py @@ -11,6 +11,7 @@ live in one place and stay unit-testable on their own. import torch import torch.nn.functional as F +from giant.config import ParticleTypeConfig from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM from giant.model.objectives import build_objective from giant.sample import sample_secondaries_ar @@ -31,7 +32,7 @@ def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) - def _type_repr( sec_type_idx: torch.Tensor, sec_cont: torch.Tensor, - particle_type_cfg: dict, + particle_type_cfg: ParticleTypeConfig, cond_enc: torch.nn.Module, emb_dim: int, ) -> torch.Tensor: @@ -46,7 +47,7 @@ def _type_repr( latter must always reflect the true physical secondary that came before, regardless of what the *current* token's own training objective is. """ - target = particle_type_cfg.get("target", "physical") + target = particle_type_cfg.target if target == "physical": return sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM] if target == "onehot": @@ -57,7 +58,7 @@ def _type_repr( def _assemble_stage2_ar_target( sec_cont: torch.Tensor, sec_type_idx: torch.Tensor, - particle_type_cfg: dict, + particle_type_cfg: ParticleTypeConfig, generator: str, cond_enc: torch.nn.Module, emb_dim: int, @@ -79,7 +80,7 @@ def _assemble_stage2_ar_target( one-hot of the true class, relaxed on the *generated* side only, by the caller; or the conditioning's own detached embedding-table row). """ - target = particle_type_cfg.get("target", "physical") + target = particle_type_cfg.target if target == "physical": return sec_cont cont = sec_cont[..., :CONT_SLOT_DIM] @@ -92,7 +93,7 @@ def _assemble_stage2_ar_target( def _assemble_stage2_real( sec_cont: torch.Tensor, sec_type_idx: torch.Tensor, - particle_type_cfg: dict, + particle_type_cfg: ParticleTypeConfig, generator: str, cond_enc: torch.nn.Module, emb_dim: int, @@ -157,7 +158,7 @@ def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tenso def _assemble_stage2_ar_inputs( sec_cont: torch.Tensor, sec_type_idx: torch.Tensor, - particle_type_cfg: dict, + particle_type_cfg: ParticleTypeConfig, cond_enc: torch.nn.Module, emb_dim: int, ) -> dict[str, torch.Tensor]: @@ -200,7 +201,7 @@ def _stage2_tf_prob(mode: str, p_start: float, p_end: float, epoch: int, total_e def _history_repr_from_ar_sample( sec_cont_pred: torch.Tensor, sec_type_pred: torch.Tensor, - particle_type_cfg: dict, + particle_type_cfg: ParticleTypeConfig, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """`(fraction, direction, type_repr)` — the same triple `_type_repr` / `_stick_fraction` derive from ground truth, but from a free-running @@ -213,7 +214,7 @@ def _history_repr_from_ar_sample( representation.""" fraction = torch.sigmoid(sec_cont_pred[..., 0]) direction = sec_cont_pred[..., 1:4] - if particle_type_cfg.get("target", "physical") == "onehot": + if particle_type_cfg.target == "onehot": type_dim = sec_type_pred.size(-1) type_repr = F.one_hot(sec_type_pred.argmax(-1), num_classes=type_dim).float() else: @@ -229,7 +230,7 @@ def _assemble_stage2_ar_inputs_scheduled( sec_cont: torch.Tensor, sec_type_idx: torch.Tensor, n_sec: torch.Tensor, - particle_type_cfg: dict, + particle_type_cfg: ParticleTypeConfig, cond_enc: torch.nn.Module, emb_dim: int, p_tf: float, diff --git a/giant/training/trainers.py b/giant/training/trainers.py index e91ff8e..70b3a05 100644 --- a/giant/training/trainers.py +++ b/giant/training/trainers.py @@ -146,7 +146,7 @@ class StageSpec: n_sec_lambda=s2_spec.n_sec.lambda_weight, particle_type=s2_spec.particle_type, particle_type_n_classes=resolve_type_n_classes( - s2_spec.particle_type.to_dict(), cfg["conditioning"]["particle"]["emb_dim"] + s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"] ), # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge # (giant/config.py), so TrainConfig.from_dict never has to fall @@ -233,7 +233,7 @@ class StageTrainer: self.router = _stage_router(self.model) self._modules = (self.model, *extra_modules) - self.particle_type_cfg = spec.particle_type.to_dict() + self.particle_type_cfg = spec.particle_type self.particle_type_n_classes = spec.particle_type_n_classes self.ema_decay = spec.ema_decay @@ -453,14 +453,14 @@ class FlowDDPMStageTrainer(StageTrainer): ) super().__init__(spec, model, device) self.objective = objective - self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0) + self.particle_type_lambda = self.particle_type_cfg.lambda_weight # Width of the type slice actually folded into x1_s2 by _sec_target, # under this trainer's objective (flow/ddpm only — see the # NotImplementedError above, neither folds the type slice): "physical" # keeps it folded in (PARTICLE_PHYS_DIM wide, unchanged from v0.2); # "onehot"/"embedding" pull it out into model.type_head instead (0 # here). - self._flow_type_dim = None if self.particle_type_cfg.get("target", "physical") == "physical" else 0 + self._flow_type_dim = None if self.particle_type_cfg.target == "physical" else 0 self.params = list(self.model.parameters()) self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay) @@ -549,7 +549,7 @@ class FlowDDPMStageTrainer(StageTrainer): type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx) mask = sec_mask.float() denom = mask.sum().clamp(min=1) - if self.particle_type_cfg.get("target") == "onehot": + if self.particle_type_cfg.target == "onehot": ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, reduction="none") l_type = (ce * mask).sum() / denom type_acc = ((type_out.argmax(-1) == sec_type_idx).float() * mask).sum() / denom @@ -726,7 +726,7 @@ class WGANStageTrainer(StageTrainer): "grad_norm_d", "grad_norm_g", ] - if self.is_stage2 and self.particle_type_cfg.get("target") == "onehot": + if self.is_stage2 and self.particle_type_cfg.target == "onehot": # Differentiability instrumentation — only meaningful when the # type slice is a straight-through Gumbel relaxation. train_keys += ["grad_norm_type_slice", "grad_norm_cont_slice"] @@ -804,7 +804,7 @@ class WGANStageTrainer(StageTrainer): global_step, device, ) - if self.particle_type_cfg.get("target", "physical") == "onehot": + if self.particle_type_cfg.target == "onehot": # Straight-through Gumbel-softmax relaxation of the type # slice only — the critic must see a hard one-hot forward # (matching what "real" data looks like) while gradient diff --git a/giant/validate.py b/giant/validate.py index acf3f87..1e68816 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -82,7 +82,7 @@ def validate_marginals( one-shot-vs-autoregressive-agnostic): n_sec distribution (+ classification accuracy), per-slot energy-fraction marginals, and a particle-type marginal whose shape depends on - `sec_decoder.particle_type_cfg["target"]` — restricted to each side's own + `sec_decoder.particle_type_cfg.target` — restricted to each side's own valid slots (real: `n_sec`; generated: the resolved `n_sec_pred`), since the two need not agree on how many slots are valid. Adds {"n_sec_real", "n_sec_pred", "n_sec_accuracy", "energy_fraction_kl"} plus, under @@ -103,7 +103,7 @@ def validate_marginals( sec_decoder.eval() k_max = sec_decoder.k_max if sec_decoder is not None else 0 - target = sec_decoder.particle_type_cfg.get("target", "physical") if sec_decoder is not None else "physical" + target = sec_decoder.particle_type_cfg.target if sec_decoder is not None else "physical" all_real, all_gen = [], [] all_n_sec_real, all_n_sec_pred = [], [] diff --git a/tests/test_flow.py b/tests/test_flow.py index aae4a93..6b0e256 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -1,11 +1,12 @@ import torch +from giant.config import ConditioningAxisConfig from giant.constants import COND_DIM from giant.model.network import Stage1Model from giant.model.schedule import CosineSchedule, flow_matching_loss from giant.sample import sample_flow, sample_ddim -PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} -MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} +PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) +MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) def _small_model(): diff --git a/tests/test_network.py b/tests/test_network.py index 73662eb..c82f77a 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -23,10 +23,10 @@ from giant.model.network import ( stage2_type_dim, ) -PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} -MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} -ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1} -ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1} +PARTICLE_CFG = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) +MATERIAL_CFG = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) +ONEHOT_PARTICLE_CFG = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=6, n_layers=1) +ONEHOT_MATERIAL_CFG = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=4, n_layers=1) def test_sinusoidal_embedding_shape(): @@ -135,28 +135,31 @@ def test_stage1_model_n_sec_head_cfg_controls_hidden_width_and_depth(): def test_stage2_type_dim_physical_is_particle_phys_dim(): - assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM + assert stage2_type_dim(gconfig.ParticleTypeConfig(target="physical"), emb_dim=16) == PARTICLE_PHYS_DIM def test_stage2_type_dim_onehot_and_embedding_are_emb_dim(): - assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16 - assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16 + assert stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim=16) == 16 + assert stage2_type_dim(gconfig.ParticleTypeConfig(target="embedding"), emb_dim=16) == 16 def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim(): k_max = 15 - assert stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM - assert stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM + physical = gconfig.ParticleTypeConfig(target="physical") + assert stage2_trunk_sec_dim(physical, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM + assert stage2_trunk_sec_dim(physical, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in(): k_max = 15 - assert stage2_trunk_sec_dim({"target": "onehot"}, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16) + onehot = gconfig.ParticleTypeConfig(target="onehot") + assert stage2_trunk_sec_dim(onehot, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16) def test_stage2_trunk_sec_dim_onehot_flow_excludes_type(): k_max = 15 - assert stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM + onehot = gconfig.ParticleTypeConfig(target="onehot") + assert stage2_trunk_sec_dim(onehot, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM # --- ConditionEncoder onehot mode ------------------------------------------- @@ -164,8 +167,8 @@ def test_stage2_trunk_sec_dim_onehot_flow_excludes_type(): def test_condition_encoder_onehot_forward_shape_and_gradients(): B = 8 - particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"]) - material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"]) + particle_emb_dim = ONEHOT_PARTICLE_CFG.emb_dim + material_emb_dim = ONEHOT_MATERIAL_CFG.emb_dim enc = ConditionEncoder( pdg_vocab=5, mat_vocab=3, @@ -196,12 +199,12 @@ def test_condition_encoder_onehot_is_a_true_one_hot_vector(): verify the concatenated input segment really is one-hot, not e.g. an accidentally-learned embedding.""" B = 4 - particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"]) + particle_emb_dim = ONEHOT_PARTICLE_CFG.emb_dim enc = ConditionEncoder( pdg_vocab=5, mat_vocab=3, particle_cfg=ONEHOT_PARTICLE_CFG, - material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1}, + material_cfg=gconfig.ConditioningAxisConfig(type="physical", emb_dim=4, n_layers=1), out_dim=16, ) cond_cont = torch.zeros(B, COND_DIM) @@ -223,13 +226,11 @@ def test_condition_encoder_onehot_is_a_true_one_hot_vector(): def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot: - particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1} - if target != "physical": - particle_cfg = dict(particle_cfg) - if target == "embedding": - particle_cfg["type"] = "embedding" + particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=emb_dim, n_layers=1) + if target == "embedding": + particle_cfg = gconfig.ConditioningAxisConfig(type="embedding", emb_dim=emb_dim, n_layers=1) k_max = 5 - sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim) + sec_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target=target), generator, k_max, emb_dim) return Stage2OneShot( pdg_vocab=5, mat_vocab=3, @@ -242,7 +243,7 @@ def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneSho sec_dim=sec_dim, generator=generator, k_max=k_max, - particle_type_cfg={"target": target, "lambda": 1.0}, + particle_type_cfg=gconfig.ParticleTypeConfig(target=target), ) @@ -293,8 +294,8 @@ def test_stage2_oneshot_predict_type_raises_when_no_type_head(): def test_stage2_oneshot_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth(): """gitea #36: n_sec_head_cfg/type_head_cfg are independently tunable.""" k_max, emb_dim = 5, 6 - particle_cfg = {"type": "onehot", "emb_dim": emb_dim, "n_layers": 1} - sec_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim) + particle_cfg = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=emb_dim, n_layers=1) + sec_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target="onehot"), "flow", k_max, emb_dim) model = Stage2OneShot( pdg_vocab=5, mat_vocab=3, @@ -307,7 +308,7 @@ def test_stage2_oneshot_n_sec_head_and_type_head_cfg_control_hidden_width_and_de sec_dim=sec_dim, generator="flow", k_max=k_max, - particle_type_cfg={"target": "onehot", "lambda": 1.0}, + particle_type_cfg=gconfig.ParticleTypeConfig(target="onehot"), n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1}, type_head_cfg={"hidden_ratio": 0.75, "depth": 2}, ) @@ -350,8 +351,8 @@ def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim() conditioning.particle.emb_dim, sizes the onehot type_head/type_dim when explicitly set — the two used to be silently the same number.""" k_max = 5 - particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1} - particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20} + particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1) + particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=20) sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", k_max, 20) model = Stage2OneShot( pdg_vocab=5, @@ -367,7 +368,7 @@ def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim() k_max=k_max, particle_type_cfg=particle_type_cfg, ) - assert model.type_dim == 20 # not particle_cfg["emb_dim"] == 6 + assert model.type_dim == 20 # not particle_cfg.emb_dim == 6 assert model.type_head is not None assert model.type_head[-1].out_features == k_max * 20 @@ -516,10 +517,9 @@ def _build_stage2_ar( k_max: int = 5, history: str = "markov", ) -> Stage2Autoregressive: - particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1} + particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=emb_dim, n_layers=1) if target == "embedding": - particle_cfg = dict(particle_cfg) - particle_cfg["type"] = "embedding" + particle_cfg = gconfig.ConditioningAxisConfig(type="embedding", emb_dim=emb_dim, n_layers=1) return Stage2Autoregressive( pdg_vocab=5, mat_vocab=3, @@ -531,7 +531,7 @@ def _build_stage2_ar( context_dim=8, generator=generator, k_max=k_max, - particle_type_cfg={"target": target, "lambda": 1.0}, + particle_type_cfg=gconfig.ParticleTypeConfig(target=target), history=history, ) @@ -552,8 +552,8 @@ def test_stage2_autoregressive_history_invalid_raises(): def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_emb_dim(): """gitea #29, Stage2Autoregressive side — see the Stage2OneShot version of this test for the full rationale.""" - particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1} - particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20} + particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1) + particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=20) model = Stage2Autoregressive( pdg_vocab=5, mat_vocab=3, @@ -567,7 +567,7 @@ def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_em k_max=5, particle_type_cfg=particle_type_cfg, ) - assert model.type_dim == 20 # not particle_cfg["emb_dim"] == 6 + assert model.type_dim == 20 # not particle_cfg.emb_dim == 6 assert model.type_head is not None assert model.type_head[-1].out_features == 20 @@ -575,8 +575,8 @@ def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_em def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth(): """gitea #36, Stage2Autoregressive side — see the Stage2OneShot version of this test for the full rationale.""" - particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1} - particle_type_cfg = {"target": "onehot", "lambda": 1.0} + particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1) + particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot") model = Stage2Autoregressive( pdg_vocab=5, mat_vocab=3, @@ -612,9 +612,9 @@ def test_stage2_autoregressive_forward_shape(target, generator, history): cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.zeros(B, 2, dtype=torch.long) stage1_out = torch.randn(B, 9) - type_dim = stage2_type_dim({"target": target}, emb_dim) + type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target=target), emb_dim) history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim) - token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim) + token_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target=target), generator, 1, emb_dim) if generator == "wgan": x_t = torch.randn(B, K, model.noise_dim) t = None @@ -651,7 +651,7 @@ def test_stage2_autoregressive_predict_type_shape(): cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.zeros(B, 2, dtype=torch.long) stage1_out = torch.randn(B, 9) - type_dim = stage2_type_dim({"target": "onehot"}, emb_dim) + type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim) history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim) out = model.predict_type( cond_cont, @@ -672,7 +672,7 @@ def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, gen cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.zeros(B, 2, dtype=torch.long) stage1_out = torch.randn(B, 9) - type_dim = stage2_type_dim({"target": target}, emb_dim) + type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target=target), emb_dim) history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim) with pytest.raises(RuntimeError): model.predict_type( @@ -692,7 +692,7 @@ def test_stage2_autoregressive_gradients_flow_wgan_onehot(): cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.zeros(B, 2, dtype=torch.long) stage1_out = torch.randn(B, 9) - type_dim = stage2_type_dim({"target": "onehot"}, emb_dim) + type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim) history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim) z = torch.randn(B, K, model.noise_dim) gen_out = model( @@ -717,9 +717,9 @@ def test_stage2_autoregressive_gradients_flow_onehot(): cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.zeros(B, 2, dtype=torch.long) stage1_out = torch.randn(B, 9) - type_dim = stage2_type_dim({"target": "onehot"}, emb_dim) + type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim) history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim) - token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim) + token_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target="onehot"), "flow", 1, emb_dim) x_t = torch.randn(B, K, token_dim) t = torch.rand(B, K) flow_out = model( @@ -759,7 +759,7 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder(): B, K, emb_dim = 3, 6, 6 model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention") model.eval() - type_dim = stage2_type_dim({"target": "physical"}, emb_dim) + type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="physical"), emb_dim) hist_in_dim = CONT_SLOT_DIM + type_dim own_feat = torch.randn(B, K, hist_in_dim) # token i's own raw feature has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1) @@ -835,6 +835,54 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete assert shared_ids <= {id(p) for p in stage2.parameters()} +def test_condition_encoder_stores_the_exact_particle_and_material_cfg_instances_passed_in(): + """gitea #38: ConditionEncoder must not round-trip particle_cfg/ + material_cfg through a dict — the exact ConditioningAxisConfig instance + passed in is what `.particle_cfg`/`.material_cfg` hold afterward.""" + particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) + material_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) + enc = ConditionEncoder(pdg_vocab=3, mat_vocab=2, particle_cfg=particle_cfg, material_cfg=material_cfg) + assert enc.particle_cfg is particle_cfg + assert enc.material_cfg is material_cfg + + +def test_stagemodel_stores_the_exact_particle_type_cfg_instance_passed_in(): + """gitea #38: a StageModel subclass must not round-trip particle_type_cfg + through a dict — the exact ParticleTypeConfig instance passed in is what + `.particle_type_cfg` holds afterward.""" + particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=11) + sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", 5, 11) + model = Stage2OneShot( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=16, + n_res_blocks=1, + k_max=5, + sec_dim=sec_dim, + particle_type_cfg=particle_type_cfg, + ) + assert model.particle_type_cfg is particle_type_cfg + + +def test_build_models_particle_type_cfg_and_conditioning_axes_are_dataclasses_not_dicts(): + """gitea #38: build_models must pass the parsed ConditioningAxisConfig/ + ParticleTypeConfig dataclasses themselves down to the model constructors, + not re-serialize them to a dict first (the inversion the issue names) — + before the fix, .particle_type_cfg was a plain dict (s2_spec.particle_type + .to_dict()) and .cond_enc.particle_cfg came from the raw, unparsed + conditioning["particle"] dict.""" + cfg = _minimal_model_config(share_stages=False) + cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11} + built = build_models(cfg) + stage1, stage2 = built["stage1"], built["stage2"] + assert stage1 is not None and stage2 is not None + assert isinstance(stage2.particle_type_cfg, gconfig.ParticleTypeConfig) + assert isinstance(stage1.cond_enc.particle_cfg, gconfig.ConditioningAxisConfig) + assert isinstance(stage1.cond_enc.material_cfg, gconfig.ConditioningAxisConfig) + + def test_build_models_particle_type_n_classes_overrides_conditioning_emb_dim(): """gitea #29 end-to-end through build_models: setting stage2_model.particle_type.n_classes independently of @@ -895,7 +943,7 @@ def _partial_model_config() -> dict: def test_build_models_omitted_decoder_and_particle_type_match_default_config(): built = build_models(_partial_model_config()) assert isinstance(built["stage2"], Stage2Autoregressive) - assert built["stage2"].particle_type_cfg["target"] == "onehot" + assert built["stage2"].particle_type_cfg.target == "onehot" def test_build_models_custom_heads_block_controls_head_shapes(): diff --git a/tests/test_objectives.py b/tests/test_objectives.py index aaf043a..74104f6 100644 --- a/tests/test_objectives.py +++ b/tests/test_objectives.py @@ -6,6 +6,7 @@ stage2_inputs.py/trainers.py.""" import pytest import torch +from giant.config import ConditioningAxisConfig from giant.constants import COND_DIM, CONT_SLOT_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM from giant.model.network import ( OBJECTIVE_REGISTRY, @@ -19,7 +20,7 @@ from giant.model.network import ( ) from giant.model.schedule import CosineSchedule, flow_matching_loss, flow_matching_loss_secondary -_PHYS_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} +_PHYS_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/tests/test_phase2.py b/tests/test_phase2.py index 7a8efec..38270b1 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -4,6 +4,7 @@ import numpy as np import pytest import torch +from giant.config import ConditioningAxisConfig from giant.constants import ( COND_DIM, CONT_SLOT_DIM, @@ -23,9 +24,9 @@ from giant.sample import sample_secondaries # ── helpers ────────────────────────────────────────────────────────────────── -def _particle_material_cfg(conditioning: str) -> tuple[dict, dict]: - cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1} - return dict(cfg), dict(cfg) +def _particle_material_cfg(conditioning: str) -> tuple[ConditioningAxisConfig, ConditioningAxisConfig]: + cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1) + return cfg, cfg def _stage1(pdg=3, mat=2, conditioning="embedding"): diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 21fa7ec..36aa83b 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -8,6 +8,7 @@ import numpy as np import pytest import torch +from giant.config import ConditioningAxisConfig, ParticleTypeConfig from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX from giant.data.loader import TopNMap from giant.data.transforms import Normalizer @@ -27,8 +28,8 @@ MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1} def _models(conditioning="embedding"): - particle_cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1} - material_cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1} + particle_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1) + material_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1) s1 = Stage1Model( pdg_vocab=3, mat_vocab=2, @@ -366,8 +367,8 @@ def _models_v3( emb_dim=4, stage2_has_n_sec_head=True, ): - particle_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1} - material_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1} + particle_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1) + material_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1) # A fresh v0.3.0 Stage1Model — no n_sec_head_k_max, unlike _models() above # (n_sec ownership moves to stage 2 by default). s1 = Stage1Model( @@ -380,7 +381,7 @@ def _models_v3( generator=generator1, noise_dim=8, ) - particle_type_cfg = {"target": target} + particle_type_cfg = ParticleTypeConfig(target=target) # Explicit kwargs rather than a shared **common dict: a dict() call whose # values have heterogeneous types (str/int/dict/bool) widens under static # analysis to dict[str, ], which then makes every constructor @@ -557,8 +558,8 @@ COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_member def _onehot_conditioning_models(): - particle_cfg = {"type": "onehot", "emb_dim": len(PDG_MAP), "n_layers": 1} - material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1} + particle_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(PDG_MAP), n_layers=1) + material_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(MAT_MAP), n_layers=1) s1 = Stage1Model( pdg_vocab=3, mat_vocab=2, @@ -574,7 +575,7 @@ def _onehot_conditioning_models(): material_cfg=material_cfg, hidden_dim=32, n_res_blocks=2, - sec_dim=stage2_trunk_sec_dim({"target": "physical"}, "flow", K_MAX, 3), + sec_dim=stage2_trunk_sec_dim(ParticleTypeConfig(target="physical"), "flow", K_MAX, 3), generator="flow", time_dim=16, ) @@ -637,9 +638,9 @@ def _run_conditioning_and_type_onehot_different_n_classes(): conditioning.particle.emb_dim (gitea #29).""" cond_emb_dim = len(PDG_MAP) # 3 type_n_classes = 5 # deliberately different from cond_emb_dim - particle_cfg = {"type": "onehot", "emb_dim": cond_emb_dim, "n_layers": 1} - material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1} - particle_type_cfg = {"target": "onehot", "n_classes": type_n_classes} + particle_cfg = ConditioningAxisConfig(type="onehot", emb_dim=cond_emb_dim, n_layers=1) + material_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(MAT_MAP), n_layers=1) + particle_type_cfg = ParticleTypeConfig(target="onehot", n_classes=type_n_classes) s1 = Stage1Model( pdg_vocab=3, diff --git a/tests/test_router.py b/tests/test_router.py index af1b403..64c7325 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -3,6 +3,7 @@ import pytest import torch +from giant.config import ConditioningAxisConfig from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( BLOCK_REGISTRY, @@ -26,8 +27,8 @@ from giant.model.network import ( build_router, ) -PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} -MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} +PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) +MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) def _cond(B=8, pdg=3, mat=2): diff --git a/tests/test_sample.py b/tests/test_sample.py index 7906069..9ea8aa0 100644 --- a/tests/test_sample.py +++ b/tests/test_sample.py @@ -5,6 +5,7 @@ for the one-shot samplers.""" import pytest import torch +from giant.config import ConditioningAxisConfig, ParticleTypeConfig from giant.constants import COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, X_DIM from giant.model.network import ( Stage1Model, @@ -20,12 +21,12 @@ from giant.sample import ( sample_wgan, ) -_PHYS_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} +_PHYS_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) -def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[dict, dict]: - cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1} - return dict(cfg), dict(cfg) +def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[ConditioningAxisConfig, ConditioningAxisConfig]: + cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1) + return cfg, cfg def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]: @@ -43,7 +44,7 @@ def _conditioning_for(target: str) -> str: def _stage2_oneshot(target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2) -> Stage2OneShot: particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim) - particle_type_cfg = {"target": target} + particle_type_cfg = ParticleTypeConfig(target=target) # build_models (giant/model/network.py) computes sec_dim this same way # before constructing Stage2OneShot — its own default (SEC_DIM, the # "physical" width) is only correct for target="physical". @@ -84,7 +85,7 @@ def _stage2_ar( time_dim=16, noise_dim=8, k_max=k_max, - particle_type_cfg={"target": target}, + particle_type_cfg=ParticleTypeConfig(target=target), history=history, attn_n_heads=2, attn_n_layers=1, diff --git a/tests/test_train.py b/tests/test_train.py index 6815309..98f7f0b 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock import pytest import torch +from giant.config import ParticleTypeConfig from giant.constants import ( COND_DIM, CONT_SLOT_DIM, @@ -158,7 +159,7 @@ def test_type_repr_shapes_and_values(target): cond_enc = torch.nn.Module() if target == "embedding": cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim) - repr_ = _type_repr(sec_type_idx, sec_cont, {"target": target}, cond_enc, emb_dim) + repr_ = _type_repr(sec_type_idx, sec_cont, ParticleTypeConfig(target=target), cond_enc, emb_dim) expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim assert repr_.shape == (B, K, expected_width) if target == "physical": @@ -187,7 +188,7 @@ def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(target cond_enc = torch.nn.Module() if target == "embedding": cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim) - particle_type_cfg = {"target": target} + particle_type_cfg = ParticleTypeConfig(target=target) flat = _assemble_stage2_real(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim) unflat = _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim) assert torch.equal(unflat.flatten(1), flat) @@ -198,7 +199,7 @@ def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width(): sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM) sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX)) cond_enc = torch.nn.Module() - out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim) + out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, ParticleTypeConfig(target="physical"), cond_enc, emb_dim) assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM) assert out["has_prev"].shape == (B, K_MAX) assert out["remaining_frac"].shape == (B, K_MAX) diff --git a/tests/test_validate.py b/tests/test_validate.py index d19a80f..928684a 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -1,17 +1,18 @@ import numpy as np import torch +from giant.config import ConditioningAxisConfig, ParticleTypeConfig from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM from giant.data.dataset import StepBatch from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim from giant.validate import validate_marginals -_PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} -_MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} +_PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) +_MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) _K_MAX = 5 -def _tiny_models(particle_type_cfg: dict | None = None): +def _tiny_models(particle_type_cfg: ParticleTypeConfig | None = None): """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head, so n_sec always comes from Stage2OneShot.""" s1 = Stage1Model( @@ -22,12 +23,13 @@ def _tiny_models(particle_type_cfg: dict | None = None): hidden_dim=16, n_res_blocks=1, ) - target = (particle_type_cfg or {}).get("target", "physical") + resolved_type_cfg = particle_type_cfg or ParticleTypeConfig(target="physical") + target = resolved_type_cfg.target sec_dim = stage2_trunk_sec_dim( - particle_type_cfg or {"target": "physical"}, + resolved_type_cfg, "flow", _K_MAX, - int(_PARTICLE_CFG["emb_dim"]), + _PARTICLE_CFG.emb_dim, ) s2 = Stage2OneShot( pdg_vocab=3, @@ -42,7 +44,7 @@ def _tiny_models(particle_type_cfg: dict | None = None): sec_dim=sec_dim, particle_type_cfg=particle_type_cfg, ) - assert s2.particle_type_cfg.get("target", "physical") == target + assert s2.particle_type_cfg.target == target return s1.eval(), s2.eval() @@ -95,7 +97,7 @@ def test_validate_marginals_physical_target_shapes(): def test_validate_marginals_onehot_type_class_marginal(): - particle_type_cfg = {"target": "onehot"} + particle_type_cfg = ParticleTypeConfig(target="onehot") s1, s2 = _tiny_models(particle_type_cfg) loader = _loader(n_sec_value=2, n_classes=s2.type_dim) diff --git a/tests/test_wgan.py b/tests/test_wgan.py index 5b17aa8..dcb2cb3 100644 --- a/tests/test_wgan.py +++ b/tests/test_wgan.py @@ -1,12 +1,13 @@ import torch +from giant.config import ConditioningAxisConfig from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM from giant.model.network import CriticModel, Stage1Model, Stage2OneShot from giant.model.wgan import critic_loss, generator_loss, gradient_penalty from giant.sample import sample_secondaries_wgan, sample_wgan -PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} -MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} +PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) +MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) def _cond(B=8): From cc37a55183cf559177280539a28129a3b1036652 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 14 Aug 2026 15:05:32 +0200 Subject: [PATCH 3/3] Bump patch version to 0.3.1 Co-Authored-By: Claude Opus 5 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c9cd7ce..3789cdc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "giant" -version = "0.3.0" +version = "0.3.1" description = "Geant4 step-function surrogate via conditional flow matching" readme = "README.md" requires-python = ">=3.12" diff --git a/uv.lock b/uv.lock index 97f81e7..c4e2c8e 100644 --- a/uv.lock +++ b/uv.lock @@ -633,7 +633,7 @@ wheels = [ [[package]] name = "giant" -version = "0.3.0" +version = "0.3.1" source = { editable = "." } dependencies = [ { name = "numpy" },