Writing to a shared array in vmapped function

[Cattaneo123]

Greetings
I was wondering if there's a mechanism by which I can write to a shared array within a vmapped function instead of creating new values. I know JAX uses immutable arrays, but it mentions in the guide that while in jitted or vmapped functions the compiler attempts to make the changes in place. I ask because I have a function that performs an action at each value in a vector, but depending on the contents of that value it may do something to the next value as well. This causes a benign race, but I don't actually mind if that is the case as the correct answer should be produced either way. My attempts to get JAX to output a copy of the input array with the changes made have so far ended in failure unfortunately. I would greatly appreciate some guidance.

[jakevdp]

From what you describe, it seems the operation you have in mind is not expressible by vmap. In particular this piece:

I have a function that performs an action at each value in a vector, but depending on the contents of that value it may do something to the next value as well

This does not sound like a batchwise operation, but rather an iterative operation, and so I suspect vmap is not the right tool.

If you could put together a minimal example of the kind of operation you have in mind, we could probably give you some more specific help.

[Cattaneo123]

If I can break down one large operation into a series of smaller non-overlapping sequential operations each of which happen in parallel, then presumably I can represent the initial reproducible example as a vmapped scan function by breaking down the initial array into a series of subarrays and vmapping over that. If I do so, is there a way to do it such that the instead of returning a new array, it performs a write operation (or a write on a copy) of the larger original array.

I've tried this but it hasn't worked.

def scanOverVector(largeVector):
    def func(init, value, index):
        largeVector.at[index].set(value)
        return init,value/2

    return jax.lax.scan(func,0,largeVector,index)
    

performForSeveralVectors = jax.vmap(scanOverVector)
largeVector = [1,2,3,...n]
vectorIndexGroups = [[1,2,3],[4,5,6],...[n-2,n-1,n]]
performForSeveralVectors(vectorIndexGroups)
print(largeVector) #this, ideally has changed or been overwritten or a copy has been made with the changes

My hope is that after calling performForSeveralVectors largeector is changed, or performForSeveralVectors (and by extension scanOverVector) is changed such that a copy with those changes made it returned.

[jakevdp]

The issue is that largeVector.at[index].set(value) is not an in-place operation: rather it returns a modified copy of largeVector. So you'll need to include that in the carry of your scanned function if you're using this approach.

[Cattaneo123]

The issue is that largeVector.at[index].set(value) is not an in-place operation: rather it returns a modified copy of largeVector. So you'll need to include that in the carry of your scanned function if you're using this approach.

How would I go about that? If I return the copy then wouldn't I end with multiple slightly different copies of largeVector and no easy way of combining them?

[jakevdp]

You would include the array to be updated within the carry of the scan, and then update it at each step, passing the updated copy to the next step. Does that make sense?

[Cattaneo123]

You would include the array to be updated within the carry of the scan, and then update it at each step, passing the updated copy to the next step. Does that make sense?

Ah, I see, yes that makes sense.
I'll give that a shot. Thank you very much for your assistance