jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def batch_get_3children_orient_svd(
    rel_pose_skeleton, rel_rest_pose, rot_mat_chain_parent, children_list, dtype
):
    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))
    det_u_v = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_v
    rot_mat = torch.bmm(torch.bmm(V, det_modify_mat), 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
):
    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))
    det_u_v = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_v
    rot_mat = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)
    return rot_mat
//  [python]
def kbash_alignment(coor1, coor2):
    '''
    coor1: point cloud 1, shape = (N, 3)
    coor2: point cloud 2, shape = (N, 3)
    '''
    import torch

    A = ( coor2 - coor2.mean(dim=0) ).t() @ ( coor1 - coor1.mean(dim=0) )
    U, S, Vt = torch.svd(A)
    rotation = ( U @ torch.diag(torch.Tensor([1, 1, torch.sign(torch.det(A))]).to(coor1.device)) ) @ Vt
    translation = coor2.mean(dim=0) - (rotation @ coor1.mean(dim=0, keepdim=True).t()).t()

    # detach the gradients if there are
    rotation = rotation.detach()
    translation = translation.detach()

    coor1 = coor1 @ rotation.t() + translation
    return coor1, rotation, translation
//  [python]
def kbash_alignment(coor1, coor2):
    '''
    coor1: point cloud 1, shape = (N, 3)
    coor2: point cloud 2, shape = (N, 3)
    '''
    import torch

    A = ( coor2 - coor2.mean(dim=0) ).t() @ ( coor1 - coor1.mean(dim=0) )
    U, S, Vt = torch.linalg.svd(A)
    rotation = ( U @ torch.diag(torch.Tensor([1, 1, torch.sign(torch.det(A))]).to(coor1.device)) ) @ Vt
    translation = coor2.mean(dim=0) - (rotation @ coor1.mean(dim=0, keepdim=True).t()).t()

    # detach the gradients if there are
    rotation = rotation.detach()
    translation = translation.detach()

    coor1 = coor1 @ rotation.t() + translation
    return coor1, rotation, translation
//  [python]
def _svd_basis(points):
        # Select 4 virtual control points with SVD
        center = points.mean(dim=-2, keepdim=True)
        translated = points - center
        u, s, vh = torch.svd(translated.mT @ translated)
        controls = center + s.sqrt().unsqueeze(-1) * vh.mT
        return torch.cat([center, controls], dim=-2)
//  [python]
def _svd_basis(points):
        # Select 4 virtual control points with SVD
        center = points.mean(dim=-2, keepdim=True)
        translated = points - center
        u, s, vh = torch.linalg.svd(translated.mT @ translated)
        controls = center + s.sqrt().unsqueeze(-1) * vh.mT
        return torch.cat([center, controls], dim=-2)
//  [python]
module_names = module_dict.keys()

# Plot Conv2d or Linear
module_visual = [i for i in module_dict.keys() if isinstance(
    module_dict[i], torch.nn.Conv2d) or isinstance(module_dict[i], torch.nn.Linear) or isinstance(module_dict[i], torch.nn.BatchNorm2d)]

df = None

max_num_neuron = 0
for module_name in module_visual:
    target_layer = module_dict[module_name]
    
    print(f'Collecting Lips {target_layer}')
    if isinstance(target_layer, torch.nn.Linear):
        channel_lips = []
        for idx in range(target_layer.weight.shape[0]):
            w = target_layer.weight[idx].reshape(target_layer.weight.shape[1], -1)
            # Just norm of weight for linear layer
            channel_lips.append(torch.svd(w)[1].max())
        channel_lips = torch.Tensor(channel_lips)

    elif isinstance(target_layer, torch.nn.BatchNorm2d):
        std = target_layer.running_var.sqrt()
        weight = target_layer.weight

        channel_lips = []
        for idx in range(weight.shape[0]):
            w = conv.weight[idx].reshape(conv.weight.shape[1], -1) * (weight[idx]/std[idx]).abs()
            channel_lips.append(torch.svd(w)[1].max())
        channel_lips = torch.Tensor(channel_lips)

        
        # Convolutional layer should be followed by a BN layer by default
    elif isinstance(target_layer, torch.nn.Conv2d):
        conv = target_layer    
            
        channel_lips = []
        for idx in range(target_layer.weight.shape[0]):
            w = target_layer.weight[idx].reshape(target_layer.weight.shape[1], -1)
            channel_lips.append(torch.svd(w)[1].max())
        channel_lips = torch.Tensor(channel_lips)
    else:
        assert False, "Unknown layer type"

    for neuron_i in range(channel_lips.shape[0]):
        base_row = {}
        base_row['layer'] = module_name
        base_row['Neuron'] = neuron_i
        base_row['Lips'] = channel_lips[neuron_i].item()
        if df is None:
            df = pd.DataFrame.from_dict([base_row])
        else:
            df.loc[df.shape[0]] = base_row
//  [python]
module_names = module_dict.keys()

# Plot Conv2d or Linear
module_visual = [i for i in module_dict.keys() if isinstance(
    module_dict[i], torch.nn.Conv2d) or isinstance(module_dict[i], torch.nn.Linear) or isinstance(module_dict[i], torch.nn.BatchNorm2d)]

df = None

