Introduce new shampoo and soap

[skyw]

Improve abstraction. clear path for more variant of Shampoo and Soap.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[skyw]

@greptileai review the draft

[skyw]

/ok to test a9b2fcf

[greptile-apps]

Greptile Summary

The PR moves the experimental SOAP implementation into a dedicated Shampoo package and introduces reusable preconditioner protocols and a ReKLS variant.

• Splits common Shampoo/SOAP interfaces and tensor-pair state into shampoo_base.py.
• Refactors SOAP initialization and updates through interchangeable preconditioner classes.
• Adds per-parameter-group beta handling and parity tests against legacy SOAP and ReKLS.
• Adds explicit precision contexts around KL-Shampoo factor updates.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
emerging_optimizers/shampoo/shampoo_base.py	Introduces shared tensor-pair and preconditioner protocol abstractions without an identified blocking defect.
emerging_optimizers/shampoo/soap_v3.py	Refactors SOAP preconditioning, adds ReKLS, and fully guards the previously reported KL factor-update precision paths.
tests/test_soap_v3.py	Adds state-layout, legacy-parity, ReKLS-parity, and smoke-test coverage for the refactored optimizers.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Loop["SoapBase.step"] --> Protocol["SoapPreconditionerProtocol"]
  Protocol --> KL["KlSoapPreconditioner"]
  Protocol --> ReKLS["ReklsPreconditioner"]
  KL --> QR["Orthogonal iteration"]
  ReKLS --> Eigh["Full eigendecomposition"]
  KL --> Scalar["Adam / MAdam scalar update"]
  ReKLS --> Scalar
  Scalar --> Param["Parameter update"]

Loading

_{Reviews (5): Last reviewed commit: "add more tests" | Re-trigger Greptile}

[skyw]

/ok to test a9b2fcf

[github-actions]

Test Results

   91 files  + 2    169 suites  +8   1m 57s ⏱️ +5s
1 254 tests +24  1 251 ✅ +21  0 💤 ±0  3 ❌ +3 
2 930 runs  +48  2 926 ✅ +44  0 💤 ±0  4 ❌ +4 

For more details on these failures, see this check.

Results for commit 4f516ff. ± Comparison against base commit 170d1b8.

♻️ This comment has been updated with latest results.

[skyw]

/ok to test 715ce11

[skyw]

/ok to test 4f516ff