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

Add support for alternative automatic differentiation backends #11

Merged
merged 25 commits into from
Oct 6, 2024

Conversation

matt-graham
Copy link
Owner

@matt-graham matt-graham commented Sep 26, 2024

Adds support for JAX and SymNum as alternative automatic differentiation backends to the current Autograd fallback.

All systems now expose a backend argument in their initializer that can be used to specify name of backend to use, defaulting to None, corresponding to not using an automatic differentiation backend.

Also adds a vjp_and_value differential operator to Autograd wrapper which was previously missing.

Todo

  • Update README with details of autodiff backends
  • Rename autodiff optional dependencies group to autograd
  • Add new optional dependencies groups for jax and symnum
  • Add tests for autodiff backends

@matt-graham matt-graham marked this pull request as ready for review September 26, 2024 21:22
@matt-graham matt-graham merged commit ea4adee into main Oct 6, 2024
4 checks passed
@matt-graham matt-graham deleted the mmg/jax-support branch October 6, 2024 12:57
@matt-graham matt-graham added the enhancement (Requests for) new features label Oct 7, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement (Requests for) new features
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant