Add jax_config.abstract_mesh_context_manager to core.get_sharding cache context.

[botev]

At the moment, if you are under a shard_map and have the following abstract eval rule of a primitive:

y = jax.core.ShapedArray(
  shape=shape,
  dtype=x.dtype,
  vma=x.vma,
)

where the goal is to have a value with sharding jax.P(*[None] * len(shape)) the above line throws an error:

  File "my_program.py", line 86, in inner_abstract_eval
    y = jax.core.ShapedArray(
        ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/root/lib/python3.12/site-packages/jax/_src/core.py", line 2220, in __init__
    self.vma = get_vma(vma, self.sharding.mesh)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/root/lib/python3.12/site-packages/jax/_src/core.py", line 2187, in get_vma
    assert not vma, vma
           ^^^^^^^
AssertionError: frozenset({'batch'})

Because the get_sharding caching context does not include the shard_map abstract mesh. This should actually be included in the caching context, and used inside.
This is obviously fixable, by me passing in the explicit sharding, but imho, at the moment it behaves quite unintuitively (and this was not an issue 0.8.0)

[yashk2810]

No, this is working as intended. ShapedArray does not read from the mesh context and shardings/mesh need to be passed explicitly to it.

In this case, you are setting the vma but not the mesh which is why you are seeing this error because vma without a mesh makes no sense.

[botev]

Ok, but then why do we have these global abstract mesh contexts, if they are not ever used for defaults?

[yashk2810]

We use the context mesh in a lot of places. One of the places that's kept explicit is ShapedArray i.e. jax's type. It is important to keep it context free otherwise a lot of things go wrong.