Update global variable on pmap function

[lucasliunju]

Hi,

I would like to ask a question about how to update the global value on pmap function. For example:

v = initialize_v() // initialize v

//training process
def train_step(train_step, learning_rate, data):

    global v

    v = train(v, batch)

    output = g(v)
    return output


pmapped_train_step = jax.pmap(
      functools.partial(
          train_step,
          learning_rate_fn=learning_rate_fn),
      axis_name='batch',
      donate_argnums=(0, 1))

output = pmapped_train_step(batch)

I would like to reuse the value of global "v" and update it when I run pmap(). However, I find I havn't update the value of "v" when I run it. Could you give me some advice on how to solve it.

Thnak you very much!

YOng

[jakevdp]

When a function changes a global value, that is known as a side effect, and JAX transforms like pmap assume functions are side-effect free (see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions). So to answer your question directly: there's no way to update a global value within a pmapped operation.

The operation you're doing looks more like a scan than a map; you can probably implement it with lax.scan. It would look something like this, assuming your batches are along the leading axis:

from jax import lax

def f(v, x):
  v = train(v, x)
  return v, g(v)

v = initialize_v()

v, outs = lax.scan(f, v, batches)