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

Adds a registered vanilla Shampoo optimizer and reorganizes the shared preconditioner protocols while preserving the existing SOAP variants.

• Introduces Shampoo factor accumulation, inverse-root preconditioning, momentum, and optimizer-state handling.
• Renames the shared protocol module from shampoo_base.py to precond_base.py and updates SOAP imports.
• Re-exports Shampoo-family optimizers through the package initializer.
• Adds Shampoo coverage and adjusts SOAP tests for the reorganized implementation.

Confidence Score: 5/5

The PR appears safe to merge because no eligible blocking failure or known outstanding prior finding remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
emerging_optimizers/shampoo/shampoo.py	Adds the vanilla Shampoo preconditioner, shared optimizer lifecycle, registered optimizer, and momentum update.
emerging_optimizers/shampoo/precond_base.py	Renames the shared preconditioner protocol module without changing its contents.
emerging_optimizers/shampoo/soap_v3.py	Updates SOAP variants to import shared protocols and tensor pairs from the renamed module.
emerging_optimizers/shampoo/init.py	Exposes Shampoo and the existing SOAP-family optimizers through the package API.
tests/test_shampoo.py	Adds unit and smoke coverage for Shampoo state, factor updates, inverse roots, validation, and optimizer steps.
tests/test_soap_v3.py	Updates SOAP test imports to follow the shared protocol-module rename.
pyproject.toml	Excludes abstract NotImplementedError branches from coverage reporting.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Training loop calls Shampoo.step] --> B[Initialize per-parameter state]
    B --> C[Update left and right Kronecker factors]
    C --> D[Compute factor inverse roots]
    D --> E[Apply momentum scalar update]
    E --> F[Precondition update on both sides]
    F --> G[Apply weight decay and parameter update]
    G --> H[Rebind state and increment step]

Loading

_{Reviews (5): Last reviewed commit: "add tikhonov and bias correction" | 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