utils

JAX Utils

clear_default_sharding()

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

Source code in jaxquantum/utils/utils.py

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

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

def get_default_sharding():
    """Return the configured default sharding, or ``None`` if unset."""
    from jaxquantum.core.settings import SETTINGS
    return SETTINGS["default_sharding"]

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

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

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

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'")