@ankit-maverick Sorry, my fault, of course. But this could be speeded up by broadcasting or np.einsum.