In [1]:

# !pip install git+https://github.com/EQuS/jaxquantum.git  # Uncomment when running in Colab.

# !pip install git+https://github.com/EQuS/jaxquantum.git # Uncomment when running in Colab.

Sharding Across Multiple Devices¶

This is an executable IPython Notebook tutorial.

jaxquantum runs out of the box on a single device. With one line of configuration, every Qarray you construct — whether Dense or SparseDIA — automatically lives sharded across multiple CPUs, GPUs, or TPUs:

jqt.set_device_mesh(shape=(2,), axis_names=('dp',))

All downstream JAX operations (@, eigh, diffrax, mesolve, …) propagate the sharding through XLA. There is no separate sharded backend, no .shard() method to call, and no need to touch .data.

This tutorial covers:

1. Spoofing multiple devices on a laptop — for development without a GPU/TPU
2. Configuring the default sharding — set_device_mesh and set_default_sharding
3. Sharded Qarrays in action — Dense and SparseDIA both pick it up automatically
4. Tensor products — sharding is re-applied to the (larger) output
5. Validation — clear errors when matrix dims don't divide the mesh
6. Solver-level sharding — running a sharded mesolve end-to-end
7. SparseBCOO is unsupported — and why
8. Manual control — explicit PartitionSpec, per-call overrides, clearing
9. Data-parallel — many systems, one device each (parameter sweeps)
10. Model-parallel — one large system, multiple devices
11. Both at once — 2D meshes for combined parallelism

How sharding works in jaxquantum¶

You set a global default once with jqt.set_device_mesh(...). Every Qarray you build afterwards routes its underlying JAX array through jax.lax.with_sharding_constraint at construction time, so dense matrices, sparse-diagonal _diags, kets, bras, and batched arrays all live on the configured mesh without any per-call boilerplate. Pure JAX ops (@, eigh, diffrax) propagate the sharding through XLA.

Two parallelism modes are supported, picked by the name of each mesh axis:

• Data-parallel (axis_names=('dp',) or 'data'): shards the batch dim of (B, N, N) arrays — useful for parameter sweeps and vmap'd Monte Carlo trajectories.
• Model-parallel (axis_names=('mp',) or 'model'): shards the matrix dim of (N, N) operators — useful for one large simulation that won't fit (or doesn't run fast enough) on a single device.
• Both at once with a 2D mesh: axis_names=('dp', 'mp'). Each device gets a slice of the batch and a slice of each operator.

There is no "sharded backend" — sharding is orthogonal to storage layout. SparseBCOO is unsupported (variable non-zeros per shard). To return to single-device mode, call jqt.clear_default_sharding().

1. Spoofing multiple CPU devices on a laptop¶

Set the XLA_FLAGS env var before importing JAX. This makes XLA expose virtual CPU devices, so you can develop sharded code on any machine and only switch to real GPUs/TPUs when you scale up. We spoof 8 here so we can demonstrate both 1D meshes (using a slice) and 2D meshes (using all 8) below.

In [2]:

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
print(jax.devices())  # 8 CPU devices

import os os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8" import jax print(jax.devices()) # 8 CPU devices

[CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]

In [3]:

import jaxquantum as jqt
import jax.numpy as jnp
import jax.scipy as jsp

import jaxquantum as jqt import jax.numpy as jnp import jax.scipy as jsp

/opt/hostedtoolcache/Python/3.11.16/x64/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

2. Configuring the default sharding¶

The set_device_mesh helper builds a Mesh and NamedSharding for you, then registers them as the global default. Every Qarray built afterwards routes its underlying jnp.ndarray through jax.lax.with_sharding_constraint.

With partition_spec=None (the default), the helper installs a rank-adaptive sharding that picks where to partition each array based on the name of each mesh axis: 'dp' (data-parallel) prefers the leading batch dim, while 'mp' (model-parallel) prefers the matrix dim. Anything else falls back to 'mp' behavior.

For now we'll use a 1D mesh ('dp',) with two devices — the same as the rest of this tutorial up to section 8. Sections 9–11 explore the full picture.

