jax.custom_remat

jax.custom_remat(f, f_fwd, f_rem, f_bwd, *, static_argnums=(), static_argnames=())

Wrap f with custom rematerialization behavior for reverse-mode AD.

Where jax.checkpoint() policies select saveable values by name, a custom_remat-wrapped function carries its own rematerialization rules, which can depend on the ambient checkpoint policy. Requires the jax_remat3 implementation.

Parameters:

• f – the function to wrap, called (or traced) for ordinary evaluation.

• f_fwd – forward-pass rule under rematerialized differentiation, of signature f_fwd(policy, *args) -> (out, res). It receives the ambient checkpoint policy along with the arguments of f, and returns the primal output paired with residuals to save (which may be None, to save nothing).

• f_rem – rematerialization rule, of signature f_rem(res, *args) -> (out, res2). On the backward pass it receives the residuals saved by f_fwd and the arguments of f, and recomputes the primal output paired with the residuals that f_bwd needs.

• f_bwd – backward-pass rule, of signature f_bwd(res2, out_ct) -> args_ct, returning a tuple of cotangents with one entry per argument of f.

• static_argnums – as in jax.jit().

• static_argnames – as in jax.jit().

Returns:

A wrapped version of f with the same call behavior, but with the given rules applied when it is differentiated in reverse mode under rematerialization (e.g. under jax.checkpoint()). Forward-mode differentiation falls back to differentiating f.