ArrayRef, pure functions and vmap

[Egiob]

Hello!
First, I discovered ArrayRef recently, I think it a pretty recent feature, and I absolutely love it. That said there is something I don't get with how it works with vmap.

The guide says

If a function only uses refs internally, it is still considered pure. Purity is in the eye of the caller. For example:

# internal refs => still pure
@jax.jit
def pure1(x):
  ref = jax.array_ref(x)
  ref[...] = ref[...] + ref[...]
  return ref[...]

Pure functions, even those that use refs internally, are familiar: for example, they work with transformations like jax.grad,
jax.vmap, jax.shard_map, and others in the usual way.

Now it seems to me that this function is also pure, as it only uses ref internally:

@jax.jit
def isitpure(x):
    ref = jax.array_ref(1)
    ref[...] += x
    return ref[...]

The difference is that it performs a computation using a traced value (x). When called with vmap this fails with the error:

performing a set/swap operation with vmapped value on an unbatched array reference of type ...

I do understand that the vmap limitations are addressed later in the guide, but I don't see why the limitations would apply here, since the function seems to be pure (in my view at least). IMO the expected behavior would be to silently replicate the array ref over the batched axis as it would be done with a normal array.

Thanks in advance to anyone that could clarify this for me 😄

[dougalm]

I think this is just a bug on our end. (ArrayRefs haven't been tested much yet so bugs are expected.) Thanks for letting us know about it.

[Egiob]

Thanks for the quick answer! I'll be waiting for a patch then 😄

[dougalm]

#31672 should fix it. Let us know when you run into more bugs!