clipping: support default unitwise_norm for 5D params #1576

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Open

staticpayload wants to merge 2 commits into google-deepmind:main from staticpayload:unitwise-norm-5d-default

staticpayload commented Jan 24, 2026

Summary

extend default unitwise_norm behavior to rank-5 tensors (Conv3D kernels)
update tests to cover the new default behavior

Testing

python3.11 -m venv .venv-test
pip install -e .
pip install pytest
pytest optax/transforms/_clipping_test.py -q

staticpayload added 2 commits

January 24, 2026 13:37


          clipping: support default unitwise_norm for 5D params

ff3afca


          tests: wrap galore shape asserts for lint

2f6f74d

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet