jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def sample_inputs_logdet(op_info, device, dtype, requires_grad, **kwargs):
    def make_nonzero_det(A, *, sign=1, min_singular_value=0.1, **kwargs):
        u, s, vh = torch.linalg.svd(A, full_matrices=False)
        s.clamp_(min=min_singular_value)
        A = (u * s.unsqueeze(-2)) @ vh
        det = A.det()
        if sign is not None:
            if A.dim() == 2:
                if (det < 0) ^ (sign < 0):
                    A[0, :].neg_()
            else:
                cond = ((det < 0) ^ (sign < 0)).nonzero()
                if cond.size(0) > 0:
                    for i in range(cond.size(0)):
                        A[list(cond[i])][0, :].neg_()
        return A

    # cases constructed using make_tensor()
    tensor_shapes = (
        (S, S),
        (1, 1),
        (3, 3, S, S),
        (3, 3, 1, 1)
    )

    for shape in tensor_shapes:
        t = make_tensor(shape, device=device, dtype=dtype)
        d = make_nonzero_det(t).requires_grad_(requires_grad)
        yield SampleInput(d)

    # cases constructed using:
    #  1) make_symmetric_matrices
    #  2) make_symmetric_pd_matrices
    #  3) make_fullrank_matrices_with_distinct_singular_values
    symmetric_shapes = (
        (S, S),
        (3, S, S),
    )


    def _helper(constructor, *shape, **kwargs):
        t = constructor(*shape, device=device, dtype=dtype)
        d = make_nonzero_det(t, **kwargs).requires_grad_(requires_grad)
        yield SampleInput(d)

    for shape in symmetric_shapes:
        _helper(make_symmetric_matrices, *shape)
        _helper(make_symmetric_pd_matrices, *shape)
        _helper(make_fullrank_matrices_with_distinct_singular_values, *shape, min_singular_value=0)
//  [python]
def __call__(self, inputs: List[torch.Tensor]):
        updates = torch.stack(inputs, dim=0)
        d = len(updates[0])

        benign_ids = []
        for i in range(self.num_iters):
            indices = torch.randperm(d)[: self.sub_dim]
            sub_updates = updates[:, indices]
            mu = sub_updates.mean(dim=0)
            centered_update = sub_updates - mu
            v = torch.svd(centered_update)[2][0, :]
            s = np.array(
                [(torch.dot(update - mu, v) ** 2).item() for update in sub_updates]
            )

            good = s.argsort()[
                : len(updates) - int(self.fliter_frac * self.num_byzantine)
            ]
            benign_ids.append(good)

        # Convert the first list to a set to start the intersection
        intersection_set = set(benign_ids[0])

        # Iterate over the rest of the lists and get the intersection
        for lst in benign_ids[1:]:
            intersection_set.intersection_update(lst)

        # Convert the set back to a list
        benign_ids = list(intersection_set)
        benign_updates = updates[benign_ids, :].mean(dim=0)
        return benign_updates
//  [python]
def __call__(self, inputs: List[torch.Tensor]):
        updates = torch.stack(inputs, dim=0)
        d = len(updates[0])

        benign_ids = []
        for i in range(self.num_iters):
            indices = torch.randperm(d)[: self.sub_dim]
            sub_updates = updates[:, indices]
            mu = sub_updates.mean(dim=0)
            centered_update = sub_updates - mu
            v = torch.linalg.svd(centered_update, full_matrices=False)[2][0, :]
            s = np.array(
                [(torch.dot(update - mu, v) ** 2).item() for update in sub_updates]
            )

            good = s.argsort()[
                : len(updates) - int(self.fliter_frac * self.num_byzantine)
            ]
            benign_ids.append(good)

        # Convert the first list to a set to start the intersection
        intersection_set = set(benign_ids[0])

        # Iterate over the rest of the lists and get the intersection
        for lst in benign_ids[1:]:
            intersection_set.intersection_update(lst)

        # Convert the set back to a list
        benign_ids = list(intersection_set)
        benign_updates = updates[benign_ids, :].mean(dim=0)
        return benign_updates
//  [python]
def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False):
        r"""
        Returns a randomly sampled matrix on the manifold by sampling a matrix according
        to ``init_`` and projecting it onto the manifold.

        The output of this method can be used to initialize a parametrized tensor
        that has been parametrized with this or any other manifold as::

            >>> layer = nn.Linear(20, 20)
            >>> M = LowRank(layer.weight.size(), rank=6)
            >>> geotorch.register_parametrization(layer, "weight", M)
            >>> layer.weight = M.sample()

        Args:
            init\_ (callable): Optional. A function that takes a tensor and fills it
                    in place according to some distribution. See
                    `torch.init <https://pytorch.org/docs/stable/nn.init.html>`_.
                    Default: ``torch.nn.init.xavier_normal_``
        """
        with torch.no_grad():
            device = self[0].base.device
            dtype = self[0].base.dtype
            X = torch.empty(
                *(self.tensorial_size + (self.n, self.k)), device=device, dtype=dtype
            )
            init_(X)
            U, S, Vt = torch.svd(X, some=False)
            U, S, Vt = U[..., : self.rank], S[..., : self.rank], Vt[..., : self.rank, :]
            if factorized:
                if self.transposed:
                    return Vt.transpose(-2, -1), S, U
                else:
                    return U, S, Vt.transpose(-2, -1)
            else:
                X = (U * S.unsqueeze(-2)) @ Vt
                if self.transposed:
                    X = X.transpose(-2, -1)
                return X
//  [python]
def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False):
        r"""
        Returns a randomly sampled matrix on the manifold by sampling a matrix according
        to ``init_`` and projecting it onto the manifold.

        The output of this method can be used to initialize a parametrized tensor
        that has been parametrized with this or any other manifold as::

            >>> layer = nn.Linear(20, 20)
            >>> M = LowRank(layer.weight.size(), rank=6)
            >>> geotorch.register_parametrization(layer, "weight", M)
            >>> layer.weight = M.sample()

        Args:
            init\_ (callable): Optional. A function that takes a tensor and fills it
                    in place according to some distribution. See
                    `torch.init <https://pytorch.org/docs/stable/nn.init.html>`_.
                    Default: ``torch.nn.init.xavier_normal_``
        """
        with torch.no_grad():
            device = self[0].base.device
            dtype = self[0].base.dtype
            X = torch.empty(
                *(self.tensorial_size + (self.n, self.k)), device=device, dtype=dtype
            )
            init_(X)
            U, S, Vt = torch.linalg.svd(X, full_matrices=False)
            U, S, Vt = U[..., : self.rank], S[..., : self.rank], Vt[..., : self.rank, :]
            if factorized:
                if self.transposed:
                    return Vt.transpose(-2, -1), S, U
                else:
                    return U, S, Vt.transpose(-2, -1)
            else:
                X = (U * S.unsqueeze(-2)) @ Vt
                if self.transposed:
                    X = X.transpose(-2, -1)
                return X
//  [python]
def sample_transformation(eigenvalues, normalize=False):
    n_dims = len(eigenvalues)
    U, _, _ = torch.svd(torch.randn(n_dims, n_dims))
    t = U @ torch.diag(eigenvalues) @ torch.transpose(U, 0, 1)
    if normalize:
        norm_subspace = torch.sum(eigenvalues**2)
        t *= math.sqrt(n_dims / norm_subspace)
    return t
//  [python]
def sample_transformation(eigenvalues, normalize=False):
    n_dims = len(eigenvalues)
    U, _, _ = torch.linalg.svd(torch.randn(n_dims, n_dims))
    t = U @ torch.diag(eigenvalues) @ torch.transpose(U, 0, 1)
    if normalize:
        norm_subspace = torch.sum(eigenvalues**2)
        t *= math.sqrt(n_dims / norm_subspace)
    return t
//  [python]
def kabsch(xyz1, xyz2, eps=1e-6):
    """Superimposes `xyz2` coordinates onto `xyz1`, returns RMSD and rotation matrix."""
    # center to CA centroid
    xyz1 = xyz1 - xyz1.mean(0)
    xyz2 = xyz2 - xyz2.mean(0)

    # Computation of the covariance matrix
    C = xyz2.T @ xyz1

    # Compute optimal rotation matrix using SVD
    V, S, W = torch.svd(C)

    # get sign to ensure right-handedness
    d = torch.ones([3,3])
    d[:,-1] = torch.sign(torch.det(V)*torch.det(W))

    # Rotation matrix U
    U = (d*V) @ W

    # Rotate xyz2
    xyz2_ = xyz2 @ U

    L = xyz2_.shape[0]

    rmsd = torch.sqrt(torch.sum((xyz2_-xyz1)*(xyz2_-xyz1), axis=(0,1)) / L + eps)

    return rmsd, U
//  [python]
def kabsch(xyz1, xyz2, eps=1e-6):
    """Superimposes `xyz2` coordinates onto `xyz1`, returns RMSD and rotation matrix."""
    # center to CA centroid
    xyz1 = xyz1 - xyz1.mean(0)
    xyz2 = xyz2 - xyz2.mean(0)

    # Computation of the covariance matrix
    C = xyz2.T @ xyz1

    # Compute optimal rotation matrix using SVD
    V, S, W = torch.linalg.svd(C)

    # get sign to ensure right-handedness
    d = torch.ones([3,3])
    d[:,-1] = torch.sign(torch.linalg.det(V)*torch.linalg.det(W))

    # Rotation matrix U
    U = (d*V) @ W

    # Rotate xyz2
    xyz2_ = xyz2 @ U

    L = xyz2_.shape[0]

    rmsd = torch.sqrt(torch.sum((xyz2_-xyz1)*(xyz2_-xyz1), axis=(0,1)) / L + eps)

    return rmsd, U
//  [python]
def low_rank_project(M, rank):
    """Supports batches of matrices as well.
    """
    U, S, Vt = torch.svd(M)
    S_sqrt = S[..., :rank].sqrt()
    U = U[..., :rank] * rearrange(S_sqrt, '... rank -> ... 1 rank')
    Vt = rearrange(S_sqrt, '... rank -> ... rank 1') * Vt[..., :rank, :]
    return U, Vt
//  [python]
def low_rank_project(M, rank):
    """Supports batches of matrices as well.
    """
    U, S, Vt = torch.linalg.svd(M)
    S_sqrt = S[..., :rank].sqrt()
    U = U[..., :rank] * rearrange(S_sqrt, '... rank -> ... 1 rank')
    Vt = rearrange(S_sqrt, '... rank -> ... rank 1') * Vt[..., :rank, :]
    return U, Vt
//  [python]
def nrqm(
    img: Tensor,
    linear_param,
    rf_param,
) -> Tensor:
    """Calculate NRQM
    Args:
        img (Tensor): Input image.
        linear_param (np.array): (4, 1) linear regression params
        rf_param: params of 3 random forest for 3 kinds of features
    """
    assert img.ndim == 4, ('Input image must be a gray or Y (of YCbCr) image with shape (b, c, h, w).')

    # crop image
    b, c, h, w = img.shape
    img = img.double()
    img_pyr = get_guass_pyramid(img / 255.)

    # DCT features
    f1 = []
    for im in img_pyr:
        f1.append(block_dct(im))
    f1 = torch.cat(f1, dim=1)

    # gsm features
    f2 = global_gsm(img)

    # svd features
    f3 = []
    for im in img_pyr:
        col = im2col(im, 5, 'distinct')
        _, s, _ = torch.svd(col)
        f3.append(s)
    f3 = torch.cat(f3, dim=1)

    # Random forest regression. Currently not differentiable and only support CPU
    preds = torch.ones(b, 1)
    for feat, rf in zip([f1, f2, f3], rf_param):
        tmp_pred = random_forest_regression(feat, *rf)
        preds = torch.cat((preds, tmp_pred), dim=1)
    quality = preds @ torch.tensor(linear_param)

    return quality.squeeze()
//  [python]
def nrqm(
    img: Tensor,
    linear_param,
    rf_param,
) -> Tensor:
    """Calculate NRQM
    Args:
        img (Tensor): Input image.
        linear_param (np.array): (4, 1) linear regression params
        rf_param: params of 3 random forest for 3 kinds of features
    """
    assert img.ndim == 4, ('Input image must be a gray or Y (of YCbCr) image with shape (b, c, h, w).')

    # crop image
    b, c, h, w = img.shape
    img = img.double()
    img_pyr = get_guass_pyramid(img / 255.)

    # DCT features
    f1 = []
    for im in img_pyr:
        f1.append(block_dct(im))
    f1 = torch.cat(f1, dim=1)

    # gsm features
    f2 = global_gsm(img)

    # svd features
    f3 = []
    for im in img_pyr:
        col = im2col(im, 5, 'distinct')
        _, s, _ = torch.linalg.svd(col, full_matrices=False)
        f3.append(s)
    f3 = torch.cat(f3, dim=1)

    # Random forest regression. Currently not differentiable and only support CPU
    preds = torch.ones(b, 1)
    for feat, rf in zip([f1, f2, f3], rf_param):
        tmp_pred = random_forest_regression(feat, *rf)
        preds = torch.cat((preds, tmp_pred), dim=1)
    quality = preds @ torch.tensor(linear_param)

    return quality.squeeze()
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()

        feat = model.forward_features(x)
        u, s, v = torch.svd(feat)
        feat = feat - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat = feat - power_iteration(feat, iter=20)
        feat = model.norm(feat)
        logits = model.head(feat[:, 0])

        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()

        feat = model.forward_features(x)
        u, s, v = torch.linalg.svd(feat)
        feat = feat - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat = feat - power_iteration(feat, iter=20)
        feat = model.norm(feat)
        logits = model.head(feat[:, 0])

        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def U_gk(self) -> torch.Tensor:
        r"""
        Principal components corresponding to eigenvalues ``L_k``.

        .. note::
           Gradients are disabled, used for inference only.
        """
        return torch.svd(self.W_kg.T).U
//  [python]
def U_gk(self) -> torch.Tensor:
        r"""
        Principal components corresponding to eigenvalues ``L_k``.

        .. note::
           Gradients are disabled, used for inference only.
        """
        return torch.linalg.svd(self.W_kg.T, full_matrices=False).U
//  [python]
def get_rmsd(a, b, eps=1e-6):
    """
    align crds b to a : always use all alphas
    expexted tensor of shape (L,3)
    jake's torch adapted version
    """
    assert a.shape == b.shape, "make sure tensors are the same size"
    L = a.shape[0]
    assert a.shape == torch.Size([L, 3]), "make sure tensors are in format [L,3]"

    # center to CA centroid
    a = a - a.mean(dim=0)
    b = b - b.mean(dim=0)

    # Computation of the covariance matrix
    C = torch.einsum(
        "kj,ji->ki", torch.transpose(b.type(torch.float32), 0, 1), a.type(torch.float32)
    )

    # Compute optimal rotation matrix using SVD
    U, S, V = torch.svd(C)

    # get sign to ensure right-handedness
    d = torch.ones([3, 3])
    d[:, -1] = torch.sign(torch.det(U) * torch.det(V))

    # Rotation matrix U
    U = torch.einsum("kj,ji->ki", (d * U), V)

    # Rotate xyz_hal
    rP = torch.einsum("kj,ji->ki", b.type(torch.float32), U.type(torch.float32))

    L = rP.shape[0]
    rmsd = torch.sqrt(torch.sum((rP - a) * (rP - a), axis=(0, 1)) / L + eps)

    return rmsd, U
//  [python]
def get_rmsd(a, b, eps=1e-6):
    """
    align crds b to a : always use all alphas
    expexted tensor of shape (L,3)
    jake's torch adapted version
    """
    assert a.shape == b.shape, "make sure tensors are the same size"
    L = a.shape[0]
    assert a.shape == torch.Size([L, 3]), "make sure tensors are in format [L,3]"

    # center to CA centroid
    a = a - a.mean(dim=0)
    b = b - b.mean(dim=0)

    # Computation of the covariance matrix
    C = torch.einsum(
        "kj,ji->ki", torch.transpose(b.type(torch.float32), 0, 1), a.type(torch.float32)
    )

    # Compute optimal rotation matrix using SVD
    V, S, W = torch.linalg.svd(C)

    # get sign to ensure right-handedness
    d = torch.ones([3, 3])
    d[:, -1] = torch.sign(torch.linalg.det(V) * torch.linalg.det(W))

    # Rotation matrix U
    U = torch.einsum("kj,ji->ki", (d * V), W)

    # Rotate xyz_hal
    rP = torch.einsum("kj,ji->ki", b.type(torch.float32), U.type(torch.float32))

    L = rP.shape[0]
    rmsd = torch.sqrt(torch.sum((rP - a) * (rP - a), axis=(0, 1)) / L + eps)

    return rmsd, U
