clipping: support default unitwise_norm for 5D params

[staticpayload]

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

[vroulet]

Can you sync with head, update the docs and add a test?
Thanks

[staticpayload]

Can you sync with head, update the docs and add a test? Thanks

Done

[vroulet]

Thanks!

[staticpayload]

🫡