jax

8a62a9b6 - block-unrolled scan primitive implementation (#3738)

block-unrolled scan primitive implementation (#3738) * block-unrolled scan implementation, via optional `_unroll` scan parameter * index statically in the inlined path of lax.scan * make `unroll` a required scan parameter, and test that it unrolls