From 4ee75d00429ceb753644ba17648a7f6d71970850 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 8 Jul 2026 10:11:17 +0200 Subject: [PATCH] Clamp n_sec classification label to K_MAX Real data has steps with up to ~37 secondaries, but the n_sec head only has K_MAX+1=16 classes. The unclamped label occasionally overflowed cross_entropy's valid range and crashed CUDA training with "unique_by_key: failed to synchronize: cudaErrorAssert". The continuous secondary targets were already truncated to K_MAX slots; only this label was missed. Co-Authored-By: Claude Sonnet 5 --- giant/data/transforms.py | 10 ++++++++-- tests/test_transforms.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index bfcc3d4..f5df8db 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -485,7 +485,13 @@ def build_features( mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64) cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2) - n_sec = data["n_sec"].astype(np.int64) # (N,) + n_sec_raw = data["n_sec"].astype(np.int64) # (N,) unclamped, for the valid-slot mask + # Clamp the classification label to K_MAX: the head only has K_MAX+1 classes + # (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already + # applied to sec_cont/sec_pdg_idx by the loader's list padding. Without this, + # a rare high-multiplicity step (real data goes up to ~37) hands + # cross_entropy an out-of-range target and CUDA asserts. + n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64) # (N,) # Secondary continuous targets sec_E_list = data.get("sec_E_list") @@ -493,7 +499,7 @@ def build_features( sec_pdg_list = data.get("sec_pdg_list") if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None: - sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None] # (N, K_MAX) + sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None] # (N, K_MAX) sec_cont = encode_secondaries( sec_E_list, sec_dir_list, sec_valid, data["e_sec"], data["pre_dir"] ) # (N, K_MAX, 4) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 69c1d1f..d92f347 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1,6 +1,8 @@ import numpy as np import pytest +from giant.constants import K_MAX from giant.data.transforms import ( + build_features, energy_simplex_decode, energy_simplex_encode, inv_local_frame_rotation, @@ -205,3 +207,35 @@ def test_normalizer_serialization(): assert norm2.std is not None and norm.std is not None np.testing.assert_allclose(norm2.mean, norm.mean) np.testing.assert_allclose(norm2.std, norm.std) + + +def test_build_features_clamps_n_sec_label_to_k_max(): + """A step with more secondaries than K_MAX must not overflow the + n_sec classifier's K_MAX+1 classes (regression test: this used to hand + cross_entropy an out-of-range target and crash CUDA training with + 'unique_by_key: failed to synchronize: cudaErrorAssert').""" + N = 3 + raw_n_sec = np.array([0, 5, K_MAX + 20], dtype=np.int32) + rng = np.random.default_rng(0) + data = { + "pdg": np.array([11, 11, 11], dtype=np.int32), + "material": np.array(["PbWO4", "PbWO4", "PbWO4"], dtype=object), + "pre_pos": rng.standard_normal((N, 3)).astype(np.float32), + "pre_E": np.full(N, 10.0, dtype=np.float32), + "pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)), + "layer_id": np.zeros(N, dtype=np.int32), + "n_sec": raw_n_sec, + "e_sec": np.full(N, 1.0, dtype=np.float32), + "step_length": np.full(N, 1.0, dtype=np.float32), + "post_E": np.full(N, 9.0, dtype=np.float32), + "edep": np.full(N, 1.0, dtype=np.float32), + "post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)), + "post_pos": rng.standard_normal((N, 3)).astype(np.float32), + } + pdg_map = {11: 0} + mat_map = {"PbWO4": 0} + + _, _, _, n_sec, _, _, _, _ = build_features(data, pdg_map, mat_map) + + assert n_sec.max() <= K_MAX + np.testing.assert_array_equal(n_sec, [0, 5, K_MAX])