flax

6ae22681 - Merge pull request #2339 from sanchit-gandhi:add-unroll-to-scan-with-axes