jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def convert_pca(model, dim, device, embed_dict):
    """
    Perform SVD to reduce the dimensionality of embeddings through the final weight matrix in 
    a resnet model
    Inputs:
    - model: resnet model
    - dim: dimension to reduce to
    - embed_dict: {'embed': array-like (n_samples, 2048), 
                 'labels': [n_samples, ], 
                 'preds': [n_samples, ], 
                 'subclass_labels': subclass_labels}
    Outputs:
    - results = {'embed': array-like (n_samples, dim), 
                 'labels': [n_samples, ], 
                 'preds': [n_samples, ], 
                 'subclass_labels': subclass_labels}
    """
    with torch.no_grad():
        # svd
        u, s, v = torch.linalg.svd(model.fc.weight.T, full_matrices=False) # 1000, 2048 -> 2048, 1000
        # print(u.shape, s.shape, v.shape)
        s[dim:] = 0
        r = u * s # reduction matrix # 2048, 1000 -> cols past dim are 0
        embed_old = embed_dict['embed'].to(device) # n_samples, 2048
        embed_red = embed_old @ r # n_samples, 1000, w/0s on end
        embed_dict['embed'] = embed_red[:, :dim].cpu()
        embed_dict['preds'] = torch.argmax(embed_red @ v, dim=1).tolist()
        return embed_dict
//  [python]
def forward(self, offset_embeddings):
        layer1_act = F.relu(self.layer1(offset_embeddings))
        layer1_mean = torch.mean(layer1_act, dim=0)
        gate = torch.sigmoid(self.layer2(layer1_mean))
        offset_embeddings = torch.unsqueeze(torch.squeeze(offset_embeddings),dim=-2)
        inv_cov1 = torch.matmul(torch.transpose(offset_embeddings[0],-2,-1),offset_embeddings[0])
        inv_cov2 = torch.matmul(torch.transpose(offset_embeddings[-1],-2,-1),offset_embeddings[-1])
        offset = torch.svd(inv_cov1 + inv_cov2)[1]
        return offset * gate
//  [python]
def forward(self, offset_embeddings):
        layer1_act = F.relu(self.layer1(offset_embeddings))
        layer1_mean = torch.mean(layer1_act, dim=0)
        gate = torch.sigmoid(self.layer2(layer1_mean))
        offset_embeddings = torch.unsqueeze(torch.squeeze(offset_embeddings),dim=-2)
        inv_cov1 = torch.matmul(torch.transpose(offset_embeddings[0],-2,-1),offset_embeddings[0])
        inv_cov2 = torch.matmul(torch.transpose(offset_embeddings[-1],-2,-1),offset_embeddings[-1])
        offset = torch.linalg.svd(inv_cov1 + inv_cov2,compute_uv=True)[1]
        return offset * gate
//  [python]
def dist_mat_to_coords(d: torch.Tensor, k: int = 3) -> CoordTensor:
    """
    Computes a set of coordinates that satisfy the (valid) Euclidean
    distance matrix ``d``.

    See: https://math.stackexchange.com/questions/156161/finding-the-coordinates-of-points-from-distance-matrix

    :param d: Valid euclidean distance matrix
    :type d: torch.Tensor
    :param k: Number of dimensions of coordinate output, defaults to ``3``.
    :type k: int, optional
    :return: Set of coordinates that satisfy the distance matrix ``d``.
    :rtype: graphein.protein.tensor.types.CoordTensor
    """
    D_ij_2 = torch.pow(d, 2)
    D_1i_2 = torch.pow(d[0, :], 2).unsqueeze(1)
    D_1j_2 = torch.pow(d[:, 0], 2).unsqueeze(0)
    Mij = (D_1i_2 + D_1j_2 - D_ij_2) / 2

    # try:
    U, S, V = torch.svd(Mij)
    # except Exception:  # torch.svd may have convergence issues for GPU and CPU.
    #    U, S, _ = torch.linalg.svd(
    #        Mij + 1e-4 * Mij.mean() * torch.rand_like(Mij)
    #    )
    X = U * torch.sqrt(S)
    X = X[:, :k]
    return X
//  [python]
def dist_mat_to_coords(d: torch.Tensor, k: int = 3) -> CoordTensor:
    """
    Computes a set of coordinates that satisfy the (valid) Euclidean
    distance matrix ``d``.

    See: https://math.stackexchange.com/questions/156161/finding-the-coordinates-of-points-from-distance-matrix

    :param d: Valid euclidean distance matrix
    :type d: torch.Tensor
    :param k: Number of dimensions of coordinate output, defaults to ``3``.
    :type k: int, optional
    :return: Set of coordinates that satisfy the distance matrix ``d``.
    :rtype: graphein.protein.tensor.types.CoordTensor
    """
    D_ij_2 = torch.pow(d, 2)
    D_1i_2 = torch.pow(d[0, :], 2).unsqueeze(1)
    D_1j_2 = torch.pow(d[:, 0], 2).unsqueeze(0)
    Mij = (D_1i_2 + D_1j_2 - D_ij_2) / 2

    # try:
    U, S, _ = torch.linalg.svd(Mij)
    # except Exception:  # torch.svd may have convergence issues for GPU and CPU.
    #    U, S, _ = torch.linalg.svd(
    #        Mij + 1e-4 * Mij.mean() * torch.rand_like(Mij)
    #    )
    X = U * torch.sqrt(S)
    X = X[:, :k]
    return X
//  [python]
def training_epoch_end(self, outputs):
        features = torch.vstack([out["features"] for out in outputs])
        acc = torch.hstack([out["is_correct"] for out in outputs]).float().mean() * 100
        norm = torch.hstack([out["norms"] for out in outputs]).float().mean()
        self.avg_train_norm = norm

        logger.info("\nTraining set acc {:.2f}%".format(acc))
        logger.info("\nTraining set avg norm {:.2f}".format(norm))

        # Calc regrets
        t0 = time()
        x_t_x = torch.matmul(features.t(), features)
        U, s, V = torch.linalg.svd(x_t_x)
        logger.info(f"Training set singular values largest: {s[:5]}")
        logger.info(f"Training set singular values smallest: {s[-5:]}")
        # self.x_t_x_inv = torch.linalg.inv(x_t_x) if s[-1] > 1e-16 else torch.linalg.pinv(x_t_x, hermitian=True)
        self.x_t_x = x_t_x
        self.x_t_x_inv = torch.linalg.pinv(
            self.x_t_x, hermitian=False, rcond=self.pinv_rcond
        )
        logger.info(f"Finish inv in {time() -t0 :.2f} sec")

        # Plot svd
        fig, ax = plt.subplots(1, 1)
        ax.plot(s.cpu(), "*")
        ax.set_xlabel("Singular value number")
        ax.set_yscale("log")
        ax.grid()
        plt.savefig(osp.join(self.out_dir, "svd.jpg"))
        plt.close(fig)

        torch.save(s, osp.join(self.out_dir, "singular_values.pt"))
//  [python]
def training_epoch_end(self, outputs):
        features = torch.vstack([out["features"] for out in outputs])
        acc = torch.hstack([out["is_correct"] for out in outputs]).float().mean() * 100
        norm = torch.hstack([out["norms"] for out in outputs]).float().mean()
        self.avg_train_norm = norm

        logger.info("\nTraining set acc {:.2f}%".format(acc))
        logger.info("\nTraining set avg norm {:.2f}".format(norm))

        # Calc regrets
        t0 = time()
        x_t_x = torch.matmul(features.t(), features)
        _, s, _ = torch.linalg.svd(x_t_x, compute_uv=False)
        logger.info(f"Training set singular values largest: {s[:5]}")
        logger.info(f"Training set singular values smallest: {s[-5:]}")
        # self.x_t_x_inv = torch.linalg.inv(x_t_x) if s[-1] > 1e-16 else torch.linalg.pinv(x_t_x, hermitian=True)
        self.x_t_x = x_t_x
        self.x_t_x_inv = torch.linalg.pinv(
            self.x_t_x, hermitian=False, rcond=self.pinv_rcond
        )
        logger.info(f"Finish inv in {time() -t0 :.2f} sec")

        # Plot svd
        fig, ax = plt.subplots(1, 1)
        ax.plot(s.cpu(), "*")
        ax.set_xlabel("Singular value number")
        ax.set_yscale("log")
        ax.grid()
        plt.savefig(osp.join(self.out_dir, "svd.jpg"))
        plt.close(fig)

        torch.save(s, osp.join(self.out_dir, "singular_values.pt"))
//  [python]
def compute_kabsch_RMSD(T_true, T_pred, weight, unit=10.0, dclamp=10.0, eps = 1e-4, reduction='mean'):
    # T_true, T_pred is (N_batch, N_max_len_of_batch, 3)
    # weight is the mask of size (N_batch, N_max_len_of_batch)

    # First center with mean coordinate of meaningful tokens, which involves scaling up the mean coordinates w.r.t mean of weight.
    # And then we zero the non-meaningful ones - zero coordinates in SVD will not affect the rotation matrix

    weight = weight.type(T_true.dtype)

    T_true_cm = (T_true * weight[:, :, None]).mean(axis=1)[:,None,:] / (weight.mean(1)[:, None, None])
    T_pred_cm = (T_pred * weight[:, :, None]).mean(axis=1)[:,None,:] / (weight.mean(1)[:, None, None])

    T_true_cen = T_true - T_true_cm
    T_pred_cen = T_pred - T_pred_cm

    # Kabsch method
    C = torch.matmul(torch.transpose(T_pred_cen * weight[:, :, None], -2, -1), T_true_cen * weight[:, :, None])
    with torch.autocast('cuda',enabled=False):
        # svd doesn't support fp16/bf16 on AMD
        U, S, V = torch.svd(C.float())
    V_prime = V * torch.tensor([1, 1, -1], device=V.device, dtype=V.dtype)[None, :, None] # Alternate version of V if d < 0
    d = (torch.det(U) * torch.det(V)) < 0.0

    V = torch.where(d[:, None, None], V_prime, V) # We have to broadcast d to select the right V's
    # Create Rotation matrix U
    U = torch.matmul(V, torch.transpose(U, -2, -1))
    T_pred_align = torch.matmul(T_pred_cen, U) # Rotate P
    diff = T_pred_align - T_true_cen
    dist = torch.sqrt((diff**2).sum(-1) + eps)     # (N_batch, N_max_len_of_batch)

    if dclamp is not None:
        dist = torch.clamp(dist, max=dclamp)
    dist = weight*dist

    loss = torch.sum(dist, dim=-1) / weight.sum(dim=-1) / unit

    # mini-batch reduction
    if reduction == 'mean':
        loss = torch.mean(loss)
    elif reduction == 'gmean':
        loss = torch.exp(torch.mean(torch.log(loss)))
    else:
        raise ValueError

    return loss
//  [python]
def compute_kabsch_RMSD(T_true, T_pred, weight, unit=10.0, dclamp=10.0, eps = 1e-4, reduction='mean'):
    # T_true, T_pred is (N_batch, N_max_len_of_batch, 3)
    # weight is the mask of size (N_batch, N_max_len_of_batch)

    # First center with mean coordinate of meaningful tokens, which involves scaling up the mean coordinates w.r.t mean of weight.
    # And then we zero the non-meaningful ones - zero coordinates in SVD will not affect the rotation matrix

    weight = weight.type(T_true.dtype)

    T_true_cm = (T_true * weight[:, :, None]).mean(axis=1)[:,None,:] / (weight.mean(1)[:, None, None])
    T_pred_cm = (T_pred * weight[:, :, None]).mean(axis=1)[:,None,:] / (weight.mean(1)[:, None, None])

    T_true_cen = T_true - T_true_cm
    T_pred_cen = T_pred - T_pred_cm

    # Kabsch method
    C = torch.matmul(torch.transpose(T_pred_cen * weight[:, :, None], -2, -1), T_true_cen * weight[:, :, None])
    with torch.autocast('cuda',enabled=False):
        # svd doesn't support fp16/bf16 on AMD
        V, S, W = torch.linalg.svd(C.float())
    V_prime = V * torch.tensor([1, 1, -1], device=V.device, dtype=V.dtype)[None, :, None] # Alternate version of V if d < 0
    d = (torch.linalg.det(V) * torch.linalg.det(W)) < 0.0

    V = torch.where(d[:, None, None], V_prime, V) # We have to broadcast d to select the right V's
    # Create Rotation matrix U
    U = torch.matmul(V, W)
    T_pred_align = torch.matmul(T_pred_cen, U) # Rotate P
    diff = T_pred_align - T_true_cen
    dist = torch.sqrt((diff**2).sum(-1) + eps)     # (N_batch, N_max_len_of_batch)

    if dclamp is not None:
        dist = torch.clamp(dist, max=dclamp)
    dist = weight*dist

    loss = torch.sum(dist, dim=-1) / weight.sum(dim=-1) / unit

    # mini-batch reduction
    if reduction == 'mean':
        loss = torch.mean(loss)
    elif reduction == 'gmean':
        loss = torch.exp(torch.mean(torch.log(loss)))
    else:
        raise ValueError

    return loss
//  [python]
def post_train(self, task_id, trn_loader):
        """Runs after training all the epochs of the task (after the train session)"""
        with torch.no_grad():
            if task_id > 0 and self.sigma_proto_update != -1:
                print("Final Computing Update Proto")
                start = time()
                new_features, old_features  =  get_old_new_features(self.model, self.old_model, 
                                                                                 trn_loader, self.device)
                
                drift = self.compute_drift(new_features, old_features,  device="cpu")
                drift = drift.cpu()
                end = time()
                print("Elapsed time {}".format(end-start))
                for i, (p, var, proto_label) in enumerate(zip(self.proto_generator.prototype, 
                                                            self.proto_generator.running_proto_variance,
                                                            self.proto_generator.class_label)):
            
                    mean = p + drift[i] 
                    self.proto_generator.update_gaussian(proto_label, mean, var)
                    # final update the mean
                    self.proto_generator.prototype[i] = mean 
                    
                self.proto_generator.running_proto = deepcopy(self.proto_generator.prototype)

        
            efm_matrix = EmpiricalFeatureMatrix(self.device, out_path=self.out_path)
            efm_matrix.compute(self.model,  deepcopy(trn_loader), task_id)
            self.previous_efm = efm_matrix.get()           
            U, S, V = torch.svd(self.previous_efm)
            matrix_rank = torch.matrix_rank(self.previous_efm)
            print("Computed Matrix Rank {}".format(matrix_rank))
            self.R = U
            self.L = S
            self.matrix_rank = matrix_rank

        # save matrix after each task for analysis
        save_efm(self.previous_efm, task_id, self.out_path)
            
        print("Computing New Task Prototypes")
        self.proto_generator.compute(self.model, deepcopy(trn_loader), task_id)
//  [python]
def post_train(self, task_id, trn_loader):
        """Runs after training all the epochs of the task (after the train session)"""
        with torch.no_grad():
            if task_id > 0 and self.sigma_proto_update != -1:
                print("Final Computing Update Proto")
                start = time()
                new_features, old_features  =  get_old_new_features(self.model, self.old_model, 
                                                                                 trn_loader, self.device)
                
                drift = self.compute_drift(new_features, old_features,  device="cpu")
                drift = drift.cpu()
                end = time()
                print("Elapsed time {}".format(end-start))
                for i, (p, var, proto_label) in enumerate(zip(self.proto_generator.prototype, 
                                                            self.proto_generator.running_proto_variance,
                                                            self.proto_generator.class_label)):
            
                    mean = p + drift[i] 
                    self.proto_generator.update_gaussian(proto_label, mean, var)
                    # final update the mean
                    self.proto_generator.prototype[i] = mean 
                    
                self.proto_generator.running_proto = deepcopy(self.proto_generator.prototype)

        
            efm_matrix = EmpiricalFeatureMatrix(self.device, out_path=self.out_path)
            efm_matrix.compute(self.model,  deepcopy(trn_loader), task_id)
            self.previous_efm = efm_matrix.get()           
            R, L, V = torch.linalg.svd(self.previous_efm)
            matrix_rank = torch.linalg.matrix_rank(self.previous_efm)
            print("Computed Matrix Rank {}".format(matrix_rank))
            self.R = R
            self.L = L
            self.matrix_rank = matrix_rank

        # save matrix after each task for analysis
        save_efm(self.previous_efm, task_id, self.out_path)
            
        print("Computing New Task Prototypes")
        self.proto_generator.compute(self.model, deepcopy(trn_loader), task_id)
//  [python]
def decompose_and_triangulate(points, models, weights=None):
    n_matches, _ = points.shape
    assert points.shape == (n_matches, 4)
    n_models, _, _ = models.shape
    assert models.shape == (n_models, 3, 3)

    Y1 = torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=models.dtype, device=models.device)
    Y2 = Y1.T
    u, d, vt = torch.svd(models)
    t1 = u[..., -1]
    t2 = -t1
    R1 = u @ Y1 @ vt
    R2 = u @ Y2 @ vt
    R_mirror = ((torch.det(R1) > 0).long() * 2 - 1).view(n_models, 1, 1)
    R1 = R1 * R_mirror
    R2 = R2 * R_mirror

    # def err(t, R):
    #     E_hat = cross_mat(t) @ R
    #     dist = torch.linalg.norm(models - (models * E_hat).sum().div(E_hat.square().sum()) * E_hat).square()
    #     return dist
    # e1, e2, e3, e4 = err(t1, R1), err(t2, R2), err(t1, R2), err(t2, R1)

    R_batch = torch.stack([R1, R2, R1, R2], dim=0).view(n_models * 4, 3, 3)
    T_batch = torch.stack([t1, t1, t2, t2], dim=0).view(n_models * 4, 3)
    u1, u2 = points[:, :2], points[:, -2:]

    # z1z2 = triangulate(u1, u2, R_batch, T_batch).view(4, n_models, n_matches, 2)
    # inline triangulation:
    c1 = -torch.nn.functional.pad(u1, (0 ,1), value=1.).view(1, -1, 3) @ R_batch.transpose(-1, -2) # (n_models * 4, n_matches, 3)
    c2 = torch.nn.functional.pad(u2, (0 ,1), value=1.).view(1, n_matches, 3).expand(n_models * 4, n_matches, 3)  # (n_models * 4, n_matches, 3)
    c1c2 = torch.stack([c1, c2], axis=-1)
    z1z2 = batch_2x2_inv(c1c2.transpose(-1, -2) @ c1c2) @ c1c2.transpose(-1, -2) @ T_batch.view(n_models * 4, 1, 3, 1)
    z1z2 = z1z2.view(4, n_models, n_matches, 2)

    chirality_num_pass = torch.all(z1z2 > 0.01, dim=-1).float()  # (4, n_models, n_matches)
    if weights is not None:
        chirality_num_pass = chirality_num_pass * weights.squeeze()
    chirality_num_pass = chirality_num_pass.sum(-1)  # (4, n_models)

    best_model = torch.argmax(chirality_num_pass, dim=0)
    model_range = torch.arange(n_models, device=models.device)
    R = R_batch.view(4, n_models, 3, 3)[best_model, model_range].squeeze(0)
    T = T_batch.view(4, n_models, 3)[best_model, model_range].squeeze(0)
    return R, T, z1z2[best_model, model_range].squeeze(0)
