pyfector.backend

Array backend abstraction: seamless NumPy / CuPy switching.

Usage: from pyfector.backend import get_backend, set_device set_device("gpu") # or "cpu" (default) xp = get_backend() # returns numpy or cupy module A = xp.random.randn(100, 50)

All pyfector internals call get_backend() to obtain the active array module, so every matrix operation automatically runs on CPU or GPU.

 1"""
 2Array backend abstraction: seamless NumPy / CuPy switching.
 3
 4Usage:
 5    from pyfector.backend import get_backend, set_device
 6    set_device("gpu")          # or "cpu" (default)
 7    xp = get_backend()         # returns numpy or cupy module
 8    A = xp.random.randn(100, 50)
 9
10All pyfector internals call ``get_backend()`` to obtain the active array
11module, so every matrix operation automatically runs on CPU or GPU.
12"""
13
14from __future__ import annotations
15
16import numpy as np
17from typing import Literal
18
19_DEVICE: Literal["cpu", "gpu"] = "cpu"
20_CUPY_AVAILABLE: bool | None = None
21
22
23def _check_cupy() -> bool:
24    global _CUPY_AVAILABLE
25    if _CUPY_AVAILABLE is None:
26        try:
27            import cupy  # noqa: F401
28            _CUPY_AVAILABLE = True
29        except ImportError:
30            _CUPY_AVAILABLE = False
31    return _CUPY_AVAILABLE
32
33
34def set_device(device: Literal["cpu", "gpu"]) -> None:
35    """Set the compute device globally."""
36    global _DEVICE
37    if device == "gpu" and not _check_cupy():
38        raise ImportError(
39            "CuPy is required for GPU support. Install it with: "
40            "pip install cupy-cuda12x  (adjust for your CUDA version)"
41        )
42    _DEVICE = device
43
44
45def get_device() -> Literal["cpu", "gpu"]:
46    """Return the current device."""
47    return _DEVICE
48
49
50def get_backend():
51    """Return the active array module (numpy or cupy)."""
52    if _DEVICE == "gpu":
53        import cupy
54        return cupy
55    return np
56
57
58def to_numpy(arr) -> np.ndarray:
59    """Ensure *arr* is a host numpy array (move from GPU if needed)."""
60    if _DEVICE == "gpu":
61        import cupy
62        if isinstance(arr, cupy.ndarray):
63            return cupy.asnumpy(arr)
64    return np.asarray(arr)
65
66
67def to_device(arr):
68    """Move a numpy array to the current device."""
69    if _DEVICE == "gpu":
70        import cupy
71        return cupy.asarray(arr)
72    return np.asarray(arr)
73
74
75def make_rng(seed: int | None = None):
76    """Create a seeded host RNG.
77
78    Cross-validation and inference use this generator only for indices
79    and seeds, so keeping it on the host avoids backend-specific RNG
80    differences while preserving reproducibility.
81    """
82    if seed is not None:
83        return np.random.default_rng(seed)
84    return np.random.default_rng()
def set_device(device: Literal['cpu', 'gpu']) -> None: 
35def set_device(device: Literal["cpu", "gpu"]) -> None:
36    """Set the compute device globally."""
37    global _DEVICE
38    if device == "gpu" and not _check_cupy():
39        raise ImportError(
40            "CuPy is required for GPU support. Install it with: "
41            "pip install cupy-cuda12x  (adjust for your CUDA version)"
42        )
43    _DEVICE = device

Set the compute device globally.

def get_device() -> Literal['cpu', 'gpu']: 
46def get_device() -> Literal["cpu", "gpu"]:
47    """Return the current device."""
48    return _DEVICE

Return the current device.

def get_backend(): 
51def get_backend():
52    """Return the active array module (numpy or cupy)."""
53    if _DEVICE == "gpu":
54        import cupy
55        return cupy
56    return np

Return the active array module (numpy or cupy).

def to_numpy(arr) -> numpy.ndarray: 
59def to_numpy(arr) -> np.ndarray:
60    """Ensure *arr* is a host numpy array (move from GPU if needed)."""
61    if _DEVICE == "gpu":
62        import cupy
63        if isinstance(arr, cupy.ndarray):
64            return cupy.asnumpy(arr)
65    return np.asarray(arr)

Ensure arr is a host numpy array (move from GPU if needed).

def to_device(arr): 
68def to_device(arr):
69    """Move a numpy array to the current device."""
70    if _DEVICE == "gpu":
71        import cupy
72        return cupy.asarray(arr)
73    return np.asarray(arr)

Move a numpy array to the current device.

def make_rng(seed: int | None = None): 
76def make_rng(seed: int | None = None):
77    """Create a seeded host RNG.
78
79    Cross-validation and inference use this generator only for indices
80    and seeds, so keeping it on the host avoids backend-specific RNG
81    differences while preserving reproducibility.
82    """
83    if seed is not None:
84        return np.random.default_rng(seed)
85    return np.random.default_rng()

Create a seeded host RNG.

Cross-validation and inference use this generator only for indices and seeds, so keeping it on the host avoids backend-specific RNG differences while preserving reproducibility.