jax.enable_custom_vjp_by_custom_transpose

jax.enable_custom_vjp_by_custom_transpose#

jax.enable_custom_vjp_by_custom_transpose = <jax._src.config._StateContextManager object>#

Context manager for jax_enable_custom_vjp_by_custom_transpose config option (transient).

Enables an internal upgrade that implements jax.custom_vjp by reduction to jax.custom_jvp and jax.custom_transpose. This will be enabled by default in future versions of JAX, at which point all uses of the flag will be considered deprecated (following the API compatibility policy).