//  [python]
def decompose_and_triangulate(points, models, weights=None):
    n_matches, _ = points.shape
    assert points.shape == (n_matches, 4)
    n_models, _, _ = models.shape
    assert models.shape == (n_models, 3, 3)

    Y1 = torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=models.dtype, device=models.device)
    Y2 = Y1.T
    u, d, vt = torch.linalg.svd(models)
    t1 = u[..., -1]
    t2 = -t1
    R1 = u @ Y1 @ vt
    R2 = u @ Y2 @ vt
    R_mirror = ((torch.linalg.det(R1) > 0).long() * 2 - 1).view(n_models, 1, 1)
    R1 = R1 * R_mirror
    R2 = R2 * R_mirror

    # def err(t, R):
    #     E_hat = cross_mat(t) @ R
    #     dist = torch.linalg.norm(models - (models * E_hat).sum().div(E_hat.square().sum()) * E_hat).square()
    #     return dist
    # e1, e2, e3, e4 = err(t1, R1), err(t2, R2), err(t1, R2), err(t2, R1)

    R_batch = torch.stack([R1, R2, R1, R2], dim=0).view(n_models * 4, 3, 3)
    T_batch = torch.stack([t1, t1, t2, t2], dim=0).view(n_models * 4, 3)
    u1, u2 = points[:, :2], points[:, -2:]

    # z1z2 = triangulate(u1, u2, R_batch, T_batch).view(4, n_models, n_matches, 2)
    # inline triangulation:
    c1 = -torch.nn.functional.pad(u1, (0 ,1), value=1.).view(1, -1, 3) @ R_batch.transpose(-1, -2) # (n_models * 4, n_matches, 3)
    c2 = torch.nn.functional.pad(u2, (0 ,1), value=1.).view(1, n_matches, 3).expand(n_models * 4, n_matches, 3)  # (n_models * 4, n_matches, 3)
    c1c2 = torch.stack([c1, c2], axis=-1)
    z1z2 = batch_2x2_inv(c1c2.transpose(-1, -2) @ c1c2) @ c1c2.transpose(-1, -2) @ T_batch.view(n_models * 4, 1, 3, 1)
    z1z2 = z1z2.view(4, n_models, n_matches, 2)

    chirality_num_pass = torch.all(z1z2 > 0.01, dim=-1).float()  # (4, n_models, n_matches)
    if weights is not None:
        chirality_num_pass = chirality_num_pass * weights.squeeze()
    chirality_num_pass = chirality_num_pass.sum(-1)  # (4, n_models)

    best_model = torch.argmax(chirality_num_pass, dim=0)
    model_range = torch.arange(n_models, device=models.device)
    R = R_batch.view(4, n_models, 3, 3)[best_model, model_range].squeeze(0)
    T = T_batch.view(4, n_models, 3)[best_model, model_range].squeeze(0)
    return R, T, z1z2[best_model, model_range].squeeze(0)
//  [python]
def project_flow_to_bases(bases, flow):
    # flow:  Bx2xHxW
    # bases: Bx2xHxWxN
    bases = bases + (torch.randn(bases.shape, device=bases.device) * 0.00001)
    bases = bases.permute(0, 4, 2, 3, 1).flatten(2).permute(0, 2, 1)
    try:
        U, S, Vh = torch.svd(bases)
        epsilon = 1e-5
        mask = S >= epsilon
        mask = mask.unsqueeze(1).repeat(1, U.shape[1], 1)
        Us = U * mask

        projected_flow = Us @ (Us.transpose(1, 2) @ flow.permute(0, 2, 3, 1).flatten(1).unsqueeze(2))
        return projected_flow.reshape(flow.permute(0, 2, 3, 1).shape).permute(0, 3, 1, 2)
    except:
        return None
//  [python]
def project_flow_to_bases(bases, flow):
    # flow:  Bx2xHxW
    # bases: Bx2xHxWxN
    bases = bases + (torch.randn(bases.shape, device=bases.device) * 0.00001)
    bases = bases.permute(0, 4, 2, 3, 1).flatten(2).permute(0, 2, 1)
    try:
        U, S, Vh = torch.linalg.svd(bases, full_matrices=False)
        epsilon = 1e-5
        mask = S >= epsilon
        mask = mask.unsqueeze(1).repeat(1, U.shape[1], 1)
        Us = U * mask

        projected_flow = Us @ (Us.transpose(1, 2) @ flow.permute(0, 2, 3, 1).flatten(1).unsqueeze(2))
        return projected_flow.reshape(flow.permute(0, 2, 3, 1).shape).permute(0, 3, 1, 2)
    except:
        return None
