flax

1b2754a9 - Split the attention softmax so that the expensive elementwise division happens on an array that's O(N) rather than O(N^2)