In [4]:

jqt.set_device_mesh(shape=(2,), axis_names=('dp',))
jqt.get_default_sharding()  # the rank-adaptive callable we just installed

jqt.set_device_mesh(shape=(2,), axis_names=('dp',)) jqt.get_default_sharding() # the rank-adaptive callable we just installed

Out[4]:

<function jaxquantum.utils.utils.set_device_mesh.<locals>._adaptive(arr)>

3. Sharded Qarrays in action¶

Just build operators as you normally would — they come out sharded.

In [5]:

N = 8
I = jqt.identity(N)
print('shape:', I.data.shape)
print('sharding:', I.data.sharding)
print('devices:', I.data.sharding.device_set)

N = 8 I = jqt.identity(N) print('shape:', I.data.shape) print('sharding:', I.data.sharding) print('devices:', I.data.sharding.device_set)

shape: (8, 8)
sharding: NamedSharding(mesh=Mesh('dp': 2, axis_types=(Auto,)), spec=P('dp', None), memory_kind=device)
devices: {CpuDevice(id=0), CpuDevice(id=1)}

In [6]:

a = jqt.destroy(N)
ad = jqt.create(N)
x = (a + ad)
print('x = a + a^dag is sharded:', x.data.sharding.device_set)

n = ad @ a
print('n = a^dag @ a is sharded:', n.data.sharding.device_set)

a = jqt.destroy(N) ad = jqt.create(N) x = (a + ad) print('x = a + a^dag is sharded:', x.data.sharding.device_set) n = ad @ a print('n = a^dag @ a is sharded:', n.data.sharding.device_set)

x = a + a^dag is sharded: {CpuDevice(id=0), CpuDevice(id=1)}
n = a^dag @ a is sharded: {CpuDevice(id=0), CpuDevice(id=1)}

SparseDIA picks it up too — only the _diags array (a regular jnp.ndarray) is sharded; the _offsets tuple is static Python metadata and lives on every device.

In [7]:

a_dia = jqt.destroy(N, implementation=jqt.QarrayImplType.SPARSE_DIA)
print('SparseDIA _diags shape:', a_dia.data.diags.shape)
print('SparseDIA _diags sharded across:', a_dia.data.diags.sharding.device_set)

a_dia = jqt.destroy(N, implementation=jqt.QarrayImplType.SPARSE_DIA) print('SparseDIA _diags shape:', a_dia.data.diags.shape) print('SparseDIA _diags sharded across:', a_dia.data.diags.sharding.device_set)

SparseDIA _diags shape: (1, 8)
SparseDIA _diags sharded across: {CpuDevice(id=0), CpuDevice(id=1)}

4. Tensor products re-shard the output¶

kron produces a matrix whose shape is (N·M, N·M) — new dimensions that the input partition spec didn't anticipate. _make re-applies the default sharding to the larger array so the invariant holds for arbitrary nested tensor products.

In [8]:

T = jqt.tensor(jqt.identity(4), jqt.identity(2))
print('tensor shape:', T.data.shape)
print('tensor sharded across:', T.data.sharding.device_set)

T = jqt.tensor(jqt.identity(4), jqt.identity(2)) print('tensor shape:', T.data.shape) print('tensor sharded across:', T.data.sharding.device_set)

tensor shape: (8, 8)
tensor sharded across: {CpuDevice(id=0), CpuDevice(id=1)}

5. Validation: matrix dims must divide the mesh¶

If you commit to a specific PartitionSpec and ask for a matrix size that can't be partitioned evenly across the named mesh axis, you get a friendly error at construction time naming the offending shape and mesh-axis size — much friendlier than a deep XLA error during a jit. (The rank-adaptive default never errors; it falls back to full replication when nothing divides — this is what makes scalars and oddly-sized arrays Just Work.)

In [9]:

from jax.sharding import PartitionSpec

# Commit to row-sharding — now an indivisible size (5) errors clearly.
jqt.set_device_mesh(shape=(2,), axis_names=('dp',), partition_spec=PartitionSpec('dp', None))
try:
    jqt.identity(5)
