Add an option to choose between flax.nn.custom_vjp and jax.custom_vjp. - Githubissues

google / aqt

Apache License 2.0

262 stars 27 forks source link

Add an option to choose between flax.nn.custom_vjp and jax.custom_vjp. #543

Closed copybara-service[bot] closed 7 months ago

copybara-service[bot] commented 7 months ago

Add an option to choose between flax.nn.custom_vjp and jax.custom_vjp.