jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, kernel, channels, img_dim, device):
        self.img_dim = img_dim
        self.channels = channels
        # build 1D conv matrix
        H_small = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel.shape[0] // 2, i + kernel.shape[0] // 2):
                if j < 0 or j >= img_dim: continue
                H_small[i, j] = kernel[j - i + kernel.shape[0] // 2]
        # torch.cholesky(H_small) # raise error if H_small is not positive definite
        self.H_small = H_small # positive definite matrix
        # get the evd of the 1D conv
        self.U_small, self.singulars_small, _ = torch.linalg.svd(H_small, full_matrices=False)
        self.V_small = self.U_small # V should equal U since H is symmetric
        
        ZERO = 3e-2
        # truncate
        self.singulars_small[self.singulars_small < ZERO] = ZERO
        # calculate the singular values of the big matrix using Kronecker product
        self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1),
                                       self.singulars_small.reshape(1, img_dim)).reshape(img_dim ** 2)
        self._singulars[self._singulars > 1] = 1
        print("contains zero?", torch.any(self._singulars == 0))
        # raise NotImplementedError
        # sort the big matrix singulars and save the permutation
        self._singulars, self._perm = self._singulars.sort(descending=True)
//  [python]
def align_compressed_coils(self):
        """Virtual Coil Alignment."""
        self.aligned_data = self.unaligned_data

        _, sy, nc = self.aligned_data.shape
        ncc = sy

        n0 = nc // 2

        A00 = self.aligned_data[:, :ncc, n0 - 1]

        A0 = A00
        for n in range(n0, 0, -1):
            A1 = self.aligned_data[:, :ncc, n - 1]
            C = torch.conj(A1).T @ A0
            u, _, v = torch.svd(C, some=False)
            P = v @ torch.conj(u).T
            self.aligned_data[:, :ncc, n - 1] = A1 @ torch.conj(P).T
            A0 = self.aligned_data[:, :ncc, n - 1]

        A0 = A00
        for n in range(n0, nc):
            A1 = self.aligned_data[:, :ncc, n]
            C = torch.conj(A1).T @ A0
            u, _, v = torch.svd(C, some=False)
            P = v @ torch.conj(u).T
            self.aligned_data[:, :ncc, n] = A1 @ torch.conj(P).T
            A0 = self.aligned_data[:, :ncc, n]
//  [python]
def align_compressed_coils(self):
        """Virtual Coil Alignment."""
        self.aligned_data = self.unaligned_data

        _, sy, nc = self.aligned_data.shape
        ncc = sy

        n0 = nc // 2

        A00 = self.aligned_data[:, :ncc, n0 - 1]

        A0 = A00
        for n in range(n0, 0, -1):
            A1 = self.aligned_data[:, :ncc, n - 1]
            C = torch.conj(A1).T @ A0
            u, _, v = torch.linalg.svd(C, full_matrices=False)
            P = v @ torch.conj(u).T
            self.aligned_data[:, :ncc, n - 1] = A1 @ torch.conj(P).T
            A0 = self.aligned_data[:, :ncc, n - 1]

        A0 = A00
        for n in range(n0, nc):
            A1 = self.aligned_data[:, :ncc, n]
            C = torch.conj(A1).T @ A0
            u, _, v = torch.linalg.svd(C, full_matrices=False)
            P = v @ torch.conj(u).T
            self.aligned_data[:, :ncc, n] = A1 @ torch.conj(P).T
            A0 = self.aligned_data[:, :ncc, n]