except ValueError as e:
    print('Got expected error:')
    print(' ', e)

jqt.clear_default_sharding()

from jax.sharding import PartitionSpec # Commit to row-sharding — now an indivisible size (5) errors clearly. jqt.set_device_mesh(shape=(2,), axis_names=('dp',), partition_spec=PartitionSpec('dp', None)) try: jqt.identity(5) except ValueError as e: print('Got expected error:') print(' ', e) jqt.clear_default_sharding()

Got expected error:
  Cannot apply default sharding NamedSharding(mesh=Mesh('dp': 2, axis_types=(Auto,)), spec=P('dp', None), memory_kind=device) to array of shape (5, 5): axis 0 (size 5) is not divisible by mesh axis ('dp',) (size 2). Pass an explicit `sharding=` per call, or call `jqt.clear_default_sharding()` to disable.

6. Solver-level sharding¶

diffrax (and therefore jqt.mesolve / jqt.sesolve) preserves sharding through ODE evolution: a sharded Hamiltonian and initial state produce a sharded trajectory.

In [10]:

# Re-establish the rank-adaptive default for the rest of this section.
jqt.set_device_mesh(shape=(2,), axis_names=('dp',))

# Build a Kerr Hamiltonian. All operators inherit the default sharding.
N = 16
a = jqt.destroy(N)
ad = jqt.create(N)
K = 0.01
H = K * (ad @ ad @ a @ a)
psi0 = jqt.basis(N, 1)

ts = jnp.linspace(0.0, 10.0, 21)
states = jqt.sesolve(H, psi0, ts)
print('trajectory shape:', states.data.shape)
print('trajectory sharded across:', states.data.sharding.device_set)

# Re-establish the rank-adaptive default for the rest of this section. jqt.set_device_mesh(shape=(2,), axis_names=('dp',)) # Build a Kerr Hamiltonian. All operators inherit the default sharding. N = 16 a = jqt.destroy(N) ad = jqt.create(N) K = 0.01 H = K * (ad @ ad @ a @ a) psi0 = jqt.basis(N, 1) ts = jnp.linspace(0.0, 10.0, 21) states = jqt.sesolve(H, psi0, ts) print('trajectory shape:', states.data.shape) print('trajectory sharded across:', states.data.sharding.device_set)

  0% |░░░░░░░░░░| [00:00<?, ?%/s]

  5% |▒░░░░░░░░░| [00:00<00:00, 3741.57%/s]

100% |██████████| [00:00<00:00, 33506.18%/s]

trajectory shape: (21, 16)
trajectory sharded across: {CpuDevice(id=0), CpuDevice(id=1)}

7. SparseBCOO is unsupported under sharding¶

BCOO stores (value, row, col) triples; the number of non-zeros per row is data-dependent, so it can't be evenly partitioned across devices the way Dense or SparseDIA can. Trying to build a SparseBCOO Qarray while a default sharding is configured raises a clear error.

In [11]:

try:
    jqt.Qarray.create(jnp.eye(4), implementation=jqt.QarrayImplType.SPARSE_BCOO)
except NotImplementedError as e:
    print('Got expected error:')
    print(' ', e)

try: jqt.Qarray.create(jnp.eye(4), implementation=jqt.QarrayImplType.SPARSE_BCOO) except NotImplementedError as e: print('Got expected error:') print(' ', e)

Got expected error:
  Sharding is not supported for SparseBCOO (variable nnz per shard). Convert to Dense or SparseDIA, or call `jqt.clear_default_sharding()` before constructing a SparseBCOO Qarray.

8. Manual control¶

For more control, pass an explicit PartitionSpec to set_device_mesh, or hand a fully-built Sharding to set_default_sharding. To return to single-device mode, call clear_default_sharding.

In [12]:

from jax.sharding import PartitionSpec

