Add chex to dependencies

[jakevdp]

Fixes #150

[Tomas542]

Also kinda fixes #120

[jakevdp]

Looking more closely, we may want to rethink this. chex currently depends on the dm-tree package which appears to be more-or-less unmaintained (no release since 2022: https://pypi.org/project/dm-tree/#history) and in particular is not installable on Python 3.13, which means chex is not compatible with Python 3.13.

I'd like some assurances that chex and its requirements will continue to be maintained before we make it an official part of the stack.

[jakevdp]

The chex python 3.13 issue is addressed in google-deepmind/chex#378

[jakevdp]

Given that the dependency issue is solved, and given that optax is unlikely to remove the chex requirement, I think it makes sense to add chex to the jax-ai-stack package. I'm going to merge this!