//  [python]
def iterate_data_rankfeatweight(data_loader, model, temperature):
    confs = []
    
    weight = model.blocks[-1].mlp.fc2.weight.data
    B, C = weight.size()
    weight_sub = power_iteration(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    model.blocks[-1].mlp.fc2.weight.data = weight
    
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()

        feat = model.forward_features(x)
        u, s, v = torch.svd(feat)
        feat = feat - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat = feat - power_iteration(feat, iter=20)
        feat = model.norm(feat)
        logits = model.head(feat[:, 0])

        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeatweight(data_loader, model, temperature):
    confs = []
    
    weight = model.blocks[-1].mlp.fc2.weight.data
    B, C = weight.size()
    weight_sub = power_iteration(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    model.blocks[-1].mlp.fc2.weight.data = weight
    
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()

        feat = model.forward_features(x)
        u, s, v = torch.linalg.svd(feat)
        feat = feat - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat = feat - power_iteration(feat, iter=20)
        feat = model.norm(feat)
        logits = model.head(feat[:, 0])

        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def overwrite_base(base_model, tuned_model, rank, clamp_quantile):
    device = base_model.device
    dtype = base_model.dtype

    for lor_base, lor_tune in zip(_iter_lora(base_model), _iter_lora(tuned_model)):

        if isinstance(lor_base, LoraInjectedLinear):
            residual = lor_tune.linear.weight.data - lor_base.linear.weight.data
            # SVD on residual
            print("Distill Linear shape ", residual.shape)
            residual = residual.float()
            U, S, Vh = torch.svd(residual)
            U = U[:, :rank]
            S = S[:rank]
            U = U @ torch.diag(S)

            Vh = Vh[:rank, :]

            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp_quantile)
            low_val = -hi_val

            U = U.clamp(low_val, hi_val)
            Vh = Vh.clamp(low_val, hi_val)

            assert lor_base.lora_up.weight.shape == U.shape
            assert lor_base.lora_down.weight.shape == Vh.shape

            lor_base.lora_up.weight.data = U.to(device=device, dtype=dtype)
            lor_base.lora_down.weight.data = Vh.to(device=device, dtype=dtype)

        if isinstance(lor_base, LoraInjectedConv2d):
            residual = lor_tune.conv.weight.data - lor_base.conv.weight.data
            print("Distill Conv shape ", residual.shape)

            residual = residual.float()
            residual = residual.flatten(start_dim=1)

            # SVD on residual
            U, S, Vh = torch.svd(residual)
            U = U[:, :rank]
            S = S[:rank]
            U = U @ torch.diag(S)

            Vh = Vh[:rank, :]

            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp_quantile)
            low_val = -hi_val

            U = U.clamp(low_val, hi_val)
            Vh = Vh.clamp(low_val, hi_val)

            # U is (out_channels, rank) with 1x1 conv. So,
            U = U.reshape(U.shape[0], U.shape[1], 1, 1)
            # V is (rank, in_channels * kernel_size1 * kernel_size2)
            # now reshape:
            Vh = Vh.reshape(
                Vh.shape[0],
                lor_base.conv.in_channels,
                lor_base.conv.kernel_size[0],
                lor_base.conv.kernel_size[1],
            )

            assert lor_base.lora_up.weight.shape == U.shape
            assert lor_base.lora_down.weight.shape == Vh.shape

            lor_base.lora_up.weight.data = U.to(device=device, dtype=dtype)
            lor_base.lora_down.weight.data = Vh.to(device=device, dtype=dtype)
//  [python]
def overwrite_base(base_model, tuned_model, rank, clamp_quantile):
    device = base_model.device
    dtype = base_model.dtype

    for lor_base, lor_tune in zip(_iter_lora(base_model), _iter_lora(tuned_model)):

        if isinstance(lor_base, LoraInjectedLinear):
            residual = lor_tune.linear.weight.data - lor_base.linear.weight.data
            # SVD on residual
            print("Distill Linear shape ", residual.shape)
            residual = residual.float()
            U, S, Vh = torch.linalg.svd(residual)
            U = U[:, :rank]
            S = S[:rank]
            U = U @ torch.diag(S)

            Vh = Vh[:rank, :]

            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp_quantile)
            low_val = -hi_val

            U = U.clamp(low_val, hi_val)
            Vh = Vh.clamp(low_val, hi_val)

            assert lor_base.lora_up.weight.shape == U.shape
            assert lor_base.lora_down.weight.shape == Vh.shape

            lor_base.lora_up.weight.data = U.to(device=device, dtype=dtype)
            lor_base.lora_down.weight.data = Vh.to(device=device, dtype=dtype)

        if isinstance(lor_base, LoraInjectedConv2d):
            residual = lor_tune.conv.weight.data - lor_base.conv.weight.data
            print("Distill Conv shape ", residual.shape)

            residual = residual.float()
            residual = residual.flatten(start_dim=1)

            # SVD on residual
            U, S, Vh = torch.linalg.svd(residual)
            U = U[:, :rank]
            S = S[:rank]
            U = U @ torch.diag(S)

            Vh = Vh[:rank, :]

            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp_quantile)
            low_val = -hi_val

            U = U.clamp(low_val, hi_val)
            Vh = Vh.clamp(low_val, hi_val)

            # U is (out_channels, rank) with 1x1 conv. So,
            U = U.reshape(U.shape[0], U.shape[1], 1, 1)
            # V is (rank, in_channels * kernel_size1 * kernel_size2)
            # now reshape:
            Vh = Vh.reshape(
                Vh.shape[0],
                lor_base.conv.in_channels,
                lor_base.conv.kernel_size[0],
                lor_base.conv.kernel_size[1],
            )

            assert lor_base.lora_up.weight.shape == U.shape
            assert lor_base.lora_down.weight.shape == Vh.shape

            lor_base.lora_up.weight.data = U.to(device=device, dtype=dtype)
            lor_base.lora_down.weight.data = Vh.to(device=device, dtype=dtype)
//  [python]
def get_orthogonal_matrix(self, weights, rank, type):
        module_params = weights

        if module_params.data.dtype != torch.float:
            float_data = False
            original_type = module_params.data.dtype
            original_device = module_params.data.device
            matrix = module_params.data.float()
        else:
            float_data = True
            matrix = module_params.data

        U, s, Vh = torch.svd(matrix)

        # make the smaller matrix always to be orthogonal matrix
        if type == 'right':
            A = U[:, :rank] @ torch.diag(s[:rank])
            B = Vh[:rank, :]

            if not float_data:
                B = B.to(original_device).type(original_type)
            return B
        elif type == 'left':
            A = U[:, :rank]
            B = torch.diag(s[:rank]) @ Vh[:rank, :]
            if not float_data:
                A = A.to(original_device).type(original_type)
            return A
        elif type == 'full':
            A = U[:, :rank]
            B = Vh[:rank, :]
            if not float_data:
                A = A.to(original_device).type(original_type)
                B = B.to(original_device).type(original_type)
            return [A, B]
        else:
            raise ValueError('type should be left, right or full')
