quantax.utils.make_mesh

quantax.utils.make_mesh() → Mesh

Build the device mesh spanning all devices in jax.devices().

Returns:

A ("process", "device") jax.sharding.Mesh whose two axes have sizes jax.process_count() and jax.local_device_count(), with both axes set to the automatic AxisType.Auto.

Note

The mesh is built directly from jax.devices() reshaped to (process_count, local_device_count) rather than via jax.make_mesh(). jax.make_mesh rejects multi-slice topologies (raising on multi-host GPU since JAX 0.10), but jax.devices is already process-major, so the reshape reproduces the same ("process", "device") layout while supporting multi-node runs.