v0.0.7 - Don't require jax[cuda12], just jax

What's Changed

• Don't require jax[cuda12] as the dep, just jax by @lebrice in #2

Full Changelog: v0.0.6...v0.0.7

v0.0.6 - Reduce python version requirement to 3.10

What's Changed

• Reduce Python version requirement from 3.11-->3.10 by @lebrice in #1

New Contributors

• @lebrice made their first contribution in #1

Full Changelog: v0.0.5...v0.0.6

v0.0.5 - Rename JaxModule->WrappedJaxFunction, add wrapper for jax scalar-valued functions

• Rename JaxModule -> WrappedJaxFunction
• Add WrappedJaxScalarFunction:
  • Offered as an alternative to WrappedJaxFunction for scalar-valued functions (although both will work).
  • This is potentially more efficient, since it uses a jax.jit-ed jax.value_and_grad to do a fused forward and backward pass compared with WrappedJaxFunction that uses jax.vjp.
• Add more tests
• Add more doctests, which serve as both documentation and unit tests.
• Add some examples in the README.

Full Changelog: v0.0.4...v0.0.5

Add JaxModule: enables using jax functions in PyTorch autograd graphs

Example: https://github.com/lebrice/torch_jax_interop/blob/754445a1ce9a473b51d126b7919512cda22a0fa6/torch_jax_interop/to_torch_test.py#L279

Full Changelog: v0.0.3...v0.0.4