Skip to content

utils

Utils

Ec_to_inv_pF(Ec)

GHz -> 1/picoFarad

Source code in jaxquantum/utils/units.py 
81
82
83
84
85
86
87
88
def Ec_to_inv_pF(Ec):
    """
    GHz -> 1/picoFarad
    """
    joule = GHz_to_joule(Ec)
    Gjoule = joule / 1e9
    inv_nFarad = Gjoule / ((constants.e) ** 2 / (2))
    return inv_nFarad * 1e-3

as_series(*arrs)

Return arguments as a list of 1-d arrays.

The returned list contains array(s) of dtype double, complex double, or object. A 1-d argument of shape (N,) is parsed into N arrays of size one; a 2-d argument of shape (M,N) is parsed into M arrays of size N (i.e., is "parsed by row"); and a higher dimensional array raises a Value Error if it is not first reshaped into either a 1-d or 2-d array.

Parameters

arrs : array_like 1- or 2-d array_like trim : boolean, optional When True, trailing zeros are removed from the inputs. When False, the inputs are passed through intact.

Returns

a1, a2,... : 1-D arrays A copy of the input data as 1-d arrays.

Source code in jaxquantum/utils/hermgauss.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def as_series(*arrs):
    """Return arguments as a list of 1-d arrays.

    The returned list contains array(s) of dtype double, complex double, or
    object.  A 1-d argument of shape ``(N,)`` is parsed into ``N`` arrays of
    size one; a 2-d argument of shape ``(M,N)`` is parsed into ``M`` arrays
    of size ``N`` (i.e., is "parsed by row"); and a higher dimensional array
    raises a Value Error if it is not first reshaped into either a 1-d or 2-d
    array.

    Parameters
    ----------
    arrs : array_like
        1- or 2-d array_like
    trim : boolean, optional
        When True, trailing zeros are removed from the inputs.
        When False, the inputs are passed through intact.

    Returns
    -------
    a1, a2,... : 1-D arrays
        A copy of the input data as 1-d arrays.

    """
    arrays = tuple(jnp.array(a, ndmin=1) for a in arrs)
    arrays = promote_dtypes_inexact(*arrays)
    if len(arrays) == 1:
        return arrays[0]
    return tuple(arrays)

benchmark_jax_function(function, *args, iterations=25, warmup=1, clear_caches=True, include_hlo=False, compare_precision=False, jit_kwargs=None, call_kwargs=None)

Collect synchronized JAX timing, memory, HLO, cost, and precision stats.

