update ai agent rules and memory (#242)

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

CLAUDE.md

@@ -60,7 +60,7 @@ torchrun --nproc_per_node=<n> --no-python python tests/test_distributed_muon_uti
 
 CI test suites are encoded as shell scripts in `tests/ci/`:
 
-- `L0_Tests_CPU.sh` — distributed CPU + a couple of CPU-only tests.
+- `L0_Tests_CPU.sh` — all `tests/test_distributed_*_cpu.py` at 8 and 4 ranks (rekls also at 1 rank) + a couple of CPU-only tests.
 - `L0_Tests_GPU.sh` — every `tests/test_*.py` (excluding `*_cpu.py`) on `--device=cuda`, run twice (random seed, then `--seed=42`), plus the `tests/convergence/*_test.py` set.
 - `L1_Tests_GPU.sh` — convergence tests at `--seed=77` plus the `moe_c4_convergence.py` runs against muon/soap with a loss target.
 
@@ -87,11 +87,12 @@ Run a single test class/method using absl conventions, e.g. `python tests/test_o
 
 - `registry.py` — global `_OPTIMIZERS` dict with a `@register_optimizer("name")` decorator; `get_optimizer_cls` / `get_configured_optimizer_cls` provide lookup. **Important: a class is only in the registry once its module has been imported.** External callers must import the optimizer module (e.g. `from emerging_optimizers.orthogonalized_optimizers import muon`) before `get_optimizer_cls("muon")` will find it.
 - `mixin.py` — `WeightDecayMixin` with four `weight_decay_method` modes: `decoupled`, `independent`, `l2`, `palm`. Optimizers that include this mixin set `self.weight_decay_method` and call `_apply_weight_decay_inplace(p, grad, lr, wd)` inside `step`.
-- `orthogonalized_optimizers/` — Muon and friends (`muon`, `adaptive_muon`, `muon_hyperball`, `mop`, `polargrad`, `scion`, `sinkhorn_muon`). All extend `OrthogonalizedOptimizer` (`orthogonalized_optimizer.py`), which implements the EMA-momentum + (optional Nesterov) + `scaled_orthogonalize_fn` pipeline. Subclasses customize behavior by overriding `orthogonalize(p, grad, **group_kwargs)` (e.g. for split-QKV preconditioning) or the `pre_weight_update_fn_inplace` / `post_weight_update_fn_inplace` hooks (used by hyperball-style updates that preserve weight norms). `muon_utils` holds the Newton–Schulz iteration coefficients and helpers; `spectral_clipping_utils` implements spectral norm clipping.
-- `soap/` — SOAP optimizer (`soap.py`) and REKLS variant (`rekls.py`), with shared `soap_utils.py`. Uses Kronecker-factored preconditioning and tracks eigenbases per-parameter.
-- `psgd/` — PSGD-Kron (`psgd.py`) plus `psgd_kron_contractions.py` (Einsum contractions for triangular factors), `procrustes_step.py`, and `psgd_utils.py`.
-- `scalar_optimizers/` — Lion (`lion.py`); `update_functions/` holds reusable update primitives.
-- `riemannian_optimizers/` — `normalized_optimizer.py` (manifold-normalized updates).
+- `orthogonalized_optimizers/` — Muon and friends (`muon`, `adaptive_muon`, `muon_hyperball`, `mop`, `polargrad`, `scion`, `sinkhorn_muon`, `spel`). All extend `OrthogonalizedOptimizer` (`orthogonalized_optimizer.py`), which implements the EMA-momentum + (optional Nesterov) + `scaled_orthogonalize_fn` pipeline. Subclasses customize behavior by overriding `orthogonalize(p, grad, **group_kwargs)` (e.g. for split-QKV preconditioning) or the `pre_weight_update_fn_inplace` / `post_weight_update_fn_inplace` hooks (used by hyperball-style updates that preserve weight norms). `muon_utils` holds the Newton–Schulz iteration coefficients and helpers, including the tensor-parallel `newton_schulz_tp`; `spectral_clipping_utils` implements spectral norm clipping.
+- `soap/` — the SOAP family, sharing `soap_utils.py`: `soap.py` (`SOAP`, plus `StackedSoap` which optimizes batched/3D params through transient 2D stacking), `rekls.py` (REKLS), `tp_rekls.py` (experimental tensor-parallel REKLS), and `moso.py` (MOSO, momentum one-sided SOAP). Uses Kronecker-factored preconditioning and tracks eigenbases per-parameter.
+- `psgd/` — `PSGDPro` (`psgd.py`, registered as `psgd_pro`) plus `psgd_kron_contractions.py` (Einsum contractions for triangular factors), `procrustes_step.py`, and `psgd_utils.py`.
+- `scalar_optimizers/` — element-wise optimizers (`lion`, `signum`, `laprop`, `sim_ademamix`), defined in `__init__.py` on top of the `base.py` hierarchy (`SingleMomentumOptimizer` / `TwoMomentsOptimizer`); `update_functions/` holds the reusable per-algorithm update primitives (adam, lion, laprop, ademamix, madam, signum).
+- `riemannian_optimizers/` — `normalized_optimizer.py` (manifold-normalized updates; `oblique_sgd`, `oblique_adam`).
+- `contrib/` — contributed/experimental optimizers kept out of the registry (`muown.py`: `Muown`, a Muon variant).
 - `triton_kernels/syrk.py` — Triton SYRK kernel used to accelerate Gram-matrix computations in preconditioner updates. Triton is an optional/test-only dep; gate imports.
 - `utils/` — `eig.py` (eigendecomposition helpers), `sinkhorn_mapper.py` (`SinkhornMapper`), `modules.py`, and the `fp32_matmul_precision` context manager (sets `torch.set_float32_matmul_precision` for the duration of orthogonalization). `FP32MatmulPrecT = Literal["highest", "high", "medium"]` is the canonical type alias.
 
@@ -124,4 +125,9 @@ Don't override `step` — override `orthogonalize` (and the pre/post hooks if yo
 
 ## CI specifics
 
-CI runs in containers built via `docker/Dockerfile.ci`. Key env flags set by the test scripts: `TORCH_COMPILE_DISABLE=1` (CPU script) and `TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0` (GPU scripts) — preserve these when reproducing CI failures locally.
+CI runs in containers built via `docker/Dockerfile.ci`. Key env flags set by the test scripts: `TORCH_COMPILE_DISABLE=1` (CPU and L0 GPU scripts) and `TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0` (GPU scripts) — preserve these when reproducing CI failures locally.
+
+## Rules for agent
+
+* Never write or modify comment including docstring unless explicitly asked.
+* Never run tests unless explicitly asked.