//  [python]
def get_orthogonal_matrix(self, weights, rank, type):
        module_params = weights

        if module_params.data.dtype != torch.float:
            float_data = False
            original_type = module_params.data.dtype
            original_device = module_params.data.device
            matrix = module_params.data.float()
        else:
            float_data = True
            matrix = module_params.data

        U, s, Vh = torch.linalg.svd(matrix, full_matrices=False)

        # make the smaller matrix always to be orthogonal matrix
        if type == 'right':
            A = U[:, :rank] @ torch.diag(s[:rank])
            B = Vh[:rank, :]

            if not float_data:
                B = B.to(original_device).type(original_type)
            return B
        elif type == 'left':
            A = U[:, :rank]
            B = torch.diag(s[:rank]) @ Vh[:rank, :]
            if not float_data:
                A = A.to(original_device).type(original_type)
            return A
        elif type == 'full':
            A = U[:, :rank]
            B = Vh[:rank, :]
            if not float_data:
                A = A.to(original_device).type(original_type)
                B = B.to(original_device).type(original_type)
            return [A, B]
        else:
            raise ValueError('type should be left, right or full')
//  [python]
def test_svd_feature_reduction():
    assert str(SVDFeatureReduction(10)) == 'SVDFeatureReduction(10)'

    x = torch.randn(4, 16)
    U, S, _ = torch.svd(x)
    data = Data(x=x)
    data = SVDFeatureReduction(10)(data)
    assert torch.allclose(data.x, torch.mm(U[:, :10], torch.diag(S[:10])))

    x = torch.randn(4, 8)
    data.x = x
    data = SVDFeatureReduction(10)(Data(x=x))
    assert torch.allclose(data.x, x)
//  [python]
def test_svd_feature_reduction():
    assert str(SVDFeatureReduction(10)) == 'SVDFeatureReduction(10)'

    x = torch.randn(4, 16)
    U, S, _ = torch.linalg.svd(x)
    data = Data(x=x)
    data = SVDFeatureReduction(10)(data)
    assert torch.allclose(data.x, torch.mm(U[:, :10], torch.diag(S[:10])))

    x = torch.randn(4, 8)
    data.x = x
    data = SVDFeatureReduction(10)(Data(x=x))
    assert torch.allclose(data.x, x)
//  [python]
def get_lowrank_tuple_torch_gpu(tensor, max_rank, eps=None):
    t       = tensor.t().float()
    u, s, v = torch.svd(t)
    u, s, v = u[:,:max_rank], s[:max_rank], v[:max_rank, :]
    us      = torch.matmul(u, torch.diag(s))
    A, B    = (v.t(), us.t()) #t ~ AB
    if(eps is not None):
        A  = A.clamp(min=-eps, max=eps)
        B   = B.clamp(min=-eps, max=eps)
    return A.to(tensor.dtype), B.to(tensor.dtype)
//  [python]
def get_lowrank_tuple_torch_gpu(tensor, max_rank, eps=None):
    t       = tensor.t().float()
    u, s, v = torch.linalg.svd(t)
    u, s, v = u[:,:max_rank], s[:max_rank], v[:max_rank, :]
    us      = torch.matmul(u, torch.diag(s))
    A, B    = (v.t(), us.t()) #t ~ AB
    if(eps is not None):
        A  = A.clamp(min=-eps, max=eps)
        B   = B.clamp(min=-eps, max=eps)
    return A.to(tensor.dtype), B.to(tensor.dtype)
//  [python]
def solve_svd(I):
    print(I.shape)
    # SVD on I using I = L' dot B model
    # I.shape = (N, P) (luminance matrix)
    # L.shape = (3, N), L'.shape = (N, 3) (light matrix)
    # B.shape = (3, P) (pseudo-normal matrix) = normal times albedo (scaling value)

    rank = 3
    U, S, VH = torch.svd(I)
    L_t = U[:, :rank] @ torch.diag(torch.sqrt(S[:rank]))
    B = torch.diag(torch.sqrt(S[:rank])) @ VH[:rank, :] # Pseudo-normal matrix
    L = L_t.T # light matrix (3, N)
    
    return B.float(), L.float()
//  [python]
def solve_svd(I):
    print(I.shape)
    # SVD on I using I = L' dot B model
    # I.shape = (N, P) (luminance matrix)
    # L.shape = (3, N), L'.shape = (N, 3) (light matrix)
    # B.shape = (3, P) (pseudo-normal matrix) = normal times albedo (scaling value)

    rank = 3
    U, S, VH = torch.linalg.svd(I, full_matrices=False)
    L_t = U[:, :rank] @ torch.diag(torch.sqrt(S[:rank]))
    B = torch.diag(torch.sqrt(S[:rank])) @ VH[:rank, :] # Pseudo-normal matrix
    L = L_t.T # light matrix (3, N)
    
    return B.float(), L.float()
//  [python]
def update_model_reduce_layer(self, layer_type, layer_number):
        layer_id = f"{layer_type}+{layer_number}"
        if layer_id in self.modified_layers:
            print(f"Layer {layer_id} has already been modified. Skipping.")
            return False

        for name, module in self.model.named_modules():
            if layer_type in name and str(layer_number) in name:
                print(f"Reconstructing layer: {name}")
                original_dtype = module.weight.dtype
                self.original_weights[name] = module.weight.detach().clone()
                weights = module.weight.double()
                U, S, V = torch.svd(weights)

                # Estimate sigma using the full IQR method
                sigma_estimated_full_iqr = self.estimate_sigma_with_full_iqr(S)

                # Calculate Marchenko-Pastur threshold
                n, m = weights.shape
                mp_threshold_full_iqr = self.marchenko_pastur_threshold(sigma_estimated_full_iqr, n, m)

                # Retain only the singular values above the MP threshold
                S_reduced = torch.zeros_like(S)
                k = (S > mp_threshold_full_iqr).sum().item()
                S_reduced[:k] = S[:k]
                print(f"Reduced from {S.shape} to {k}")

                # Reconstruct the matrix using the thresholded singular values
                reconstructed_weights = U @ torch.diag(S_reduced) @ V.t()
                reconstructed_weights = reconstructed_weights.to(original_dtype)
                module.weight = torch.nn.Parameter(reconstructed_weights)
                self.modified_layers.add(layer_id)
                return True
