Why *does* this work?

[robertmaxton42]

As part of a project I have ended up needing to roll my own wrapper around jnp.vectorize. Part of my code looks like this:

class Quad:
    def vectorize(cls, f: Callable):
        out_treedef: jtree.PyTreeDef
        def _quadified(treedef, *args):
            nonlocal out_treedef

            qargs = jtree.tree_unflatten(treedef, args)
            out = f(*qargs, **kwargs)
            flatout, out_treedef = jtree.tree_flatten(out)
            return tuple(flatout)
        @wraps(f)
        def _vectorized(*args):
            flatargs, treedef = jtree.tree_flatten(args)            
            _vec_quadified = jnp.vectorize(_quadified, excluded=frozenset({0}))
            flatout = _vec_quadified(treedef, *flatargs)
            return jtree.tree_unflatten(out_treedef, flatout)  

(Yes, I know, I'm calling jnp.vectorize on an internal local function every time the vectorized function is called; I'm not terribly happy about that, but in the final result the above also includes a signature argument I have to modify on the fly, including the output result structure, and I won't in general know the output structure until I run it once. Presumably I could avoid this by exploiting JAX's existing tracing capabilities but they aren't documented and I don't have time to really dig into that code. Fortunately, in my use case every vectorized function is only used once or twice, just on large inputs.)

Anyway, I've used a nonlocal here because I don't see any other way to communicate out_treedef back to _vectorized, as jnp.vectorize will always try and batch outputs into arrays and then promptly fail because JAX arrays don't support objects like PyTreeDef. And it works!

Except... why does it work? This is sort of the archetypal case of depending on a side effect, after all; I've defined a variable in Quad.vectorize's closure, I've not included it as an argument, and now I'm am passing a function that modifies that variable into a JAX transformation (jnp.vectorize calls vmap under the hood). Presumably, since JAX has to run _quadified once to trace it, the tracers happen to provide the output structure that I need, but I'm still not entirely clear why it doesn't throw an error...

(Not that I'm complaining, lol; I need this code to work!)

[jakevdp]

I think you've got it right: the reason this works is that even though out_treedef is a variable modified as a side effect, two things combine to save you:

1. every time you modify out_treedef by tracing _quadified you immediately consume it, in the next line, and never refer to it in other contexts
2. the only thing necessary to create this value is the abstract value of the input structure, so it is safe to write while tracing: i.e. it doesn't lead to tracer leaks because it doesn't store traced values, only static attributes.

If you're going to continue with this approach, I might actually define out_treedef and _quadified within _vectorize to ensure that they can't be called elsewhere, just for an added degree of safety.

[robertmaxton42]

Thanks for the answer! I'm glad to hear my understanding of JAX isn't too far off!

I actually did it that way originally, but then I'm worried that I really will be compiling _quadified every time I call the vectorized function. This way, the compiler helps me check that the only dependency that could prompt a compilation cache miss is if I changed the signature between calls of _vectorized, which shouldn't happen ever and and definitely would only happen rarely. Is there a reason I can be fairly confident that I won't be prompting JAX to JIT _quadified every time _vectorized is called, even if _quadified becomes a true local function?

[jakevdp]

As currently written, I don't think the location that you define _quadified will make any difference at runtime. JIT is not coming into play at all, since you've not used jit anywhere in your implementation, and if there's an outer JIT it will compile the same sequence of operations no matter where the definition is.

[robertmaxton42]

Huh. jnp.vectorize doesn't compile its input? I had assumed that since it involved vmap (and seemed in general faster than a for loop even without GPU/parallelization), it would get compiled into XLA.

... Huh. Upon checking, vmap doesn't say it compiles its inputs either. Interesting.

[jakevdp]

No, neither vectorize nor vmap involves jit. But if you want to JIT-compile your vmapped code, you can always compose the two.