You're thinking of `jax.closure_convert`. :) (Although technically that works by tracing and extracting all constants from the jaxpr, rather than introspecting

You're thinking of `jax.closure_convert`. :) (Although technically that works by tracing and extracting all constants from the jaxpr, rather than introspecting the function's closure cells -- it sounds like your trick is the latter.) When you discuss dynamic allocation, I'm guessing you're mainly referring to not being able to backprop through `jax.lax.while_loop`. If so, you might find `equinox.internal.while_loop` interesting, which is an unbounded while loop that you can backprop through! The secret sauce is to use a treeverse-style checkpointing scheme. https://github.com/patrick-kidger/equinox/blob/f95a8ba13fb35cc1160102168f2ecdae6119580a/equinox/internal/_loop/loop.py#L21 https://github.com/patrick-kidger/equinox/blob/f95a8ba13fb35...