Use of jax.device_put vs jax.lax.with_sharding_contraints

[miltonllera]

As the question says, it is unclear what the different uses of these two are based on the documentation. First some context. The question stems from @mjo22's effort to rewrite the auto-parallelisation tutorial for equinox and my own attempt at implementing the code in the tutorial, which lead to me noticing a duplicate use of equinox's filter_shard function.

Essentially, the issue is that with 0.7.0, Jax recommends using device_put to shard arrays in eager mode and lax.with_sharding_constraints being used inside jitted functions. equinox.filter_shard uses the latter and while testing I found that this works outside jitted functions just fine. The first question is, is there any downside to using a lax function such as this one in eager mode? My understanding is that there isn't in this case and that the only difference between the two is that device_put forces a particular sharding layout on the arrays which the compiler must stick to while with_sharding_constraints only provides a suggestion which the compiler only needs to consider but is otherwise free to optimize the sharding layout as wanted.

If this is true, then how could we achieve the same result in eager mode, that is, before passing the arrays to jitted functions? Is the only way to use in_sharding with Layout.AUTO?

This is important as it will influence how the equinox filter functions are re-written, especially eqx.filter_jit.

[yashk2810]

with_sharding_constraint can be used in eager mode too but it's just an identity jit i.e. jax.jit(lambda x: x, out_shardings=sharding)(x).

Under a jit, with_sharding_constraint is a strict constraint which compiler has to respect.

device_put has broader semantics than eager with_sharding_constraint and it is a "transfer" API. While with_sharding_constraint is a compiler tickler API which tells the XLA how to constrain sharding of an intermediate.

Usually the recommendation is to use device_put outside jit and if you want to constraint shardings of an intermediate inside jit, then with_sharding_constraint is the API. But if you want that same code to run under jax.disable_jit, then with_sharding_constraint will still work as it works in eager mode too.

[michael-0brien]

What is the recommended API long-term for constraining the outputs of a jit? Is it also with_sharding_constraint? Or is it jax.jit(..., out_shardings=...)? Further, will these two approaches act in the same way?

Though if I understand correctly, it seems for most use-cases it suffices to run jax.device_put on the inputs.

[yashk2810]

There is no recommended API as such. It depends on your use case.

Further, will these two approaches act in the same way?

Yeah they should.

[michael-0brien]

Okay, great thanks.

[miltonllera]

Thanks for the reply, but some things are still not clear to me. When using eqx.filter_shard in eager mode, which just calls with_sharding_constraints, I do get my arrays to shard correctly when I feed them to a jitted function. So clearly it is not an identity. Or am I misunderstanding something? This is also confusing when you say that it is a "tickler" API. I am not a compiler expert but this seems like non-standard terminology.

My understanding is that wcs gives compiler hints, but doesn't shard the array itself when used in eager mode, while device_put does shard the array itself. It is not clear then what is the difference between the two in eager mode, and why would device_put be needed.