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

Merge pull request 'Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base (gitea #39)' (#56) from fix/issue-39 into master

Browse Source 
Reviewed-on: #56
This commit is contained in:
Lars Bogner
2026-08-14 15:16:02 +02:00
parent c71210f006 cc37a55183
commit 5982ddcd49
22 changed files with 620 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/model/builders.py
+11 -13
View File
@@ -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,
giant/model/encoders.py
+19 -18
View File
@@ -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:
giant/model/models.py
+256 -169
View File
@@ -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,20 +58,23 @@ 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)
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`
@@ -83,8 +85,130 @@ class Stage1Model(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: 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
# `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.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,
mat_vocab: int,
particle_cfg: ConditioningAxisConfig,
material_cfg: ConditioningAxisConfig,
hidden_dim: int = 256,
n_res_blocks: int = 6,
cond_out_dim: int = 128,
@@ -100,27 +224,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 +261,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).
@@ -156,7 +291,7 @@ class Stage2OneShot(nn.Module):
(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
@@ -169,17 +304,14 @@ 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__(
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,
@@ -195,59 +327,45 @@ class Stage2OneShot(nn.Module):
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,
) -> 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,
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,
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,
)
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
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 +390,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 +405,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,17 +432,14 @@ 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__(
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,
@@ -349,7 +454,7 @@ class Stage2Autoregressive(nn.Module):
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,
@@ -357,20 +462,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,36 +495,29 @@ 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.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.target
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,
)
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:
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)
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)
ctx = self.context_adapter(stage1_out)
@@ -511,12 +608,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 +623,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,
@@ -562,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,
giant/model/network.py
+2
View File
@@ -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",
giant/pipeline.py
+3 -3
View File
@@ -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
giant/rollout.py
+2 -2
View File
@@ -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']"
giant/sample.py
+2 -2
View File
@@ -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
giant/training/stage2_inputs.py
+10 -9
View File
@@ -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,
giant/training/trainers.py
+7 -7
View File
@@ -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
giant/validate.py
+2 -2
View File
@@ -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 = [], []
pyproject.toml
+1 -1
View File
@@ -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"
tests/test_flow.py
+3 -2
View File
@@ -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():
tests/test_network.py
+256 -46
View File
@@ -14,17 +14,19 @@ from giant.model.network import (
Stage1Model,
Stage2Autoregressive,
Stage2OneShot,
StageModel,
build_critics,
build_history,
build_models,
build_objective,
stage2_trunk_sec_dim,
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():
@@ -133,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 -------------------------------------------
@@ -162,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,
@@ -194,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)
@@ -221,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,
@@ -240,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),
)
@@ -291,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,
@@ -305,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},
)
@@ -348,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,
@@ -365,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
@@ -514,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,
@@ -529,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,
)
@@ -550,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,
@@ -565,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
@@ -573,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,
@@ -610,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
@@ -649,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,
@@ -670,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(
@@ -690,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(
@@ -715,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(
@@ -757,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)
@@ -833,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
@@ -893,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():
@@ -987,3 +1037,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
tests/test_objectives.py
+2 -1
View File
@@ -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]:
tests/test_phase2.py
+4 -3
View File
@@ -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"):
tests/test_rollout.py
+12 -11
View File
@@ -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, <big union>], 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,
tests/test_router.py
+3 -2
View File
@@ -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):
tests/test_sample.py
+7 -6
View File
@@ -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,
tests/test_train.py
+4 -3
View File
@@ -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)
tests/test_validate.py
+10 -8
View File
@@ -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)
tests/test_wgan.py
+3 -2
View File
@@ -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):
uv.lock
Generated
+1 -1
View File
@@ -633,7 +633,7 @@ wheels = [
[[package]]
name = "giant"
version = "0.3.0"
version = "0.3.1"
source = { editable = "." }
dependencies = [
{ name = "numpy" },