Uncommit a committed array

[tehret]

By default, arrays are created uncommitted. This is a very practical state as we can use them whatever the sharding of the other arrays. Is there a way to go back to this state once an array has been committed?
I looked across the documentation and I didn't find any hint that this could be possible (I couldn't even find a way to create an uncommitted version of the array directly in JAX, as soon as I try to copy the values it also copies the committed state).

The workaround I found was to use numpy as an intermediate step (committed array -> numpy array -> jax numpy array). It's kind of a device_put to the cpu (which I'm fine with) except that after using jax.numpy.asarray my array is now back to this very practical uncommitted state. It works but it's not very elegant as it requires using numpy as proxy.

A simple (useless) example where this could be useful:

import jax.numpy as jnp
from jax.sharding import PartitionSpec as P, NamedSharding

mesh = jax.make_mesh((2,), ('batch',))

@jax.jit
def parallel_heavy_computation(x, y):
    return x + y   

a = jnp.ones((4,4))
b = jax.device_put(a, NamedSharding(mesh, P(('batch',))))

c = parallel_heavy_computation(a,b)

In this example, c is now committed. This means that if I want to reuse it in another sharded setting like

d = jax.device_put(a, jax.devices()[0])
e = c + d # Error

I get the obvious error Received incompatible devices for jitted computation. Therefore, I need to make sure manually that the sharding is compatible. If I were able to "uncommit" c then I could go back to the ideal state where JAX takes care of the sharding so that it works at all time (at the cost of some memory transfer obviously).

[yashk2810]

Going to numpy and then jnp.asarray is the best way to do this.

[tehret]

Thanks for the confirmation !