Coverage for jaxquantum/utils/utils.py: 94%

1""" 

2JAX Utils 

3""" 

4 

5from numbers import Number 

6from typing import Dict 

7 

8from jax import lax, Array, device_put, config 

9from jax._src.scipy.special import gammaln 

10import jax.numpy as jnp 

11import numpy as np 

12 

13from typing import Literal 

14 

15config.update("jax_enable_x64", True) 

16 

17 

18def device_put_params(params: Dict, non_device_params=None): 

19 non_device_params = [] if non_device_params is None else non_device_params 

20 for param, value in params.items(): 

21 if param in non_device_params: 

22 continue 

23 if isinstance(value, (Number, np.ndarray)): 

24 params[param] = device_put(value) 

25 return params 

26 

27 

28def comb(N, k): 

29 """ 

30 NCk 

31 

32 #TODO: replace with jsp.special.comb once issue is closed: 

33 https://github.com/google/jax/issues/9709 

34 

35 Args: 

36 N: total items 

37 k: # of items to choose 

38 

39 Returns: 

40 NCk: N choose k 

41 """ 

42 one = 1 

43 N_plus_1 = lax.add(N, one) 

44 k_plus_1 = lax.add(k, one) 

45 return lax.exp( 

46 lax.sub( 

47 gammaln(N_plus_1), lax.add(gammaln(k_plus_1), gammaln(lax.sub(N_plus_1, k))) 

48 ) 

49 ) 

50 

51 

52def complex_to_real_iso_matrix(A): 

53 return jnp.block([[jnp.real(A), -jnp.imag(A)], [jnp.imag(A), jnp.real(A)]]) 

54 

55 

56def real_to_complex_iso_matrix(A): 

57 N = A.shape[0] 