//  [python]
def calculate_gcc(self):
        """Calculates Geometric Coil-Compression."""
        ws = (self.virtual_coils // 2) * 2 + 1

        Nx, Ny, Nz, Nc = self.data.shape

        im = torch.view_as_complex(
            fft.ifft2(torch.view_as_real(self.data), self.fft_centered, self.fft_normalization, spatial_dims=0)
        )

        s = torch.as_tensor([Nx + ws - 1, Ny, Nz, Nc])
        idx = [
            torch.arange(
                abs(int(im.shape[n] // 2 + torch.ceil((-s[n] / 2).clone().detach()))),
                abs(int(im.shape[n] // 2 + torch.ceil((s[n] / 2).clone().detach())) + 1),
            )
            for n in range(len(s))
        ]

        zpim = torch.zeros((Nx + ws - 1, Ny, Nz, Nc)).type(im.dtype)
        zpim[idx[0][0] : idx[0][-1], idx[1][0] : idx[1][-1], idx[2][0] : idx[2][-1], idx[3][0] : idx[3][-1]] = im

        self.unaligned_data = torch.zeros((Nc, min(Nc, ws * Ny * Nz), Nx)).type(im.dtype)
        for n in range(Nx):
            tmpc = utils.reshape_fortran(zpim[n : n + ws, :, :, :], (ws * Ny * Nz, Nc))
            _, _, v = torch.svd(tmpc, some=False)
            self.unaligned_data[:, :, n] = v

        self.unaligned_data = self.unaligned_data[:, : self.virtual_coils, :]
//  [python]
def calculate_gcc(self):
        """Calculates Geometric Coil-Compression."""
        ws = (self.virtual_coils // 2) * 2 + 1

        Nx, Ny, Nz, Nc = self.data.shape

        im = torch.view_as_complex(
            fft.ifft2(torch.view_as_real(self.data), self.fft_centered, self.fft_normalization, spatial_dims=0)
        )

        s = torch.as_tensor([Nx + ws - 1, Ny, Nz, Nc])
        idx = [
            torch.arange(
                abs(int(im.shape[n] // 2 + torch.ceil((-s[n] / 2).clone().detach()))),
                abs(int(im.shape[n] // 2 + torch.ceil((s[n] / 2).clone().detach())) + 1),
            )
            for n in range(len(s))
        ]

        zpim = torch.zeros((Nx + ws - 1, Ny, Nz, Nc)).type(im.dtype)
        zpim[idx[0][0] : idx[0][-1], idx[1][0] : idx[1][-1], idx[2][0] : idx[2][-1], idx[3][0] : idx[3][-1]] = im

        self.unaligned_data = torch.zeros((Nc, min(Nc, ws * Ny * Nz), Nx)).type(im.dtype)
        for n in range(Nx):
            tmpc = utils.reshape_fortran(zpim[n : n + ws, :, :, :], (ws * Ny * Nz, Nc))
            U, S, V = torch.linalg.svd(tmpc, full_matrices=False)
            self.unaligned_data[:, :, n] = V

        self.unaligned_data = self.unaligned_data[:, : self.virtual_coils, :]
//  [python]
def kabsch_torch(X, Y, cpu=True):
    # Assumes X,Y are both (N_points x Dims). See below for wrapper.
    device = X.device
    X, Y = X.transpose(0, 1), Y.transpose(0, 1)

    #  center X and Y to the origin
    X_ = X - X.mean(dim=-1, keepdim=True)
    Y_ = Y - Y.mean(dim=-1, keepdim=True)
    # calculate convariance matrix (for each prot in the batch)
    C = torch.matmul(X_, Y_.t()).detach()
    if cpu:
        C = C.cpu()

    # Optimal rotation matrix via SVD
    if int(torch.__version__.split(".")[1]) < 8:
        # warning! int torch 1.<8 : W must be transposed
        V, S, W = torch.svd(C)
        W = W.t()
    else:
        V, S, W = torch.linalg.svd(C)

    # determinant sign for direction correction
    d = (torch.det(V) * torch.det(W)) < 0.0
    if d:
        S[-1]    = S[-1] * (-1)
        V[:, -1] = V[:, -1] * (-1)
    # Create Rotation matrix U
    U = torch.matmul(V, W).to(device)
    # calculate rotations
    X_ = torch.matmul(X_.t(), U).t()
    # return centered and aligned
    return X_.transpose(0, 1), Y_.transpose(0, 1)
//  [python]
def kabsch_torch(X, Y, cpu=True):
    # Assumes X,Y are both (N_points x Dims). See below for wrapper.
    device = X.device
    X, Y = X.transpose(0, 1), Y.transpose(0, 1)

    #  center X and Y to the origin
    X_ = X - X.mean(dim=-1, keepdim=True)
    Y_ = Y - Y.mean(dim=-1, keepdim=True)
    # calculate convariance matrix (for each prot in the batch)
    C = torch.matmul(X_, Y_.t()).detach()
    if cpu:
        C = C.cpu()

    # Optimal rotation matrix via SVD
    V, S, W = torch.linalg.svd(C)

    # determinant sign for direction correction
    d = (torch.det(V) * torch.det(W)) < 0.0
    if d:
        S[-1]    = S[-1] * (-1)
        V[:, -1] = V[:, -1] * (-1)
    # Create Rotation matrix U
    U = torch.matmul(V, W).to(device)
    # calculate rotations
    X_ = torch.matmul(X_.t(), U).t()
    # return centered and aligned
    return X_.transpose(0, 1), Y_.transpose(0, 1)
//  [python]
def BNM(src, tar):
    """ Batch nuclear-norm maximization, CVPR 2020.
    tar: a tensor, softmax target output.
    NOTE: this does not require source domain data.
    """
    _, out, _ = torch.svd(tar)
    loss = -torch.mean(out)
    return loss
//  [python]
def BNM(src, tar):
    """ Batch nuclear-norm maximization, CVPR 2020.
    tar: a tensor, softmax target output.
    NOTE: this does not require source domain data.
    """
    _, out, _ = torch.linalg.svd(tar)
    loss = -torch.mean(out)
    return loss
//  [python]
def gmm_register(pi_s, mu_s, mu_t, sigma_t):
    '''
    Inputs:
        pi: B x J
        mu: B x J x 3
        sigma: B x J x 3 x 3
    '''
    c_s = pi_s.unsqueeze(1) @ mu_s
    c_t = pi_s.unsqueeze(1) @ mu_t
    Ms = torch.sum((pi_s.unsqueeze(2) * (mu_s - c_s)).unsqueeze(3) @
                   (mu_t - c_t).unsqueeze(2) @ sigma_t.inverse(), dim=1)
    U, _, V = torch.svd(Ms.cpu())
    U = U.cuda()
    V = V.cuda()
    S = torch.eye(3).unsqueeze(0).repeat(U.shape[0], 1, 1).to(U.device)
    S[:, 2, 2] = torch.det(V @ U.transpose(1, 2))
    R = V @ S @ U.transpose(1, 2)
    t = c_t.transpose(1, 2) - R @ c_s.transpose(1, 2)
    bot_row = torch.Tensor([[[0, 0, 0, 1]]]).repeat(R.shape[0], 1, 1).to(R.device)
    T = torch.cat([torch.cat([R, t], dim=2), bot_row], dim=1)
    return T
//  [python]
def gmm_register(pi_s, mu_s, mu_t, sigma_t):
    '''
    Inputs:
        pi: B x J
        mu: B x J x 3
        sigma: B x J x 3 x 3
    '''
    c_s = pi_s.unsqueeze(1) @ mu_s
    c_t = pi_s.unsqueeze(1) @ mu_t
    Ms = torch.sum((pi_s.unsqueeze(2) * (mu_s - c_s)).unsqueeze(3) @
                   (mu_t - c_t).unsqueeze(2) @ sigma_t.inverse(), dim=1)
    U, S, V = torch.linalg.svd(Ms.cpu())
    U = U.cuda()
    V = V.cuda()
    S = torch.eye(3).unsqueeze(0).repeat(U.shape[0], 1, 1).to(U.device)
    S[:, 2, 2] = torch.det(V @ U.transpose(1, 2))
    R = V @ S @ U.transpose(1, 2)
    t = c_t.transpose(1, 2) - R @ c_s.transpose(1, 2)
    bot_row = torch.Tensor([[[0, 0, 0, 1]]]).repeat(R.shape[0], 1, 1).to(R.device)
    T = torch.cat([torch.cat([R, t], dim=2), bot_row], dim=1)
    return T
//  [python]
def is_minimum(x, energy, threshold):
    x = x.clone()
    x = x.requires_grad_(True)
    e = energy(x)
    J = torch.autograd.grad(e.sum(), x, create_graph=True, retain_graph=True)[0]
    jtol = J.pow(2).sum(dim=-1)
    necessary = jtol < threshold
    outs =[J[:, i].sum() for i in range(J.shape[-1])]
    H = [torch.autograd.grad(out, x, create_graph=True, retain_graph=True)[0] for out in outs]
    H = torch.stack(H, dim=-1)
    _, s, _ = torch.svd(H, compute_uv=False)
    sel = s.min(-1)[0] > 1e-9
    sufficient = torch.zeros_like(necessary).bool()
    eig = torch.symeig(H[sel])[0]
    sufficient[sel] = torch.all((eig > 0), -1)
    return necessary & sufficient
//  [python]
def is_minimum(x, energy, threshold):
    x = x.clone()
    x = x.requires_grad_(True)
    e = energy(x)
    J = torch.autograd.grad(e.sum(), x, create_graph=True, retain_graph=True)[0]
    jtol = J.pow(2).sum(dim=-1)
    necessary = jtol < threshold
    outs =[J[:, i].sum() for i in range(J.shape[-1])]
    H = [torch.autograd.grad(out, x, create_graph=True, retain_graph=True)[0] for out in outs]
    H = torch.stack(H, dim=-1)
    _, s, _ = torch.linalg.svd(H)
    sel = s.min(-1)[0] > 1e-9
    sufficient = torch.zeros_like(necessary).bool()
    eig = torch.symeig(H[sel])[0]
    sufficient[sel] = torch.all((eig > 0), -1)
    return necessary & sufficient
//  [python]
def compute_rigid_transform(a: torch.Tensor, b: torch.Tensor, weights: torch.Tensor):
	"""Compute rigid transforms between two point sets

	Args:
		a (torch.Tensor): (B, M, 3) points
		b (torch.Tensor): (B, N, 3) points
		weights (torch.Tensor): (B, M)

	Returns:
		Transform T (B, 3, 4) to get from a to b, i.e. T*a = b
	"""

	weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
	centroid_a = torch.sum(a * weights_normalized, dim=1)
	centroid_b = torch.sum(b * weights_normalized, dim=1)
	a_centered = a - centroid_a[:, None, :]
	b_centered = b - centroid_b[:, None, :]
	cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)

	# Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]
	# and choose based on determinant to avoid flips
	u, s, v = torch.svd(cov, some=False, compute_uv=True)
	rot_mat_pos = v @ u.transpose(-1, -2)
	v_neg = v.clone()
	v_neg[:, :, 2] *= -1
	rot_mat_neg = v_neg @ u.transpose(-1, -2)
	rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)
	assert torch.all(torch.det(rot_mat) > 0)

	# Compute translation (uncenter centroid)
	translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]

	transform = torch.cat((rot_mat, translation), dim=2)
	return transform
//  [python]
def compute_rigid_transform(a: torch.Tensor, b: torch.Tensor, weights: torch.Tensor):
	"""Compute rigid transforms between two point sets

	Args:
		a (torch.Tensor): (B, M, 3) points
		b (torch.Tensor): (B, N, 3) points
		weights (torch.Tensor): (B, M)

	Returns:
		Transform T (B, 3, 4) to get from a to b, i.e. T*a = b
	"""

	weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
	centroid_a = torch.sum(a * weights_normalized, dim=1)
	centroid_b = torch.sum(b * weights_normalized, dim=1)
	a_centered = a - centroid_a[:, None, :]
	b_centered = b - centroid_b[:, None, :]
	cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)

	# Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]
	# and choose based on determinant to avoid flips
	u, s, v = torch.linalg.svd(cov)
	rot_mat_pos = v @ u.transpose(-1, -2)
	v_neg = v.clone()
	v_neg[:, :, 2] *= -1
	rot_mat_neg = v_neg @ u.transpose(-1, -2)
	rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)
	assert torch.all(torch.det(rot_mat) > 0)

	# Compute translation (uncenter centroid)
	translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]

	transform = torch.cat((rot_mat, translation), dim=2)
	return transform
//  [python]
def rank(m, tol=1e-14):
    """
    Compute the rank of a matrix
    :param m: Matrix
    """
    # SVD on M
    Um, Sm, _ = torch.svd(m)

    # How many sv above threshold
    return int(torch.sum(1.0 * (Sm > tol)))
//  [python]
def rank(m, tol=1e-14):
    """
    Compute the rank of a matrix
    :param m: Matrix
    """
    # SVD on M
    Um, Sm, _ = torch.linalg.svd(m)

    # How many sv above threshold
    return int(torch.sum(1.0 * (Sm > tol)))
//  [python]
def svd_orthogonalization(lyr):
	r"""Applies regularization to the training by performing the
	orthogonalization technique described in the paper "An Analysis and Implementation of
	the FFDNet Image Denoising Method." Tassano et al. (2019).
	For each Conv layer in the model, the method replaces the matrix whose columns
	are the filters of the layer by new filters which are orthogonal to each other.
	This is achieved by setting the singular values of a SVD decomposition to 1.

	This function is to be called by the torch.nn.Module.apply() method,
	which applies svd_orthogonalization() to every layer of the model.
	"""
	classname = lyr.__class__.__name__
	if classname.find('Conv') != -1:
		weights = lyr.weight.data.clone()
		c_out, c_in, f1, f2 = weights.size()
		dtype = lyr.weight.data.type()

		# Reshape filters to columns
		# From (c_out, c_in, f1, f2)  to (f1*f2*c_in, c_out)
		weights = weights.permute(2, 3, 1, 0).contiguous().view(f1*f2*c_in, c_out)

		try:
			# SVD decomposition and orthogonalization
			mat_u, _, mat_v = torch.svd(weights)
			weights = torch.mm(mat_u, mat_v.t())

			lyr.weight.data = weights.view(f1, f2, c_in, c_out).permute(3, 2, 0, 1).contiguous().type(dtype)
		except:
			pass
	else:
		pass
//  [python]
def svd_orthogonalization(lyr):
	r"""Applies regularization to the training by performing the
	orthogonalization technique described in the paper "An Analysis and Implementation of
	the FFDNet Image Denoising Method." Tassano et al. (2019).
	For each Conv layer in the model, the method replaces the matrix whose columns
	are the filters of the layer by new filters which are orthogonal to each other.
	This is achieved by setting the singular values of a SVD decomposition to 1.

	This function is to be called by the torch.nn.Module.apply() method,
	which applies svd_orthogonalization() to every layer of the model.
	"""
	classname = lyr.__class__.__name__
	if classname.find('Conv') != -1:
		weights = lyr.weight.data.clone()
		c_out, c_in, f1, f2 = weights.size()
		dtype = lyr.weight.data.type()

		# Reshape filters to columns
		# From (c_out, c_in, f1, f2)  to (f1*f2*c_in, c_out)
		weights = weights.permute(2, 3, 1, 0).contiguous().view(f1*f2*c_in, c_out)

		try:
			# SVD decomposition and orthogonalization
			mat_u, _, mat_v = torch.linalg.svd(weights)
			weights = torch.mm(mat_u, mat_v.t())

			lyr.weight.data = weights.view(f1, f2, c_in, c_out).permute(3, 2, 0, 1).contiguous().type(dtype)
		except:
			pass
	else:
		pass
//  [python]
def fit_latent(self, Y, init='pca', num_iters=100, tolerance=1e-5, optimizer=None, verbose=False):
        """Fit the given data by maximizing the marginal likelihood.

        Args:
            Y (torch.Tensor, np.array): observational data.
            init (str, torch.nn.Module): initialization scheme or back-constraint function. If a string, it is the
                initialization scheme. Currently, you can choose between 'pca' and 'random'. If torch.nn.Module, it
                is the back-constraint function which has some parameters to optimize.
            num_iters (int): number of iterations to optimize the marginal likelihood.
            tolerance (float): termination tolerance on function value/parameter changes (default: 1e-5).
            optimizer (None, torch.optim.Optimizer): optimizer to use to optimize the marginal likelihood. By default,
                it will use L-BFGS.
            verbose (bool): if True, it will print information during the optimization.
        """
        # convert observation data to torch.tensor
        Y = self._convert_to_torch(Y)

        # initialize the latent space
        if isinstance(init, torch.nn.Module):  # back-constraint function
            raise NotImplementedError
        elif isinstance(init, str):
            init = init.lower()
            if init == 'random':
                X = torch.rand(Y.shape[0], self.latent_dim)
            elif init == 'pca':
                X = Y - torch.mean(Y, dim=0).expand_as(Y)
                U, S, V = torch.svd(X)
                X = torch.mm(X, V[:, :self.latent_dim])
            else:
                raise NotImplementedError("Currently, only the 'random' or 'pca' initialization schemes have been "
                                          "implemented")
            X.requires_grad = True
            X = torch.nn.Parameter(X)

        self.X = X

        # fit the data
        self.fit(X, Y, num_iters=num_iters, tolerance=tolerance, optimizer=optimizer, verbose=verbose)
//  [python]
def fit_latent(self, Y, init='pca', num_iters=100, tolerance=1e-5, optimizer=None, verbose=False):
        """Fit the given data by maximizing the marginal likelihood.

        Args:
            Y (torch.Tensor, np.array): observational data.
            init (str, torch.nn.Module): initialization scheme or back-constraint function. If a string, it is the
                initialization scheme. Currently, you can choose between 'pca' and 'random'. If torch.nn.Module, it
                is the back-constraint function which has some parameters to optimize.
            num_iters (int): number of iterations to optimize the marginal likelihood.
            tolerance (float): termination tolerance on function value/parameter changes (default: 1e-5).
            optimizer (None, torch.optim.Optimizer): optimizer to use to optimize the marginal likelihood. By default,
                it will use L-BFGS.
            verbose (bool): if True, it will print information during the optimization.
        """
        # convert observation data to torch.tensor
        Y = self._convert_to_torch(Y)

        # initialize the latent space
        if isinstance(init, torch.nn.Module):  # back-constraint function
            raise NotImplementedError
        elif isinstance(init, str):
            init = init.lower()
            if init == 'random':
                X = torch.rand(Y.shape[0], self.latent_dim)
            elif init == 'pca':
                X = Y - torch.mean(Y, dim=0).expand_as(Y)
                U, S, V = torch.linalg.svd(X)
                X = torch.mm(X, V[:, :self.latent_dim])
            else:
                raise NotImplementedError("Currently, only the 'random' or 'pca' initialization schemes have been "
                                          "implemented")
            X.requires_grad = True
            X = torch.nn.Parameter(X)

        self.X = X

        # fit the data
        self.fit(X, Y, num_iters=num_iters, tolerance=tolerance, optimizer=optimizer, verbose=verbose)
//  [python]
def EPnP(x_img,x_w,K):
    M = x_img.shape[0]
    N = x_img.shape[2]
    f = K[0,0]
    px = K[0,2]
    py = K[1,2]

    # get control points
    c_w = getControlPoints(x_w)

    # solve alphas
    alphas = solveAlphas(x_w,c_w)
    if x_img.is_cuda:
        alphas = alphas.cuda()
    Matrix = setupM(alphas,x_img.permute(0,2,1),px,py,f)

    if ~torch.all(Matrix == Matrix):
        print(Matrix)
        print(f)
        print(x_w)
        return

    u,d,v = torch.svd(Matrix)
    km = v[:,:,-4:]

    # solve N=1
    beta_n1 = torch.zeros((M,4))
    if x_img.is_cuda:
        beta_n1 = beta_n1.cuda()
    beta_n1[:,3] = 1
    c_c_n1 = km[:,:,-1].reshape((M,4,3)).permute(0,2,1)
    _, x_c_n1,s1 = scaleControlPoints(c_c_n1,c_w[:3,:],alphas,x_w)
    mask1 = s1 == s1
    Rn1, Tn1 = getExtrinsics(x_c_n1[mask1],x_w)
    reproj_error2_n1 = getReprojError2(x_img[mask1],x_w,Rn1,Tn1,K,loss='l2')

    return km, c_w, beta_n1 / s1.unsqueeze(1), alphas

    # solve N=2
    d12, d13, d14, d23, d24, d34 = getDistances(c_w)
    distances = torch.stack([d12,d13,d14,d23,d24,d34])**2
    betasq_n2 = getBetaN2(km[:,:,-2:],distances).squeeze()
    b1_n2 = torch.sqrt(torch.abs(betasq_n2[:,0]))
    b2_n2 = torch.sqrt(torch.abs(betasq_n2[:,2])) * torch.sign(betasq_n2[:,1]) * torch.sign(betasq_n2[:,0]).detach()
    beta_n2 = torch.zeros((M,4))
    if x_img.is_cuda:
        beta_n2 = beta_n2.cuda()
    beta_n2[:,2] = b1_n2
    beta_n2[:,3] = b2_n2
    c_c_n2 = getControlPointsN2(km[:,:,-2:],beta_n2)
    _,x_c_n2,s2 = scaleControlPoints(c_c_n2,c_w[:3,:],alphas,x_w)
    mask2 = s2 == s2
    Rn2,Tn2 = getExtrinsics(x_c_n2[mask2],x_w)
    reproj_error2_n2 = getReprojError2(x_img[mask2],x_w,Rn2,Tn2,K,loss='l2')

    if torch.any(mask1 != mask2):
        error1 = torch.zeros(M).to(reproj_error2_n1.device) + 10000
        error2 = torch.zeros(M).to(reproj_error2_n2.device) + 10000
        error1[mask1] = reproj_error2_n1
        error2[mask2] = reproj_error2_n2
    else:
        error1 = reproj_error2_n1
        error2 = reproj_error2_n2


    s = torch.stack((s1,s2))

    # determine best solution in terms of 2d reprojection
    mask = error1 > error2
    mask = mask.long()
    betas = torch.stack((beta_n1,beta_n2))
    best_betas = betas.gather(0,torch.stack(4*[mask],dim=1).unsqueeze(0)).squeeze()
    best_scale = s.gather(0,mask.unsqueeze(0)).squeeze()
    scaled_betas = best_betas / best_scale.unsqueeze(1)

    return km, c_w, scaled_betas, alphas
//  [python]
def EPnP(x_img,x_w,K):
    M = x_img.shape[0]
    N = x_img.shape[2]
    f = K[0,0]
    px = K[0,2]
    py = K[1,2]

    # get control points
    c_w = getControlPoints(x_w)

    # solve alphas
    alphas = solveAlphas(x_w,c_w)
    if x_img.is_cuda:
        alphas = alphas.cuda()
    Matrix = setupM(alphas,x_img.permute(0,2,1),px,py,f)

    if ~torch.all(Matrix == Matrix):
        print(Matrix)
        print(f)
        print(x_w)
        return

    u,s,v = torch.linalg.svd(Matrix)
    km = v[:,:,-4:]

    # solve N=1
    beta_n1 = torch.zeros((M,4))
    if x_img.is_cuda:
        beta_n1 = beta_n1.cuda()
    beta_n1[:,3] = 1
    c_c_n1 = km[:,:,-1].reshape((M,4,3)).permute(0,2,1)
    _, x_c_n1,s1 = scaleControlPoints(c_c_n1,c_w[:3,:],alphas,x_w)
    mask1 = s1 == s1
    Rn1, Tn1 = getExtrinsics(x_c_n1[mask1],x_w)
    reproj_error2_n1 = getReprojError2(x_img[mask1],x_w,Rn1,Tn1,K,loss='l2')

    return km, c_w, beta_n1 / s1.unsqueeze(1), alphas

    # solve N=2
    d12, d13, d14, d23, d24, d34 = getDistances(c_w)
    distances = torch.stack([d12,d13,d14,d23,d24,d34])**2
    betasq_n2 = getBetaN2(km[:,:,-2:],distances).squeeze()
    b1_n2 = torch.sqrt(torch.abs(betasq_n2[:,0]))
    b2_n2 = torch.sqrt(torch.abs(betasq_n2[:,2])) * torch.sign(betasq_n2[:,1]) * torch.sign(betasq_n2[:,0]).detach()
    beta_n2 = torch.zeros((M,4))
    if x_img.is_cuda:
        beta_n2 = beta_n2.cuda()
    beta_n2[:,2] = b1_n2
    beta_n2[:,3] = b2_n2
    c_c_n2 = getControlPointsN2(km[:,:,-2:],beta_n2)
    _,x_c_n2,s2 = scaleControlPoints(c_c_n2,c_w[:3,:],alphas,x_w)
    mask2 = s2 == s2
    Rn2,Tn2 = getExtrinsics(x_c_n2[mask2],x_w)
    reproj_error2_n2 = getReprojError2(x_img[mask2],x_w,Rn2,Tn2,K,loss='l2')

    if torch.any(mask1 != mask2):
        error1 = torch.zeros(M).to(reproj_error2_n1.device) + 10000
        error2 = torch.zeros(M).to(reproj_error2_n2.device) + 10000
        error1[mask1] = reproj_error2_n1
        error2[mask2] = reproj_error2_n2
    else:
        error1 = reproj_error2_n1
        error2 = reproj_error2_n2


    s = torch.stack((s1,s2))

    # determine best solution in terms of 2d reprojection
    mask = error1 > error2
    mask = mask.long()
    betas = torch.stack((beta_n1,beta_n2))
    best_betas = betas.gather(0,torch.stack(4*[mask],dim=1).unsqueeze(0)).squeeze()
    best_scale = s.gather(0,mask.unsqueeze(0)).squeeze()
    scaled_betas = best_betas / best_scale.unsqueeze(1)

    return km, c_w, scaled_betas, alphas
//  [python]
def quota(conceptor_matrix):
    """
    Compute quota of a conceptor matrix
    """
    _, Se, _ = torch.svd(conceptor_matrix)
    return float(torch.sum(Se).item() / conceptor_matrix.size(0))
//  [python]
def quota(conceptor_matrix):
    """
    Compute quota of a conceptor matrix
    """
    _, Se, _ = torch.linalg.svd(conceptor_matrix)
    return float(torch.sum(Se).item() / conceptor_matrix.size(0))
//  [python]
def compute_singular_values(stats):
    """
    Compute singular values
    :param states:
    :return:
    """
    # Compute R (correlation matrix)
    R = stats.t().mm(stats) / stats.shape[0]

    # Compute singular values
    return torch.svd(R)
//  [python]
def compute_singular_values(stats):
    """
    Compute singular values
    :param states:
    :return:
    """
    # Compute R (correlation matrix)
    R = stats.t().mm(stats) / stats.shape[0]

    # Compute singular values
    return torch.linalg.svd(R)
//  [python]
def apply_sefa(generator, backbone, z, fake_label, num_semantic_axis, maximum_variations, num_cols):
    generator = misc.peel_model(generator)
    w = generator.linear0.weight
    if backbone == "big_resnet":
        zs = z
        z = torch.split(zs, generator.chunk_size, 0)[0]
    eigen_vectors = torch.svd(w).V.to(z.device)[:, :num_semantic_axis]

    z_dim = len(z)
    zs_start = z.repeat(num_semantic_axis).view(-1, 1, z_dim)
    zs_end = (z.unsqueeze(1) + maximum_variations * eigen_vectors).T.view(-1, 1, z_dim)
    if backbone == "big_resnet":
        zs_shard = zs[z_dim:].expand([1, 1, -1]).repeat(num_semantic_axis, 1, 1)
        zs_start = torch.cat([zs_start, zs_shard], axis=2)
        zs_end = torch.cat([zs_end, zs_shard], axis=2)
    zs_canvas = misc.interpolate(x0=zs_start, x1=zs_end, num_midpoints=num_cols - 2).view(-1, zs_start.shape[-1])
    images_canvas = generator(zs_canvas, fake_label.repeat(len(zs_canvas)), eval=True)
    return images_canvas
//  [python]
def apply_sefa(generator, backbone, z, fake_label, num_semantic_axis, maximum_variations, num_cols):
    generator = misc.peel_model(generator)
    w = generator.linear0.weight
    if backbone == "big_resnet":
        zs = z
        z = torch.split(zs, generator.chunk_size, 0)[0]
    eigen_vectors = torch.linalg.svd(w).V.to(z.device)[:, :num_semantic_axis]

    z_dim = len(z)
    zs_start = z.repeat(num_semantic_axis).view(-1, 1, z_dim)
    zs_end = (z.unsqueeze(1) + maximum_variations * eigen_vectors).T.view(-1, 1, z_dim)
    if backbone == "big_resnet":
        zs_shard = zs[z_dim:].expand([1, 1, -1]).repeat(num_semantic_axis, 1, 1)
        zs_start = torch.cat([zs_start, zs_shard], axis=2)
        zs_end = torch.cat([zs_end, zs_shard], axis=2)
    zs_canvas = misc.interpolate(x0=zs_start, x1=zs_end, num_midpoints=num_cols - 2).view(-1, zs_start.shape[-1])
    images_canvas = generator(zs_canvas, fake_label.repeat(len(zs_canvas)), eval=True)
    return images_canvas
//  [python]
def get_matrix_kernel(A, eps = 1e-10):
    '''
    Compute an orthonormal basis of the kernel (x_1, x_2, ...)
    A x_i = 0
    scalar_product(x_i, x_j) = delta_ij

    :param A: matrix
    :return: matrix where each row is a basis vector of the kernel of A
    '''
    A = rearrange(A, '... d -> (...) d')
    _u, s, v = torch.svd(A)
    kernel = v.t()[s < eps]
    return kernel
//  [python]
def get_matrix_kernel(A, eps = 1e-10):
    '''
    Compute an orthonormal basis of the kernel (x_1, x_2, ...)
    A x_i = 0
    scalar_product(x_i, x_j) = delta_ij

    :param A: matrix
    :return: matrix where each row is a basis vector of the kernel of A
    '''
    A = rearrange(A, '... d -> (...) d')
    _u, s, v = torch.linalg.svd(A)
    kernel = v.t()[s < eps]
    return kernel
//  [python]
def test_svd_no_singularvectors(self, device):
        for size in [(5, 5), (5, 20), (20, 5)]:
            a = torch.randn(*size, device=device)
            u, s_expect, v = torch.svd(a)
            u, s_actual, v = torch.svd(a, compute_uv=False)
            self.assertEqual(s_expect, s_actual, msg="Singular values don't match")
//  [python]
def test_svd_no_singularvectors(self, device):
        for size in [(5, 5), (5, 20), (20, 5)]:
            a = torch.randn(*size, device=device)
            u, s_expect, v = torch.linalg.svd(a)
            u, s_actual, v = torch.linalg.svd(a, compute_uv=False)
            self.assertEqual(s_expect, s_actual, msg="Singular values don't match")
//  [python]
def svd(
        self,
    ) -> Tuple[
        TT[TH.leading_dims : ..., TH.ldim, TH.mdim],
        TT[TH.leading_dims : ..., TH.mdim],
        TT[TH.leading_dims : ..., TH.rdim, TH.mdim],
    ]:
        """
        Efficient algorithm for finding Singular Value Decomposition, a tuple (U, S, Vh) for matrix M st S is a vector and U, Vh are orthogonal matrices, and U @ S.diag() @ Vh.T == M

        (Note that Vh is given as the transpose of the obvious thing)
        """
        Ua, Sa, Vha = torch.svd(self.A)
        Ub, Sb, Vhb = torch.svd(self.B)
        middle = Sa[..., :, None] * utils.transpose(Vha) @ Ub * Sb[..., None, :]
        Um, Sm, Vhm = torch.svd(middle)
        U = Ua @ Um
        Vh = Vhb @ Vhm
        S = Sm
        return U, S, Vh
//  [python]
def svd(
        self,
    ) -> Tuple[
        TT[TH.leading_dims : ..., TH.ldim, TH.mdim],
        TT[TH.leading_dims : ..., TH.mdim],
        TT[TH.leading_dims : ..., TH.rdim, TH.mdim],
    ]:
        """
        Efficient algorithm for finding Singular Value Decomposition, a tuple (U, S, Vh) for matrix M st S is a vector and U, Vh are orthogonal matrices, and U @ S.diag() @ Vh.T == M

        (Note that Vh is given as the transpose of the obvious thing)
        """
        Ua, Sa, Vha = torch.linalg.svd(self.A)
        Ub, Sb, Vhb = torch.linalg.svd(self.B)
        middle = Sa[..., :, None] * utils.transpose(Vha) @ Ub * Sb[..., None, :]
        Um, Sm, Vhm = torch.linalg.svd(middle)
        U = Ua @ Um
        Vh = Vhb @ Vhm
        S = Sm
        return U, S, Vh
//  [python]
def rot9d_to_rotmat(x):
    '''
    Converts 9D rotation output to valid 3x3 rotation amtrix.
    Based on Levinson et al., An Analysis of SVD for Deep Rotation Estimation.

    Input:
        (B, 9)
    Output:
        (B, 9)
    '''
    B = x.size()[0]
    x = x.reshape((B, 3, 3))
    u, s, v = torch.svd(x)

    v_T = v.transpose(-2, -1)
    s_p = torch.eye(3).to(x).reshape((1, 3, 3)).expand_as(x).clone()
    s_p[:, 2, 2] = torch.det(torch.matmul(u, v_T))
    x_out = torch.matmul(torch.matmul(u, s_p), v_T)

    return x_out.reshape((B, 9))
//  [python]
def rot9d_to_rotmat(x):
    '''
    Converts 9D rotation output to valid 3x3 rotation amtrix.
    Based on Levinson et al., An Analysis of SVD for Deep Rotation Estimation.

    Input:
        (B, 9)
    Output:
        (B, 9)
    '''
    B = x.size()[0]
    x = x.reshape((B, 3, 3))
    u, s, v = torch.linalg.svd(x)

    v_T = v.transpose(-2, -1)
    s_p = torch.eye(3).to(x).reshape((1, 3, 3)).expand_as(x).clone()
    s_p[:, 2, 2] = torch.det(torch.matmul(u, v_T))
    x_out = torch.matmul(torch.matmul(u, s_p), v_T)

    return x_out.reshape((B, 9))
//  [python]
def svd(covar, fast=False):
    if fast:
        from torch_batch_svd import svd as fast_svd
        return fast_svd(covar)
    else:
        u, s, v = torch.svd(covar.cpu())
        s = s.to(covar.device)
        u = u.to(covar.device)
        v = v.to(covar.device)
        return u, s, v
//  [python]
def svd(covar, fast=False):
    if fast:
        from torch_batch_svd import svd as fast_svd
        return fast_svd(covar)
    else:
        u, s, v = torch.linalg.svd(covar.cpu())
        s = s.to(covar.device)
        u = u.to(covar.device)
        v = v.to(covar.device)
        return u, s, v
//  [python]
def embed_SVD(model: HookedTransformer):
    '''
    Returns vector S, where W_E = U @ diag(S) @ V.T in singular value decomp.
    '''
    U, S, V = t.svd(model.W_E[:, :-1])
    return S
//  [python]
def embed_SVD(model: HookedTransformer):
    '''
    Returns vector S, where W_E = U @ diag(S) @ V.T in singular value decomp.
    '''
    U, S, V = t.linalg.svd(model.W_E[:, :-1])
    return S
//  [python]
def svd(feat, iden=False, device='cuda:0'):
	size = feat.size()
	mean = torch.mean(feat, 1)
	mean = mean.unsqueeze(1).expand_as(feat)
	_feat = feat.clone()
	_feat -= mean
	if size[1] > 1:
		conv = torch.mm(_feat, _feat.t()).div(size[1] - 1)
	else:
		conv = torch.mm(_feat, _feat.t())
	if iden:
		conv += torch.eye(size[0]).to(device)
	#pdb.set_trace()
	u, e, v = torch.svd(conv, some=False)
	return u, e, v
//  [python]
def svd(feat, iden=False, device='cuda:0'):
	size = feat.size()
	mean = torch.mean(feat, 1)
	mean = mean.unsqueeze(1).expand_as(feat)
	_feat = feat.clone()
	_feat -= mean
	if size[1] > 1:
		conv = torch.mm(_feat, _feat.t()).div(size[1] - 1)
	else:
		conv = torch.mm(_feat, _feat.t())
	if iden:
		conv += torch.eye(size[0]).to(device)
	#pdb.set_trace()
	u, e, v = torch.linalg.svd(conv)
	return u, e, v
//  [python]
def batch_get_3children_orient_svd(rel_pose_skeleton, rel_rest_pose,
                                   rot_mat_chain_parent, children_list, dtype):
    """Get pelvis orientation for batch data.

    Args:
        rel_pose_skeleton (torch.tensor):
            Locations of root-normalized pose skeleton with shape (Bx29x3)
        rel_rest_pose (torch.tensor):
            Locations of rest/ template pose with shape (Bx29x3)
        rot_mat_chain_parents (torch.tensor):
            parent's rotation matrix with shape (Bx3x3)
        children (List[int]): list of indexes of kinematic children with len 29
        dtype (torch.dtype, optional):
            Data type of the created tensors, the default is torch.float32

    Returns:
        rot_mat (torch.tensor):
            Child's rotation matrix with shape (Bx3x3)
    """
    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(rot_mat_chain_parent.transpose(1, 2), target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    U, _, V = torch.svd(S)

    rot_mat = torch.bmm(V, U.transpose(1, 2))
    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)
    return rot_mat
//  [python]
def batch_get_3children_orient_svd(rel_pose_skeleton, rel_rest_pose,
                                   rot_mat_chain_parent, children_list, dtype):
    """Get pelvis orientation for batch data.

    Args:
        rel_pose_skeleton (torch.tensor):
            Locations of root-normalized pose skeleton with shape (Bx29x3)
        rel_rest_pose (torch.tensor):
            Locations of rest/ template pose with shape (Bx29x3)
        rot_mat_chain_parents (torch.tensor):
            parent's rotation matrix with shape (Bx3x3)
        children (List[int]): list of indexes of kinematic children with len 29
        dtype (torch.dtype, optional):
            Data type of the created tensors, the default is torch.float32

    Returns:
        rot_mat (torch.tensor):
            Child's rotation matrix with shape (Bx3x3)
    """
    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(rot_mat_chain_parent.transpose(1, 2), target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    U, S, V = torch.linalg.svd(S)

    rot_mat = torch.bmm(V, U.transpose(1, 2))
    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)
    return rot_mat
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents,
                                children, dtype):
    """Get pelvis orientation svd for batch data.

    Args:
        rel_pose_skeleton (torch.tensor):
            Locations of root-normalized pose skeleton with shape (Bx29x3)
        rel_rest_pose (torch.tensor):
            Locations of rest/ template pose with shape (Bx29x3)
        parents (List[int]): list of indexes of kinematic parents with len 29
        children (List[int]): list of indexes of kinematic children with len 29
        dtype (torch.dtype, optional):
            Data type of the created tensors, the default is torch.float32

    Returns:
        rot_mat (torch.tensor):
            Rotation matrix of pelvis with shape (Bx3x3)
    """
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    U, _, V = torch.svd(S_non_zero)

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    rot_mat_non_zero = torch.bmm(V, U.transpose(1, 2))
    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents,
                                children, dtype):
    """Get pelvis orientation svd for batch data.

    Args:
        rel_pose_skeleton (torch.tensor):
            Locations of root-normalized pose skeleton with shape (Bx29x3)
        rel_rest_pose (torch.tensor):
            Locations of rest/ template pose with shape (Bx29x3)
        parents (List[int]): list of indexes of kinematic parents with len 29
        children (List[int]): list of indexes of kinematic children with len 29
        dtype (torch.dtype, optional):
            Data type of the created tensors, the default is torch.float32

    Returns:
        rot_mat (torch.tensor):
            Rotation matrix of pelvis with shape (Bx3x3)
    """
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    U, S, V = torch.linalg.svd(S_non_zero)

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    rot_mat_non_zero = torch.bmm(V, U.transpose(1, 2))
    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat
//  [python]
def _compute_nullspace(self):
        with torch.no_grad():
            u,s,vh = torch.svd(self.W)
            N = vh.mm(vh.transpose(1,0))
        return N
//  [python]
def _compute_nullspace(self):
        with torch.no_grad():
            u,s,vh = torch.linalg.svd(self.W)
            N = vh.mm(vh.transpose(1,0))
        return N
//  [python]
def lyapunov_svd(A, C, rtol=1e-4, use_svd=False):
    """Solve AX+XA=C"""

    assert A.shape[0] == A.shape[1]
    assert len(A.shape) == 2
    if use_svd:
        U, S, V = torch.svd(A)
    else:
        S, U = torch.symeig(A, eigenvectors=True)
    S = S.diag() @ torch.ones(A.shape)
    X = U @ ((U.t() @ C @ U) / (S + S.t())) @ U.t()
    error = A @ X + X @ A - C
    relative_error = torch.max(torch.abs(error)) / torch.max(torch.abs(A))
    if relative_error > rtol:
        print(f"Warning, error {relative_error} encountered in lyapunov_svd")

    return X
//  [python]
def lyapunov_svd(A, C, rtol=1e-4, use_svd=False):
    """Solve AX+XA=C"""

    assert A.shape[0] == A.shape[1]
    assert len(A.shape) == 2
    if use_svd:
        U, S, V = torch.linalg.svd(A)
    else:
        S, U = torch.symeig(A, eigenvectors=True)
    S = S.diag() @ torch.ones(A.shape)
    X = U @ ((U.t() @ C @ U) / (S + S.t())) @ U.t()
    error = A @ X + X @ A - C
    relative_error = torch.max(torch.abs(error)) / torch.max(torch.abs(A))
    if relative_error > rtol:
        print(f"Warning, error {relative_error} encountered in lyapunov_svd")

    return X
//  [python]
def test_torch_svd(N0=3, N1=23):
    np0 = np.random.randn(N0, N1, N1)
    ret_ = np.stack([np.linalg.svd(x, compute_uv=False) for x in np0])
    ret_U, ret_S, ret_V = torch.svd(torch.tensor(np0))
    ret_U, ret_S, ret_V = ret_U.numpy(), ret_S.numpy(), ret_V.numpy()
    assert hfe(ret_, ret_S) < 1e-7
    assert max(hfe(np.eye(N1), x@x.T) for x in ret_U) < 1e-7
    assert max(hfe(np.eye(N1), x@x.T) for x in ret_V) < 1e-7
    assert hfe(np0, np.stack([(x*y)@z.T for x,y,z in zip(ret_U,ret_S,ret_V)])) < 1e-7
//  [python]
def test_torch_svd(N0=3, N1=23):
    np0 = np.random.randn(N0, N1, N1)
    ret_ = np.stack([np.linalg.svd(x, compute_uv=False) for x in np0])
    ret_U, ret_S, ret_V = torch.linalg.svd(torch.tensor(np0))
    ret_U, ret_S, ret_V = ret_U.numpy(), ret_S.numpy(), ret_V.numpy()
    assert hfe(ret_, ret_S) < 1e-7
    assert max(hfe(np.eye(N1), x@x.T) for x in ret_U) < 1e-7
    assert max(hfe(np.eye(N1), x@x.T) for x in ret_V) < 1e-7
    assert hfe(np0, np.stack([(x*y)@z.T for x,y,z in zip(ret_U,ret_S,ret_V)])) < 1e-7
//  [python]
def _estimate_nullspace(self):

        if self.last_transform < self.n_steps_recompute:
            self.last_transform += 1
            return self.proj

        with torch.no_grad():
            U,S,Vh = torch.svd(self.model.classifier.weight, some=True)
            self.proj = Vh.mm(Vh.transpose(1,0))

        self.last_transform = 0
        return self.proj
//  [python]
def _estimate_nullspace(self):

        if self.last_transform < self.n_steps_recompute:
            self.last_transform += 1
            return self.proj

        with torch.no_grad():
            U, S, Vh = torch.linalg.svd(self.model.classifier.weight, full_matrices=True)
            self.proj = Vh.mm(Vh.transpose(1,0))

        self.last_transform = 0
        return self.proj
//  [python]
def test_svd_lowrank(self, device):
        import torch
        from torch.testing._internal.common_utils import random_lowrank_matrix, random_sparse_matrix

        dtype = torch.double

        def run_subtest(actual_rank, matrix_size, batches, device, svd_lowrank, **options):
            density = options.pop('density', 1)
            if isinstance(matrix_size, int):
                rows = columns = matrix_size
            else:
                rows, columns = matrix_size
            if density == 1:
                a_input = random_lowrank_matrix(actual_rank, rows, columns, *batches, device=device, dtype=dtype)
                a = a_input
            else:
                assert batches == ()
                a_input = random_sparse_matrix(rows, columns, density, device=device, dtype=dtype)
                a = a_input.to_dense()

            q = min(*size)
            u, s, v = svd_lowrank(a_input, q=q, **options)

            # check if u, s, v is a SVD
            u, s, v = u[..., :q], s[..., :q], v[..., :q]
            A = u.matmul(s.diag_embed()).matmul(v.transpose(-2, -1))
            self.assertEqual(A, a)

            # check if svd_lowrank produces same singular values as torch.svd
            U, S, V = torch.svd(a)
            self.assertEqual(s.shape, S.shape)
            self.assertEqual(u.shape, U.shape)
            self.assertEqual(v.shape, V.shape)
            self.assertEqual(s, S)

            if density == 1:
                # actual_rank is known only for dense inputs
                #
                # check if pairs (u, U) and (v, V) span the same
                # subspaces, respectively
                u, s, v = u[..., :actual_rank], s[..., :actual_rank], v[..., :actual_rank]
                U, S, V = U[..., :actual_rank], S[..., :actual_rank], V[..., :actual_rank]
                self.assertEqual(u.transpose(-2, -1).matmul(U).det().abs(), torch.ones(batches, device=device, dtype=dtype))
                self.assertEqual(v.transpose(-2, -1).matmul(V).det().abs(), torch.ones(batches, device=device, dtype=dtype))

        all_batches = [(), (1,), (3,), (2, 3)]
        for actual_rank, size, all_batches in [
                (2, (17, 4), all_batches),
                (4, (17, 4), all_batches),
                (4, (17, 17), all_batches),
                (10, (100, 40), all_batches),
                (7, (1000, 1000), [()]),
        ]:
            # dense input
            for batches in all_batches:
                run_subtest(actual_rank, size, batches, device, torch.svd_lowrank)
                if size != size[::-1]:
                    run_subtest(actual_rank, size[::-1], batches, device, torch.svd_lowrank)

        # sparse input
        for size in [(17, 4), (4, 17), (17, 17), (100, 40), (40, 100), (1000, 1000)]:
            for density in [0.005, 0.1]:
                run_subtest(None, size, (), device, torch.svd_lowrank, density=density)

        # jitting support
        jitted = torch.jit.script(torch.svd_lowrank)
        actual_rank, size, batches = 2, (17, 4), ()
        run_subtest(actual_rank, size, batches, device, jitted)