Skip to content

backend

sorix.backend

Backend utilities for Sorix

Provides a thin abstraction over NumPy and CuPy, selecting the appropriate module based on a Tensor's device and CuPy availability.

get_xp 

get_xp(*args)

Return the array module (NumPy or CuPy) to use for args.

This is the single array-module selector used across the library: it is resolved once per operation instead of scattering device == 'cuda' conditionals through every call site.

Parameters:

  • *args (Any, default: () ) –

    Objects to inspect, typically :class:sorix.tensor.Tensor instances. Anything without a CUDA device is ignored, so raw NumPy arrays and scalars can be passed freely.

Returns:

  • Any ( Any ) –

    cp if any argument lives on a CUDA device and CuPy is

  • Any –

    available; otherwise np.

Source code in sorix/backend.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def get_xp(*args: Any) -> Any:
    """Return the array module (NumPy or CuPy) to use for ``args``.

    This is the single array-module selector used across the library: it is
    resolved once per operation instead of scattering ``device == 'cuda'``
    conditionals through every call site.

    Args:
        *args: Objects to inspect, typically :class:`sorix.tensor.Tensor`
            instances. Anything without a CUDA ``device`` is ignored, so raw
            NumPy arrays and scalars can be passed freely.

    Returns:
        Any: ``cp`` if any argument lives on a CUDA device and CuPy is
        available; otherwise ``np``.
    """
    if _cupy_available:
        for arg in args:
            device = getattr(arg, "device", None)
            if getattr(device, "type", None) == "cuda":
                return cp
    return np