Skip to content

Conversation

@NicolasKorjahn
Copy link
Contributor

Add Shampoo optimizer support for neural network training in SystemDS.
Includes full-matrix and diagonal preconditioning, momentum updates, and a heuristic variant with delayed preconditioner updates / infrequent root recomputation. Extend the existing NN training scripts with tests to validate correctness and convergence.

Tests: src/test/scripts/applications/nn/component/shampoo_test.dml
src/test/scripts/applications/nn/component/shampoo_test2.dml
src/test/scripts/applications/nn/component/shampoo_test.py

Experiments in staging/shampoo_optimizer

@codecov
Copy link

codecov bot commented Jan 21, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.50%. Comparing base (b394e32) to head (1f54461).
⚠️ Report is 6 commits behind head on main.

Additional details and impacted files
@@             Coverage Diff              @@
##               main    #2401      +/-   ##
============================================
- Coverage     71.51%   71.50%   -0.01%     
+ Complexity    47441    47433       -8     
============================================
  Files          1539     1539              
  Lines        182605   182605              
  Branches      35916    35916              
============================================
- Hits         130585   130573      -12     
- Misses        42028    42032       +4     
- Partials       9992    10000       +8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@christinadionysio
Copy link
Contributor

Thank you for the patch @NicolasKorjahn! Could you add the missing licenses and then I would merge it in. For your follow-up experiments etc. you can then create a new PR.

@christinadionysio christinadionysio merged commit 2cbf63a into apache:main Feb 2, 2026
120 of 124 checks passed
@github-project-automation github-project-automation bot moved this from In Progress to Done in SystemDS PR Queue Feb 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

2 participants