def svd_lstsq(AA, BB, tol=1e-5):
    U, S, Vh = torch.linalg.svd(AA, full_matrices=False)
    Spinv = torch.zeros_like(S)
    Spinv[S>tol] = 1/S[S>tol]
    UhBB = U.adjoint() @ BB
    if Spinv.ndim!=UhBB.ndim:
      Spinv = Spinv.unsqueeze(-1)
    SpinvUhBB = Spinv * UhBB
    return Vh.adjoint() @ SpinvUhBB

import torch
print(f'Running PyTorch version: {torch.__version__}')

torchdevice = torch.device('cpu')
if torch.cuda.is_available():
  torchdevice = torch.device('cuda')
  print('Default GPU is ' + torch.cuda.get_device_name(torch.device('cuda')))
print('Running on ' + str(torchdevice))

b = 2
r = 5
c = 3
k = 1

if b==1:
  A = torch.randn(r, c, device=torchdevice)
  if k==1:
    B = torch.randn(r, device=torchdevice)
  else:
    B = torch.randn(r, k, device=torchdevice)
else:
  A = torch.randn(b, r, c, device=torchdevice)
  B = torch.randn(b, r, k, device=torchdevice)

# degrade rank
A[...,-1] = 0
print("A",A)

try:
  X_lstsq = torch.linalg.lstsq(A, B).solution
  print("X_lstsq",X_lstsq)
except Exception as error:
  print("An error occurred:", type(error).__name__, "–", error)

X_pinv=torch.linalg.pinv(A) @ B
print("X_pinv",X_pinv)

def svd_lstsq(AA, BB, tol=1e-5):
    U, S, Vh = torch.linalg.svd(AA, full_matrices=False)
    Spinv = torch.zeros_like(S)
    Spinv[S>tol] = 1/S[S>tol]
    UhBB = U.adjoint() @ BB
    if Spinv.ndim!=UhBB.ndim:
      Spinv = Spinv.unsqueeze(-1)
    SpinvUhBB = Spinv * UhBB
    return Vh.adjoint() @ SpinvUhBB

X_svd= svd_lstsq(A, B)
print("X_svd",X_svd)