jbuckman

[BUG] async_reset crashes when called multiple times

3

## Describe the bug When we call async_reset multiple times, we get a crash. The use case here is that I want to run many parallel episodes with the async...

[Feature Request] Turn off auto-reset

1

## Motivation The documentation says auto-reset is enabled by default, but for certain applications it is better to have it turned off. ## Solution When auto-reset is off, terminated environments...

Make bars stay in fixed order, replacing themselves on resume

1

I'm running several nested bars. The inner bars may be run many times for each step of the outer bars. When an inner bar completes, it "stays around", and a...

process can exit before jax.debug.prints, needs jax.effects_barrier()

3

### Description This does not print anything: ```python import jax import jax.numpy as jnp def scanner(carry, x): jax.debug.print("carry={carry} x={x}", carry=carry, x=x) return carry + x, x out = jax.lax.scan(scanner, jnp.array(0.),...

Pallas jax.lax.fori_loop over long inputs slows down

1

### Description Inside Pallas kernels, we often want a loop, and to speed up compilation, we typically use a scan function such as jax.lax.fori_loop. (For example, in the attention kernel...