//  [python]
def decompose_and_triangulate_minimal(points, models):
    n_models, n_minimal, _ = points.shape
    # assert n_minimal == 5
    assert points.shape == (n_models, n_minimal, 4)
    n_models, _, _ = models.shape
    assert models.shape == (n_models, 3, 3)

    Y1 = torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=models.dtype, device=models.device)
    Y2 = Y1.T
    u, d, vt = torch.svd(models)
    t1 = u[..., -1]
    t2 = -t1
    R1 = u @ Y1 @ vt
    R2 = u @ Y2 @ vt
    R_mirror = ((torch.det(R1) > 0).long() * 2 - 1).view(-1, 1, 1)
    R1 = R1 * R_mirror
    R2 = R2 * R_mirror

    # def err(t, R):
    #     E_hat = cross_mat(t) @ R
    #     dist = torch.linalg.norm(models - (models * E_hat).sum().div(E_hat.square().sum()) * E_hat).square()
    #     return dist
    # e1, e2, e3, e4 = err(t1, R1), err(t2, R2), err(t1, R2), err(t2, R1)

    R_batch = torch.stack([R1, R2, R1, R2], dim=0).view(-1, 3, 3)
    T_batch = torch.stack([t1, t1, t2, t2], dim=0).view(-1, 3)
    u1, u2 = points[..., :2], points[..., -2:]

    # inline triangulation:
    c1 = -torch.nn.functional.pad(u1, (0 ,1), value=1.).repeat(4, 1, 1) @ R_batch.transpose(-1, -2) # (n_models * 4, n_minimal, 3)
    c2 = torch.nn.functional.pad(u2, (0 ,1), value=1.).repeat(4, 1, 1)  # (n_models * 4, n_minimal, 3)
    c1c2 = torch.stack([c1, c2], axis=-1)
    z1z2 = batch_2x2_inv(c1c2.transpose(-1, -2) @ c1c2) @ c1c2.transpose(-1, -2) @ T_batch.view(-1, 1, 3, 1)
    # z1z2 = z1z2.view(4, -1, n_minimal, 2)

    # chirality_num_pass = torch.all(z1z2 > 0, dim=-1).long().sum(dim=-1)  # (4, n_models)
    # best_model = torch.argmax(chirality_num_pass, dim=0)
    # model_range = torch.arange(n_models, device=models.device)
    # R = R_batch.view(4, n_models, 3, 3)[best_model, model_range].squeeze(0)
    # T = T_batch.view(4, n_models, 3)[best_model, model_range].squeeze(0)
    # z1z2_ = z1z2[best_model, model_range].squeeze(0)
    # for i in range(n_models):
    #     print(i)
    #     print(z1z2[:, i])
    #     print(chirality_num_pass[:, i])
    # return R, T, z1z2_
    z1z2 = z1z2.view(-1, n_minimal * 2)
    chirality_pass = torch.all(z1z2 > 0.01, dim=-1)  # (4 * num_models)
    origin_index = chirality_pass.view(4, -1).nonzero(as_tuple=True)[1]
    R = R_batch[chirality_pass]
    T = T_batch[chirality_pass]
    z1z2 = z1z2[chirality_pass].view(-1, n_minimal, 2)
    return R, T, z1z2, origin_index
//  [python]
def decompose_and_triangulate_minimal(points, models):
    n_models, n_minimal, _ = points.shape
    # assert n_minimal == 5
    assert points.shape == (n_models, n_minimal, 4)
    n_models, _, _ = models.shape
    assert models.shape == (n_models, 3, 3)

    Y1 = torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=models.dtype, device=models.device)
    Y2 = Y1.T
    u, d, vt = torch.linalg.svd(models)
    t1 = u[..., -1]
    t2 = -t1
    R1 = u @ Y1 @ vt
    R2 = u @ Y2 @ vt
    R_mirror = ((torch.linalg.det(R1) > 0).long() * 2 - 1).view(-1, 1, 1)
    R1 = R1 * R_mirror
    R2 = R2 * R_mirror

    # def err(t, R):
    #     E_hat = cross_mat(t) @ R
    #     dist = torch.linalg.norm(models - (models * E_hat).sum().div(E_hat.square().sum()) * E_hat).square()
    #     return dist
    # e1, e2, e3, e4 = err(t1, R1), err(t2, R2), err(t1, R2), err(t2, R1)

    R_batch = torch.stack([R1, R2, R1, R2], dim=0).view(-1, 3, 3)
    T_batch = torch.stack([t1, t1, t2, t2], dim=0).view(-1, 3)
    u1, u2 = points[..., :2], points[..., -2:]

    # inline triangulation:
    c1 = -torch.nn.functional.pad(u1, (0 ,1), value=1.).repeat(4, 1, 1) @ R_batch.transpose(-1, -2) # (n_models * 4, n_minimal, 3)
    c2 = torch.nn.functional.pad(u2, (0 ,1), value=1.).repeat(4, 1, 1)  # (n_models * 4, n_minimal, 3)
    c1c2 = torch.stack([c1, c2], axis=-1)
    z1z2 = batch_2x2_inv(c1c2.transpose(-1, -2) @ c1c2) @ c1c2.transpose(-1, -2) @ T_batch.view(-1, 1, 3, 1)
    # z1z2 = z1z2.view(4, -1, n_minimal, 2)

    # chirality_num_pass = torch.all(z1z2 > 0, dim=-1).long().sum(dim=-1)  # (4, n_models)
    # best_model = torch.argmax(chirality_num_pass, dim=0)
    # model_range = torch.arange(n_models, device=models.device)
    # R = R_batch.view(4, n_models, 3, 3)[best_model, model_range].squeeze(0)
    # T = T_batch.view(4, n_models, 3)[best_model, model_range].squeeze(0)
    # z1z2_ = z1z2[best_model, model_range].squeeze(0)
    # for i in range(n_models):
    #     print(i)
    #     print(z1z2[:, i])
    #     print(chirality_num_pass[:, i])
    # return R, T, z1z2_
    z1z2 = z1z2.view(-1, n_minimal * 2)
    chirality_pass = torch.all(z1z2 > 0.01, dim=-1)  # (4 * num_models)
    origin_index = chirality_pass.view(4, -1).nonzero(as_tuple=True)[1]
    R = R_batch[chirality_pass]
    T = T_batch[chirality_pass]
    z1z2 = z1z2[chirality_pass].view(-1, n_minimal, 2)
    return R, T, z1z2, origin_index
//  [python]
def forward(
        self,
        query_feat,
        support_feat,
        way_num,
        shot_num,
        normalize=True,
        discriminative=False,
    ):
        e, ws, d = support_feat.size()

        support_feat = support_feat.reshape(e, way_num, shot_num, -1)
        # add way dim
        query_feat = query_feat.unsqueeze(1)

        UU, _, _ = torch.svd(
                support_feat.permute(0, 1, 3, 2).double()
            )  # [episode_size, way_num, dim, dim]
        UU = UU.float()
        subspace = UU[:, :, :, : shot_num - 1].permute(
            0, 1, 3, 2
        )  # [episode_size, way_num, subspace_dim, dim]

        projection = (
            subspace.permute(0, 1, 3, 2)
            .matmul(subspace.matmul(query_feat.permute(0, 1, 3, 2)))
            .permute(0, 1, 3, 2)
        )  # [episode_size, way_num, shot_num, dim]

        dist = torch.sum((query_feat - projection) ** 2, dim=-1).permute(0, 2, 1)

        logits = -dist
        if normalize:
            logits /= d

        disc_loss = 0
        if discriminative:
            # P is [dim, subspace_dim] while subspace here is [subspace_dim, dim]
            subspace_metric = torch.norm(
                torch.matmul(
                    subspace.unsqueeze(1), subspace.unsqueeze(2).transpose(-2, -1)
                ),
                p="fro",
                dim=[-2, -1],
            )  # [episode_size, way_num, way_num]
            mask = torch.eye(way_num).bool()
            subspace_metric = subspace_metric[:, ~mask]
            disc_loss = torch.sum(subspace_metric**2)

        return logits, disc_loss
