Remove loop for grad_lnlikelihood calculation

The result (using einsum) is cleaner and perhaps slightly faster.