58 return A[: N // 2, : N // 2] + 1j * A[N // 2 :, : N // 2] 

59 

60 

61def complex_to_real_iso_vector(v): 

62 return jnp.block([[jnp.real(v)], [jnp.imag(v)]]) 

63 

64 

65def real_to_complex_iso_vector(v): 

66 N = v.shape[0] 

67 return v[: N // 2, :] + 1j * v[N // 2 :, :] 

68 

69 

70def imag_times_iso_vector(v): 

71 N = v.shape[0] 

72 return jnp.block([[-v[N // 2 :, :]], [v[: N // 2, :]]]) 

73 

74 

75def imag_times_iso_matrix(A): 

76 N = A.shape[0] 

77 Ar = A[: N // 2, : N // 2] 

78 Ai = A[N // 2 :, : N // 2] 

79 return jnp.block([[-Ai, -Ar], [Ar, -Ai]]) 

80 

81 

82def conj_transpose_iso_matrix(A): 

83 N = A.shape[0] 

84 Ar = A[: N // 2, : N // 2].T 

85 Ai = A[N // 2 :, : N // 2].T 

86 return jnp.block([[Ar, Ai], [-Ai, Ar]]) 

87 

88 

89def robust_isscalar(val): 

90 is_scalar = isinstance(val, Number) or np.isscalar(val) 

91 if isinstance(val, Array): 

92 is_scalar = is_scalar or (len(val.shape) == 0) 

93 return is_scalar 

94 

95 

96# ===================================================== 

97 

98# Precision 

99 

100def set_precision(precision: Literal["single", "double"]): 

101 """ 

102 Set the precision of JAX operations. 

103 

104 Args: 

105 precision: 'single' or 'double' 

106 

107 Raises: 

108 ValueError: if precision is not 'single' or 'double' 

109 """ 

110 if precision == "single": 

111 config.update("jax_enable_x64", False) 

112 elif precision == "double": 

113 config.update("jax_enable_x64", True) 

114 else: 

115 raise ValueError("precision must be 'single' or 'double'") 

116 

117 

118# ===================================================== 

119 

120# Sharding 

121 

122 

123def set_default_sharding(sharding): 

124 """Configure the global default ``Sharding`` applied to every Qarray. 

125 

126 Once set, every ``DenseImpl`` and ``SparseDiaImpl`` construction routes 

127 its underlying ``jnp.ndarray`` through ``jax.lax.with_sharding_constraint`` 

128 using *sharding*. ``SparseBCOOImpl`` is unsupported under sharding and 

129 will raise from ``Qarray.create(..., implementation=SPARSE_BCOO)``. 

130 

131 Args: 

132 sharding: Either a ``jax.sharding.Sharding`` (typically 

133 ``NamedSharding(mesh, PartitionSpec(...))``) applied to every 

134 array regardless of rank, or a callable ``(arr) -> Sharding`` 

135 for rank-adaptive partitioning. Pass ``None`` to disable 

136 (equivalent to ``clear_default_sharding()``). 

137 """ 

138 from jaxquantum.core.settings import SETTINGS 

139 SETTINGS["default_sharding"] = sharding 

140 

141 

142def get_default_sharding(): 

143 """Return the configured default sharding, or ``None`` if unset.""" 

144 from jaxquantum.core.settings import SETTINGS 

145 return SETTINGS["default_sharding"] 

146 

147 

148def clear_default_sharding(): 

149 """Disable default sharding (return to single-device behaviour).""" 

150 from jaxquantum.core.settings import SETTINGS 

151 SETTINGS["default_sharding"] = None 

152 

153 

154def set_device_mesh(shape, axis_names, partition_spec=None, devices=None): 

155 """Configure default sharding from a high-level mesh description. 

156 

157 Convenience wrapper around :func:`set_default_sharding` that builds a 

158 ``Mesh`` and ``NamedSharding`` for you. Mirrors the pattern used in 

159 ``experiments/distributed/1-demo.ipynb``. 

160 

161 Args: 

162 shape: Tuple of mesh dimensions, e.g. ``(2,)`` for a 2-device 1D mesh 

163 or ``(2, 4)`` for a 2x4 2D mesh. 

164 axis_names: Tuple of mesh axis names, same length as *shape*, e.g. 

165 ``('dp',)`` or ``('dp', 'mp')``. 

166 partition_spec: Optional ``jax.sharding.PartitionSpec``. If ``None``, 

167 stores a rank-adaptive callable that picks the partition for 

168 each array based on the *name* of each mesh axis: 

169 

170 * Names starting with ``'dp'`` / ``'data'`` → data-parallel: 

171 prefer the leading batch axes ``[0, 1, ..., rank-3]``, then 

172 fall through to matrix axes ``[-2, -1]``. 

173 * Names starting with ``'mp'`` / ``'model'`` → model-parallel: 

174 prefer matrix axes ``[-2, -1]``, then fall through to leading 

175 batch axes. 

176 * Anything else behaves like ``'mp'``. 

177 

178 Each mesh axis is greedy-bound to the first un-claimed array 

179 axis (in priority order) whose size is divisible by the mesh 

180 axis size. Mesh axes that find no binding are unused (the array 

181 replicates along them). This produces e.g. 

182 

183 * ``('dp',)`` on ``(B, N, N)`` → ``P('dp', None, None)`` 

184 (parameter sweep — each device gets a slice of B). 

185 * ``('mp',)`` on ``(N, N)`` → ``P('mp', None)`` (single large 

186 system, matrix-row sharded). 

187 * ``('dp', 'mp')`` on ``(B, N, N)`` → ``P('dp', 'mp', None)`` 

188 (both modes simultaneously). 

189 devices: Optional explicit list of devices. Defaults to 

190 ``jax.devices()``. 

191 

192 Raises: 

193 ValueError: if ``len(shape) != len(axis_names)``. 

194 """ 

195 if len(shape) != len(axis_names): 

196 raise ValueError( 

197 f"shape ({shape}) and axis_names ({axis_names}) must have the " 

198 "same length" 

199 ) 

200 

201 import jax 

202 from math import prod 

203 from jax.experimental import mesh_utils 

204 from jax.sharding import Mesh, NamedSharding, PartitionSpec 

205 

206 if devices is None: 

207 # Slice to the first prod(shape) devices so a 1D mesh works on a 

208 # host with extra devices (e.g. tests run with XLA_FLAGS spoofing 8 

209 # CPUs but the user wants a (2,) mesh). 

210 all_devices = jax.devices() 

211 needed = prod(shape) 

212 if len(all_devices) < needed: 

213 raise ValueError( 

214 f"set_device_mesh(shape={shape}) needs {needed} devices but " 

215 f"only {len(all_devices)} are available." 

216 ) 

217 devices = all_devices[:needed] 

218 

219 mesh_devices = mesh_utils.create_device_mesh(shape, devices=devices) 

220 mesh = Mesh(mesh_devices, axis_names) 

221 

222 if partition_spec is None: 

223 # Rank-adaptive: pick partition per array based on each mesh axis's 

224 # name. 'dp'/'data' prefers leading batch axes; 'mp'/'model' (and 

225 # anything else) prefers matrix axes. Greedy first-fit binding; 

226 # unbound mesh axes leave the array replicated along them. See the 

227 # docstring for the full priority table. 

228 mesh_axis_priorities = [ 

229 (name, _array_axis_priority(name)) for name in axis_names 

230 ] 

231 

232 def _adaptive(arr): 

233 rank = arr.ndim 

234 if rank == 0: 

235 return NamedSharding(mesh, PartitionSpec()) 

236 spec_parts = [None] * rank 

237 used_array_axes: set = set() 

238 for mesh_axis, priority_fn in mesh_axis_priorities: 

239 mesh_size = mesh.shape[mesh_axis] 

240 for array_axis in priority_fn(arr.shape): 

241 if array_axis in used_array_axes: 

242 continue 

243 if arr.shape[array_axis] % mesh_size != 0: 

244 continue 

245 spec_parts[array_axis] = mesh_axis 

246 used_array_axes.add(array_axis) 

247 break 

248 return NamedSharding(mesh, PartitionSpec(*spec_parts)) 

249 

250 set_default_sharding(_adaptive) 

251 else: 

252 set_default_sharding(NamedSharding(mesh, partition_spec)) 

253 

254 

255def _array_axis_priority(mesh_axis_name: str): 

256 """Return a function ``shape -> list[int]`` giving the array-axis 

257 priority order for the given mesh-axis name. 

258 

259 'dp'/'data' → batch dims first (leading axes), then space dims. 

260 'mp'/'model' (or anything else) → space dims first, then batch. 

261 

262 The space-vs-batch split is shape-aware so that 1‑D ``(N,)`` and batched 

263 ``(B, N)`` state vectors shard their Hilbert-space axis under model-parallel 

264 rather than their batch axis: 

265 

266 * rank ≤ 1 → the sole axis is the space axis. 

267 * ``shape[-2] == shape[-1]`` → operator-like; the last two axes are the 

268 matrix (space) dims, the rest are batch. 

269 * otherwise → batched vector; the last axis is the space axis, the rest are 

270 batch. 

271 

272 The only residual ambiguity is a *square* batched vector (``B == N``), which 

273 is indistinguishable from an operator by shape alone and is treated as an 

274 operator. This is rare and never incorrect (sharding stays valid). 

275 """ 

276 name = mesh_axis_name.lower() 

277 is_dp = name.startswith(("dp", "data")) 

278 

279 def priority(shape) -> list: 

280 rank = len(shape) 

281 if rank <= 1: 

282 space = list(range(rank)) 

283 batch = [] 

284 elif shape[-2] == shape[-1]: 

285 space = [rank - 2, rank - 1] 

286 batch = list(range(rank - 2)) 

287 else: 

288 space = [rank - 1] 

289 batch = list(range(rank - 1)) 

290 return (batch + space) if is_dp else (space + batch) 

291 

292 return priority