[x] I've formatted the new code by running poetry run pre-commit run --all-files --show-diff-on-failure before committing.
[X] I've added tests for new code.
[X] I've added docstrings for the new code.
Description
This PR updates our codebase, docs, and unit tests to be compatible with the latest version of NNX. Due to a bug in TFP's sample_halton_sequence fn., I've had to upper bound our JAX/jaxlib version.
Type of changes
Checklist
poetry run pre-commit run --all-files --show-diff-on-failurebefore committing.Description
This PR updates our codebase, docs, and unit tests to be compatible with the latest version of NNX. Due to a bug in TFP's
sample_halton_sequencefn., I've had to upper bound our JAX/jaxlib version.Issue Number: N/A