4 ms·
One curious thing I discovered a few months ago: you can sort of hack higher-order functions into JAX by defining “Pytree closures” which introspect on normal c
by mccoyb 3y ago
One curious thing I discovered a few months ago: you can sort of hack higher-order functions into JAX by defining “Pytree closures” which introspect on normal closures, and pull out the JAX tracer data from the closure environment (and put it back in, when tracing is required) —- and this works! You can pass these Pytree closures in and out of JIT boundaries, etc.
I believe JAX has a utility for this somewhere, can’t quite remember what this is called.
I typically think of JAX as quite restrictive — but I think the reality is that the only real limit on expressivity is that you can’t dynamically allocate inside of unbounded control flow (e.g. creating new allocations inside of a while loop).
- patrickkidger 3y agoYou'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...