# Shard rows instead of columns (matches the pattern used in the demo notebook).
jqt.set_device_mesh(shape=(2,), axis_names=('dp',), partition_spec=PartitionSpec('dp', None))
row_sharded = jqt.identity(8)
print('row-sharded:', row_sharded.data.sharding.spec)

jqt.clear_default_sharding()
single = jqt.identity(8)
print('after clear:', getattr(single.data, 'sharding', None))

from jax.sharding import PartitionSpec # Shard rows instead of columns (matches the pattern used in the demo notebook). jqt.set_device_mesh(shape=(2,), axis_names=('dp',), partition_spec=PartitionSpec('dp', None)) row_sharded = jqt.identity(8) print('row-sharded:', row_sharded.data.sharding.spec) jqt.clear_default_sharding() single = jqt.identity(8) print('after clear:', getattr(single.data, 'sharding', None))

row-sharded: P('dp', None)
after clear: SingleDeviceSharding(device=CpuDevice(id=0), memory_kind=device)

9. Data-parallel: many systems, one device each (parameter sweeps)¶

A parameter sweep is the canonical data-parallel workload — you have many independent Hamiltonians (one per parameter value) and want each device to take a slice of the sweep. Name your mesh axis 'dp', batch your operators with vmap, and the default sharding will partition the batch dim — each device gets full operators for B/dp items.

In [13]:

jqt.clear_default_sharding()
jqt.set_device_mesh(shape=(2,), axis_names=('dp',))

N = 8
a = jqt.destroy(N)
ad = jqt.create(N)

# Sweep over four Kerr coefficients.
Ks = jnp.array([0.005, 0.01, 0.02, 0.05])
H_data = jax.vmap(lambda K: K * (ad @ ad @ a @ a).data)(Ks)  # (4, N, N)
H = jqt.Qarray.create(H_data, bdims=(4,))

print('Hamiltonian shape:', H.data.shape)
print('partition spec:', H.data.sharding.spec)  # batch dim sharded
print('devices:', H.data.sharding.device_set)

# Eigvalsh across the sweep — each device computes 2 eigenproblems.
spectra = jax.vmap(jnp.linalg.eigvalsh)(H.data)
print('spectra shape:', spectra.shape, '— sharded across', len(spectra.sharding.device_set), 'devices')

jqt.clear_default_sharding() jqt.set_device_mesh(shape=(2,), axis_names=('dp',)) N = 8 a = jqt.destroy(N) ad = jqt.create(N) # Sweep over four Kerr coefficients. Ks = jnp.array([0.005, 0.01, 0.02, 0.05]) H_data = jax.vmap(lambda K: K * (ad @ ad @ a @ a).data)(Ks) # (4, N, N) H = jqt.Qarray.create(H_data, bdims=(4,)) print('Hamiltonian shape:', H.data.shape) print('partition spec:', H.data.sharding.spec) # batch dim sharded print('devices:', H.data.sharding.device_set) # Eigvalsh across the sweep — each device computes 2 eigenproblems. spectra = jax.vmap(jnp.linalg.eigvalsh)(H.data) print('spectra shape:', spectra.shape, '— sharded across', len(spectra.sharding.device_set), 'devices')

Hamiltonian shape: (4, 8, 8)
partition spec: P('dp', None, None)
devices: {CpuDevice(id=0), CpuDevice(id=1)}
spectra shape: (4, 8) — sharded across 2 devices

10. Model-parallel: one large system, multiple devices¶

A single, large (N, N) Hamiltonian that doesn't fit (or doesn't run fast enough) on one device. Name your mesh axis 'mp' and the default sharding will partition the matrix-row dim — each device holds a row-block of every operator.

In [14]:

jqt.clear_default_sharding()
jqt.set_device_mesh(shape=(2,), axis_names=('mp',))

N = 32  # large enough that intra-matrix parallelism is worth it
a = jqt.destroy(N)
ad = jqt.create(N)
H = (ad @ a) + 0.1 * (a + ad) + 0.01 * (ad @ ad @ a @ a)

print('partition spec:', H.data.sharding.spec)  # P('mp', None) — row-sharded
print('devices:', H.data.sharding.device_set)

