Fixed a memory leak by using Einstein summation.

Still memory leak problems somewhere, but it is working
better now. I suspect that there's a broadcasting that
can be replaced by np.einsum instead.