Metadata-Version: 2.4
Name: acai-control
Version: 0.1.0a1
Summary: Backend-independent array control-flow primitives
Author: sequince-dev
License-Expression: MIT
Classifier: Development Status :: 3 - Alpha
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: array-api-compat>=1.9
Requires-Dist: optree>=0.14
Dynamic: license-file

# 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:

```console
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](https://github.com/metaopt/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:

```python
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`.
