Add NorMuon optimizer implementation and related tests

This commit introduces the NorMuon optimizer, which enhances the Muon algorithm by incorporating neuron-wise normalization of update magnitudes. The implementation includes the core functionality in `_normuon.py`, along with a dedicated test suite in `_normuon_test.py` to validate its behavior. Additionally, references to the new optimizer have been added in the documentation and the main `__init__.py` file.

Key changes:
- New optimizer: NorMuon with associated scaling and state management.
- Tests for NorMuon functionality and behavior.
- Documentation updates to include NorMuon in the API reference.

PiperOrigin-RevId: [insert-rev-id-here]

Author: tinker495

docs/api/contrib.rst

@@ -34,6 +34,8 @@ are not supported by the main library.
     MomoAdamState
     muon
     MuonState
+    normuon
+    NorMuonState
     prodigy
     ProdigyState
     sam
@@ -127,6 +129,12 @@ Muon
 .. autofunction:: scale_by_muon
 .. autoclass:: MuonState
 
+NorMuon
+~~~~~~~
+.. autofunction:: normuon
+.. autofunction:: scale_by_normuon
+.. autoclass:: NorMuonState
+
 Prodigy
 ~~~~~~~
 .. autofunction:: prodigy

optax/contrib/__init__.py

@@ -54,6 +54,9 @@
 from optax.contrib._muon import MuonDimensionNumbers
 from optax.contrib._muon import MuonState
 from optax.contrib._muon import scale_by_muon
+from optax.contrib._normuon import normuon
+from optax.contrib._normuon import NorMuonState
+from optax.contrib._normuon import scale_by_normuon
 from optax.contrib._privacy import differentially_private_aggregate
 from optax.contrib._privacy import DifferentiallyPrivateAggregateState
 from optax.contrib._privacy import dpsgd

optax/contrib/_common_test.py

@@ -53,6 +53,7 @@
     {'opt_name': 'momo', 'opt_kwargs': {'learning_rate': 1e-1}},
     {'opt_name': 'momo_adam', 'opt_kwargs': {'learning_rate': 1e-1}},
     {'opt_name': 'muon', 'opt_kwargs': {'learning_rate': 1e-2}},
+    {'opt_name': 'normuon', 'opt_kwargs': {'learning_rate': 1e-2}},
     {'opt_name': 'prodigy', 'opt_kwargs': {'learning_rate': 1e-1}},
     {
         'opt_name': 'schedule_free_sgd',