Jeremie Coullon jeremiecoullon

## Jax_progress_bar_tqdm.py
from tqdm.auto import tqdm
from jax import jit, vmap, grad, random, lax, partial
import jax.numpy as jnp
from jax.experimental import host_callback


def progress_bar_scan(num_samples, message=None):
    "Progress bar for a JAX scan"
    if message is None:
            message = f"Running for {num_samples:,} iterations"

## jax_progress_bar_basic.py

import jax.numpy as jnp
from jax import jit, vmap, grad, random, lax, partial
from jax.experimental import host_callback

# ========
# define progress bar

def _print_consumer(arg, transform):
    iter_num, num_samples = arg
	from tqdm.auto import tqdm
	from jax import jit, vmap, grad, random, lax, partial
	import jax.numpy as jnp
	from jax.experimental import host_callback


	def progress_bar_scan(num_samples, message=None):
	"Progress bar for a JAX scan"
	if message is None:
	message = f"Running for {num_samples:,} iterations"