//  [python]
def forward(
        self,
        query_feat,
        support_feat,
        way_num,
        shot_num,
        normalize=True,
        discriminative=False,
    ):
        e, ws, d = support_feat.size()

        support_feat = support_feat.reshape(e, way_num, shot_num, -1)
        # add way dim
        query_feat = query_feat.unsqueeze(1)

        try:
            # faster but need pytorch>=1.8.0
            UU, _, _ = torch.linalg.svd(
                support_feat.permute(0, 1, 3, 2).double()
            )  # [episode_size, way_num, dim, dim]
        except AttributeError:
            UU, _, _ = torch.svd(
                support_feat.permute(0, 1, 3, 2).double()
            )  # [episode_size, way_num, dim, dim]
        UU = UU.float()
        subspace = UU[:, :, :, : shot_num - 1].permute(
            0, 1, 3, 2
        )  # [episode_size, way_num, subspace_dim, dim]

        projection = (
            subspace.permute(0, 1, 3, 2)
            .matmul(subspace.matmul(query_feat.permute(0, 1, 3, 2)))
            .permute(0, 1, 3, 2)
        )  # [episode_size, way_num, shot_num, dim]

        dist = torch.sum((query_feat - projection) ** 2, dim=-1).permute(0, 2, 1)

        logits = -dist
        if normalize:
            logits /= d

        disc_loss = 0
        if discriminative:
            # P is [dim, subspace_dim] while subspace here is [subspace_dim, dim]
            subspace_metric = torch.norm(
                torch.matmul(
                    subspace.unsqueeze(1), subspace.unsqueeze(2).transpose(-2, -1)
                ),
                p="fro",
                dim=[-2, -1],
            )  # [episode_size, way_num, way_num]
            mask = torch.eye(way_num).bool()
            subspace_metric = subspace_metric[:, ~mask]
            disc_loss = torch.sum(subspace_metric**2)

        return logits, disc_loss
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []
    m = torch.nn.Softmax(dim=-1).cuda()
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()
        with torch.no_grad():
            
            for i in range(6):
                x = model.features[i](x)
            
            feat = x
            if b == 0:
                print(feat.size())
            B, window_size, window_size, feat_dim = feat.size()
            feat = feat.view(B, window_size*window_size, feat_dim)

            u, s, v = torch.svd(feat1)
            feat1 = feat1 - 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_plus(feat, iter=20)

            feat = feat.view(B, window_size, window_size, feat_dim)
            
            feat = model.features[6](feat)
            feat = model.features[7](feat)

            feat = model.norm(feat)
            feat = model.permute(feat)
            feat = model.avgpool(feat)
            feat = model.flatten(feat)
            logits = model.head(feat)

            # conf = temperature * torch.logsumexp(logits / temperature, dim=1)
            conf, _ = torch.max(m(logits), dim=-1)

            confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []
    m = torch.nn.Softmax(dim=-1).cuda()
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()
        with torch.no_grad():
            
            for i in range(6):
                x = model.features[i](x)
            
            feat = x
            if b == 0:
                print(feat.size())
            B, window_size, window_size, feat_dim = feat.size()
            feat = feat.view(B, window_size*window_size, feat_dim)

            u, s, v = torch.linalg.svd(feat1)
            feat1 = feat1 - 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_plus(feat, iter=20)

            feat = feat.view(B, window_size, window_size, feat_dim)
            
            feat = model.features[6](feat)
            feat = model.features[7](feat)

            feat = model.norm(feat)
            feat = model.permute(feat)
            feat = model.avgpool(feat)
            feat = model.flatten(feat)
            logits = model.head(feat)

            # conf = temperature * torch.logsumexp(logits / temperature, dim=1)
            conf, _ = torch.max(m(logits), dim=-1)

            confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def __init__(self, rot_mats, trans, debug=False):
        """
        Args:
            rot_mats: (N, 3, 3)
            trans: (N, 3)
        """
        super().__init__()
        # Init as above
        dtype = rot_mats.dtype
        device = rot_mats.device
        Q = avg_matrix_approx(rot_mats)  # (3, 3)
        Q_6d = matrix_to_rotation_6d(Q)
        t2 = trans.mean(0)
        P_6d = matrix_to_rotation_6d(torch.eye(3, dtype=dtype, device=device))
        t1 = torch.zeros(3, dtype=dtype, device=device)

        d_rot_mats = torch.matmul(Q.view(1, 3, 3).permute(0, 2, 1), rot_mats)  # rot_mat(v) = drot x Q
        U, S, V = torch.svd(d_rot_mats.mean(0))
        w = U[:, 0]  # (3,)
        w = w / w.norm(dim=0)
        w = -w if w[w.abs().argmax()] < 0 else w
        Rw = make_basis(w)  # [u, v, w]; (3, 3)
        if debug:
            print('w (most invariant)', w)
            print('Rw (basis)', Rw)
            self.d_rot_mats = d_rot_mats  # debug
            self.dRxy = dRxy  # debug, the 3rd column should be [0, 0, 1]
            self.init_cos = init_cos # debug
            self.init_ths = init_ths # debug

        # dRxy = torch.matmul(Rw.view(1, 3, 3).permute(0, 2, 1), d_rot_mats)  # (N, 3, 3)
        dRxy = Rw.view(1, 3, 3).matmul(d_rot_mats).matmul(Rw.view(1, 3, 3).permute(0, 2, 1))  # (N, 3, 3)
        # init_cos = so3_rotation_angle(dRxy, eps=1e-6, cos_angle=True)  # (N,) the trace trick
        # init_ths = torch.atanh(init_cos)  # (N,), tanh
        init_sin = so3_rotation_angle(dRxy, eps=1e-6, cos_angle=True)  # (N,) the trace trick
        init_ths = torch.atanh(init_sin)  # (N,), tanh
        Rw_6d = matrix_to_rotation_6d(Rw)

        self.Q_6d = nn.Parameter(Q_6d.detach().clone(), requires_grad=True)
        self.t2 = nn.Parameter(t2.detach().clone(), requires_grad=True)
        self.P_6d = nn.Parameter(P_6d.detach().clone(), requires_grad=True)
        self.t1 = nn.Parameter(t1.detach().clone(), requires_grad=True)
        self.Rw_6d = nn.Parameter(Rw_6d.detach().clone(), requires_grad=True)
        self.ths = nn.Parameter(init_ths.detach().clone(), requires_grad=True)
