lars
4692cee699
CI / Lint (ruff check) (push) Successful in 30s
CI / Format (ruff format) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 43s
CI / Lint (ruff check) (pull_request) Successful in 40s
CI / Format (ruff format) (pull_request) Successful in 43s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 45s
CI / Tests (pull_request) Successful in 3m24s
CI / Tests (push) Successful in 3m33s
The conditioning arrays' column order was written down three times — twice
in giant/data/transforms.py (build_cond_features and build_features each
built cond_cont and cond_cat from scratch) and again in
giant/model/encoders.py (cat_col_layout, plus hand-written
COND_DIM_BASE + PARTICLE_PHYS_DIM slicing in ConditionEncoder). The three
were held in sync only by parallel comments, so a wrong column order
produced silently mis-indexed features rather than an exception.
The drift had already happened, twice, both times in build_features:
- 5b63dfd added per-axis vocab-lookup strictness (an out-of-vocab
pdg/material must not KeyError under "physical"/"onehot", where the
index is never read) to build_cond_features only.
- _cond_normalizer_transform's legacy-normalizer padding, which keeps a
pre-physical-conditioning 8-wide cond normalizer loadable, was likewise
only wired into build_cond_features — so `giant predict` on such a
checkpoint died with a broadcast error.
New giant/cond_layout.py holds a frozen CondLayout built from the
(particle, material) mode pair, exposing named cond_cont slices
(base/particle_phys/material_phys) and cond_cat columns
(PDG_COL/MAT_COL/particle_topn_col/material_topn_col/cat_dim). Both
builders now share one _build_cond_arrays, ConditionEncoder reads its
slices off the same object, and PdgRouter/ProcessRouter use the named
dense-vocab columns instead of literal 0/1. CondLayout also absorbs the
two duplicated axis-type validations, keeping their message text verbatim.
Decisions taken while planning:
- Scope is CondLayout only. The issue's second half — a
CONDITIONING_AXIS_REGISTRY registering (feature_columns, encoder_module)
as a pair — is deferred: it would force ConditioningConfig's fixed
particle/material fields into a dynamic axis map and ripple through
pipeline.py, checkpoint_io.py and rollout.py, i.e. a config-schema break
with no consumer yet.
- The two divergences above are unified onto build_cond_features'
behaviour rather than preserved as parameters, so the new single source
of truth doesn't carry the old split forward. Each gets a regression
test that fails before this commit.
- cat_col_layout is replaced outright (deleted, dropped from network.py's
__all__, its four tests rewritten against CondLayout) rather than kept
as a wrapper — two spellings of the same fact is the defect itself.
cond_cat's width is now the layout's call rather than "did the caller pass
a map", so an "onehot" axis without its top-N map raises instead of
yielding a narrower array that ConditionEncoder would index out of bounds.
pipeline.py's normalizer-fitting pass reads only cond_cont but had to be
handed the maps to satisfy that.
No parameter, buffer or state_dict change; existing checkpoints load
unchanged, and the protected migration surfaces are untouched.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
88 lines
3.9 KiB
Python
88 lines
3.9 KiB
Python
import pytest
|
|
from giant.cond_layout import AXIS_TYPES, CondLayout
|
|
from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
|
|
|
|
# ── cond_cat column layout ───────────────────────────────────────────────────
|
|
|
|
|
|
def test_topn_cols_neither_onehot():
|
|
layout = CondLayout.from_types("physical", "embedding")
|
|
assert (layout.particle_topn_col, layout.material_topn_col) == (None, None)
|
|
assert layout.cat_dim == 2
|
|
|
|
|
|
def test_topn_cols_particle_only():
|
|
layout = CondLayout.from_types("onehot", "physical")
|
|
assert (layout.particle_topn_col, layout.material_topn_col) == (2, None)
|
|
assert layout.cat_dim == 3
|
|
|
|
|
|
def test_topn_cols_material_only():
|
|
layout = CondLayout.from_types("physical", "onehot")
|
|
assert (layout.particle_topn_col, layout.material_topn_col) == (None, 2)
|
|
assert layout.cat_dim == 3
|
|
|
|
|
|
def test_topn_cols_both_onehot_particle_then_material():
|
|
layout = CondLayout.from_types("onehot", "onehot")
|
|
assert (layout.particle_topn_col, layout.material_topn_col) == (2, 3)
|
|
assert layout.cat_dim == 4
|
|
|
|
|
|
def test_dense_vocab_cols_are_mode_independent():
|
|
"""Columns 0/1 are always the dense pdg/material index — giant.model.routers
|
|
reads them without knowing the conditioning mode."""
|
|
assert (CondLayout.PDG_COL, CondLayout.MAT_COL) == (0, 1)
|
|
for particle in AXIS_TYPES:
|
|
for material in AXIS_TYPES:
|
|
layout = CondLayout.from_types(particle, material)
|
|
assert layout.particle_topn_col not in (layout.PDG_COL, layout.MAT_COL)
|
|
assert layout.material_topn_col not in (layout.PDG_COL, layout.MAT_COL)
|
|
|
|
|
|
# ── cond_cont slice layout ───────────────────────────────────────────────────
|
|
|
|
|
|
def test_cont_slices_tile_cond_cont_exactly():
|
|
"""base / particle_phys / material_phys must partition cond_cont with no
|
|
gap and no overlap — a gap or overlap is exactly the silent
|
|
mis-indexing this object exists to prevent."""
|
|
layout = CondLayout.from_types("physical", "physical")
|
|
covered = list(range(*layout.base.indices(COND_DIM)))
|
|
covered += list(range(*layout.particle_phys.indices(COND_DIM)))
|
|
covered += list(range(*layout.material_phys.indices(COND_DIM)))
|
|
assert covered == list(range(COND_DIM))
|
|
|
|
|
|
def test_cont_slice_widths_match_constants():
|
|
layout = CondLayout.from_types("embedding", "embedding")
|
|
assert layout.base == slice(0, COND_DIM_BASE)
|
|
assert layout.particle_phys.stop - layout.particle_phys.start == PARTICLE_PHYS_DIM
|
|
assert layout.material_phys.stop - layout.material_phys.start == MATERIAL_PHYS_DIM
|
|
assert layout.cont_dim == COND_DIM
|
|
|
|
|
|
def test_cont_slices_are_mode_independent():
|
|
"""cond_cont is COND_DIM wide in every mode — a non-"physical" axis gets
|
|
its block zero-filled rather than dropped, so the slices never move."""
|
|
physical = CondLayout.from_types("physical", "physical")
|
|
for particle in AXIS_TYPES:
|
|
for material in AXIS_TYPES:
|
|
layout = CondLayout.from_types(particle, material)
|
|
assert layout.base == physical.base
|
|
assert layout.particle_phys == physical.particle_phys
|
|
assert layout.material_phys == physical.material_phys
|
|
|
|
|
|
# ── validation ───────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_unknown_particle_type_raises():
|
|
with pytest.raises(ValueError, match="unknown conditioning.particle.type 'bogus'"):
|
|
CondLayout.from_types("bogus", "physical")
|
|
|
|
|
|
def test_unknown_material_type_raises():
|
|
with pytest.raises(ValueError, match="unknown conditioning.material.type 'bogus'"):
|
|
CondLayout.from_types("physical", "bogus")
|