-
Notifications
You must be signed in to change notification settings - Fork 96
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Refactor with best python practices; update to latest JAX #8
Open
MilesCranmer
wants to merge
42
commits into
master
Choose a base branch
from
refactor
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Conversation
This file contains 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
This was referenced Jun 25, 2024
Open
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Finally getting around to rewrite a lot of the repo to follow best practices. Also updates to the latest JAX version (changed behavior for stax, optimizers, odeint), and fixes the new tracer issue.
Note this makes the repo a proper python project. So install it with
pip install .
This also deletes the notebooks and uses jupytext-compatible scripts in their place. Use:
to get a notebook back.
I also turn on ruff formatting, as well as isort + pycln.
TODO:
There is some issue in the double pendulum example pointed about by a student. I can't seem to figure it out. Perhaps something deep inside JAX changed and is breaking things now? We might have to leave this as-is, because at least now the code runs. (Just seems to get there much slower for some reason)
base-nn-double-pendulum.ipynb
#3 (I think)