//  [python]
def __init__(self, rot_mats, trans, debug=False):
        """
        Args:
            rot_mats: (N, 3, 3)
            trans: (N, 3)
        """
        super().__init__()
        # Init as above
        dtype = rot_mats.dtype
        device = rot_mats.device
        Q = avg_matrix_approx(rot_mats)  # (3, 3)
        Q_6d = matrix_to_rotation_6d(Q)
        t2 = trans.mean(0)
        P_6d = matrix_to_rotation_6d(torch.eye(3, dtype=dtype, device=device))
        t1 = torch.zeros(3, dtype=dtype, device=device)

        d_rot_mats = torch.matmul(Q.view(1, 3, 3).permute(0, 2, 1), rot_mats)  # rot_mat(v) = drot x Q
        w = torch.linalg.svd(d_rot_mats.mean(0)).U[:, 0]  # (3,)
        w = w / w.norm(dim=0)
        w = -w if w[w.abs().argmax()] < 0 else w
        Rw = make_basis(w)  # [u, v, w]; (3, 3)
        if debug:
            print('w (most invariant)', w)
            print('Rw (basis)', Rw)
            self.d_rot_mats = d_rot_mats  # debug
            self.dRxy = dRxy  # debug, the 3rd column should be [0, 0, 1]
            self.init_cos = init_cos # debug
            self.init_ths = init_ths # debug

        # dRxy = torch.matmul(Rw.view(1, 3, 3).permute(0, 2, 1), d_rot_mats)  # (N, 3, 3)
        dRxy = Rw.view(1, 3, 3).matmul(d_rot_mats).matmul(Rw.view(1, 3, 3).permute(0, 2, 1))  # (N, 3, 3)
        # init_cos = so3_rotation_angle(dRxy, eps=1e-6, cos_angle=True)  # (N,) the trace trick
        # init_ths = torch.atanh(init_cos)  # (N,), tanh
        init_sin = so3_rotation_angle(dRxy, eps=1e-6, cos_angle=True)  # (N,) the trace trick
        init_ths = torch.atanh(init_sin)  # (N,), tanh
        Rw_6d = matrix_to_rotation_6d(Rw)

        self.Q_6d = nn.Parameter(Q_6d.detach().clone(), requires_grad=True)
        self.t2 = nn.Parameter(t2.detach().clone(), requires_grad=True)
        self.P_6d = nn.Parameter(P_6d.detach().clone(), requires_grad=True)
        self.t1 = nn.Parameter(t1.detach().clone(), requires_grad=True)
        self.Rw_6d = nn.Parameter(Rw_6d.detach().clone(), requires_grad=True)
        self.ths = nn.Parameter(init_ths.detach().clone(), requires_grad=True)
