Implement SOAP optimizer

### **Description:**
I’d like to propose adding the SOAP (Second-order Optimization with Alternating Projections) optimizer to optax.contrib. SOAP is a high-performance second-order method that improves upon Shampoo by running Adam steps in the preconditioner's eigenbasis.

Paper (ICLR 2025): SOAP: [Improving and Stabilizing Shampoo using Adam](https://openreview.net/forum?id=IDxZhXrpNf)

### **Idea**
* Performance: ~40% fewer iterations and ~35% faster wall-clock time vs AdamW on LLMs.
* Stability: Much less performance drop when eigendecomposition is infrequent.
* Continuous Updates: Second-moment updates happen in the current basis, avoiding stale statistics.
Proposed Implementation:
* Core: Add scale_by_soap and soap alias to optax/contrib/_soap.py.
* State: Use a JAX-idiomatic NamedTuple to manage basis and momentum buffers.
* JIT: Use jax.lax.cond for periodic eigendecomposition to remain JIT-friendly.
* Generalization: Support nD parameter shapes via dimension numbers (Muon/Shampoo style).
* Tests: Full unit tests in optax/contrib/_soap_test.py.

### **Motivation:**
I will be referencing [haydn-jones/SOAP_JAX](https://github.com/haydn-jones/SOAP_JAX) for the industrialization.

Would you be open to a PR for this?