JAX-tqdm
Add a tqdm progress bar to your JAX scans and loops.
The code is explained in this blog post.
Example usage
in jax.lax.scan
from jax_tqdm import scan_tqdm
from jax import lax
import jax.numpy as jnp
n = 10_000
@scan_tqdm(n)
def step(carry, x):
return carry + 1, carry + 1
last_number, all_numbers = lax.scan(step, 0, jnp.arange(n))
in jax.lax.fori_loop
from jax_tqdm import loop_tqdm
from jax import lax
n = 10_000
@loop_tqdm(n)
def step(i, val):
return val + 1
last_number = lax.fori_loop(0, n, step, 0)
1 maintainer
Metadata
1 maintainer
Release files for jax-tqdm 0.1.0
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_tqdm-0.1.0.tar.gz | 2.7 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| jax_tqdm-0.1.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 5.4 kB
Release files / jax_tqdm-0.1.0.tar.gz
| Download URL | jax_tqdm-0.1.0.tar.gz |
|---|---|
| Size | 2.7 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
d293e13e68abdf66b0c593cb514298946262b04b530149ebfd309f10df5450ff
|
|
BLAKE2b-256 checksum How to use checksums |
a686728cec6500eeb18e799ecf6002add603a77859c3ac5589aa30f50b4952b2
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
poetry/1.3.2 CPython/3.10.6 Linux/5.15.0-58-generic
|
Release files / jax_tqdm-0.1.0-py3-none-any.whl
| Download URL | jax_tqdm-0.1.0-py3-none-any.whl |
|---|---|
| Size | 2.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
bb027e1bfc006391d8ca2648e40b5bd891f1544f88a09c006bb9347e66b8ed21
|
|
BLAKE2b-256 checksum How to use checksums |
dd38b9dfece50bcefe25c39dfee63454252f8b65abe47d08d4e0838b05c591d5
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
poetry/1.3.2 CPython/3.10.6 Linux/5.15.0-58-generic
|
