Autograd -> Jax conversion #433
Open
Conversation
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
|
Wavebot is working now when converted to jax. Working on updating the code so that only the necessary functions use jax.numpy while the rest use numpy. |
* bug bix : DC and Nyquist frequency should not be devided by two before ifft * Changed td_to_fd to scale single sided frequency components rather than TD signal * minor bug fix from issue332 sandialabs#332
* added initial file changes based on sphinx_multiversion docs and WEC-Sim implementation * removed sphinx-multiversion since it is no longer supported and made manual multiversion * now uses absolute paths, commented out linkcheck for debugging * fixed docstring errors in utilities module * updating files again that somehow got reverted * fixing path in conf.py * don't run tutorials (will revert later) * handle file moves correctly, fixed if statement to make other versions appear * fixed two bugs in versions template * reverted temp changes, changes latest to main * switched latest to main * main branch now in root directory of pages * fixed URLs with change from last commit * make other branches visible before building * switched main branch tag for more testing * fixed typo * switched dev branch to an existing branch * renamed main to latest, changed version.html file name to avoid confusion * added prints about moving files so Sphinx output isn't misleading * fixed typo with quotations * changed versions.html name back because that broke things I guess * modified contributing documentation to reflect changes * add logic to remove duplicate 'latest' branch * Fixed pathing when already on latest * remove typo * Troubleshooting complete, switching back to correct branches for deployment * Removed extra word in docstring * removed redundant function * fixed pathing so returns to same file (and fixes tutorial/API docs) * changed latest branch for demonstration * switched back latest branch for deployment
* removed conda environment from workflows since newer capytaine/wavespectra work with Windows * fixed unnecessary capitalization * still create CI conda environment to fix Mac environment failures * added conda env fully back in, push workflow deploys docs, split PR workflow * conda environment activates again * mambaforge instead of miniforge * manual cache reset * reset to older version of setup-miniconda to troubleshoot
…supported versions (sandialabs#390) Co-authored-by: jtgrasb <[email protected]>
* Try specifying subversion * Test new cache * revert to 3.12 * Revert comment back to normal
jtgrasb
force-pushed
the
jax_conversion
branch
from
0c0e42b to
7ecc18d
Compare
|
To do:
|
|
I added just-in-time compilation to the optimization (objective function, constraints, and relative gradients) using jax.jit which should speed up the code. I also had to change the call to Here is the computation time with and without jit for the AquaHarmonics parameter sweep cell. For some reason, I'm not seeing that much of a speedup so need to look into this further. To do:
|
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.
Description
Convert autograd to jax.
Wavebot tutorial is working for the first optimization.