transforms: pad legacy cond normalizers for pre-physical-conditioning checkpoints
CI / Lint (ruff check) (push) Successful in 58s
CI / Format (ruff format) (push) Failing after 1m6s
CI / Type check (ty) (push) Successful in 1m6s
CI / Tests (push) Successful in 1m45s
CI / Lint (ruff check) (pull_request) Successful in 1m1s
CI / Format (ruff format) (pull_request) Failing after 1m1s
CI / Type check (ty) (pull_request) Successful in 1m6s
CI / Tests (pull_request) Successful in 1m50s
CI / Bump version, build & publish wheel (push) Has been skipped
CI / Bump version, build & publish wheel (pull_request) Has been skipped

Checkpoints trained before commit 68fb99b (physical-property
conditioning, COND_DIM 8->15) saved a COND_DIM_BASE-wide cond
normalizer, fit before build_cond_features grew the extra physical
columns. Any inference against such a checkpoint under current code
(predict/rollout/router_gating) crashed broadcasting a 15-wide
cond_cont against an 8-wide mean/std.

In "embedding" mode those physical columns are never read by
ConditionEncoder, so padding the missing entries with mean=0/std=1 is
a safe no-op. "physical" mode reads them directly, so a mismatch there
still raises instead of silently normalizing garbage.
This commit is contained in:
2026-07-27 11:06:48 +02:00
parent 4d6101dcd7
commit eb751d968d
2 changed files with 77 additions and 1 deletions
+34 -1
View File
@@ -530,11 +530,44 @@ def build_cond_features(
cond_cat = np.column_stack([pdg_idx, mat_idx])
if cond_normalizer is not None:
cond_cont = cond_normalizer.transform(cond_cont)
cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning)
return cond_cont, cond_cat
def _cond_normalizer_transform(
cond_cont: np.ndarray, cond_normalizer: "Normalizer", conditioning: str
) -> np.ndarray:
"""Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed.
Checkpoints trained before physical-property conditioning (``COND_DIM``
8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond
normalizer, fit before ``build_cond_features`` grew the extra physical
columns. In "embedding" mode those columns are never read by
``ConditionEncoder`` (``giant/model/network.py``), so padding the missing
entries with mean=0/std=1 is a safe no-op that keeps such checkpoints
usable under the current, always-``COND_DIM``-wide contract. In
"physical" mode the physical columns are load-bearing, so a mismatch
there is a real incompatibility, not something to paper over.
"""
mean, std = cond_normalizer.mean, cond_normalizer.std
assert mean is not None and std is not None, "Normalizer not fitted"
width = cond_cont.shape[-1]
if mean.shape[-1] < width:
if conditioning != "embedding":
raise ValueError(
f"cond normalizer has {mean.shape[-1]} columns, expected "
f"{width}, and conditioning={conditioning!r} reads the "
"physical columns directly — this checkpoint predates "
"physical-property conditioning and can't be safely padded; "
"retrain it under the current code."
)
pad = width - mean.shape[-1]
mean = np.concatenate([mean, np.zeros(pad, dtype=mean.dtype)])
std = np.concatenate([std, np.ones(pad, dtype=std.dtype)])
return ((cond_cont - mean) / std).astype(np.float32)
def build_features(
data: dict[str, np.ndarray],
pdg_map: dict[int, int],
+43
View File
@@ -397,3 +397,46 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
)
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0])
def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode():
"""A pre-physical-conditioning checkpoint's cond normalizer is COND_DIM_BASE
(8) wide, fit before build_cond_features grew the extra physical columns.
In "embedding" mode those columns are never read downstream, so a legacy
normalizer should be usable as-is (padded, not rejected)."""
data = _minimal_step_data(3)
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
legacy_norm = Normalizer()
legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32)
legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
cond_cont, _ = build_cond_features(
data, pdg_map, mat_map, cond_normalizer=legacy_norm, conditioning="embedding"
)
assert cond_cont.shape[-1] == COND_DIM
# padded physical columns are zero-filled pre-normalization and
# mean=0/std=1 post-normalization, so they should come out as zero
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)
def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
fake_material_props,
):
"""Unlike "embedding" mode, "physical" mode actually reads the physical
columns, so a legacy 8-wide normalizer can't be silently padded — that
would silently feed the network un-normalized physical properties."""
data = _minimal_step_data(3)
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
legacy_norm = Normalizer()
legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32)
legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
with pytest.raises(ValueError, match="predates physical-property conditioning"):
build_cond_features(
data,
pdg_map,
mat_map,
cond_normalizer=legacy_norm,
conditioning="physical",
)