max_num_neuron = 0
for module_name in module_visual:
    target_layer = module_dict[module_name]
    
    print(f'Collecting Lips {target_layer}')
    if isinstance(target_layer, torch.nn.Linear):
        channel_lips = []
        for idx in range(target_layer.weight.shape[0]):
            w = target_layer.weight[idx].reshape(target_layer.weight.shape[1], -1)
            # Just norm of weight for linear layer
            channel_lips.append(torch.linalg.svd(w)[1].max())
        channel_lips = torch.Tensor(channel_lips)

    elif isinstance(target_layer, torch.nn.BatchNorm2d):
        std = target_layer.running_var.sqrt()
        weight = target_layer.weight

        channel_lips = []
        for idx in range(weight.shape[0]):
            w = conv.weight[idx].reshape(conv.weight.shape[1], -1) * (weight[idx]/std[idx]).abs()
            channel_lips.append(torch.linalg.svd(w)[1].max())
        channel_lips = torch.Tensor(channel_lips)

        
        # Convolutional layer should be followed by a BN layer by default
    elif isinstance(target_layer, torch.nn.Conv2d):
        conv = target_layer    
            
        channel_lips = []
        for idx in range(target_layer.weight.shape[0]):
            w = target_layer.weight[idx].reshape(target_layer.weight.shape[1], -1)
            channel_lips.append(torch.linalg.svd(w)[1].max())
        channel_lips = torch.Tensor(channel_lips)
    else:
        assert False, "Unknown layer type"

    for neuron_i in range(channel_lips.shape[0]):
        base_row = {}
        base_row['layer'] = module_name
        base_row['Neuron'] = neuron_i
        base_row['Lips'] = channel_lips[neuron_i].item()
        if df is None:
            df = pd.DataFrame.from_dict([base_row])
        else:
            df.loc[df.shape[0]] = base_row
//  [python]
def svd_sigma(tensor):
    # returns the singular values of a tensor
    # working with different torch versions
    if version.parse(torch.__version__) > version.parse('1.8.1'):
        sigma = torch.svd(tensor, compute_uv=True).S
    else:
        sigma = torch.svd(tensor, compute_uv=True).S
    
    return sigma
//  [python]
def svd_sigma(tensor):
    # returns the singular values of a tensor
    # working with different torch versions
    if version.parse(torch.__version__) > version.parse('1.8.1'):
        sigma = torch.linalg.svd(tensor).S
    else:
        sigma = torch.svd(tensor, compute_uv=True).S
    
    return sigma
//  [python]
def svd_cpu(x, eps=1e-4, require_backward=False):
    device = x.device
    x = x.cpu()

    _perturb = lambda _y, _eps: _y + _eps * _y.mean() * torch.rand_like(_y)

    try:
        U, S, Vh = torch.svd(x)
    except:
        U, S, Vh = torch.svd(_perturb(x, eps))

    while require_backward and len(torch.unique(S)) != len(S):
        U, S, Vh = torch.svd(_perturb(x, eps))
        eps *= 2

    return U.to(device), S.to(device), torch.t(Vh).to(device)
//  [python]
def svd_cpu(x, eps=1e-4, require_backward=False):
    device = x.device
    x = x.cpu()

    _perturb = lambda _y, _eps: _y + _eps * _y.mean() * torch.rand_like(_y)

    try:
        U, S, Vh = torch.linalg.svd(x)
    except:
        U, S, Vh = torch.linalg.svd(_perturb(x, eps))

    while require_backward and len(torch.unique(S)) != len(S):
        U, S, Vh = torch.linalg.svd(_perturb(x, eps))
        eps *= 2

    return U.to(device), S.to(device), torch.t(Vh).to(device)
//  [python]
def _svd_lowrank(A: Tensor, q: Optional[int] = 6, niter: Optional[int] = 2,
                 M: Optional[Tensor] = None) -> Tuple[Tensor, Tensor, Tensor]:
    q = 6 if q is None else q
    m, n = A.shape[-2:]
    matmul = _utils.matmul
    if M is None:
        M_t = None
    else:
        M_t = _utils.transpose(M)
    A_t = _utils.transpose(A)

    # Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
    # the number conjugate and transpose operations
    if m < n or n > q:
        # computing the SVD approximation of a transpose in
        # order to keep B shape minimal (the m < n case) or the V
        # shape small (the n > q case)
        Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B_t = matmul(A, Q_c)
        else:
            B_t = matmul(A, Q_c) - matmul(M, Q_c)
        assert B_t.shape[-2] == m, (B_t.shape, m)
        assert B_t.shape[-1] == q, (B_t.shape, q)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.svd(B_t)
        V = Vh.mH
        V = Q.matmul(V)
    else:
        Q = get_approximate_basis(A, q, niter=niter, M=M)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B = matmul(A_t, Q_c)
        else:
            B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
        B_t = _utils.transpose(B)
        assert B_t.shape[-2] == q, (B_t.shape, q)
        assert B_t.shape[-1] == n, (B_t.shape, n)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.svd(B_t)
        V = Vh.mH
        U = Q.matmul(U)

    return U, S, V
//  [python]
def _svd_lowrank(A: Tensor, q: Optional[int] = 6, niter: Optional[int] = 2,
                 M: Optional[Tensor] = None) -> Tuple[Tensor, Tensor, Tensor]:
    q = 6 if q is None else q
    m, n = A.shape[-2:]
    matmul = _utils.matmul
    if M is None:
        M_t = None
    else:
        M_t = _utils.transpose(M)
    A_t = _utils.transpose(A)

    # Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
    # the number conjugate and transpose operations
    if m < n or n > q:
        # computing the SVD approximation of a transpose in
        # order to keep B shape minimal (the m < n case) or the V
        # shape small (the n > q case)
        Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B_t = matmul(A, Q_c)
        else:
            B_t = matmul(A, Q_c) - matmul(M, Q_c)
        assert B_t.shape[-2] == m, (B_t.shape, m)
        assert B_t.shape[-1] == q, (B_t.shape, q)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.linalg.svd(B_t, full_matrices=False)
        V = Vh.mH
        V = Q.matmul(V)
    else:
        Q = get_approximate_basis(A, q, niter=niter, M=M)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B = matmul(A_t, Q_c)
        else:
            B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
        B_t = _utils.transpose(B)
        assert B_t.shape[-2] == q, (B_t.shape, q)
        assert B_t.shape[-1] == n, (B_t.shape, n)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.linalg.svd(B_t, full_matrices=False)
        V = Vh.mH
        U = Q.matmul(U)

    return U, S, V
//  [python]
def normalizeSE3_inplace(T):
    R = T[...,:3,:3]
    U, S, Vh = torch.svd(R)
    T[...,:3,:3] = torch.matmul(U, Vh)
