Re: [ot][spam][crazy][data] transformer model 'attention' improvement

[k]

I've got my local code working with the mgrid data.  I made at least
two bugs: an incorrect einsum, and dotting with the weights rather
than the exponent.  The mgrid data doesn't test the softmax since each
vector has the same maximum.  Time to figure out how to make random
tensors in jax.