Skip to content
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
wants to merge 42 commits into
base: master
Choose a base branch
from

Conversation

MilesCranmer
Copy link
Owner

@MilesCranmer MilesCranmer commented Jun 25, 2024

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:

jupytext --to ipynb SpecialRelativity.py

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)


Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
1 participant