Add `unroll` argument to functions that use `lax.scan`

**Feature request:** Add an `unroll` argument (default `1`) to functions that use [`lax.scan`](https://docs.jax.dev/en/latest/_autosummary/jax.lax.scan.html).

`scan` is known to be slow on GPUs because it requires a kernel launch on each iteration, unless multiple iterations are unrolled.

This change would allow users to do that, and thus potentially speed up computation.

I can submit a PR for this.