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 <noreply@anthropic.com>

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)

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])