Fix: Modernize sharding detection and error messages for JAX 0.8.2 compatibility (#422)#423

Open

Ashutosh0x wants to merge 2 commits intogoogle-deepmind:mainfrom

Ashutosh0x:fix/jax-0.8.2-compat

Commits on Jan 4, 2026

Fix: Modernize sharding detection and error messages for JAX 0.8.2 compatibility (#422)
Ashutosh0x
committed

Commits on Jan 6, 2026

fix: address review comments (compatibility and style adjustments)
Ashutosh0x
committed