Public API for mesh axis index size

[gabbard]

You can use jax.lax.axis_index(axis) to get your index along a mesh axis without needing access to the mesh object. However, there does not seem to be a way to get access to the size of a mesh axis through the public API.

We have been doing:

        tensor_parallel_size = jax.interpreters.pxla.thread_resources.env.shape["TP"]

but jax.interpreters is not part of the public API.

Should such a function be added? Or are we doing something wrong if we think we sometimes need it?

[yashk2810]

No, there is no public API yet but we are working on it.

One thing you can do is create your own context manager and thread the mesh around and access it wherever you want.

[gabbard]

Thanks, that's a good idea for a workaround.