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.