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

[Ashutosh0x]

Resolves #422. This PR modernizes the _check_sharding function to use the leaf.sharding attribute directly for jax.Array objects, bypassing jax.typeof(leaf) which could be unreliable or inconsistent in newer JAX versions (0.4.35+ and the upcoming 0.8.x). It also refines sharding-related error messages to be more accurate while maintaining backward compatibility with existing test regexes.

[hbq1]

Hi @Ashutosh0x thanks for the PR! Happy to accept it if you could address my comment + fix the style warnings from the CI.

[Ashutosh0x]

Thanks @hbq1! I've addressed the feedback:

• Added the else branch in _check_sharding to restore backward compatibility using jax.typeof(x).sharding.
• Fixed the style warnings by cleaning up the redundant parentheses in assert_tree_is_on_device.

Ready for another look!

[Ashutosh0x]

@hbq1 Since the PR is approved and CI is green, could you please merge it from your side? Thanks!

[Ashutosh0x]

@hbq1 gentle follow-up on this one, with a status check rather than just a ping.

I re-verified against main today, and the fix still has not landed - _check_sharding in chex/_src/asserts.py is still the pre-fix version:

if not jax.typeof(x).sharding.is_fully_replicated:
    return True
else:
    return len(x.sharding.device_set) > 1

asserts.py has changed upstream since this PR was opened, but for unrelated reasons - so this is genuinely still outstanding rather than something that landed by another route. #422 is also still open.

Current state, for whoever picks it up:

• Approved by you on 2026-01-06
• import/copybara reported the change imported to the internal review system (cl/852788807) and feedback/copybara reported internal checks passing
• mergeable is still true - despite the upstream drift there is no conflict, so no rebase is needed and I have deliberately not pushed anything, since a new commit would only reset the import state

One caveat worth flagging: the green CI on this PR is from January, including the jax=nightly job. That result is seven months old and says nothing about current JAX. If it would help, I am happy to rebase purely to re-trigger CI against present-day JAX so you have a fresh signal before merging - just say the word and I will do it. Otherwise I will leave the branch untouched.