//  [python]
def iterate_data_rankfeatweight(data_loader, model, temperature):
    confs = []
    
    weight = model.features[5][17].mlp[3].weight.data
    B, C = weight.size()
    weight_sub = power_iteration_plus(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    weight = model.features[5][17].mlp[3].weight.data = weight
    
    m = torch.nn.Softmax(dim=-1).cuda()
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()

        with torch.no_grad():
            
            for i in range(6):
                x = model.features[i](x)
            feat = x
            B, window_size, window_size, feat_dim = feat.size()
            feat = feat.view(B, window_size*window_size, feat_dim)

            u, s, v = torch.svd(feat1)
            feat1 = feat1 - 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_plus(feat, iter=20)

            feat = feat.view(B, window_size, window_size, feat_dim)
            
            feat = model.features[6](feat)
            feat = model.features[7](feat)

            feat = model.norm(feat)
            feat = model.permute(feat)
            feat = model.avgpool(feat)
            feat = model.flatten(feat)
            logits = model.head(feat)

            # conf = temperature * torch.logsumexp(logits / temperature, dim=1)
            conf, _ = torch.max(m(logits), dim=-1)

            confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeatweight(data_loader, model, temperature):
    confs = []
    
    weight = model.features[5][17].mlp[3].weight.data
    B, C = weight.size()
    weight_sub = power_iteration_plus(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    weight = model.features[5][17].mlp[3].weight.data = weight
    
    m = torch.nn.Softmax(dim=-1).cuda()
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()

        with torch.no_grad():
            
            for i in range(6):
                x = model.features[i](x)
            feat = x
            B, window_size, window_size, feat_dim = feat.size()
            feat = feat.view(B, window_size*window_size, feat_dim)

            u, s, v = torch.linalg.svd(feat1)
            feat1 = feat1 - 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_plus(feat, iter=20)

            feat = feat.view(B, window_size, window_size, feat_dim)
            
            feat = model.features[6](feat)
            feat = model.features[7](feat)

            feat = model.norm(feat)
            feat = model.permute(feat)
            feat = model.avgpool(feat)
            feat = model.flatten(feat)
            logits = model.head(feat)

            # conf = temperature * torch.logsumexp(logits / temperature, dim=1)
            conf, _ = torch.max(m(logits), dim=-1)

            confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def random_square_matrix_of_rank(l, rank, dtype=torch.double, device="cpu"):
    assert rank <= l
    A = torch.randn(l, l, dtype=dtype, device=device)
    u, s, vh = torch.svd(A)
    for i in range(l):
        if i >= rank:
            s[i] = 0
        elif s[i] == 0:
            s[i] = 1
    return (u * s.to(dtype).unsqueeze(-2)) @ vh
//  [python]
def random_square_matrix_of_rank(l, rank, dtype=torch.double, device="cpu"):
    assert rank <= l
    A = torch.randn(l, l, dtype=dtype, device=device)
    u, s, vh = torch.linalg.svd(A, full_matrices=False)
    for i in range(l):
        if i >= rank:
            s[i] = 0
        elif s[i] == 0:
            s[i] = 1
    return (u * s.to(dtype).unsqueeze(-2)) @ vh
//  [python]
def make_fullrank_matrices_with_distinct_singular_values(*shape, device, dtype):
    assert shape[-1] == shape[-2]
    t = make_tensor(shape, device=device, dtype=dtype)
    u, _, vh = t.svd()
    # TODO: improve the handling of complex tensors here
    real_dtype = t.real.dtype if t.dtype.is_complex else t.dtype
    s = torch.arange(1.0, shape[-1] + 1, dtype=real_dtype, device=device).mul_(
        1.0 / (shape[-1] + 1)
    )
    return (u * s.to(dtype)) @ vh
//  [python]
def make_fullrank_matrices_with_distinct_singular_values(*shape, device, dtype):
    assert shape[-1] == shape[-2]
    t = make_tensor(shape, device=device, dtype=dtype)
    u, _, vh = torch.linalg.svd(t, full_matrices=False)
    # TODO: improve the handling of complex tensors here
    real_dtype = t.real.dtype if t.dtype.is_complex else t.dtype
    s = torch.arange(1.0, shape[-1] + 1, dtype=real_dtype, device=device).mul_(
        1.0 / (shape[-1] + 1)
    )
    return (u * s.to(dtype)) @ vh
//  [python]
def random_fullrank_matrix_distinct_singular_value(matrix_size, *batch_dims, **kwargs):
    dtype = kwargs.get("dtype", torch.double)
    device = kwargs.get("device", "cpu")
    silent = kwargs.get("silent", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(matrix_size, matrix_size, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (matrix_size, matrix_size), dtype=dtype, device=device)
    u, s, vh = torch.svd(A)
    real_dtype = A.real.dtype if A.dtype.is_complex else A.dtype
    s = torch.arange(1.0, matrix_size + 1, dtype=real_dtype, device=device).mul_(
        1.0 / (matrix_size + 1)
    )
    return (u * s.to(A.dtype)) @ vh
//  [python]
def random_fullrank_matrix_distinct_singular_value(matrix_size, *batch_dims, **kwargs):
    dtype = kwargs.get("dtype", torch.double)
    device = kwargs.get("device", "cpu")
    silent = kwargs.get("silent", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(matrix_size, matrix_size, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (matrix_size, matrix_size), dtype=dtype, device=device)
    u, _, vh = torch.linalg.svd(A, full_matrices=False)
    real_dtype = A.real.dtype if A.dtype.is_complex else A.dtype
    s = torch.arange(1.0, matrix_size + 1, dtype=real_dtype, device=device).mul_(
        1.0 / (matrix_size + 1)
    )
    return (u * s.to(A.dtype)) @ vh
//  [python]
def random_well_conditioned_matrix(*shape, dtype, device, mean=1.0, sigma=0.001):
    """
    Returns a random rectangular matrix (batch of matrices)
    with singular values sampled from a Gaussian with
    mean `mean` and standard deviation `sigma`.
    The smaller the `sigma`, the better conditioned
    the output matrix is.
    """
    primitive_dtype = {
        torch.float: torch.float,
        torch.double: torch.double,
        torch.cfloat: torch.float,
        torch.cdouble: torch.double,
    }
    x = torch.rand(shape, dtype=dtype, device=device)
    m = x.size(-2)
    n = x.size(-1)
    u, s, vh = torch.svd(x)
    s = (
        (
            torch.randn(
                *(shape[:-2] + (min(m, n),)),
                dtype=primitive_dtype[dtype],
                device=device,
            )
            * sigma
            + mean
        )
        .sort(-1, descending=True)
        .values.to(dtype)
    )
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def random_well_conditioned_matrix(*shape, dtype, device, mean=1.0, sigma=0.001):
    """
    Returns a random rectangular matrix (batch of matrices)
    with singular values sampled from a Gaussian with
    mean `mean` and standard deviation `sigma`.
    The smaller the `sigma`, the better conditioned
    the output matrix is.
    """
    primitive_dtype = {
        torch.float: torch.float,
        torch.double: torch.double,
        torch.cfloat: torch.float,
        torch.cdouble: torch.double,
    }
    x = torch.rand(shape, dtype=dtype, device=device)
    m = x.size(-2)
    n = x.size(-1)
    u, _, vh = torch.linalg.svd(x, full_matrices=False)
    s = (
        (
            torch.randn(
                *(shape[:-2] + (min(m, n),)),
                dtype=primitive_dtype[dtype],
                device=device,
            )
            * sigma
            + mean
        )
        .sort(-1, descending=True)
        .values.to(dtype)
    )
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def dist(self, x, y):
        """from https://pymanopt.org/docs/latest/_modules/pymanopt/manifolds/grassmann.html"""
        _, s, _ = torch.svd(torch.bmm(torch.transpose(x, dim0=-2, dim1=-1), y))
        s[s > 1] = 1
        s = torch.arccos(s)
        return torch.linalg.norm(s, dim=1)
//  [python]
def dist(self, x, y):
        """from https://pymanopt.org/docs/latest/_modules/pymanopt/manifolds/grassmann.html"""
        _, s, _ = torch.linalg.svd(torch.bmm(torch.transpose(x, dim0=-2, dim1=-1), y))
        s[s > 1] = 1
        s = torch.arccos(s)
        return torch.linalg.norm(s, dim=1)
//  [python]
def random_matrix(rows, columns, *batch_dims, **kwargs):
    """Return rectangular matrix or batches of rectangular matrices.

    Parameters:
      dtype - the data type
      device - the device kind
      singular - when True, the output will be singular
    """
    dtype = kwargs.get("dtype", torch.double)
    device = kwargs.get("device", "cpu")
    silent = kwargs.get("silent", False)
    singular = kwargs.get("singular", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(rows, columns, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (rows, columns), dtype=dtype, device=device)
    u, s, vh = torch.svd(A)
    k = min(rows, columns)
    s = torch.linspace(1 / (k + 1), 1, k, dtype=dtype, device=device)
    if singular:
        # make matrix singular
        s[k - 1] = 0
        if k > 2:
            # increase the order of singularity so that the pivoting
            # in LU factorization will be non-trivial
            s[0] = 0
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def random_matrix(rows, columns, *batch_dims, **kwargs):
    """Return rectangular matrix or batches of rectangular matrices.

    Parameters:
      dtype - the data type
      device - the device kind
      singular - when True, the output will be singular
    """
    dtype = kwargs.get("dtype", torch.double)
    device = kwargs.get("device", "cpu")
    silent = kwargs.get("silent", False)
    singular = kwargs.get("singular", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(rows, columns, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (rows, columns), dtype=dtype, device=device)
    u, _, vh = torch.linalg.svd(A, full_matrices=False)
    k = min(rows, columns)
    s = torch.linspace(1 / (k + 1), 1, k, dtype=dtype, device=device)
    if singular:
        # make matrix singular
        s[k - 1] = 0
        if k > 2:
            # increase the order of singularity so that the pivoting
            # in LU factorization will be non-trivial
            s[0] = 0
    return (u * s.unsqueeze(-2)) @ vh