Differences of pjit and jit after interface merge

[zhaopku]

Now that pjit and jit have been merged and have the same interface, what are the differences between pjit and jit? Can I use jit at places where I previously used pjit?

[yashk2810]

The only difference is that pjit respects the mesh context manager i.e. with jax.sharding.Mesh(...) and jax.jit does not do that.

The next natural question is why doesn't jax.jit respect it? One of the main reasons is that even simple jnp operations under a multihost mesh will become a blocking operation all of a sudden which would be a major shift in the semantics of jit.

Example:

# mesh spans across hosts which means that jnp.arange should be executed on all hosts at the same time.
with jax.sharding.Mesh(...):
  jnp.arange(16)

So what's the recommendation? The recommendation is that use jax.jit everywhere and remove the usage of with mesh from your codebase since jax.jit has sharding follows data semantics and accepts in_shardings/out_shardings arguments which you can use instead.