Source code in jaxquantum/utils/benchmarking.py 
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
def benchmark_jax_function(
    function: Callable,
    *args,
    iterations: int = 25,
    warmup: int = 1,
    clear_caches: bool = True,
    include_hlo: bool = False,
    compare_precision: bool = False,
    jit_kwargs: Mapping[str, Any] | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Collect synchronized JAX timing, memory, HLO, cost, and precision stats."""
    if compare_precision:
        return benchmark_precision(
            function,
            *args,
            iterations=iterations,
            warmup=warmup,
            clear_caches=clear_caches,
            include_hlo=include_hlo,
            jit_kwargs=jit_kwargs,
            call_kwargs=call_kwargs,
        )
    if iterations < 1 or warmup < 0:
        raise ValueError("iterations must be positive and warmup non-negative")
    return _benchmark_once(
        function,
        args,
        iterations,
        warmup,
        clear_caches,
        include_hlo,
        jit_kwargs,
        dict(call_kwargs or {}),
    )[0]

benchmark_precision(function, *args, iterations=25, warmup=1, clear_caches=True, include_hlo=False, jit_kwargs=None, call_kwargs=None)

Compare double and single precision performance and output accuracy.

Source code in jaxquantum/utils/benchmarking.py 
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
def benchmark_precision(
    function: Callable,
    *args,
    iterations: int = 25,
    warmup: int = 1,
    clear_caches: bool = True,
    include_hlo: bool = False,
    jit_kwargs: Mapping[str, Any] | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Compare double and single precision performance and output accuracy."""
    if iterations < 1 or warmup < 0:
        raise ValueError("iterations must be positive and warmup non-negative")
    original_x64 = jax.config.x64_enabled
    reports = {}
    outputs = {}
    try:
        for name, enabled, real_dtype, complex_dtype in (
            ("double", True, np.float64, np.complex128),
            ("single", False, np.float32, np.complex64),
        ):
            jax.config.update("jax_enable_x64", enabled)
            precision_args = _cast_precision(args, real_dtype, complex_dtype)
            precision_kwargs = _cast_precision(
                dict(call_kwargs or {}),
                real_dtype,
                complex_dtype,
            )
            reports[name], output = _benchmark_once(
                function,
                precision_args,
                iterations,
                warmup,
                clear_caches,
                include_hlo,
                jit_kwargs,
                precision_kwargs,
            )
            outputs[name] = jax.device_get(output)
            del output
    finally:
        jax.config.update("jax_enable_x64", original_x64)
        if clear_caches:
            jax.clear_caches()

    double_timing = reports["double"]["timings_s"]
    single_timing = reports["single"]["timings_s"]
    double_memory = reports["double"]["memory_bytes"]
    single_memory = reports["single"]["memory_bytes"]
    return {
        "double": reports["double"],
        "single": reports["single"],
        "accuracy": _accuracy_stats(outputs["double"], outputs["single"]),
        "single_vs_double": {
            "cold_speedup": _ratio(
                double_timing["cold_total"],
                single_timing["cold_total"],
            ),
            "warm_speedup": _ratio(
                double_timing["warm_median"],
                single_timing["warm_median"],
            ),
            "temporary_memory_ratio": _ratio(
                double_memory["temp_size_in_bytes"],
                single_memory["temp_size_in_bytes"],
            ),
            "temporary_bytes_saved": (
                _difference(
                    double_memory["temp_size_in_bytes"],
                    single_memory["temp_size_in_bytes"],
                )
            ),
            "peak_memory_ratio": _ratio(
                double_memory["peak_memory_in_bytes"],
                single_memory["peak_memory_in_bytes"],
            ),
            "peak_bytes_saved": (
                _difference(
                    double_memory["peak_memory_in_bytes"],
                    single_memory["peak_memory_in_bytes"],
                )
            ),
        },
    }

block_until_ready(tree)

Synchronize every array leaf in a PyTree.

Source code in jaxquantum/utils/benchmarking.py 
28
29
30
31
32
def block_until_ready(tree) -> None:
    """Synchronize every array leaf in a PyTree."""
    for leaf in jax.tree.leaves(tree):
        if hasattr(leaf, "block_until_ready"):
            leaf.block_until_ready()

clear_default_sharding()

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

Source code in jaxquantum/utils/utils.py 
148
149
150
151
def clear_default_sharding():
    """Disable default sharding (return to single-device behaviour)."""
    from jaxquantum.core.settings import SETTINGS
    SETTINGS["default_sharding"] = None

comb(N, k)

NCk

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

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

Parameters:

Name Type Description Default
N

total items

 required
k

of items to choose

 required

Returns:

Name Type Description
NCk

N choose k
Source code in jaxquantum/utils/utils.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def comb(N, k):
    """
    NCk

    #TODO: replace with jsp.special.comb once issue is closed:
    https://github.com/google/jax/issues/9709

    Args:
        N: total items
        k: # of items to choose

    Returns:
        NCk: N choose k
    """
    one = 1
    N_plus_1 = lax.add(N, one)
    k_plus_1 = lax.add(k, one)
    return lax.exp(
        lax.sub(
            gammaln(N_plus_1), lax.add(gammaln(k_plus_1), gammaln(lax.sub(N_plus_1, k)))
        )
    )

get_default_sharding()

Return the configured default sharding, or None if unset.

Source code in jaxquantum/utils/utils.py 
142
143
144
145
def get_default_sharding():
    """Return the configured default sharding, or ``None`` if unset."""
    from jaxquantum.core.settings import SETTINGS
    return SETTINGS["default_sharding"]

hermcompanion(c)

Return the scaled companion matrix of c.

The basis polynomials are scaled so that the companion matrix is symmetric when c is an Hermite basis polynomial. This provides better eigenvalue estimates than the unscaled case and for basis polynomials the eigenvalues are guaranteed to be real if jax.numpy.linalg.eigvalsh is used to obtain them.

Parameters

c : array_like 1-D array of Hermite series coefficients ordered from low to high degree.

Returns

mat : ndarray Scaled companion matrix of dimensions (deg, deg).

Source code in jaxquantum/utils/hermgauss.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
@jit
def hermcompanion(c):
    """Return the scaled companion matrix of c.

    The basis polynomials are scaled so that the companion matrix is
    symmetric when `c` is an Hermite basis polynomial. This provides
    better eigenvalue estimates than the unscaled case and for basis
    polynomials the eigenvalues are guaranteed to be real if
    `jax.numpy.linalg.eigvalsh` is used to obtain them.

    Parameters
    ----------
    c : array_like
        1-D array of Hermite series coefficients ordered from low to high
        degree.

    Returns
    -------
    mat : ndarray
        Scaled companion matrix of dimensions (deg, deg).

    """
    c = as_series(c)
    if len(c) < 2:
        raise ValueError("Series must have maximum degree of at least 1.")
    if len(c) == 2:
        return jnp.array([[-0.5 * c[0] / c[1]]])

    n = len(c) - 1
    mat = jnp.zeros((n, n), dtype=c.dtype)
    scl = jnp.hstack((1.0, 1.0 / jnp.sqrt(2.0 * jnp.arange(n - 1, 0, -1))))
    scl = jnp.cumprod(scl)[::-1]
    shp = mat.shape
    mat = mat.flatten()
    mat = mat.at[1 :: n + 1].set(jnp.sqrt(0.5 * jnp.arange(1, n)))
    mat = mat.at[n :: n + 1].set(jnp.sqrt(0.5 * jnp.arange(1, n)))
    mat = mat.reshape(shp)
    mat = mat.at[:, -1].add(-scl * c[:-1] / (2.0 * c[-1]))
    return mat

inductance_to_inductive_energy(L)

Convert inductance to inductive energy E_L.

Parameters:

Name Type Description Default
L float

Inductance in nH.

 required

Returns:

Name Type Description
float

Inductive energy in GHz.
Source code in jaxquantum/utils/units.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
def inductance_to_inductive_energy(L):
    """Convert inductance to inductive energy E_L.

    Args:
        L (float): Inductance in nH.

    Returns:
        float: Inductive energy in GHz.
    """

    inv_L = 1e9 / L
    El_joules = inv_L * (FLUX_QUANTUM**2) / (2 * np.pi) ** 2
    return joule_to_GHz(El_joules)

inductive_energy_to_inductance(El)

Convert inductive energy E_L to inductance.

Parameters:

Name Type Description Default
El float

inductive energy in GHz.

 required

Returns:

Name Type Description
float

Inductance in nH.
Source code in jaxquantum/utils/units.py 
44
45
46
47
48
49
50
51
52
53
54
55
def inductive_energy_to_inductance(El):
    """Convert inductive energy E_L to inductance.

    Args:
        El (float): inductive energy in GHz.

    Returns:
        float: Inductance in nH.
    """

    inv_L = GHz_to_joule(El) * (2 * np.pi) ** 2 / (FLUX_QUANTUM**2)
    return 1e9 / inv_L

inv_pF_to_Ec(inv_pfarad)

1/picoFarad -> GHz

Source code in jaxquantum/utils/units.py 
72
73
74
75
76
77
78
def inv_pF_to_Ec(inv_pfarad):
    """
    1/picoFarad -> GHz
    """
    inv_nFarad = inv_pfarad * 1e3
    Gjoule = (constants.e) ** 2 / (2) * inv_nFarad
    return joule_to_GHz(Gjoule * 1e9)

jax_device_memory_stats()

Return allocator statistics reported by each JAX device.

Source code in jaxquantum/utils/benchmarking.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def jax_device_memory_stats() -> dict[str, dict[str, int] | None]:
    """Return allocator statistics reported by each JAX device."""
    output = {}
    for device in jax.devices():
        stats = device.memory_stats()
        output[str(device)] = (
            None
            if stats is None
            else {
                key: int(value)
                for key, value in stats.items()
                if isinstance(value, (int, np.integer))
            }
        )
    return output

jax_hlo(function, *args, jit_kwargs=None, call_kwargs=None)

Return lowered StableHLO text for a function call.

Source code in jaxquantum/utils/benchmarking.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
def jax_hlo(
    function: Callable,
    *args,
    jit_kwargs: Mapping[str, Any] | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
) -> str:
    """Return lowered StableHLO text for a function call."""
    return lower_jax_function(
        function,
        *args,
        jit_kwargs=jit_kwargs,
        call_kwargs=call_kwargs,
    ).as_text()

jax_memory_stats(compiled)

Return XLA's compiled buffer-size estimates.

Source code in jaxquantum/utils/benchmarking.py 
48
49
50
51
52
53
def jax_memory_stats(compiled) -> dict[str, int | None]:
    """Return XLA's compiled buffer-size estimates."""
    memory = compiled.memory_analysis()
    if memory is None:
        return dict.fromkeys(_MEMORY_FIELDS)
    return {field: getattr(memory, field, None) for field in _MEMORY_FIELDS}

lower_jax_function(function, *args, jit_kwargs=None, call_kwargs=None)

Lower a function with the supplied JIT and call arguments.

Source code in jaxquantum/utils/benchmarking.py 
35
36
37
38
39
40
41
42
43
44
45
def lower_jax_function(
    function: Callable,
    *args,
    jit_kwargs: Mapping[str, Any] | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
):
    """Lower a function with the supplied JIT and call arguments."""
    return jax.jit(function, **dict(jit_kwargs or {})).lower(
        *args,
        **dict(call_kwargs or {}),
    )

n_thermal(frequency, temperature)

Calculate the average thermal photon number for a given frequency and temperature.

Parameters:

Name Type Description Default
frequency float

Frequency in GHz.

 required
temperature float

Temperature in Kelvin.

 required

Returns:

Name Type Description
float float

Average thermal photon number.
Source code in jaxquantum/utils/units.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def n_thermal(frequency: float, temperature: float) -> float:
    """Calculate the average thermal photon number for a given frequency and temperature.

    Args:
        frequency (float): Frequency in GHz.
        temperature (float): Temperature in Kelvin.

    Returns:
        float: Average thermal photon number.
    """
    k_B = constants.k  # Boltzmann constant in J/K
    h = constants.h  # Planck constant in J·s

    exponent = h * (frequency * 1e9) / (k_B * temperature)
    n_avg = 1 / (np.exp(exponent) - 1)
    return n_avg

set_default_sharding(sharding)

Configure the global default Sharding applied to every Qarray.

Once set, every DenseImpl and SparseDiaImpl construction routes its underlying jnp.ndarray through jax.lax.with_sharding_constraint using sharding. SparseBCOOImpl is unsupported under sharding and will raise from Qarray.create(..., implementation=SPARSE_BCOO).

Parameters:

Name Type Description Default
sharding

Either a jax.sharding.Sharding (typically NamedSharding(mesh, PartitionSpec(...))) applied to every array regardless of rank, or a callable (arr) -> Sharding for rank-adaptive partitioning. Pass None to disable (equivalent to clear_default_sharding()).

 required
Source code in jaxquantum/utils/utils.py 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def set_default_sharding(sharding):
    """Configure the global default ``Sharding`` applied to every Qarray.

    Once set, every ``DenseImpl`` and ``SparseDiaImpl`` construction routes
    its underlying ``jnp.ndarray`` through ``jax.lax.with_sharding_constraint``
    using *sharding*. ``SparseBCOOImpl`` is unsupported under sharding and
    will raise from ``Qarray.create(..., implementation=SPARSE_BCOO)``.

    Args:
        sharding: Either a ``jax.sharding.Sharding`` (typically
            ``NamedSharding(mesh, PartitionSpec(...))``) applied to every
            array regardless of rank, or a callable ``(arr) -> Sharding``
            for rank-adaptive partitioning. Pass ``None`` to disable
            (equivalent to ``clear_default_sharding()``).
    """
    from jaxquantum.core.settings import SETTINGS
    SETTINGS["default_sharding"] = sharding

set_device_mesh(shape, axis_names, partition_spec=None, devices=None)

Configure default sharding from a high-level mesh description.

Convenience wrapper around :func:set_default_sharding that builds a Mesh and NamedSharding for you. Mirrors the pattern used in experiments/distributed/1-demo.ipynb.

Parameters:

Name Type Description Default
shape

Tuple of mesh dimensions, e.g. (2,) for a 2-device 1D mesh or (2, 4) for a 2x4 2D mesh.

 required
axis_names

Tuple of mesh axis names, same length as shape, e.g. ('dp',) or ('dp', 'mp').

 required
partition_spec

Optional jax.sharding.PartitionSpec. If None, stores a rank-adaptive callable that picks the partition for each array based on the name of each mesh axis:

  • Names starting with 'dp' / 'data' → data-parallel: prefer the leading batch axes [0, 1, ..., rank-3], then fall through to matrix axes [-2, -1].
  • Names starting with 'mp' / 'model' → model-parallel: prefer matrix axes [-2, -1], then fall through to leading batch axes.
  • Anything else behaves like 'mp'.

Each mesh axis is greedy-bound to the first un-claimed array axis (in priority order) whose size is divisible by the mesh axis size. Mesh axes that find no binding are unused (the array replicates along them). This produces e.g.

  • ('dp',) on (B, N, N) → P('dp', None, None) (parameter sweep — each device gets a slice of B).
  • ('mp',) on (N, N) → P('mp', None) (single large system, matrix-row sharded).
  • ('dp', 'mp') on (B, N, N) → P('dp', 'mp', None) (both modes simultaneously).
 None
devices

Optional explicit list of devices. Defaults to jax.devices().

 None

Raises:

Type Description
ValueError

if len(shape) != len(axis_names).
Source code in jaxquantum/utils/utils.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def set_device_mesh(shape, axis_names, partition_spec=None, devices=None):
    """Configure default sharding from a high-level mesh description.

    Convenience wrapper around :func:`set_default_sharding` that builds a
    ``Mesh`` and ``NamedSharding`` for you. Mirrors the pattern used in
    ``experiments/distributed/1-demo.ipynb``.

    Args:
        shape: Tuple of mesh dimensions, e.g. ``(2,)`` for a 2-device 1D mesh
            or ``(2, 4)`` for a 2x4 2D mesh.
        axis_names: Tuple of mesh axis names, same length as *shape*, e.g.
            ``('dp',)`` or ``('dp', 'mp')``.
        partition_spec: Optional ``jax.sharding.PartitionSpec``. If ``None``,
            stores a rank-adaptive callable that picks the partition for
            each array based on the *name* of each mesh axis:

            * Names starting with ``'dp'`` / ``'data'`` → data-parallel:
              prefer the leading batch axes ``[0, 1, ..., rank-3]``, then
              fall through to matrix axes ``[-2, -1]``.
            * Names starting with ``'mp'`` / ``'model'`` → model-parallel:
              prefer matrix axes ``[-2, -1]``, then fall through to leading
              batch axes.
            * Anything else behaves like ``'mp'``.

            Each mesh axis is greedy-bound to the first un-claimed array
            axis (in priority order) whose size is divisible by the mesh
            axis size. Mesh axes that find no binding are unused (the array
            replicates along them). This produces e.g.

            * ``('dp',)`` on ``(B, N, N)`` → ``P('dp', None, None)``
              (parameter sweep — each device gets a slice of B).
            * ``('mp',)`` on ``(N, N)`` → ``P('mp', None)`` (single large
              system, matrix-row sharded).
            * ``('dp', 'mp')`` on ``(B, N, N)`` → ``P('dp', 'mp', None)``
              (both modes simultaneously).
        devices: Optional explicit list of devices. Defaults to
            ``jax.devices()``.

    Raises:
        ValueError: if ``len(shape) != len(axis_names)``.
    """
    if len(shape) != len(axis_names):
        raise ValueError(
            f"shape ({shape}) and axis_names ({axis_names}) must have the "
            "same length"
        )

    import jax
    from math import prod
    from jax.experimental import mesh_utils
    from jax.sharding import Mesh, NamedSharding, PartitionSpec

    if devices is None:
        # Slice to the first prod(shape) devices so a 1D mesh works on a
        # host with extra devices (e.g. tests run with XLA_FLAGS spoofing 8
        # CPUs but the user wants a (2,) mesh).
        all_devices = jax.devices()
        needed = prod(shape)
        if len(all_devices) < needed:
            raise ValueError(
                f"set_device_mesh(shape={shape}) needs {needed} devices but "
                f"only {len(all_devices)} are available."
            )
        devices = all_devices[:needed]

    mesh_devices = mesh_utils.create_device_mesh(shape, devices=devices)
    mesh = Mesh(mesh_devices, axis_names)

    if partition_spec is None:
        # Rank-adaptive: pick partition per array based on each mesh axis's
        # name. 'dp'/'data' prefers leading batch axes; 'mp'/'model' (and
        # anything else) prefers matrix axes. Greedy first-fit binding;
        # unbound mesh axes leave the array replicated along them. See the
        # docstring for the full priority table.
        mesh_axis_priorities = [
            (name, _array_axis_priority(name)) for name in axis_names
        ]

        def _adaptive(arr):
            rank = arr.ndim
            if rank == 0:
                return NamedSharding(mesh, PartitionSpec())
            spec_parts = [None] * rank
            used_array_axes: set = set()
            for mesh_axis, priority_fn in mesh_axis_priorities:
                mesh_size = mesh.shape[mesh_axis]
                for array_axis in priority_fn(arr.shape):
                    if array_axis in used_array_axes:
                        continue
                    if arr.shape[array_axis] % mesh_size != 0:
                        continue
                    spec_parts[array_axis] = mesh_axis
                    used_array_axes.add(array_axis)
                    break
            return NamedSharding(mesh, PartitionSpec(*spec_parts))

        set_default_sharding(_adaptive)
    else:
        set_default_sharding(NamedSharding(mesh, partition_spec))

set_precision(precision)

Set the precision of JAX operations.

Parameters:

Name Type Description Default
precision Literal['single', 'double']

'single' or 'double'

 required

Raises:

Type Description
ValueError

if precision is not 'single' or 'double'
Source code in jaxquantum/utils/utils.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def set_precision(precision: Literal["single", "double"]):
    """
    Set the precision of JAX operations.

    Args:
        precision: 'single' or 'double'

    Raises:
        ValueError: if precision is not 'single' or 'double'
    """
    if precision == "single":
        config.update("jax_enable_x64", False)
    elif precision == "double":
        config.update("jax_enable_x64", True)
    else:
        raise ValueError("precision must be 'single' or 'double'")