Unpin JAX version after next jax_tpu_embedding release

[jeffcarp]

Pinned in #106 to fix import errors in tests.

[jshin1394]

should we unpin JAX versions now that jax_tpou_embedding has had a new release?

https://pypi.org/project/jax-tpu-embedding/#history

[jshin1394]

Updated pyproject to

dependencies = [
  "clu>=0.0.12",
  "flax>=0.11.1",
  "jax>=0.6.2",
  "numpy>=2.1.3",
  "tensorboardX>=2.6.4",
]