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
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:
@@ -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],
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user