Add KL Shampoo's correction inside SOAP #24
Description
Is your feature request related to a problem? Please describe.
KL Shampoo corrects kronecker factor matrix accumulation by removing bias. Read Understanding and Improving the Shampoo Optimizer via Kullback–Leibler Minimization (https://arxiv.org/pdf/2509.03378) for full mathematical derivation.
Describe the solution you'd like
A modification or alternative update_kronecker_factors function for soap.py
Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.
Additional context
Add any other context or screenshots about the feature request here.