//  [python]
def normalizeSE3_inplace(T):
    R = T[...,:3,:3]
    U, S, Vh = torch.linalg.svd(R)
    T[...,:3,:3] = torch.matmul(U, Vh)
//  [python]
def svdtf(source, target):
    r'''
    Computes the rigid transformation ( :math:`SE(3)` ) between two sets of associated
    point clouds (source and target) using Singular Value Decomposition (SVD).

    Args:
        source (``torch.Tensor``): the coordinates of the source point cloud.
            The shape has to be (..., N, 3).
        target (``torch.Tensor``): the coordinates of the target point cloud.
            The shape has to be (..., N, 3).

    Returns:
        ``LieTensor``: The rigid transformation matrix in ``SE3Type``  that
        minimizes the mean squared error between the input point sets.

    Warning:
        The number of points N has to be the same for both point clouds.

    Example:
        >>> import torch, pypose as pp
        >>> source = torch.tensor([[0., 0., 0.],
        ...                     [1., 0., 0.],
        ...                     [0., 1., 0.]])
        >>> target = torch.tensor([[1., 1., 1.],
        ...                     [2., 1., 1.],
        ...                     [1., 2., 1.]])
        >>> pp.svdtf(source, target)
        SE3Type LieTensor:
        LieTensor([1., 1., 1., 0., 0., 0., 1.])
    '''
    assert source.size(-2) == target.size(-2), {
        "The number of points N has to be the same for both point clouds."}
    ctnsource = source.mean(dim=-2, keepdim=True)
    ctntarget = target.mean(dim=-2, keepdim=True)
    source = source - ctnsource
    target = target - ctntarget
    M = torch.einsum('...Na, ...Nb -> ...ab', target, source)
    U, S, Vh = torch.svd(M)
    R = U @ torch.diag_embed(torch.sign(torch.diagonal(U @ Vh, dim1=-2, dim2=-1))) @ Vh
    t = ctntarget.mT - R @ ctnsource.mT
    T = torch.cat((R, t), dim=-1)
    return mat2SE3(T, check=False)
//  [python]
def svdtf(source, target):
    r'''
    Computes the rigid transformation ( :math:`SE(3)` ) between two sets of associated
    point clouds (source and target) using Singular Value Decomposition (SVD).

    Args:
        source (``torch.Tensor``): the coordinates of the source point cloud.
            The shape has to be (..., N, 3).
        target (``torch.Tensor``): the coordinates of the target point cloud.
            The shape has to be (..., N, 3).

    Returns:
        ``LieTensor``: The rigid transformation matrix in ``SE3Type``  that
        minimizes the mean squared error between the input point sets.

    Warning:
        The number of points N has to be the same for both point clouds.

    Example:
        >>> import torch, pypose as pp
        >>> source = torch.tensor([[0., 0., 0.],
        ...                     [1., 0., 0.],
        ...                     [0., 1., 0.]])
        >>> target = torch.tensor([[1., 1., 1.],
        ...                     [2., 1., 1.],
        ...                     [1., 2., 1.]])
        >>> pp.svdtf(source, target)
        SE3Type LieTensor:
        LieTensor([1., 1., 1., 0., 0., 0., 1.])
    '''
    assert source.size(-2) == target.size(-2), {
        "The number of points N has to be the same for both point clouds."}
    ctnsource = source.mean(dim=-2, keepdim=True)
    ctntarget = target.mean(dim=-2, keepdim=True)
    source = source - ctnsource
    target = target - ctntarget
    M = torch.einsum('...Na, ...Nb -> ...ab', target, source)
    U, S, Vh = torch.linalg.svd(M)
    R = U @ Vh
    mask = (R.det() + 1).abs() < 1e-6
    R[mask] = - R[mask]
    t = ctntarget.mT - R @ ctnsource.mT
    T = torch.cat((R, t), dim=-1)
    return mat2SE3(T, check=False)
//  [python]
def __init__(self, config, data_feature):
        super().__init__(config, data_feature)

        self.adj_mx = data_feature.get('adj_mx')
        self.num_nodes = data_feature.get('num_nodes', 1)
        self.feature_dim = data_feature.get('feature_dim', 2)
        self.output_dim = data_feature.get('output_dim', 2)

        self.hidden_size = config.get('hidden_size', 25)
        self.n_dim = config.get('n_dim', 50)
        self.n_supports = config.get('n_supports', 1)
        self.k_hop = config.get('k_hop', 3)
        self.n_rnn_layers = config.get('n_rnn_layers', 1)
        self.n_gconv_layers = config.get('n_gconv_layers', 3)
        self.cl_decay_steps = config.get('cl_decay_steps', 300)
        self.graph_category = config.get('graph_category', 'gau')
        self.normalized_category = config.get('normalized_category', 'randomwalk')

        self.input_window = config.get('input_window', 1)
        self.output_window = config.get('output_window', 1)
        self.device = config.get('device', torch.device('cpu'))
        self._logger = getLogger()
        self._scaler = self.data_feature.get('scaler')

        self.adj_mx = graph_preprocess(self.adj_mx, self.normalized_category)
        self.adj_mx = torch.from_numpy(self.adj_mx).float().to(self.device)
        n, k = self.adj_mx.shape
        if n == k:
            self.method = 'big'
            m, p, n = torch.svd(self.adj_mx)
            initemb1 = torch.mm(m[:, :self.n_dim], torch.diag(p[:self.n_dim] ** 0.5))
            initemb2 = torch.mm(torch.diag(p[:self.n_dim] ** 0.5), n[:, :self.n_dim].t())
            self.nodevec1 = nn.Parameter(initemb1, requires_grad=True)
            self.nodevec2 = nn.Parameter(initemb2, requires_grad=True)
        else:
            self.method = 'small'
            self.w, self.m = self._delta_cal(self.adj_mx)
            self.w = self.w.to(self.device)
            self.cov = nn.Parameter(self.m, requires_grad=True)

        self.encoder = DCRNNEncoder(self.feature_dim, self.hidden_size, self.num_nodes, self.n_supports,
                                    self.k_hop, self.n_rnn_layers, self.n_gconv_layers, self.n_dim)
        self.decoder = DCRNNDecoder(self.output_dim, self.hidden_size, self.num_nodes, self.n_supports,
                                    self.k_hop, self.n_rnn_layers, self.output_window,
                                    self.n_gconv_layers, self.n_dim)

        self.w1 = nn.Parameter(torch.eye(self.n_dim), requires_grad=True)
        self.w2 = nn.Parameter(torch.eye(self.n_dim), requires_grad=True)
        self.b1 = nn.Parameter(torch.zeros(self.n_dim), requires_grad=True)
        self.b2 = nn.Parameter(torch.zeros(self.n_dim), requires_grad=True)
        self.graph0 = None
        self.graph1 = None
        self.graph2 = None