# Eigendecompose the full operator — XLA splits the work across devices.
w, v = jnp.linalg.eigh(H.data)
print('eigenvalues shape:', w.shape)
print('eigenvectors sharded across:', len(v.sharding.device_set), 'devices')

# Time-evolve a single ket under the large Hamiltonian.
psi0 = jqt.basis(N, 1)
ts = jnp.linspace(0.0, 5.0, 11)
states = jqt.sesolve(H, psi0, ts)
print('trajectory shape:', states.data.shape)
print('trajectory sharded across:', states.data.sharding.device_set)

jqt.clear_default_sharding() jqt.set_device_mesh(shape=(2,), axis_names=('mp',)) N = 32 # large enough that intra-matrix parallelism is worth it a = jqt.destroy(N) ad = jqt.create(N) H = (ad @ a) + 0.1 * (a + ad) + 0.01 * (ad @ ad @ a @ a) print('partition spec:', H.data.sharding.spec) # P('mp', None) — row-sharded print('devices:', H.data.sharding.device_set) # Eigendecompose the full operator — XLA splits the work across devices. w, v = jnp.linalg.eigh(H.data) print('eigenvalues shape:', w.shape) print('eigenvectors sharded across:', len(v.sharding.device_set), 'devices') # Time-evolve a single ket under the large Hamiltonian. psi0 = jqt.basis(N, 1) ts = jnp.linspace(0.0, 5.0, 11) states = jqt.sesolve(H, psi0, ts) print('trajectory shape:', states.data.shape) print('trajectory sharded across:', states.data.sharding.device_set)

partition spec: P('mp', None)
devices: {CpuDevice(id=0), CpuDevice(id=1)}
eigenvalues shape: (32,)
eigenvectors sharded across: 2 devices

  0% |░░░░░░░░░░| [00:00<?, ?%/s]

  0% |░░░░░░░░░░| [00:00<?, ?%/s]

 32% |███░░░░░░░| [00:00<00:00, 5245.92%/s]

 69% |██████▒░░░| [00:00<00:00, 5964.35%/s]

100% |██████████| [00:00<00:00, 6305.52%/s]

trajectory shape: (11, 32)
trajectory sharded across: {CpuDevice(id=0), CpuDevice(id=1)}

11. Both modes at once: a 2D mesh sweeps large systems¶

When you have both axes of parallelism — a parameter sweep over reasonably large systems — a 2D mesh gives each device a slice of the batch and a slice of each operator.

In [15]:

jqt.clear_default_sharding()
jqt.set_device_mesh(shape=(2, 4), axis_names=('dp', 'mp'))  # 8 devices total

N = 16
a = jqt.destroy(N)
ad = jqt.create(N)
Ks = jnp.array([0.005, 0.02])  # 2-element sweep along 'dp'

H_data = jax.vmap(lambda K: (K * (ad @ ad @ a @ a)).data)(Ks)  # (2, 16, 16)
H = jqt.Qarray.create(H_data, bdims=(2,))
print('H spec:', H.data.sharding.spec)  # P('dp', 'mp', None)

# Initial state, broadcast over the sweep.
psi0_data = jnp.broadcast_to(jqt.basis(N, 1).data, (2, N))
psi0 = jqt.Qarray.create(psi0_data, bdims=(2,), qtype='ket')
print('psi0 spec:', psi0.data.sharding.spec)  # P('dp', 'mp')

# Evolve all systems in lockstep across the 2D mesh.
ts = jnp.linspace(0.0, 1.0, 6)
states = jax.vmap(lambda H_, psi_: jqt.sesolve(
    jqt.Qarray.create(H_), jqt.Qarray.create(psi_), ts
).data)(H.data, psi0.data)
print('trajectories shape:', states.shape)         # (2, 6, 16)
print('trajectories sharded across:', len(states.sharding.device_set), 'devices')

