Array backends

Choosing an array backend

Sequince uses the Python Array API through array-api-compat to work with NumPy, JAX, PyTorch, and CuPy arrays. The initial particle array selects the backend; there is no separate global backend setting. Install the array library you intend to use:

• NumPy is included in the core installation.

• For JAX, install sequince[jax]. See the JAX installation guide for accelerator-specific packages.

• For PyTorch, follow the PyTorch installation selector for your platform and compute device.

• For CuPy, follow the CuPy installation guide to select a package compatible with your CUDA environment.

PyTorch and CuPy are installed separately; Sequince does not provide extras for them. Backend support also depends on the steps and mutation algorithms in your pipeline. In particular, selecting JAX arrays alone does not compile a run; see JAX and functional execution for functional execution and compilation requirements.

Enabling SciPy Array API support

SciPy’s experimental Array API dispatch must be enabled before SciPy is imported. Set the environment variable before starting Python:

export SCIPY_ARRAY_API=1
python your_script.py

Alternatively, set it at the beginning of your script, before importing SciPy or packages that import SciPy:

import os

os.environ["SCIPY_ARRAY_API"] = "1"

import sequince

Sequince sets SCIPY_ARRAY_API=1 during import if it is unset. It preserves an existing value and warns if that value differs from "1". Setting it explicitly avoids import-order problems. If SciPy has already been imported in a notebook, restart the kernel and set the variable before importing it. See SciPy’s Array API guide for details.

Writing backend-compatible targets

Create particles with your chosen library, and keep target evaluations in the same array namespace. The following target works with each of the backends above:

from sequince.array_api import namespace


def log_likelihood(particles, context):
    xp = namespace(particles)
    return -0.5 * xp.sum(particles**2, axis=-1)

Keep related arrays on compatible devices and use floating-point particles. Avoid converting particles to NumPy inside a target when using another backend. For JAX runs requiring 64-bit precision, call jax.config.update("jax_enable_x64", True) before creating arrays.