I was doing a float64 cumsum over the outer dimension of 16k rows and it took 8.5ms per call, which nearly doubled training time. When I changed it to do it per row in float32 over the inner dimension, that part's time dropped to about 1/7.
Translated from Japanese
commit ee7d4c3