//  [python]
def __init__(self, config, data_feature):
        super().__init__(config, data_feature)

        self.adj_mx = data_feature.get('adj_mx')
        self.num_nodes = data_feature.get('num_nodes', 1)
        self.feature_dim = data_feature.get('feature_dim', 2)
        self.output_dim = data_feature.get('output_dim', 2)

        self.hidden_size = config.get('hidden_size', 25)
        self.n_dim = config.get('n_dim', 50)
        self.n_supports = config.get('n_supports', 1)
        self.k_hop = config.get('k_hop', 3)
        self.n_rnn_layers = config.get('n_rnn_layers', 1)
        self.n_gconv_layers = config.get('n_gconv_layers', 3)
        self.cl_decay_steps = config.get('cl_decay_steps', 300)
        self.graph_category = config.get('graph_category', 'gau')
        self.normalized_category = config.get('normalized_category', 'randomwalk')

        self.input_window = config.get('input_window', 1)
        self.output_window = config.get('output_window', 1)
        self.device = config.get('device', torch.device('cpu'))
        self._logger = getLogger()
        self._scaler = self.data_feature.get('scaler')

        self.adj_mx = graph_preprocess(self.adj_mx, self.normalized_category)
        self.adj_mx = torch.from_numpy(self.adj_mx).float().to(self.device)
        n, k = self.adj_mx.shape
        if n == k:
            self.method = 'big'
            m, s, n = torch.linalg.svd(self.adj_mx)
            initemb1 = torch.mm(m[:, :self.n_dim], torch.diag(s[:self.n_dim] ** 0.5))
            initemb2 = torch.mm(torch.diag(s[:self.n_dim] ** 0.5), n[:, :self.n_dim].t())
            self.nodevec1 = nn.Parameter(initemb1, requires_grad=True)
            self.nodevec2 = nn.Parameter(initemb2, requires_grad=True)
        else:
            self.method = 'small'
            self.w, self.m = self._delta_cal(self.adj_mx)
            self.w = self.w.to(self.device)
            self.cov = nn.Parameter(self.m, requires_grad=True)

        self.encoder = DCRNNEncoder(self.feature_dim, self.hidden_size, self.num_nodes, self.n_supports,
                                    self.k_hop, self.n_rnn_layers, self.n_gconv_layers, self.n_dim)
        self.decoder = DCRNNDecoder(self.output_dim, self.hidden_size, self.num_nodes, self.n_supports,
                                    self.k_hop, self.n_rnn_layers, self.output_window,
                                    self.n_gconv_layers, self.n_dim)

        self.w1 = nn.Parameter(torch.eye(self.n_dim), requires_grad=True)
        self.w2 = nn.Parameter(torch.eye(self.n_dim), requires_grad=True)
        self.b1 = nn.Parameter(torch.zeros(self.n_dim), requires_grad=True)
        self.b2 = nn.Parameter(torch.zeros(self.n_dim), requires_grad=True)
        self.graph0 = None
        self.graph1 = None
        self.graph2 = None
//  [python]
def reset_parameters(self, input):
        with torch.no_grad():
            U, _, _ = torch.svd(input.T)
            self.weight.data = U[:, :self.k]
        self.weight.data = self.weight.data.to(self.device)  # TODO: necessary?
        self._init_done = True
//  [python]
def reset_parameters(self, input):
        with torch.no_grad():
            U, _, _ = torch.linalg.svd(input.T, full_matrices=False)
            self.weight.data = U[:, :self.k]
        self.weight.data = self.weight.data.to(self.device)  # TODO: necessary?
        self._init_done = True
//  [python]
def Grassmann_distance(P1, P2):
    proj = P1.T@P2
    if type(P1) == type(P2) == torch.Tensor:
        _, s, _ = torch.svd(proj)
        s[s > 1] = 1
        s = torch.arccos(s)
        return torch.norm(s)
    if type(P1) == type(P2) == np.ndarray:
        _, s, _ = np.linalg.svd(proj)
        s[s > 1] = 1
        s = np.arccos(s)
        return np.linalg.norm(s)
    return TypeError('Error: should be np.ndarray or torch.Tensor...')
//  [python]
def Grassmann_distance(P1, P2):
    proj = P1.T@P2
    if type(P1) == type(P2) == torch.Tensor:
        _, s, _ = torch.linalg.svd(proj)
        s[s > 1] = 1
        s = torch.arccos(s)
        return torch.linalg.norm(s)
    if type(P1) == type(P2) == np.ndarray:
        _, s, _ = np.linalg.svd(proj)
        s[s > 1] = 1
        s = np.arccos(s)
        return np.linalg.norm(s)
    return TypeError('Error: should be np.ndarray or torch.Tensor...')
