Feature request: Add an unroll argument (default 1) to functions that use lax.scan.
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.
Feature request: Add an
unrollargument (default1) to functions that uselax.scan.scanis 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.