CuPy backend

CuPy backend#

CuPy is a near drop-in replacement for NumPy that runs on NVIDIA GPUs via CUDA. Use it when you want GPU acceleration for array-API operations and you don’t need JAX autodiff/JIT.

Installation#

pip install brainunit[cupy]

Requires a working CUDA toolkit; the cupy-cuda12x wheel is pulled in by the extra. If you have CUDA 11, install cupy-cuda11x manually instead.

Graceful import#

If CuPy isn’t installed (most CI runners and laptops without an NVIDIA GPU), the snippets below skip cleanly rather than crashing.

import brainunit as u

try:
    import cupy
    HAVE_CUPY = True
    runtime = cupy.cuda.runtime
    device_id = runtime.getDevice()
    device_name = runtime.getDeviceProperties(device_id)['name']
    if isinstance(device_name, bytes):
        device_name = device_name.decode()
except ImportError:
    HAVE_CUPY = False
    print('cupy not installed; install with: pip install brainunit[cupy]')

if HAVE_CUPY:
    print('CuPy version:', cupy.__version__)
    print('CUDA devices:', runtime.getDeviceCount())
    print('CUDA driver/runtime:', runtime.driverGetVersion(), '/', runtime.runtimeGetVersion())
    print('GPU:', device_name)

print('is_cupy_array on a non-cupy object:', u.is_cupy_array([1, 2, 3]))

CuPy version: 14.1.1
CUDA devices: 1
CUDA driver/runtime: 13030 / 12090
GPU: NVIDIA GeForce RTX 3060 Laptop GPU
is_cupy_array on a non-cupy object: False

Quick start#

if HAVE_CUPY:
    q = u.Quantity(cupy.array([1.0, 2.0, 3.0]), unit=u.meter)
    print(q)
    print('backend =', q.backend)
    print('(q + q).backend =', (q + q).backend)

[1. 2. 3.] m
backend = cupy
(q + q).backend = cupy

Math operations#

brainunit.math dispatches to array_api_compat.cupy, executing on the GPU.

if HAVE_CUPY:
    x = u.Quantity(cupy.linspace(0.0, cupy.pi, 5), unit=u.UNITLESS)
    print(u.math.sin(x))

[0.00000000e+00 7.07106781e-01 1.00000000e+00 7.07106781e-01
 1.22464680e-16]

Converting between backends#

Quantity.to_cupy(device=...) moves the mantissa to the chosen GPU.

if HAVE_CUPY:
    import numpy as np
    q_cpu = u.Quantity(np.array([1.0, 2.0]), unit=u.meter)
    q_gpu = q_cpu.to_cupy(device=0)
    print('mantissa lives on device', q_gpu.mantissa.device)
    # round-trip back to NumPy
    print(q_gpu.to_numpy())

mantissa lives on device <CUDA Device 0>
[1. 2.] m

Requesting the backend explicitly#

If you ask for the CuPy backend when CuPy isn’t installed, brainunit raises BackendError (not a bare ImportError) with the install hint.

When CuPy or a usable CUDA device is unavailable, explicitly selecting the backend raises a BackendError with the installation hint:

from brainunit import BackendError

try:
    with u.using_backend('cupy'):
        u.Quantity([1.0, 2.0], unit=u.meter)
except BackendError as exc:
    print(exc)

Limitations#

CuPy has no autograd. brainunit.autograd is JAX-only.
brainunit.lax and brainunit.sparse are JAX-only.
Move data to NumPy or JAX with .to_numpy() / .to_jax() for those.

CuPy backend

Contents