//  [python]
def update_model_reduce_layer(self, layer_type, layer_number):
        layer_id = f"{layer_type}+{layer_number}"
        if layer_id in self.modified_layers:
            print(f"Layer {layer_id} has already been modified. Skipping.")
            return False

        for name, module in self.model.named_modules():
            if layer_type in name and str(layer_number) in name:
                print(f"Reconstructing layer: {name}")
                original_dtype = module.weight.dtype
                self.original_weights[name] = module.weight.detach().clone()
                weights = module.weight.double()
                U, S, V = torch.linalg.svd(weights, full_matrices=False)

                # Estimate sigma using the full IQR method
                sigma_estimated_full_iqr = self.estimate_sigma_with_full_iqr(S)

                # Calculate Marchenko-Pastur threshold
                n, m = weights.shape
                mp_threshold_full_iqr = self.marchenko_pastur_threshold(sigma_estimated_full_iqr, n, m)

                # Retain only the singular values above the MP threshold
                S_reduced = torch.zeros_like(S)
                k = (S > mp_threshold_full_iqr).sum().item()
                S_reduced[:k] = S[:k]
                print(f"Reduced from {S.shape} to {k}")

                # Reconstruct the matrix using the thresholded singular values
                reconstructed_weights = U @ torch.diag(S_reduced) @ V
                reconstructed_weights = reconstructed_weights.to(original_dtype)
                module.weight = torch.nn.Parameter(reconstructed_weights)
                self.modified_layers.add(layer_id)
                return True
//  [python]
def solver(self, source, target):
        '''
        Given source pts, and target pts, we solve the optimal rotation matrix, R and transform matrix T
        min ||R * source +t - tareget||2

        where t^* = mean(P_taregt) - R^* mean(P_source)


        For R^*
        we have, R^{*} = argmax(sum_{1}^{N}\hat{P_t}R\hat{P_s})
                        = trace(\hat{P_t}R\hat{P_s})

        argmax trace{\hat{P_t}R\hat{P_s}}

        R = v @ u.T

        but avoid reflection rotaiton
        we add a constrain
        let
        R = v diag[1, 1, det(v@u.t)] @ u.t
        return
            R, t
        '''
        mean_source = source.mean(dim=0, keepdim =True)
        mean_target= target.mean(dim=0, keepdim= True)

        source -=mean_source
        target -=mean_target
        # source [N,3]
        # target [N,3]
        pspt= source.T @ target
        # svd solve
        # u @ diag(s) @ v = pspt
        # NOTE v is equal to V.T in the tranditional svd
        u,s,v_t = torch.svd(pspt)
        # s = torch.diag(s)
        # print((u @ s @ v_t) - pspt)

        v = v_t.T
        s = torch.eye(3).to(source.device)
        s[2,2] = torch.det(v @ u.T)
        solve_R = v @ s @ u.T

        solve_t = mean_target - (solve_R @ mean_source.T).T

        return solve_R, solve_t
//  [python]
def solver(self, source, target):
        '''
        Given source pts, and target pts, we solve the optimal rotation matrix, R and transform matrix T
        min ||R * source +t - tareget||2

        where t^* = mean(P_taregt) - R^* mean(P_source)


        For R^*
        we have, R^{*} = argmax(sum_{1}^{N}\hat{P_t}R\hat{P_s})
                        = trace(\hat{P_t}R\hat{P_s})

        argmax trace{\hat{P_t}R\hat{P_s}}

        R = v @ u.T

        but avoid reflection rotaiton
        we add a constrain
        let
        R = v diag[1, 1, det(v@u.t)] @ u.t
        return
            R, t
        '''
        mean_source = source.mean(dim=0, keepdim =True)
        mean_target= target.mean(dim=0, keepdim= True)

        source -=mean_source
        target -=mean_target
        # source [N,3]
        # target [N,3]
        pspt= source.T @ target
        # svd solve
        # u @ diag(s) @ v = pspt
        # NOTE v is equal to V.T in the tranditional svd
        u,s,v_t = torch.linalg.svd(pspt)
        # s = torch.diag(s)
        # print((u @ s @ v_t) - pspt)

        v = v_t.T
        s = torch.eye(3).to(source.device)
        s[2,2] = torch.det(v @ u.T)
        solve_R = v @ s @ u.T

        solve_t = mean_target - (solve_R @ mean_source.T).T

        return solve_R, solve_t
//  [python]
def test_linalg_svd_compute_uv(self, device, dtype):
        """
        Test the default case, compute_uv=True. Here we have the very same behavior as
        numpy
        """
        t = torch.randn((10, 11), device=device, dtype=dtype)
        np_t = t.cpu().numpy()
        for full_matrices in (True, False):
            # check linalg.svd vs numpy
            expected = np.linalg.svd(np_t, full_matrices, compute_uv=True)
            actual = t.svd(full_matrices, compute_uv=True)
            # sign/phase of the singular vectors is not unique and therefore absolute values are compared
            self.assertEqual(abs(actual[0]), abs(expected[0]))
            self.assertEqual(actual[1], expected[1])
            self.assertEqual(abs(actual[2]), abs(expected[2]))
            # check linalg.svd vs linalg.svd(out=...)
            out = (torch.empty_like(actual[0]),
                   torch.empty_like(actual[1]),
                   torch.empty_like(actual[2]))
            out2 = t.svd(full_matrices, compute_uv=True, out=out)
            self.assertEqual(actual, out)
            self.assertEqual(actual, out2)
//  [python]
def test_linalg_svd_compute_uv(self, device, dtype):
        """
        Test the default case, compute_uv=True. Here we have the very same behavior as
        numpy
        """
        t = torch.randn((10, 11), device=device, dtype=dtype)
        np_t = t.cpu().numpy()
        for full_matrices in (True, False):
            # check linalg.svd vs numpy
            expected = np.linalg.svd(np_t, full_matrices, compute_uv=True)
            actual = torch.linalg.svd(t, full_matrices, compute_uv=True)
            # sign/phase of the singular vectors is not unique and therefore absolute values are compared
            self.assertEqual(abs(actual[0]), abs(expected[0]))
            self.assertEqual(actual[1], expected[1])
            self.assertEqual(abs(actual[2]), abs(expected[2]))
            # check linalg.svd vs linalg.svd(out=...)
            out = (torch.empty_like(actual[0]),
                   torch.empty_like(actual[1]),
                   torch.empty_like(actual[2]))
            out2 = torch.linalg.svd(t, full_matrices, compute_uv=True, out=out)
            self.assertEqual(actual, out)
            self.assertEqual(actual, out2)
