GPU backend

AMBER’s GPU support has two layers:

1. Native placement (0.4.4, preferred for a single large run) — model.gpu().run() on a vectorized Model (step_vectorized or legacy step). Numeric columns stay device-resident for the run. Switch back with model.cpu(mode="vectorized").run(). Private model-specific loops require approve_fast_path(evidence) and contract="off" (see Snapshot-view contract).

2. Array-module helpers — get_array_module, to_device, to_host, scatter_add for code that writes against xp (CuPy when available, else NumPy).

Requires NVIDIA GPU + CuPy (not Apple Metal/MPS). Install CuPy matching your CUDA toolkit (see Installation and Going faster (lanes)). Polars Lazy engine="gpu" is not AMBER’s agent GPU runtime.

Native placement

import ambr as am

class WealthModel(am.Model):
    def setup(self):
        n = int(self.p.get("n", 10_000))
        self.add_agents(n, wealth=self.rng.integers(1, 10, size=n))

    def step_vectorized(self):
        xp = self.xp
        wealth = self.agents.array("wealth")
        donors = xp.nonzero(wealth > 0)[0]
        if int(donors.size) == 0:
            return
        wealth[donors] -= 1
        recipients = self.rng.choice(
            self.agents.array("id"), size=int(donors.size)
        )
        self.agents.at[recipients].scatter_add(wealth=1)

# Vectorized lane on CPU or GPU
results = WealthModel({"n": 100_000, "steps": 50, "seed": 0}).gpu().run()
# results = WealthModel(...).cpu(mode="vectorized").run()

For Python Agent objects, implement step_oop() and run with model.cpu(mode="oop"). GPU execution does not support the OOP lane.

Mode defaults to vectorized. You can also write model.gpu(mode="vectorized").run() or pass mode= / device= to run() (run overrides fluent placement when both are given).

For many short replicate runs (calibration), use the ensemble path instead: GPU batched ensemble.

Array module

An array-module abstraction over CuPy with a NumPy fallback, so device code is portable: get_array_module, to_device, and to_host resolve to CuPy when a GPU is present and fall back to NumPy when it is not.

Optional GPU backend support for AMBER (CuPy / CUDA only).

AMBER’s array kernels can run on the GPU when CuPy + NVIDIA CUDA are available. Apple Metal/MPS is not used — on Mac, prefer the vectorized CPU path with Numba (see ambr.lanes / ambr.performance).

Design notes

• Device-resident buffers (like the tensor lane, but on GPU): host↔device transfer is the bottleneck (PCIe), so state stays on device across steps.

• CuPy is NumPy-compatible: write kernels against xp = get_array_module().

• This module never hard-requires CuPy; branch on GPU_AVAILABLE.

class ambr.gpu.DeviceRNG(rng)

Bases: object

CuPy RNG with NumPy-compatible choice() for model step() code.

choice(a, size=None, replace=True, p=None)

ambr.gpu.device_name() → str

Human-readable name of the active CUDA device, or ‘cpu’.

ambr.gpu.get_array_module(prefer_gpu: bool = True)

Return the array module to write backend-agnostic kernels against.

xp = get_array_module() yields cupy when a GPU is available, else numpy – so the same kernel code runs on either backend.

ambr.gpu.make_device_rng(seed=None) → DeviceRNG

Create a GPU random generator with NumPy-shaped helpers.

ambr.gpu.require_gpu() → None

Raise a clear error if CuPy/CUDA is not available.

ambr.gpu.scatter_add(base, indices, values)

Scatter-add values into base at indices (1-D NumPy or CuPy arrays).

NumPy arrays use apply_scatter_add() (Numba when installed); CuPy arrays use cupy.add.at(). Returns the array that holds the result (may be a new buffer after a NumPy dtype upcast).

ambr.gpu.synchronize() → None

Block until all queued GPU work has finished (for honest timing).

ambr.gpu.to_device(array)

Move a host array onto the GPU (no-op when no GPU is available).

ambr.gpu.to_host(array)

Move a device array back to host as a NumPy array.