Bump jax#14
Open
jjalko wants to merge 9 commits into
Open
Conversation
… gpu with jax.devices instead
…eeds some clarification. disabled so far.
…alling function. adopted code for that.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Various small updates that make the d3p work with the newer version of jax (==0.10.2).
All the tests pass, getting some errors from the
fourier_accountantwhich are somewhat normal. I commented out one uniformity test that actually failed. See my comment intests/test_random.pystarting on line 50. Maybe you @lumip could remind me what this test is, and I can check if the failure is just bad luck with the seed.I managed to only test this on CPU, and only with the smaller models in the
examplesfolder.I updated the
setup.py, which also includes installation of thejax-chacha-prnglibrary from the current master branch of its git repository. We will probably want to change this as soon as the updatedjax-chacha-prngis on pypi.