Add vanilla shampoo

[skyw]

No description provided.

[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.

[greptile-apps]

Greptile Summary

The PR adds a registered vanilla Shampoo optimizer and reorganizes the shared preconditioner protocols, while preserving the protocol-based SOAP variants.

• Adds Shampoo factor accumulation, inverse-root preconditioning, momentum, and optimizer state management.
• Renames the shared preconditioner module and updates SOAP imports and type annotations.
• Adds Shampoo tests and adjusts SOAP parity coverage.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains within the eligible follow-up-review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
emerging_optimizers/shampoo/shampoo.py	Adds the vanilla Shampoo optimizer, factor preconditioner, shared step lifecycle, and registry entry.
emerging_optimizers/shampoo/precond_base.py	Renames the shared preconditioner protocol module without changing its contents.
emerging_optimizers/shampoo/soap_v3.py	Updates SOAP implementations to import shared protocol types directly from the renamed module.
tests/test_shampoo.py	Adds coverage for state initialization, factor updates, inverse roots, preconditioning, momentum, and optimizer steps.
tests/test_soap_v3.py	Aligns exception expectations and shortens the legacy parity test to three steps.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Train["Training loop"] --> Step["Shampoo.step()"]
  Step --> State["Per-parameter state"]
  State --> Factors["Update L and R factors"]
  Factors --> Roots["Compute inverse roots"]
  Step --> Momentum["Update momentum"]
  Momentum --> Precondition["Apply left and right inverse roots"]
  Roots --> Precondition
  Precondition --> Parameter["Update parameter"]
  Parameter --> State

Loading

_{Reviews (6): Last reviewed commit: "add test for tikhonov" | Re-trigger Greptile}

[skyw]

/ok to test d9d0a23

[github-actions]

Test Results

   93 files  +  2    175 suites  +6   1m 56s ⏱️ +3s
1 309 tests + 52  1 306 ✅ +49  0 💤 ±0  3 ❌ +3 
3 040 runs  +104  3 034 ✅ +98  0 💤 ±0  6 ❌ +6 

For more details on these failures, see this check.

Results for commit 7a7a6db. ± Comparison against base commit 65e08e6.

This pull request removes 3 and adds 55 tests. Note that renamed tests count towards both.

__main__.SoapV3AgainstLegacyTest ‑ test_5steps_close_to_legacy0 (m=4, n=4, atol=1e-05, rtol=1e-05)
__main__.SoapV3AgainstLegacyTest ‑ test_5steps_close_to_legacy1 (m=8, n=4, atol=0.0001, rtol=0.0001)
__main__.SoapV3AgainstLegacyTest ‑ test_5steps_close_to_legacy2 (m=33, n=17, atol=0.002, rtol=0.002)

__main__.ShampooBaseTest ‑ test_negative_lr_raises
__main__.ShampooBaseTest ‑ test_rejects_non_2d
__main__.ShampooBaseTest ‑ test_step_3steps_close_to_sgd0 (lr=0.25, momentum=0.0, weight_decay=0.125)
__main__.ShampooBaseTest ‑ test_step_3steps_close_to_sgd1 (lr=0.25, momentum=0.0, weight_decay=0.05)
__main__.ShampooBaseTest ‑ test_step_3steps_close_to_sgd2 (lr=0.25, momentum=75, weight_decay=0.125)
__main__.ShampooBaseTest ‑ test_step_3steps_close_to_sgd3 (lr=0.25, momentum=75, weight_decay=0.05)
__main__.ShampooBaseTest ‑ test_step_3steps_close_to_sgd4 (lr=0.125, momentum=0.0, weight_decay=0.125)
__main__.ShampooBaseTest ‑ test_step_3steps_close_to_sgd5 (lr=0.125, momentum=0.0, weight_decay=0.05)
__main__.ShampooBaseTest ‑ test_step_3steps_close_to_sgd6 (lr=0.125, momentum=75, weight_decay=0.125)
__main__.ShampooBaseTest ‑ test_step_3steps_close_to_sgd7 (lr=0.125, momentum=75, weight_decay=0.05)
…

♻️ This comment has been updated with latest results.

[skyw]

/ok to test e1f2ce9

[codecov]

Codecov Report

❌ Patch coverage is 98.18182% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
emerging_optimizers/shampoo/shampoo.py	97.89%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[skyw]

/ok to test 97e72d8

[skyw]

/ok to test deed3dc

[skyw]

/ok to test 7a7a6db