Numba - guvectorize для многомерных массивов

Question

Я новичок в программировании numba и gpu. Я пытаюсь понять, как выполнить функцию lstsq с помощью cuda. Однако я не знаю, как писать для него код с учетом размеров массивов.

@guvectorize(['float64[:,:,:],float64[:,:,:,:],float64[:,:,:]'],'(l,n,m),(k,l,n,m)->(l,n,m)', target='cuda') 
def lstsq(a,b,c):
    k,l,n,m = b.shape
    for i in range(n):
        for j in range(m):
            c[:,i,j] = cp.linalg.lstsq(b[:,:,i,j],a[:,i,j])[0]