jax-bounded-while
Bounded while loop in JAX.
This is a micro-package, containing the single function bounded_while_loop.
Reverse-mode-friendly, bounded while_loop implemented via lax.scan.
Note: This library is being renamed to
jaxmoreand expanded in scope. In addition tobounded_while_loop, it will include more JAX-related functionality — such as avmapthat supports keyword arguments. This will be the last release forjax-bounded-while.
Installation
pip install jax-bounded-while
Examples
Simple loop over a scalar:
import jax.numpy as jnp
from jax_bounded_while import bounded_while_loop
def cond_fn(x):
return x < 5
def body_fn(x):
return x + 1
result = bounded_while_loop(cond_fn, body_fn, jnp.asarray(0), max_steps=10)
print(result) # Array(5, dtype=int32)
PyTree carry (tuple):
import jax.numpy as jnp
from jax_bounded_while import bounded_while_loop
def cond_fn(state):
x, _ = state
return x < 3
def body_fn(state):
x, y = state
return x + 1, y * 2
result = bounded_while_loop(
cond_fn, body_fn, (jnp.asarray(0), jnp.asarray(1)), max_steps=5
)
print(result) # (Array(3, dtype=int32), Array(8, dtype=int32))
1 maintainer
Metadata
1 maintainer
Release files for jax-bounded-while 0.1.2
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| jax_bounded_while-0.1.2.tar.gz | 149.3 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| jax_bounded_while-0.1.2-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 156.1 kB
Release files / jax_bounded_while-0.1.2.tar.gz
| Download URL | jax_bounded_while-0.1.2.tar.gz |
|---|---|
| Size | 149.3 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
4e09c733cdfc6d904d166dc4014bcd768373aa6680d7d568906721df35dde767
|
|
BLAKE2b-256 checksum How to use checksums |
83cd9b6b9b68d0e134774fd9e6127b482586e96d057d728d5a64668a27379b66
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.7
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Feb 26, 2026.
Transparency logRelease files / jax_bounded_while-0.1.2-py3-none-any.whl
| Download URL | jax_bounded_while-0.1.2-py3-none-any.whl |
|---|---|
| Size | 6.9 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
c911bc9c18622ebab7575c5b946270a8b709de99151730ff9d0d93a235a9ec82
|
|
BLAKE2b-256 checksum How to use checksums |
49f9e0aff648bd870bf252e4b913bf0aeaa98a0bdee1622309b6066927b4c4bd
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.7
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Feb 26, 2026.
Transparency log
