jax.custom_vjp.defremat

custom_vjp.defremat(fwd, rem, bwd)

Define custom rematerialization rules for this function.

The rules apply when this function is differentiated in reverse mode under jax.remat(), where by default the fwd rule given to defvjp() is rematerialized according to the ambient checkpoint policy. Outside of jax.remat(), unless a VJP is also defined with defvjp, the rules define the VJP too, by running rem right after fwd on the forward pass. Requires the jax_custom_vjp3 and jax_remat3 implementations.

Parameters:

• fwd (Callable[..., tuple[ReturnValue, Any]]) – forward-pass rule, with the same signature as the fwd rule of defvjp(): it takes the arguments of the primal function, and returns the primal output paired with residuals to save for the backward pass.

• rem (Callable[..., tuple[ReturnValue, Any]]) – rematerialization rule, of signature rem(res, *args) -> (out, res2). On the backward pass it receives the residuals saved by fwd and the arguments of the primal function, and recomputes the primal output paired with the residuals res2 for bwd.

• bwd (Callable[..., tuple[Any, ...]]) – backward-pass rule, as in defvjp(), receiving res2 as its residuals.

Returns:

None.

Return type:

None