lars b51eafcfa5 Add opt-in straight-through Gumbel-softmax combine weights to MoE router ...

Trains the routed trunk's forward combination as a hard one-hot sample
(matching eval-time top-1 dispatch exactly) while keeping a smooth gradient
on the backward pass, targeting the train/eval mismatch identified as a
likely contributor to experts overlapping instead of partitioning in the
first energy-router rollout benchmark. Off by default (model.router.gumbel);
existing routed configs/checkpoints are unaffected.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

2026-07-30 16:29:20 +02:00

..

test_analysis_reduce.py

Rewrite analysis as streaming rollout-vs-reference plotting pipeline

2026-07-23 17:38:09 +02:00

test_bump_dataset_version.py

Fix ruff, ty, and pytest failures; apply ruff format

2026-07-02 16:59:22 +02:00

test_catalog.py

analyze: chunk per-plot aggregation across HTCondor jobs

2026-07-27 09:24:19 +02:00

test_cli_new_run.py

Add giant new-run to scaffold a config.toml + run dir ahead of training

2026-07-30 12:56:22 +02:00

test_cli_predict.py

Add --comment option to predict, recorded in YAML sidecar

2026-07-06 09:37:48 +02:00

test_condor.py

analyze: default run directory to <repo>/analysis_runs, gitignored

2026-07-27 11:29:35 +02:00

test_config.py

Add opt-in straight-through Gumbel-softmax combine weights to MoE router

2026-07-30 16:29:20 +02:00

test_create_root_files.py

Support --energy-gev in dwarf make-root for the new minicalosim energy arg

2026-07-13 13:59:09 +02:00

test_dataset.py

Offset event_id per file to avoid cross-file collisions

2026-07-30 11:36:48 +02:00

test_dwarf.py

Add dwarf warm-cache to precompute the setup-stage sidecar

2026-07-30 10:53:31 +02:00

test_flow.py

Implement Phase 2: secondary particle prediction

2026-06-29 11:34:31 +02:00

test_geometry.py

Add fast slab lookup for the GeometryOracle, replacing knn as the default

2026-07-08 11:42:12 +02:00

test_loader.py

Offset event_id per file to avoid cross-file collisions

2026-07-30 11:36:48 +02:00

test_materials.py

Condition on material/particle physical properties instead of learned embeddings

2026-07-17 15:12:54 +02:00

test_network.py

Implement Phase 2: secondary particle prediction

2026-06-29 11:34:31 +02:00

test_particles.py

Fix crashes in physical-property conditioning edge cases

2026-07-20 10:36:25 +02:00

test_phase2.py

Speed up giant train's setup stage

2026-07-29 13:34:32 +02:00

test_pipeline.py

Cache giant train's setup stage in a sidecar file

2026-07-30 10:28:37 +02:00

test_render.py

Rewrite analysis as streaming rollout-vs-reference plotting pipeline

2026-07-23 17:38:09 +02:00

test_rollout.py

Fix crashes in physical-property conditioning edge cases

2026-07-20 10:36:25 +02:00

test_router_gating.py

analyze: add MoE router gating/share diagnostic plots

2026-07-24 13:32:25 +02:00

test_router.py

Add opt-in straight-through Gumbel-softmax combine weights to MoE router

2026-07-30 16:29:20 +02:00

test_setup_cache.py

Store a quantile grid instead of a raw reservoir sample in the setup cache

2026-07-30 13:32:46 +02:00

test_steps_to_parquet_parallel.py

Fix ruff, ty, and pytest failures; apply ruff format

2026-07-02 16:59:22 +02:00

test_steps_to_parquet.py

Apply ruff format after merging phase2-secondary-prediction

2026-07-15 10:00:36 +02:00

test_train.py

Add opt-in straight-through Gumbel-softmax combine weights to MoE router

2026-07-30 16:29:20 +02:00

test_transforms.py

Speed up _WelfordAccumulator's per-chunk update

2026-07-29 13:55:07 +02:00

test_validate.py

Fix crashes in physical-property conditioning edge cases

2026-07-20 10:36:25 +02:00

test_wgan.py

Add WGAN-GP mode as a throwaway fast-eval experiment

2026-07-22 11:20:28 +02:00