//  [python]
def _svd_weighted_adapter(
        self,
        adapters,
        weights,
        new_rank,
        target,
        target_lora_A,
        target_lora_B,
        clamp=None,
        full_matrices=True,
        driver=None,
    ):
        delta_weight = weights[0] * target.get_delta_weight(adapters[0])
        for adapter, weight in zip(adapters[1:], weights[1:]):
            delta_weight += weight * target.get_delta_weight(adapter)
        conv2d = isinstance(target, Conv2d)
        if conv2d:
            conv2d_1x1 = target.weight.size()[2:4] == (1, 1)
            if not conv2d_1x1:
                delta_weight = delta_weight.flatten(start_dim=1)
            else:
                delta_weight = delta_weight.squeeze()
        if hasattr(target, "fan_in_fan_out") and target.fan_in_fan_out:
            delta_weight = delta_weight.T

        # based on https://github.com/kohya-ss/sd-scripts/blob/main/networks/svd_merge_lora.py#L114-L131
        U, S, Vh = torch.svd(delta_weight, some=full_matrices)
        U = U[:, :new_rank]
        S = S[:new_rank]
        U = U @ torch.diag(S)
        Vh = Vh[:new_rank, :]
        if clamp is not None:
            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp)
            low_val = -hi_val
            U = U.clamp(low_val, hi_val)
            Vh = Vh.clamp(low_val, hi_val)
        if conv2d:
            U = U.reshape(target_lora_B.data.shape)
            Vh = Vh.reshape(target_lora_A.data.shape)
        return Vh, U
//  [python]
def _svd_weighted_adapter(
        self,
        adapters,
        weights,
        new_rank,
        target,
        target_lora_A,
        target_lora_B,
        clamp=None,
        full_matrices=True,
        driver=None,
    ):
        delta_weight = weights[0] * target.get_delta_weight(adapters[0])
        for adapter, weight in zip(adapters[1:], weights[1:]):
            delta_weight += weight * target.get_delta_weight(adapter)
        conv2d = isinstance(target, Conv2d)
        if conv2d:
            conv2d_1x1 = target.weight.size()[2:4] == (1, 1)
            if not conv2d_1x1:
                delta_weight = delta_weight.flatten(start_dim=1)
            else:
                delta_weight = delta_weight.squeeze()
        if hasattr(target, "fan_in_fan_out") and target.fan_in_fan_out:
            delta_weight = delta_weight.T

        # based on https://github.com/kohya-ss/sd-scripts/blob/main/networks/svd_merge_lora.py#L114-L131
        U, S, Vh = torch.linalg.svd(delta_weight, full_matrices=full_matrices, driver=driver)
        U = U[:, :new_rank]
        S = S[:new_rank]
        U = U @ torch.diag(S)
        Vh = Vh[:new_rank, :]
        if clamp is not None:
            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp)
            low_val = -hi_val
            U = U.clamp(low_val, hi_val)
            Vh = Vh.clamp(low_val, hi_val)
        if conv2d:
            U = U.reshape(target_lora_B.data.shape)
            Vh = Vh.reshape(target_lora_A.data.shape)
        return Vh, U
//  [python]
def proj_pca(xs_f: torch.Tensor):
    """
    xs_f: (B, T, D)
    ===
    proj_xs_f: (B, T, 2)
    V: (D, 2) s.t. proj_xs = xs @ V
    """
    # xs_f: (batch, T, nx)
    # Only use final timestep of xs_f for PCA.
    batch, T, nx = xs_f.shape

    # (batch * T, nx)
    flat_xsf = xs_f.reshape(-1, *xs_f.shape[2:])

    # Center by subtract mean.
    # (batch, nx)
    final_xs_f = xs_f[:, -1, :]

    mean_pca_xs = torch.mean(final_xs_f, dim=0, keepdim=True)
    final_xs_f -= mean_pca_xs

    # if batch is too large, it will run out of memory.
    if batch > 200:
        rand_idxs = torch.randperm(batch)[:200]
        final_xs_f = final_xs_f[rand_idxs]

    # U: (batch, k)
    # S: (k, k)
    # VT: (k, nx)
    U, S, VT = torch.svd(final_xs_f)

    # log.info("Singular values of xs_f at final timestep:")
    # log.info(S)

    # Keep the first and last directions.
    VT = VT[:2, :]
    # VT = VT[[0, -1], :]

    assert VT.shape == (2, nx)
    V = VT.T

    # Project both xs_f and xs_b onto V.
    flat_xsf -= mean_pca_xs

    proj_xs_f = flat_xsf @ V
    proj_xs_f = proj_xs_f.reshape(batch, T, *proj_xs_f.shape[1:])

    return proj_xs_f, V
//  [python]
def proj_pca(xs_f: torch.Tensor):
    """
    xs_f: (B, T, D)
    ===
    proj_xs_f: (B, T, 2)
    V: (D, 2) s.t. proj_xs = xs @ V
    """
    # xs_f: (batch, T, nx)
    # Only use final timestep of xs_f for PCA.
    batch, T, nx = xs_f.shape

    # (batch * T, nx)
    flat_xsf = xs_f.reshape(-1, *xs_f.shape[2:])

    # Center by subtract mean.
    # (batch, nx)
    final_xs_f = xs_f[:, -1, :]

    mean_pca_xs = torch.mean(final_xs_f, dim=0, keepdim=True)
    final_xs_f -= mean_pca_xs

    # if batch is too large, it will run out of memory.
    if batch > 200:
        rand_idxs = torch.randperm(batch)[:200]
        final_xs_f = final_xs_f[rand_idxs]

    # U: (batch, k)
    # S: (k, k)
    # VT: (k, nx)
    U, S, VT = torch.linalg.svd(final_xs_f)

    # log.info("Singular values of xs_f at final timestep:")
    # log.info(S)

    # Keep the first and last directions.
    VT = VT[:2, :]
    # VT = VT[[0, -1], :]

    assert VT.shape == (2, nx)
    V = VT.T

    # Project both xs_f and xs_b onto V.
    flat_xsf -= mean_pca_xs

    proj_xs_f = flat_xsf @ V
    proj_xs_f = proj_xs_f.reshape(batch, T, *proj_xs_f.shape[1:])

    return proj_xs_f, V
