OKLS optimizer

[mkhona-nvidia]

https://blog.tilderesearch.com/blog/online-kl-shampoo

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

Zero epsilon produces NaN state

When ridge_eps=0 and the first gradient is all zero, initialization computes sqrt(rows / 0) and multiplies the zero covariance by infinity, causing NaN factors, inverse roots, and parameter values on the first step.

Suggested change

	if ridge_eps < 0.0:
	raise ValueError(f"Invalid ridge epsilon: {ridge_eps}")
	if ridge_eps <= 0.0:
	raise ValueError(f"Invalid ridge epsilon: {ridge_eps}")

[greptile-apps]

Factor matmuls inherit global precision

The covariance and final preconditioning matmuls execute outside fp32_matmul_precision, so a process-wide medium setting silently reduces their precision and makes persistent optimizer factors depend on unrelated global configuration.

Knowledge Base Used: SOAP: Shampoo-style Preconditioning

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Greptile Summary

Adds the OKLS optimizer and its scaled-CANS inverse-root implementation.

• Registers and publicly exports OKLS.
• Adds API documentation for OKLS and the matrix inverse-root utility.
• Adds optimizer, numerical utility, and registry tests.

Confidence Score: 3/5

The PR is not yet safe to merge because zero epsilon can corrupt optimizer state with NaNs and configured matmul precision is not applied throughout the OKLS update.

The current implementation still permits a zero epsilon that makes zero-gradient initialization evaluate zero times infinity, while factor and final preconditioning matmuls remain controlled by unrelated process-wide precision.

Files Needing Attention: emerging_optimizers/soap/okls.py

Important Files Changed

Filename	Overview
emerging_optimizers/soap/okls.py	Implements OKLS state initialization, factor updates, inverse-root refreshes, momentum, and parameter updates.
emerging_optimizers/soap/matrix_root_inverse_utils.py	Implements the fixed-schedule scaled-CANS approximation for FP32 matrix inverse square roots.
tests/test_okls.py	Adds CUDA optimizer smoke, state-initialization, and dimensionality-validation coverage.
tests/test_matrix_root_inverse_utils.py	Adds shape, precision-restoration, accuracy, and dtype-validation coverage for inverse roots.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    G[Gradient] --> I[Initialize or update Kronecker factors]
    I --> C[Scaled CANS inverse roots]
    G --> M[Nesterov momentum]
    C --> P[Two-sided preconditioning]
    M --> P
    P --> U[Parameter update]

Loading

_{Reviews (2): Last reviewed commit: "fix: use decoupled weight decay in OKLS" | Re-trigger Greptile}