//  [python]
def singular_value_clip(w):
    with torch.no_grad():
        dim = w.shape
        if len(dim) > 2:
            w = w.reshape(dim[0], -1)
        u, s, v = torch.svd(w)
        s[s > 1] = 1
        return (u @ torch.diag(s) @ v).view(dim)
//  [python]
def singular_value_clip(w):
    with torch.no_grad():
        dim = w.shape
        if len(dim) > 2:
            w = w.reshape(dim[0], -1)
        u, s, v = torch.linalg.svd(w, full_matrices=False)
        s[s > 1] = 1
        return (u @ torch.diag(s) @ v).view(dim)
//  [python]
def test_linalg_svd_out_different_device(self, device, dtype):
        t = torch.randn(5, 7, device=device, dtype=dtype)  # this is on cuda
        u = torch.empty((5, 5), device='cpu', dtype=dtype)
        s = torch.empty((5,), device='cpu', dtype=dtype)
        v = torch.empty((7, 7), device='cpu', dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, 'svd output tensor U is on the wrong device: expected cuda:.* got cpu'):
            torch.svd(t, out=(u, s, v))
//  [python]
def test_linalg_svd_out_different_device(self, device, dtype):
        t = torch.randn(5, 7, device=device, dtype=dtype)  # this is on cuda
        u = torch.empty((5, 5), device='cpu', dtype=dtype)
        s = torch.empty((5,), device='cpu', dtype=dtype)
        v = torch.empty((7, 7), device='cpu', dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, 'svd output tensor U is on the wrong device: expected cuda:.* got cpu'):
            torch.linalg.svd(t, out=(u, s, v))
//  [python]
def integratibility_normalization(B, L, h, w, sigma=2):

    B = B.reshape(-1, h, w) # (3, h, w)
    
    B_blurred = gaussian_2d(B, sigma)
    B_gradx = F.pad(torch.diff(B_blurred, dim=2), (0, 1, 0, 0, 0, 0)).to(B.device)
    B_grady = F.pad(torch.diff(B_blurred, dim=1), (0, 0, 0, 1, 0, 0)).to(B.device)

    B_gradx = B_gradx.reshape(B.shape[0], -1) # (3, H*W)
    B_grady = B_grady.reshape(B.shape[0], -1) # (3, H*W)

    B = B.reshape(B.shape[0], -1) # (3, H*W)
    
    A1 = B[0] * B_gradx[1] - B[1] * B_gradx[0]
    A2 = B[0] * B_gradx[2] - B[2] * B_gradx[0]
    A3 = B[1] * B_gradx[2] - B[2] * B_gradx[1]
    A4 = -B[0] * B_grady[1] + B[1] * B_grady[0]
    A5 = -B[0] * B_grady[2] + B[2] * B_grady[0]
    A6 = -B[1] * B_grady[2] + B[2] * B_grady[1]
    A = torch.stack([A1, A2, A3, A4, A5, A6]).T

    U, S, VH = torch.svd(A)
    # which one is x?
    x = VH[-1] # right singular vector corresponding to smallest singular value
    D = torch.tensor([
        [-x[2], x[5], 1],
        [x[1], -x[4], 0],
        [-x[0], x[3], 0]
    ]).float().to(B.device)
    B = torch.linalg.inv(D) @ B
    L = D.T @ L
    return B, L
//  [python]
def integratibility_normalization(B, L, h, w, sigma=2):

    B = B.reshape(-1, h, w) # (3, h, w)
    
    B_blurred = gaussian_2d(B, sigma)
    B_gradx = F.pad(torch.diff(B_blurred, dim=2), (0, 1, 0, 0, 0, 0)).to(B.device)
    B_grady = F.pad(torch.diff(B_blurred, dim=1), (0, 0, 0, 1, 0, 0)).to(B.device)

    B_gradx = B_gradx.reshape(B.shape[0], -1) # (3, H*W)
    B_grady = B_grady.reshape(B.shape[0], -1) # (3, H*W)

    B = B.reshape(B.shape[0], -1) # (3, H*W)
    
    A1 = B[0] * B_gradx[1] - B[1] * B_gradx[0]
    A2 = B[0] * B_gradx[2] - B[2] * B_gradx[0]
    A3 = B[1] * B_gradx[2] - B[2] * B_gradx[1]
    A4 = -B[0] * B_grady[1] + B[1] * B_grady[0]
    A5 = -B[0] * B_grady[2] + B[2] * B_grady[0]
    A6 = -B[1] * B_grady[2] + B[2] * B_grady[1]
    A = torch.stack([A1, A2, A3, A4, A5, A6]).T

    U, S, VH = torch.linalg.svd(A, full_matrices=False)
    # which one is x?
    x = VH[-1] # right singular vector corresponding to smallest singular value
    D = torch.tensor([
        [-x[2], x[5], 1],
        [x[1], -x[4], 0],
        [-x[0], x[3], 0]
    ]).float().to(B.device)
    B = torch.linalg.inv(D) @ B
    L = D.T @ L
    return B, L
//  [python]
def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
    out_size, in_size, kernel_size, _ = weight.size()
    U, S, Vh = torch.svd(weight.reshape(out_size, -1).to(device))
    
    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
    lora_rank = param_dict["new_rank"]

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]

    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size, kernel_size, kernel_size).cpu()
    param_dict["lora_up"] = U.reshape(out_size, lora_rank, 1, 1).cpu()
    del U, S, Vh, weight
    return param_dict
//  [python]
def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
    out_size, in_size, kernel_size, _ = weight.size()
    U, S, Vh = torch.linalg.svd(weight.reshape(out_size, -1).to(device))
    
    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
    lora_rank = param_dict["new_rank"]

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]

    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size, kernel_size, kernel_size).cpu()
    param_dict["lora_up"] = U.reshape(out_size, lora_rank, 1, 1).cpu()
    del U, S, Vh, weight
    return param_dict
//  [python]
def pca(x, k, center=True):
    if center:
        m = x.mean(0, keepdim=True)
        s = x.std(0, unbiased=False, keepdim=True)
        x -= m
        x /= s
    U, S, V = torch.svd(x)
    reduced = torch.mm(x, V[:k].T)

    return reduced
