Skip to content

_backend

Array-module backend selection for CPU (numpy) and GPU (cupy).

Kept intentionally minimal: a single resolver returns the appropriate array module, a PRNG factory, and a host-transfer helper. All hot-loop code paths receive an xp module and call xp.linalg.* / xp.asarray directly, so backend selection adds zero per-op overhead.

Functions#

default_rng #

default_rng(seed: int | None) -> np.random.Generator

Return a seeded NumPy Generator.

Always uses NumPy so that the same seed produces identical draws regardless of the device, and avoids CuPy Generator API differences (e.g. missing .normal()).

Source code in src/src_method/utils/_backend.py 
44
45
46
47
48
49
50
51
52
53
def default_rng(
    seed: int | None,
) -> np.random.Generator:
    """Return a seeded NumPy ``Generator``.

    Always uses NumPy so that the same seed produces identical draws
    regardless of the device, and avoids CuPy ``Generator`` API
    differences (e.g. missing ``.normal()``).
    """
    return np.random.default_rng(seed)

get_xp #

get_xp(device: str) -> ModuleType

Return the array module for the requested device.

PARAMETER DESCRIPTION
device

"cpu" for numpy or "gpu" for cupy.

TYPE: str

RETURNS DESCRIPTION
ModuleType

The numpy or cupy module.
RAISES DESCRIPTION
ValueError

If device is not recognised.
ImportError

If device="gpu" but cupy is not installed.
Source code in src/src_method/utils/_backend.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def get_xp(device: str) -> ModuleType:
    """Return the array module for the requested device.

    Args:
        device: ``"cpu"`` for numpy or ``"gpu"`` for cupy.

    Returns:
        The numpy or cupy module.

    Raises:
        ValueError: If ``device`` is not recognised.
        ImportError: If ``device="gpu"`` but cupy is not installed.
    """
    if device == "cpu":
        return np
    if device == "gpu":
        import cupy  # noqa: PLC0415  (lazy: optional dependency)

        return cupy
    msg = f"Unknown device {device!r}; expected 'cpu' or 'gpu'."
    raise ValueError(msg)

to_numpy #

to_numpy(arr: NDArray) -> np.ndarray

Bring an array onto the host as a numpy array (no-op for numpy).

Source code in src/src_method/utils/_backend.py 
56
57
58
59
60
61
62
def to_numpy(arr: NDArray) -> np.ndarray:
    """Bring an array onto the host as a numpy array (no-op for numpy)."""
    if isinstance(arr, np.ndarray):
        return arr
    # cupy.ndarray exposes .get(); fall back to np.asarray for other dispatchers.
    get = getattr(arr, "get", None)
    return get() if callable(get) else np.asarray(arr)