+1 much support

+1

How much code does this generate in the JAXpr/XLA? I am a little worried that this could generate surprisingly large amounts of code.

@shoyer the amount of code is a small constant factor times the logarithm of the sequence length. Months ago, when Dan first showed this to me, he demoed it on a sequence of length 1M with no problem IIRC.

Very nice, I notice now that it only calls _scan recursively exactly once.

Internal tests pass!

Hi,
It would maybe be good to precise that this is the work-efficient implementation of the prefix sum that you implemented, at least if I read the algorithm correctly.
I can do it, but just wanted to check that I wasn't misreading the recursion.
Adrien