jqt.clear_default_sharding() jqt.set_device_mesh(shape=(2, 4), axis_names=('dp', 'mp')) # 8 devices total N = 16 a = jqt.destroy(N) ad = jqt.create(N) Ks = jnp.array([0.005, 0.02]) # 2-element sweep along 'dp' H_data = jax.vmap(lambda K: (K * (ad @ ad @ a @ a)).data)(Ks) # (2, 16, 16) H = jqt.Qarray.create(H_data, bdims=(2,)) print('H spec:', H.data.sharding.spec) # P('dp', 'mp', None) # Initial state, broadcast over the sweep. psi0_data = jnp.broadcast_to(jqt.basis(N, 1).data, (2, N)) psi0 = jqt.Qarray.create(psi0_data, bdims=(2,), qtype='ket') print('psi0 spec:', psi0.data.sharding.spec) # P('dp', 'mp') # Evolve all systems in lockstep across the 2D mesh. ts = jnp.linspace(0.0, 1.0, 6) states = jax.vmap(lambda H_, psi_: jqt.sesolve( jqt.Qarray.create(H_), jqt.Qarray.create(psi_), ts ).data)(H.data, psi0.data) print('trajectories shape:', states.shape) # (2, 6, 16) print('trajectories sharded across:', len(states.sharding.device_set), 'devices')

H spec: P('dp', 'mp', None)

psi0 spec: P('dp', 'mp')

W0827 22:21:22.926569    2797 spmd_partitioner.cc:656] [SPMD] Involuntary full rematerialization. The compiler cannot go from sharding {maximal device=0} to {devices=[1,2,4]<=[8] last_tile_dim_replicate} efficiently for HLO operation %get-tuple-element = c128[6,16]{1,0} get-tuple-element(%conditional), index=3, sharding={maximal device=0}, metadata={op_name="jit(diffeqsolve)/jit(branched_error_if_impl)/cond/branch_1_fun/pure_callback" stack_frame_id=8}. As the last resort, SPMD will replicate the tensor and then partition it to obtain the target sharding, which is inefficient. This issue will be fixed by Shardy partitioner in the future, which is tracked in b/433785288. Contact Shardy or XLA team for help.

  0% |░░░░░░░░░░| [00:00<?, ?%/s]

 20% |██░░░░░░░░| [00:00<00:00, 6303.91%/s]

100% |██████████| [00:00<00:00, 16511.71%/s]

trajectories shape: (2, 6, 16)
trajectories sharded across: 8 devices

In [16]:

H.data.sharding

H.data.sharding

Out[16]:

NamedSharding(mesh=Mesh('dp': 2, 'mp': 4, axis_types=(Auto, Auto)), spec=P('dp', 'mp', None), memory_kind=device)

About axis names¶

'dp' (data-parallel) and 'mp' (model-parallel) are the SPMD/ML conventions for "shard the data" vs "shard the model". jaxquantum recognises these prefixes and routes the mesh axis accordingly. Any other name falls back to model-parallel behaviour. Override by passing partition_spec= to set_device_mesh or by handing a fully-built Sharding to set_default_sharding.

Small arrays on big meshes¶

A (N, 1) ket on a 2D mesh has no axis left for one of the two mesh dims — the ket replicates across that mesh axis. This is correct and cheap: replicating N floats costs far less memory than the N × N operator the ket interacts with.

Summary¶

Action	One-liner
Set up a 2-device mesh	jqt.set_device_mesh(shape=(2,), axis_names=('dp',))
Custom PartitionSpec	jqt.set_device_mesh(..., partition_spec=PartitionSpec(...))
Pass an existing Sharding	jqt.set_default_sharding(sharding)
Inspect the current default	jqt.get_default_sharding()
Return to single-device	jqt.clear_default_sharding()

Supported backends: Dense and SparseDIA. Unsupported: SparseBCOO (variable nnz per shard).

Where the sharding is applied: every QarrayImpl._make classmethod (called from from_data, arithmetic ops, kron, conversions, …) routes the underlying array through jax.lax.with_sharding_constraint. This is a no-op when the array is already correctly sharded, so the runtime cost is essentially zero.