How to avoid triggering compilation every time with pmap(closure) ? (vmap ok)

[marius311]

I've got the following code pattern:

class MyProblem():

    @partial(jax.jit, static_argnums=0)
    def f(self, x, y):
        print("compiling")
        return x + y

    def solve(self, xs, y):
        def calls_f(x):
            return self.f(x, y)
        return jax.pmap(calls_f)(xs)

prob = MyProblem()

prob.solve(jnp.array(range(4)), 0)

This seems to trigger compilation every time I run the last line. If I replace with vmap it only compiles once. A simplified version without the calls_f closure, I'm able to fix it with Equinox's filter_pmap, but that doesn't work here.

I'm wondering how I can rewrite this to avoid the recompilation? Is pmap just fundamentally incompatible with closures? Thanks!

[YouJiacheng]

Hi. - In fact pmap will compile the function!
see the note here.
You can try this:

class MyProblem():
    # jit here is unnecessary
    def f(self, x, y):
        print("compiling")
        return x + y

    def solve(self, xs, y):
        return jax.pmap(self.f, in_axes=(0, None))(xs, y)

prob = MyProblem()

prob.solve(jnp.array(range(4)), 0)

[marius311]

Thanks, yea saw that and figured that was the cause. But is there not any way to do it without rewriting the structure of my code? (which in my real case would hurt the logic / code-clarity by a lot more than here). I guess it goes back to my question about is pmap (or I guess really jit) just incompatible with closures?

[YouJiacheng]

I try to minimize the modification for complex code, see my new answer.

[YouJiacheng]

@marius311 I think jit is "incompatible" with closure(with new captured value every time).
But I believe there is a workaround:

from jax.tree_util import Partial # note: not functools.partial
def call_fun(f, *args, **kwargs):
    return f(*args, **kwargs)
class MyProblem():

    @partial(jax.jit, static_argnums=0)
    def f(self, x, y):
        print("compiling")
        return x + y

    def calls_f_open(self, x, y):
        return self.f(x, y)

    def solve(self, xs, y):
        calls_f = Partial(self.calls_f_open, y=y)
        # now the "closure" is a Pytree, with self.calls_f_open as aux_data which has same hash for same self
        return jax.pmap(call_fun, in_axes=(None, 0))(calls_f, xs)

prob = MyProblem()

prob.solve(jnp.array(range(4)), 0)

[YouJiacheng]

Actually there is an option inline for jit, but I think it is not available since JAX always inline(actually trace through) jit-ted sub functions.

[marius311]

I'd be curious if you could say more about "tracing can be avoided"? In my mental model, tracing is synonymous with "did the print statement get shown" and it seems like we're unable to remove the repeated print statement while also keeping the closure in my example?

[sharadmv]

Check out Matt's comments in this thread: #7155

There's also exploratory work in avoiding the retracing (#9181) but at the end of the day it's an optimization that won't drastically change the compiled program (it may inline some functions).

[YouJiacheng]

I mean if we modify JAX current implementation(not too much), we can use cached jaxpr of f and avoid tracing it again. But it is not compelling and could be misleading since compiling usually takes more time than tracing.

[marius311]

Got it, thanks both for all the replies, I think that answers everything for me.

[marius311]

To summarize the answer from the very helpful discussion:

"Nested jit-ted functions are just traced through", i.e. the inner jit is ignored, at least for now.

And this comment also explains it. Workarounds here include dont use a closure, or JIT the entire solve function if you can.