//  [python]
def spectral_gap(matrix):
    _, s, _ = torch.svd(matrix)
    abs_eigenvalues = torch.sqrt(s**2)
    return abs_eigenvalues[0] - abs_eigenvalues[1]
//  [python]
def spectral_gap(matrix):
    _, s, _ = torch.linalg.svd(matrix)
    abs_eigenvalues = torch.sqrt(s**2)
    return abs_eigenvalues[0] - abs_eigenvalues[1]
//  [python]
def get_zca_matrix(X_):
    """
    Function to compute ZCA whitening matrix (aka Mahalanobis whitening).
    INPUT:  X : [n x C x P x Q]
    OUTPUT: ZCAMatrix: [d x d] matrix
    """
    
    n, C, P, Q = X_.shape
    eps = 0.1*50000/n

    d = C*P*Q
    
    X = X_.reshape(n,d).double().cuda()
    
    D = torch.zeros(d).cuda()

    _, D[:min(n,d)], UT = torch.svd(X/sqrt(n))

    D[D < 0] = 0
    D = D**2
    D = D + eps*torch.sum(D)/d

    D = 1/D.sqrt()
    D[min(n,d):] = 0

    W = (UT.T*D) @ UT

    return W.float().cpu()
//  [python]
def get_zca_matrix(X_):
    """
    Function to compute ZCA whitening matrix (aka Mahalanobis whitening).
    INPUT:  X : [n x C x P x Q]
    OUTPUT: ZCAMatrix: [d x d] matrix
    """
    
    n, C, P, Q = X_.shape
    eps = 0.1*50000/n

    d = C*P*Q
    
    X = X_.reshape(n,d).double().cuda()
    
    D = torch.zeros(d).cuda()

    _, D[:min(n,d)], UT = torch.linalg.svd(X/sqrt(n))

    D[D < 0] = 0
    D = D**2
    D = D + eps*torch.sum(D)/d

    D = 1/D.sqrt()
    D[min(n,d):] = 0

    W = (UT.T*D) @ UT

    return W.float().cpu()
//  [python]
def train(epoch):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = net(inputs)

        loss = criterion(outputs, targets)

        loss.backward()
        # Apply Nearest Orthogonal Gradient Mapping
        if args.nog:
            if args.ow:
                weight_grad = net.module.conv1.random_conv_filter.grad
            else:
                weight_grad = net.module.conv1.weight.grad
            B, C, H, W = weight_grad.size()
            u, s, vh = torch.svd(weight_grad.view(B, C * H * W))
            s[:] = 1
            if args.ow:
                net.module.conv1.random_conv_filter.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
            else:
                net.module.conv1.weight.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
        # Apply Optimal Learnig Rate
        if args.olr:
            if args.ow:
                weight_grad = net.module.conv1.random_conv_filter.grad.view(-1, 1)
                weight = net.module.conv1.random_conv_filter.view(-1, 1)
            else:
                weight_grad = net.module.conv1.weight.grad.view(-1, 1)
                weight = net.module.conv1.weight.view(-1, 1)
            wtw = weight.t().mm(weight)
            ltw = weight_grad.t().mm(weight)
            ltl = weight_grad.t().mm(weight_grad)
            lr = (wtw * ltw) / (wtw * ltl + 2 * ltw * ltw)
            if lr > 0 and lr < optimizer.param_groups[1]['lr']:
                optimizer.param_groups[0]['lr'] = lr
            else:
                optimizer.param_groups[0]['lr'] = optimizer.param_groups[1]['lr']
        optimizer.step()

        if batch_idx % 500 == 0:
            writer.add_scalar('loss/train_loss', loss.item(), epoch*len(trainloader)+batch_idx+1)
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    acc = 100. * correct / total
    writer.add_scalar('loss/train_loss_avg', train_loss / len(trainloader), epoch)
    writer.add_scalar('train/accuracy', acc, epoch)
    writer.add_scalar('train/error', 100 - acc, epoch)
    print('current training error {}'.format(100-acc))
//  [python]
def train(epoch):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = net(inputs)

        loss = criterion(outputs, targets)

        loss.backward()
        # Apply Nearest Orthogonal Gradient Mapping
        if args.nog:
            if args.ow:
                weight_grad = net.module.conv1.random_conv_filter.grad
            else:
                weight_grad = net.module.conv1.weight.grad
            B, C, H, W = weight_grad.size()
            u, s, vh = torch.linalg.svd(weight_grad.view(B, C * H * W), full_matrices=False)
            s[:] = 1
            if args.ow:
                net.module.conv1.random_conv_filter.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
            else:
                net.module.conv1.weight.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
        # Apply Optimal Learnig Rate
        if args.olr:
            if args.ow:
                weight_grad = net.module.conv1.random_conv_filter.grad.view(-1, 1)
                weight = net.module.conv1.random_conv_filter.view(-1, 1)
            else:
                weight_grad = net.module.conv1.weight.grad.view(-1, 1)
                weight = net.module.conv1.weight.view(-1, 1)
            wtw = weight.t().mm(weight)
            ltw = weight_grad.t().mm(weight)
            ltl = weight_grad.t().mm(weight_grad)
            lr = (wtw * ltw) / (wtw * ltl + 2 * ltw * ltw)
            if lr > 0 and lr < optimizer.param_groups[1]['lr']:
                optimizer.param_groups[0]['lr'] = lr
            else:
                optimizer.param_groups[0]['lr'] = optimizer.param_groups[1]['lr']
        optimizer.step()

        if batch_idx % 500 == 0:
            writer.add_scalar('loss/train_loss', loss.item(), epoch*len(trainloader)+batch_idx+1)
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    acc = 100. * correct / total
    writer.add_scalar('loss/train_loss_avg', train_loss / len(trainloader), epoch)
    writer.add_scalar('train/accuracy', acc, epoch)
    writer.add_scalar('train/error', 100 - acc, epoch)
    print('current training error {}'.format(100-acc))
