Change to new `jax.custom_vjp` API

gehring / fax

MIT License

78 stars 9 forks source link

Change to new `jax.custom_vjp` API #10

Closed gehring closed 4 years ago

gehring commented 4 years ago

Our implementation for implicit differentiation should move away from the recently deprecated jax.custom_transform in favor of the new jax.custom_vjp API.

To do:

[ ] update version requirements
[x] replace all usage of jax.custom_transform

pierrelux commented 4 years ago

@gehring I believe this can be closed