Deactivate a mesh context

[wittenator]

Some libraries have trouble running under the automatic mesh context set by jax.set_mesh such as diffrax or optimistix. Is it possible to temporarily deactivate a mesh context without exiting it?

[yashk2810]

You can set the context mesh to an empty mesh by:

empty_concrete_mesh = Mesh(np.empty((), dtype=object), ())
jax.set_mesh(empty_concrete_mesh)