//  [python]
def rmsdalign(a, b, weights=None): # alignes B to A  # [*, N, 3]
    B = a.shape[:-2]
    N = a.shape[-2]
    if weights == None:
        weights = a.new_ones(*B, N)
    weights = weights.unsqueeze(-1)
    a_mean = (a * weights).sum(-2, keepdims=True) / weights.sum(-2, keepdims=True)
    a = a - a_mean
    b_mean = (b * weights).sum(-2, keepdims=True) / weights.sum(-2, keepdims=True)
    b = b - b_mean
    B = torch.einsum('...ji,...jk->...ik', weights * a, b)
    u, s, vh = torch.svd(B)

    # Correct improper rotation if necessary (as in Kabsch algorithm)
    '''
    if torch.linalg.det(u @ vh) < 0:
        s[-1] = -s[-1]
        u[:, -1] = -u[:, -1]
    '''
    sgn = torch.sign(torch.det(u @ vh))
    s[...,-1] *= sgn
    u[...,:,-1] *= sgn.unsqueeze(-1)
    C = u @ vh # c rotates B to A
    return b @ C.mT + a_mean
//  [python]
def rmsdalign(a, b, weights=None): # alignes B to A  # [*, N, 3]
    B = a.shape[:-2]
    N = a.shape[-2]
    if weights == None:
        weights = a.new_ones(*B, N)
    weights = weights.unsqueeze(-1)
    a_mean = (a * weights).sum(-2, keepdims=True) / weights.sum(-2, keepdims=True)
    a = a - a_mean
    b_mean = (b * weights).sum(-2, keepdims=True) / weights.sum(-2, keepdims=True)
    b = b - b_mean
    B = torch.einsum('...ji,...jk->...ik', weights * a, b)
    u, s, vh = torch.linalg.svd(B)

    # Correct improper rotation if necessary (as in Kabsch algorithm)
    '''
    if torch.linalg.det(u @ vh) < 0:
        s[-1] = -s[-1]
        u[:, -1] = -u[:, -1]
    '''
    sgn = torch.sign(torch.linalg.det(u @ vh))
    s[...,-1] *= sgn
    u[...,:,-1] *= sgn.unsqueeze(-1)
    C = u @ vh # c rotates B to A
    return b @ C.mT + a_mean
//  [python]
def check_diff_singular_value(a: torch.TensorType, b: torch.TensorType):
    assert a.shape == b.shape

    diff = a - b
    U, S, Vh = torch.svd(diff)

    return S
//  [python]
def check_diff_singular_value(a: torch.TensorType, b: torch.TensorType):
    assert a.shape == b.shape

    diff = a - b
    U, S, Vh = torch.linalg.svd(diff)

    return S
//  [python]
def kabsch_rotation(P, Q):
    """
    Calculate the best rotation that minimises the RMSD between P and Q.

    The optimal rotation matrix was calculated using Kabsch algorithm:
    https://en.wikipedia.org/wiki/Kabsch_algorithm

    Args:
    P: [N * 3] Nres is the number of atoms and each row corresponds to the atom's x,y,z coordinates
    Q: [N * 3] the same dimension as P

    return:
    A 3*3 rotation matrix
    """
    assert P.shape == torch.Size([Q.shape[0], Q.shape[1]])

    # Firstly, compute SVD of P.T * Q
    u, _, vt = torch.svd(torch.matmul(P.to(torch.float32).T,
                                             Q.to(torch.float32)))
    # Then construct s matrix
    s = torch.eye(P.shape[1], device=P.device)
    # correct the rotation matrix to ensure a right-handed coordinate
    s[-1, -1] = torch.sign(torch.det(torch.matmul(u, vt)))
    # finally compute the rotation matrix
    r_opt = torch.matmul(torch.matmul(u, s), vt)
    assert r_opt.shape == torch.Size([3,3])
    return r_opt.to(device=P.device, dtype=P.dtype)
//  [python]
def kabsch_rotation(P, Q):
    """
    Calculate the best rotation that minimises the RMSD between P and Q.

    The optimal rotation matrix was calculated using Kabsch algorithm:
    https://en.wikipedia.org/wiki/Kabsch_algorithm

    Args:
    P: [N * 3] Nres is the number of atoms and each row corresponds to the atom's x,y,z coordinates
    Q: [N * 3] the same dimension as P

    return:
    A 3*3 rotation matrix
    """
    assert P.shape == torch.Size([Q.shape[0], Q.shape[1]])

    # Firstly, compute SVD of P.T * Q
    u, _, vt = torch.linalg.svd(torch.matmul(P.to(torch.float32).T,
                                             Q.to(torch.float32)))
    # Then construct s matrix
    s = torch.eye(P.shape[1], device=P.device)
    # correct the rotation matrix to ensure a right-handed coordinate
    s[-1, -1] = torch.sign(torch.linalg.det(torch.matmul(u, vt)))
    # finally compute the rotation matrix
    r_opt = torch.matmul(torch.matmul(u, s), vt)
    assert r_opt.shape == torch.Size([3,3])
    return r_opt.to(device=P.device, dtype=P.dtype)
//  [python]
def PCA_rotation(
    X: torch.Tensor,
    components: int = -1,
    center: bool = True,
    full: bool = True,
    niter: int = 10,
) -> torch.Tensor:
    """Calculate the rotation matrix from PCA.

    If the PCA fails (e.g. if original matrix is not full rank) then this shows a warning instead of throwing an error (returns a dummy rotation).

    Args:
        X: The original matrix.
        components: The maximum number of components in the embedding. Defaults to `min(*X.shape)`.
        center: Center the matrix before calculating the PCA.
        full: Use a full SVD for the PCA (slower). Defaults to True.
        niter: The number of iterations when a randomised approach is used. Defaults to 10.

    Returns:
        Rotation matrix that turns the original matrix into the embedded space.
    """
    try:
        components = min(*X.shape, components) if components > 0 else min(*X.shape)
        if full:
            if center:
                X = X - X.mean(dim=(-2,), keepdim=True)
            U, S, V = torch.svd(X)
            return V.T[:, :components]
        else:
            return torch.pca_lowrank(X, components, center=center, niter=niter)[2]
    except Exception:
        _warn("Could not perform PCA", PCA_rotation)
        z = torch.zeros((X.shape[1], components), dtype=X.dtype, device=X.device)
        z.fill_diagonal_(1.0, True)
        return z
