heinsen_sequence

This looks really nice. Thank you for sharing it on HN! In case you didn't know, you can parallelize the slow Python loop in selective_scan that computes all the x's: x = torch.zeros((b, d_in, n)) for i in range(l): x = deltaA[:, :, i] * x + deltaB_u[:, :, i] ⋮ with only two calls to the PyTorch API. See the examples here: https://github.com/glassroom/heinsen_sequence/blob/main/READ... .[a] You can then compute all the y's with one einsum, instead of l sequential einsums. --- [a] Previous discussion on HN: https://news.ycombinator.com/item?id=38556669

It's worth noting that you can solve these linear recurrences, `x(t) = a(t)x(t-1) + b(t)`, using a single parallel prefix sum where the elements are the input tuples `(a(t), b(t))`. The following operator called `combine` works. It's associative over these tuples; the final result will be the final value of the second component. Altogether, this gives you O(n) work and O(log n) span, but using just a single parallel prefix kernel, which may be more efficient in practice. combine ((a1, b1), (a2, b2)) = (a1 * a2, b1 * a2 + b2) For example, here's a C++ implementation: https://github.com/MPLLang/parallel-ml-bench/blob/main/cpp/l... And, here's an implementation in a functional language: https://github.com/MPLLang/parallel-ml-bench/blob/main/mpl/b... I'm pretty sure this generalizes, too, to abstract multiplications and additions in a

Yes, but this uses only two plan-vanilla parallel scans, shown as a* and b* in the pdf. In the sample code, that's two calls to PyTorch's built-in API. That looks very efficient to me, almost like it shouldn't be possible. But I tested the sample code[a] and it works. Not an expert on this, but I don't think I've ever seen that before. Have you? [a] https://github.com/glassroom/heinsen_sequence