Coverage for jaxquantum/core/settings.py: 94%

1"""Core settings. 

2 

3Holds global runtime configuration read by the rest of ``jaxquantum``. 

4 

5How sharding works 

6------------------ 

7 

8A single global ``SETTINGS["default_sharding"]`` (either a 

9``jax.sharding.Sharding`` *or* a callable ``arr -> Sharding``) is consumed 

10by every ``QarrayImpl._make`` classmethod via :func:`_maybe_shard`. When 

11the entry is ``None`` (the default), every ``_make`` call is a pure 

12pass-through and behaviour matches the single-device codebase exactly. 

13 

14Why a callable: lets the rank-adaptive default produced by 

15``set_device_mesh`` return a different ``NamedSharding`` per array shape 

16(matrices vs kets vs batched arrays vs SparseDIA ``_diags``). 

17 

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

19 

20* **Data-parallel** — name a mesh axis ``'dp'`` or ``'data'``. Shards the 

21 leading batch dim of ``(B, *)`` arrays (parameter sweeps, vmap'd 

22 trajectories). 

23* **Model-parallel** — name a mesh axis ``'mp'`` or ``'model'`` (or 

24 anything else). Shards the matrix dim of ``(*, n, n)`` operators. 

25* **Both** — use a 2D mesh ``axis_names=('dp', 'mp')``. 

26 

27Chokepoint: every impl construction goes through ``_make`` → 

28``_maybe_shard`` → ``jax.lax.with_sharding_constraint``. The constraint is 

29a no-op when the array is already correctly sharded, so it's safe to apply 

30on every op (matmul, kron, conversions). Pure JAX ops between calls 

31propagate the sharding through XLA without further help. 

32 

33Not sharded: ``SparseBCOO`` (variable nnz per shard — raises in 

34``from_data``); ``_offsets`` on ``SparseDIA`` (static Python metadata, 

35not a JAX array). 

36 

37User knobs (in :mod:`jaxquantum.utils.utils`, next to ``set_precision``): 

38``set_device_mesh``, ``set_default_sharding``, ``get_default_sharding``, 

39``clear_default_sharding``. 

40""" 

41 

42from __future__ import annotations 

43 

44from typing import Any, Optional 

45 

46 

47SETTINGS: dict = { 

48 "auto_tidyup_atol": 1e-14, 

49 # Optional[jax.sharding.Sharding | Callable[[Array], jax.sharding.Sharding]] 

50 "default_sharding": None, 

51} 

52 

53 

54def _maybe_shard(arr: Any) -> Any: 

55 """Apply ``SETTINGS['default_sharding']`` to *arr* if one is configured. 

56 

57 Uses ``jax.lax.with_sharding_constraint`` (works inside and outside 

58 ``jit``; near-free no-op when the array is already correctly sharded). 

59 Short-circuits and returns *arr* unchanged when no default sharding is 

60 set — this is the common single-device path and must stay zero-cost. 

61 """ 

62 s = SETTINGS["default_sharding"] 

63 if s is None: 

64 return arr 

65 

66 # Lazy imports keep this module import-cheap. 

67 from jax.lax import with_sharding_constraint 

68 

69 sharding = s(arr) if callable(s) else s 

70 _validate_shape(arr, sharding) 

71 return with_sharding_constraint(arr, sharding) 

72 

73 

74def _validate_shape(arr: Any, sharding: Any) -> None: 

75 """Raise ``ValueError`` if *arr*'s shape can't be partitioned per *sharding*. 

76 

77 Walks the ``PartitionSpec`` against ``arr.shape`` and checks that each 

78 sharded axis is divisible by the product of the named mesh-axis sizes. 

79 Pure-Python on shape tuples — runs only when sharding is enabled. 

80 

81 Only attempts validation for ``NamedSharding`` (the common case from 

82 ``set_device_mesh``). Other ``Sharding`` types are forwarded to JAX 

83 untouched and any mismatch surfaces from JAX itself. 

84 """ 

85 try: 

86 from jax.sharding import NamedSharding 

87 except Exception: # noqa: BLE001 # pragma: no cover 

88 return 

89 

90 if not isinstance(sharding, NamedSharding): 

91 return 

92 

93 spec = sharding.spec 

94 mesh = sharding.mesh 

95 shape = getattr(arr, "shape", None) 

96 if shape is None: 

97 return 

98 

99 for axis_idx, axis_spec in enumerate(spec): 

100 if axis_spec is None or axis_idx >= len(shape): 

101 continue 

102 names = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) 

103 mesh_size = 1 

104 for name in names: 

105 mesh_size *= mesh.shape[name] 

106 dim = shape[axis_idx] 

107 if dim % mesh_size != 0: 

108 raise ValueError( 

109 f"Cannot apply default sharding {sharding} to array of " 

110 f"shape {tuple(shape)}: axis {axis_idx} (size {dim}) is not " 

111 f"divisible by mesh axis {names} (size {mesh_size}). " 

112 "Pass an explicit `sharding=` per call, or call " 

113 "`jqt.clear_default_sharding()` to disable." 

114 )