Out of memory allocating 0 bytes on CPU

[amilsted]

Does anybody have an idea why we might see a ValueError: RESOURCE_EXHAUSTED: Out of memory allocating 0 bytes. when doing anything sharded on CPU?

The context is perhaps a bit obscure: using JAX with Julia linked in via juliacall, including a Julia "system image", on Mac OS on Apple Silicon (problem not seen on Linux/x86_64). A lot of JAX works fine with this config, but lax.cond and jax.jacobian can trigger some sharded stuff (e.g. execute_sharded) that leads to these strange exceptions.

I am looking for any hints I could use to debug further or work around the issue. Many thanks!