What does jax.lax.with_sharding_constraint do in eager mode?

[mjo22]

Hi, in #30869 I got some helpful advice understanding jax.device_put. I have similar questions about jax.lax.with_sharding_constraint. Outside jitted regions, what does this function do, if anything? Is there a notion of hints to the compiler, will it act just as jax.device_put, or is it not totally clear what happens?

Testing this function, I observe that the following fails:

import jax
import numpy as np
import jax.numpy as jnp

sharding = jax.sharding.SingleDeviceSharding(jax.devices()[0])

x = np.arange(10)  # numpy array
y_wsc = jax.lax.with_sharding_constraint(x, sharding)
y_put = jax.device_put(x, sharding)
assert y_wsc.sharding == y_put.sharding == sharding  # This passes
assert y_wsc.committed == y_put.committed == True  # This passes

x = jnp.arange(10)  # jax array
y_wsc = jax.lax.with_sharding_constraint(x, sharding)
y_put = jax.device_put(x, sharding)
assert y_wsc.sharding == y_put.sharding == sharding  # This passes
assert y_wsc.committed == y_put.committed == True  # This fails; y_wsc.committed = False

However, this block of code passes!

import jax
import jax.sharding as jshard
import numpy as np
import jax.numpy as jnp

jax.config.update("jax_num_cpu_devices", 2)

num_devices = len(jax.devices())
mesh = jax.make_mesh((num_devices,), ("batch",))
sharding = jshard.NamedSharding(mesh, jshard.PartitionSpec("batch"))

x = np.arange(10)  # numpy array
y_wsc = jax.lax.with_sharding_constraint(x, sharding)
y_put = jax.device_put(x, sharding)
assert y_wsc.sharding == y_put.sharding == sharding  # This passes
assert y_wsc.committed == y_put.committed == True  # This passes

x = jnp.arange(10)  # jax array
y_wsc = jax.lax.with_sharding_constraint(x, sharding)
y_put = jax.device_put(x, sharding)
assert y_wsc.sharding == y_put.sharding == sharding  # This passes
assert y_wsc.committed == y_put.committed == True  # This passes too!

I do not understand this behavior, but of course I realize that with_sharding_constraint is documented to only be used inside jitted regions.

[yashk2810]

Inside jax.jit, with_sharding_constraint is a strict constraint to the compiler i.e. the compiler must respect the sharding specified.

Outside jax.jit, it is just an identity function i.e. jax.jit(lambda x: x, out_sharding=sharding)(x).

[mjo22]

Thanks so much!