Do you think this is this true even if the function to be applied to the data chunks is very expensive and uses jax.jit and jax.vmap?

In my experience yes. The multithreading inside of jax jit performs less well than the MPI model. Of course this is only true if your algorithm scales well.

I tried to do this using MPI "map-by" and "bind-to" options, but it didn't work. Do you have any examples of how to do this?

I never do it because on clusters I work on, it's automatically done by slurm.

Finally, for the single program multiple data problem, do you think that jax.sharding and mpi4jax generally have similar performance or not? Which of the two is more efficient?

On GPUs, in my codes, sharding outperforms by 10% mpi4jax direct gpu operation. I have no idea why. It's probably because NCCL is better tuned than MPI. Still, it's what we see.

On CPUs, well, jax may use MPI or GLOO as a communication backed for sharding. MPI as a sharding backend works, but it's largely undocumented and hard to setup. I am partly responsible for getting the support there, so we have tiny docs in here https://netket.readthedocs.io/en/latest/docs/parallelization.html#mpitrampoline-backend-very-experimental . You need to compile MPITrampoline and launch jax through it. It's a mess. But performance is identical to using mpi4jax and you get the niceties of sharding (modulo some unsupported operations).
GLOO Is terrible and it's only handy for locally testing some Sharded code. You should never use it in production.

So in essence mpi4jax performs identical to MPI sharding backend, though it's easier to setup and write code for. Sometimes sharding breaks down and starts to replicate the calculations unless you use shard map. mpi4jax is equivalent to putting a shard map on the whole of your code so it's a bit more 'reliable'.