[nnx] Improve spmd

[cgarciae]

What does this PR do?

• Adds get_named_sharding function
• Fixes a bugs in with_sharding_constraint and sharding_hook.

[codecov-commenter]

Codecov Report

Attention: 5 lines in your changes are missing coverage. Please review.

Comparison is base (008c3cc) 56.32% compared to head (3b77280) 56.31%.

Files	Patch %	Lines
flax/experimental/nnx/nnx/spmd.py	16.66%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3580      +/-   ##
==========================================
- Coverage   56.32%   56.31%   -0.02%     
==========================================
  Files         100      100              
  Lines       11951    11957       +6     
==========================================
+ Hits         6731     6733       +2     
- Misses       5220     5224       +4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[superbobry]

How hard would it be to add a test covering the changed code?