pmap + jit best practices

[cgarciae]

Hey, I just thinking about how to tackle a problem and came up with this simple toy example involving pmap:

import jax
import jax.numpy as jnp

@jax.jit
def f(w, x):
    @jax.pmap
    def g(x):
        return w * x

    return g(x)

x = jnp.ones([8, 3])

w = jnp.array([2])
print(f(w, x)) # array with 2's

w = jnp.array([4])
print(f(w, x)) # array with 4's

My initial intuition was that this wouldn't work since w is captured by g it would be treated as a constant in the eyes of pmap, however, the jit around f somehow resolves this issue. This patterns "works" but seems a bit wrong somehow and I have a lot of doubts about how w is actually being treated:

• How is w replicated across devices?
• What are the semantics about using pmap + jit?
• What are the best practices around pmap?