//  [python]
def pca(x, k, center=True):
    if center:
        m = x.mean(0, keepdim=True)
        s = x.std(0, unbiased=False, keepdim=True)
        x -= m
        x /= s
    # why pca related to svd? https://www.cs.cmu.edu/~elaw/papers/pca.pdf chap VI
    U, S, V = torch.linalg.svd(x)
    reduced = torch.mm(x, V[:k].T)

    return reduced
//  [python]
def null_space(A: Tensor) -> Tensor:
    '''
    Solve for x so that A @ x = null
    '''
    U, S, Vh = torch.svd(A)
    # get the last row of Vh
    # (the last column of V)
    return Vh[..., -1, :]
//  [python]
def null_space(A: Tensor) -> Tensor:
    '''
    Solve for x so that A @ x = null
    '''
    U, S, Vh = torch.linalg.svd(A)
    # get the last row of Vh
    # (the last column of V)
    return Vh[..., -1, :]
//  [python]
def align_pcl(Y, X, weight=None, fixed_scale=False):
    """align similarity transform to align X with Y using umeyama method
    X' = s * R * X + t is aligned with Y
    :param Y (*, N, 3) first trajectory
    :param X (*, N, 3) second trajectory
    :param weight (*, N, 1) optional weight of valid correspondences
    :returns s (*, 1), R (*, 3, 3), t (*, 3)
    """
    X = torch.as_tensor(X).clone()
    Y = torch.as_tensor(Y).clone()
    *dims, N, _ = Y.shape
    N = torch.ones(*dims, 1, 1) * N

    if weight is not None:
        Y = Y * weight
        X = X * weight
        N = weight.sum(dim=-2, keepdim=True)  # (*, 1, 1)

    # subtract mean
    my = Y.sum(dim=-2) / N[..., 0]  # (*, 3)
    mx = X.sum(dim=-2) / N[..., 0]
    y0 = Y - my[..., None, :]  # (*, N, 3)
    x0 = X - mx[..., None, :]

    if weight is not None:
        y0 = y0 * weight
        x0 = x0 * weight

    # correlation
    C = torch.matmul(y0.transpose(-1, -2), x0) / N  # (*, 3, 3)
    U, S, V = torch.svd(C)  # (*, 3, 3), (*, 3), (*, 3, 3)

    S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
    neg = torch.det(U) * torch.det(V.transpose(-1, -2)) < 0
    S[neg, 2, 2] = -1

    R = torch.matmul(U, torch.matmul(S, V.transpose(-1, -2)))  # (*, 3, 3)

    D = torch.diag_embed(S)  # (*, 3, 3)
    if fixed_scale:
        s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
    else:
        var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N  # (*, 1, 1)
        s = (
            torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(
                dim=-1, keepdim=True
            )
            / var[..., 0]
        )  # (*, 1)

    t = my - s * torch.matmul(R, mx[..., None])[..., 0]  # (*, 3)

    return s, R, t
//  [python]
def align_pcl(Y, X, weight=None, fixed_scale=False):
    """align similarity transform to align X with Y using umeyama method
    X' = s * R * X + t is aligned with Y
    :param Y (*, N, 3) first trajectory
    :param X (*, N, 3) second trajectory
    :param weight (*, N, 1) optional weight of valid correspondences
    :returns s (*, 1), R (*, 3, 3), t (*, 3)
    """
    X = torch.as_tensor(X).clone()
    Y = torch.as_tensor(Y).clone()
    *dims, N, _ = Y.shape
    N = torch.ones(*dims, 1, 1) * N

    if weight is not None:
        Y = Y * weight
        X = X * weight
        N = weight.sum(dim=-2, keepdim=True)  # (*, 1, 1)

    # subtract mean
    my = Y.sum(dim=-2) / N[..., 0]  # (*, 3)
    mx = X.sum(dim=-2) / N[..., 0]
    y0 = Y - my[..., None, :]  # (*, N, 3)
    x0 = X - mx[..., None, :]

    if weight is not None:
        y0 = y0 * weight
        x0 = x0 * weight

    # correlation
    C = torch.matmul(y0.transpose(-1, -2), x0) / N  # (*, 3, 3)
    U, D, Vh = torch.linalg.svd(C)  # (*, 3, 3), (*, 3), (*, 3, 3)

    S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
    neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0
    S[neg, 2, 2] = -1

    R = torch.matmul(U, torch.matmul(S, Vh))  # (*, 3, 3)

    D = torch.diag_embed(D)  # (*, 3, 3)
    if fixed_scale:
        s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
    else:
        var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N  # (*, 1, 1)
        s = (
            torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(
                dim=-1, keepdim=True
            )
            / var[..., 0]
        )  # (*, 1)

    t = my - s * torch.matmul(R, mx[..., None])[..., 0]  # (*, 3)

    return s, R, t
//  [python]
def test_linalg_svd_no_compute_uv(self, device, dtype):
        """
        Test the compute_uv=False case. Here we have a different return type than
        numpy: numpy returns S, we return (empty, S, empty)
        """
        t = torch.randn((10, 11), device=device, dtype=dtype)
        np_t = t.cpu().numpy()

        def is_empty(x):
            return x.numel() == 0 and x.dtype == t.dtype and x.device == t.device

        for full_matrices in (True, False):
            # check linalg.svd vs numpy
            np_s = np.linalg.svd(np_t, full_matrices, compute_uv=False)
            U, S, V = t.svd(full_matrices, compute_uv=False)
            USV = (U, S, V)
            assert is_empty(USV[0])
            self.assertEqual(USV[1], np_s)
            assert is_empty(USV[2])
            # check linalg.svd vs linalg.svd(out=...)
            out = (torch.empty_like(USV[0]), torch.empty_like(USV[1]), torch.empty_like(USV[2]))
            U, S, V = t.svd(full_matrices, compute_uv=False, out=out)
            USV = (U, S, V)
            assert USV[0] is out[0]
            assert USV[1] is out[1]
            assert USV[2] is out[2]
            self.assertEqual(USV[1], np_s)