//  [python]
def PCA_rotation(
    X: torch.Tensor,
    components: int = -1,
    center: bool = True,
    full: bool = True,
    niter: int = 10,
) -> torch.Tensor:
    """Calculate the rotation matrix from PCA.

    If the PCA fails (e.g. if original matrix is not full rank) then this shows a warning instead of throwing an error (returns a dummy rotation).

    Args:
        X: The original matrix.
        components: The maximum number of components in the embedding. Defaults to `min(*X.shape)`.
        center: Center the matrix before calculating the PCA.
        full: Use a full SVD for the PCA (slower). Defaults to True.
        niter: The number of iterations when a randomised approach is used. Defaults to 10.

    Returns:
        Rotation matrix that turns the original matrix into the embedded space.
    """
    try:
        components = min(*X.shape, components) if components > 0 else min(*X.shape)
        if full:
            if center:
                X = X - X.mean(dim=(-2,), keepdim=True)
            return torch.linalg.svd(X, full_matrices=False)[2].T[:, :components]
        else:
            return torch.pca_lowrank(X, components, center=center, niter=niter)[2]
    except Exception:
        _warn("Could not perform PCA", PCA_rotation)
        z = torch.zeros((X.shape[1], components), dtype=X.dtype, device=X.device)
        z.fill_diagonal_(1.0, True)
        return z
//  [python]
def submersion_inv(self, X, check_in_manifold=True):
        U, S, V = torch.svd(X, some=False)
        if check_in_manifold and not self.in_manifold_singular_values(S):
            raise InManifoldError(X, self)
        return U[..., : self.rank], S[..., : self.rank], V[..., : self.rank]
//  [python]
def submersion_inv(self, X, check_in_manifold=True):
        U, S, Vt = torch.linalg.svd(X, full_matrices=False)
        V = Vt.transpose(-2, -1)
        if check_in_manifold and not self.in_manifold_singular_values(S):
            raise InManifoldError(X, self)
        return U[..., : self.rank], S[..., : self.rank], V[..., : self.rank]
//  [python]
def svd_denoise(x, cut):
    x_ = x.clone().detach()
    U, S, V = torch.svd(x_)
    S[:, cut:] = 0

    return U @ torch.diag(S[0, :]) @ V
//  [python]
def svd_denoise(x, cut):
    x_ = x.clone().detach()
    U, S, V = torch.linalg.svd(x_, full_matrices=False)
    S[:, cut:] = 0

    return U @ torch.diag(S[0, :]) @ V
//  [python]
def align_vectors_torch(a, b, return_4x4=False):
    """
    Calculate the rotation matrix to align two batches of vectors
    Modified from: trimesh.geometry.align_vectors
    a and b contain batches of vectors wth shape (n, 3)
    """
    if a.shape[-1] != 3 or b.shape[-1] != 3:
        raise ValueError('vectors must be (n,3)!')
    assert a.shape[0] == b.shape[0]
    batch_size = a.shape[0]

    ar = a.reshape((-1, 3, 1))
    br = b.reshape((-1, 3, 1))

    # find the SVD of the two vectors
    au, _, _ = torch.svd(ar)
    bu, _, _ = torch.svd(br)

    au[torch.det(au) < 0, :, -1] *= -1.0
    bu[torch.det(bu) < 0, :, -1] *= -1.0

    # Transpose au along dim 1 and 2 (not batch dim 0)
    au_t = torch.transpose(au, 1, 2)
    # Batch matmul along dim 1 and 2 (not batch dim 0)?
    mat = torch.matmul(bu, au_t)

    if return_4x4:
        mat_4x4 = torch.tile(torch.eye(4), (batch_size, 1)).view((batch_size, 4, 4))
        mat_4x4[:, :3, :3] = mat
        return mat_4x4
    else:
        return mat
//  [python]
def align_vectors_torch(a, b, return_4x4=False):
    """
    Calculate the rotation matrix to align two batches of vectors
    Modified from: trimesh.geometry.align_vectors
    a and b contain batches of vectors wth shape (n, 3)
    """
    if a.shape[-1] != 3 or b.shape[-1] != 3:
        raise ValueError('vectors must be (n,3)!')
    assert a.shape[0] == b.shape[0]
    batch_size = a.shape[0]

    ar = a.reshape((-1, 3, 1))
    br = b.reshape((-1, 3, 1))

    # find the SVD of the two vectors
    au = torch.linalg.svd(ar)[0]
    bu = torch.linalg.svd(br)[0]

    au[torch.linalg.det(au) < 0, :, -1] *= -1.0
    bu[torch.linalg.det(bu) < 0, :, -1] *= -1.0

    # Transpose au along dim 1 and 2 (not batch dim 0)
    au_t = torch.transpose(au, 1, 2)
    # Batch matmul along dim 1 and 2 (not batch dim 0)?
    mat = torch.matmul(bu, au_t)

    if return_4x4:
        mat_4x4 = torch.tile(torch.eye(4), (batch_size, 1)).view((batch_size, 4, 4))
        mat_4x4[:, :3, :3] = mat
        return mat_4x4
    else:
        return mat
//  [python]
def cacluate_U(C):
    #(evals, evecs) = torch.eig(C, eigenvectors=True)
    U, _, _ = torch.svd(C.T)
    #U = U[:, 0:args.n_clusters]
    U = U[:, 0:args.n_clusters]
    return U
//  [python]
def cacluate_U(C):
    #(evals, evecs) = torch.eig(C, eigenvectors=True)
    U, _, _ = torch.linalg.svd(C.T)
    #U = U[:, 0:args.n_clusters]
    U = U[:, 0:args.n_clusters]
    return U
//  [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.svd(A)
        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)