Решение линейных систем уравнений с разложением SVD

Question

Я хочу написать функцию, которая использует декомпозицию SVD для решения системы уравнений ax = b, где a - квадратная матрица, а b - вектор значений. Функция scipy scipy.linalg.svd () должна превратить a в матрицы UW V. U и VI могут просто взять транспонирование, чтобы найти их обратное. Но для W функция дает мне одномерный массив значений, который мне нужен, чтобы сложить диагональ матрицы, а затем ввести одно значение над значением.

def solveSVD(a,b):

    U,s,V=sp.svd(a,compute_uv=True)

    Ui=np.transpose(a)
    Vi=np.transpose(V)

    W=np.diag(s)

    Wi=np.empty(np.shape(W)[0],np.shape(W)[1])
    for i in range(np.shape(Wi)[0]):
        if W[i,i]!=0:
            Wi[i,i]=1/W[i,i]

    ai=np.matmul(Ui,np.matmul(Wi,Vi))
    x=np.matmul(ai,b)

    return(x)

Однако я получаю ошибку «Ошибка типа: тип данных не понят». Я думаю, что часть проблемы в том, что

W=np.diag(s) 

не производит квадратную диагональную матрицу.

Я впервые работаю с этой библиотекой, поэтому извиняюсь, что сделал что-то очень глупое, но не могу понять, почему эта строка не сработала. Спасибо всем!