acai

Array Control-flow And Iteration: small, backend-independent control-flow primitives for code written with the Python Array API ecosystem.

Installation

Install acai with:

pip install acai-control

Array backends such as NumPy, CuPy, JAX, and PyTorch are not installed by acai-control. array-api-compat provides namespace inference and backend detection, while Optree provides pytree traversal.

Usage

The acai namespace provides four main functions:

• cond: Apply one of two functions according to a scalar predicate.
• fori_loop: Apply a function over an integer range.
• scan: Scan a function over arrays while carrying state.
• while_loop: Apply a function repeatedly while a condition is true.

The basic usage is:

import numpy as np
from array_api_compat import array_namespace

from acai import cond, fori_loop, scan, while_loop

x = np.asarray(1)
xp = array_namespace(x)

# 1 * 2 * 3 * 4 * 5
factorial = fori_loop(1, 6, lambda i, value: value * i, x, xp=xp)

# The namespace may also be inferred from an array carry.
power_of_two = while_loop(lambda value: value < 16, lambda value: value * 2, x)

# Select one branch without losing staged control flow on JAX.
magnitude = cond(x >= 0, lambda value: value, lambda value: -value, x, xp=xp)

# Accumulate values and collect every intermediate total.
total, cumulative = scan(
    lambda carry, value: (carry + value, carry + value),
    np.asarray(0),
    np.asarray([1, 2, 3]),
)

Implementation details

The loop functions use ordinary Python loops for NumPy, CuPy, and PyTorch. For JAX, they dispatch to jax.lax.while_loop and jax.lax.fori_loop, so loops remain staged under jax.jit. Likewise, cond uses ordinary Python branching except on JAX, where it dispatches to jax.lax.cond. scan uses eager iteration and stacking except on JAX, where it dispatches directly to jax.lax.scan.

PyTorch

PyTorch arrays are supported through eager Python control flow. PyTorch's native structured control-flow primitives are currently prototype features, so acai does not dispatch to them. Consequently, its control flow is not guaranteed to remain structured under torch.compile.

Dask and ndonnx

Dask arrays and ndonnx symbolic arrays are not currently supported. All primitives raise NotImplementedError when given either namespace.

Python Array API standard

acai is not part of the Python Array API standard and is not affiliated with array-api-compat or array-api-extra.