Use JAX's AbstractMesh in distribution lib

[hertschuh]

In the JAX distribution lib, use AbstractMesh instead of Mesh since it doesn't result in a JIT cache misses when the devices change. It may also simplify the distribution API.

[vedantag17]

Is anyone working on this issue?

[hertschuh]

@vedantag17

Is anyone working on this issue?

Not yet, no.

[vedantag17]

Oh, I'm a newbie in Keras, but here's my understanding:
The current function builds a concrete device array. In contrast, AbstractMesh requires an immutable tuple of (axis_name, axis_size) pairs.
Is this correct?

[hertschuh]

@vedantag17 , yes, that is correct.

[vedantag17]

@hertschuh can you review the PR, there are some failing checks.

[hertschuh]

@vedantag17

The truth is that it is more complicated than this. At some point you need to create a concrete Mesh object to actually do the distribution. The AbstractMesh part should only come into play in contexts where you use the Mesh.

Also, I'm realizing we're not ready to work on this right now because:

• the JAX APIs related to AbstractMesh are still in flux and changing right now, so we at least need to wait until the next JAX release.
• we're not using shard_map in Keras right now, which is one of the main reason for preferring AbstractMesh.