Skip to content

settings

Core settings.

Holds global runtime configuration read by the rest of jaxquantum.

How sharding works

A single global SETTINGS["default_sharding"] (either a jax.sharding.Sharding or a callable arr -> Sharding) is consumed by every QarrayImpl._make classmethod via :func:_maybe_shard. When the entry is None (the default), every _make call is a pure pass-through and behaviour matches the single-device codebase exactly.

Why a callable: lets the rank-adaptive default produced by set_device_mesh return a different NamedSharding per array shape (matrices vs kets vs batched arrays vs SparseDIA _diags).

Two parallelism modes are first-class, picked by mesh-axis name:

  • Data-parallel — name a mesh axis 'dp' or 'data'. Shards the leading batch dim of (B, *) arrays (parameter sweeps, vmap'd trajectories).
  • Model-parallel — name a mesh axis 'mp' or 'model' (or anything else). Shards the matrix dim of (*, n, n) operators.
  • Both — use a 2D mesh axis_names=('dp', 'mp').

Chokepoint: every impl construction goes through _make → _maybe_shard → jax.lax.with_sharding_constraint. The constraint is a no-op when the array is already correctly sharded, so it's safe to apply on every op (matmul, kron, conversions). Pure JAX ops between calls propagate the sharding through XLA without further help.

Not sharded: SparseBCOO (variable nnz per shard — raises in from_data); _offsets on SparseDIA (static Python metadata, not a JAX array).

User knobs (in :mod:jaxquantum.utils.utils, next to set_precision): set_device_mesh, set_default_sharding, get_default_sharding, clear_default_sharding.