Jitted function and sharding side-effect

[sh0416]

I am just wondering that the jitted function contains side effect when we use "with sharding constrint" function in their implementation.

I think there needs some process for sharding external source such as static_argname for static values in jit.

Is there any resources to handle this topic?

My code is somehow this one.

@jax.jit
def f(x, y, mesh):
  z = x + y
  jax.lax.with_sharding_constraint(z, NamedSharding(mesh, PartitionSpec('dp', 'mp')))

# better for the above???
@functools(jax.jit, static_argname=('mesh'))
def f(x, y, mesh):
  z = x + y
  jax.lax.with_sharding_constraint(z, NamedSharding(mesh, PartitionSpec('dp', 'mp')))

[jakevdp]

I'm not entirely sure I understand the question, but you'll need to mark Mesh arguments as static within JIT-compiled functions, because non-static arguments must either be arrays or pytrees of arrays. So your second example (with static_argnames=['mesh']) is probably the way to go.