singular vectors will be given by the columns of U and the
      rows of Vh respectively.
Examples:
  >>> A = torch.randn(5, 3)
  >>> U, S, Vh = torch.linalg.svd(A, full_matrices=False)
  >>> U.shape, S.shape, Vh.shape
  (torch.Size([5, 3]), torch.Size([3]), torch.Size([3, 3]))
  >>> torch.dist(A, U @ torch.diag(S) @ Vh)
  tensor(1.0486e-06)

  >>> U, S, Vh = torch.linalg.svd(A)
  >>> U.shape, S.shape, Vh.shape
  (torch.Size([5, 5]), torch.Size([3]), torch.Size([3, 3]))
  >>> torch.dist(A, U[:, :3] @ torch.diag(S) @ Vh)
  tensor(1.0486e-06)

  >>> A = torch.randn(7, 5, 3)
  >>> U, S, Vh = torch.linalg.svd(A, full_matrices=False)
  >>> torch.dist(A, U @ torch.diag_embed(S) @ Vh)
  tensor(3.0957e-06)
