This release is a pre-release and may not be stable for production use.
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.
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file acai_control-0.1.0a1.tar.gz.
File metadata
- Download URL: acai_control-0.1.0a1.tar.gz
- Upload date:
- Size: 15.5 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
5a21a79ffa9fef65e4645d5b0c75b96991693b0dfc3709039dadef29358eac7c
|
|
| MD5 |
dc9512820fddf4d9f245b7e7128d74ad
|
|
| BLAKE2b-256 |
3b6dd61e9160049b1d49e3ff3f0de4259ec52efc4dcc9cf011a91e7478b5439a
|
Provenance
The following attestation bundles were made for acai_control-0.1.0a1.tar.gz:
Publisher:
publish.yml on sequince-dev/acai
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
acai_control-0.1.0a1.tar.gz -
Subject digest:
5a21a79ffa9fef65e4645d5b0c75b96991693b0dfc3709039dadef29358eac7c - Sigstore transparency entry: 2420756874
- Sigstore integration time:
-
Permalink:
sequince-dev/acai@a115b96ead5d3a2019600612ef50b291ea7506fc -
Branch / Tag:
refs/tags/v0.1.0a1 - Owner: https://github.com/sequince-dev
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@a115b96ead5d3a2019600612ef50b291ea7506fc -
Trigger Event:
release
-
Statement type:
File details
Details for the file acai_control-0.1.0a1-py3-none-any.whl.
File metadata
- Download URL: acai_control-0.1.0a1-py3-none-any.whl
- Upload date:
- Size: 8.3 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
cd59a4630dcbe4fd2f8a4f3cba5840617f17dbad5014852f82afc0390ae42ece
|
|
| MD5 |
e12789289bad1eb27ae18733311eb69d
|
|
| BLAKE2b-256 |
70d8bdf5a6a0f2d7eac1dcb5cbdc0bbb27a9b91c9ad5a2904e24ccb32154c9f4
|
Provenance
The following attestation bundles were made for acai_control-0.1.0a1-py3-none-any.whl:
Publisher:
publish.yml on sequince-dev/acai
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
acai_control-0.1.0a1-py3-none-any.whl -
Subject digest:
cd59a4630dcbe4fd2f8a4f3cba5840617f17dbad5014852f82afc0390ae42ece - Sigstore transparency entry: 2420756881
- Sigstore integration time:
-
Permalink:
sequince-dev/acai@a115b96ead5d3a2019600612ef50b291ea7506fc -
Branch / Tag:
refs/tags/v0.1.0a1 - Owner: https://github.com/sequince-dev
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@a115b96ead5d3a2019600612ef50b291ea7506fc -
Trigger Event:
release
-
Statement type:
