jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def linear(
    points: torch.Tensor,
    K_inv: torch.Tensor,
    covariances: Optional[torch.Tensor] = None,
) -> Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
    """Does a linear projection from the image plane onto the unit sphere
    bearing vector, using
    p_3d = K_inv * point
    bv = p_3d / ||p_3d||
    \Sigma_{bv} = J K_inv \Sigma K_inv^\top J^\top

    J = (||p_3d||^2 * I - p_3d * P-3d^\top) / ||p_3d||^3

    Args:
        point (torch.Tensor): image point in homogenous coordinates (...xNx3)
        K_inv (torch.Tensor): inverse camera matrix (...x3x3)
        covariance (Optional[torch.Tensor]): covariances corresponding to the points (...xNx3x3), default = None

    Returns:
        Tuple[torch.Tensor, Optional[torch.Tensor]]: projected points (...xNx3) and covariance matrices (...xNx3x3) (if given)
    """

    transformed_points = torch.einsum("...ij,...nj->...ni", K_inv, points)
    norms = torch.norm(transformed_points, dim=-1)

    if covariances is not None:
        jacobians = (1 / torch.pow(norms, 3))[..., None, None] * (
            torch.pow(norms, 2)[..., None, None] * torch.eye(3, device=transformed_points.device, dtype=points.dtype)[None, ...]
            - torch.einsum("...i,...j->...ij", transformed_points, transformed_points)
        )
        covariances = torch.einsum(
            "...nij,...jk,...nkl,...ml,...npm->...nip",
            jacobians,
            K_inv,
            covariances,
            K_inv,
            jacobians,
        )
        transformed_points = transformed_points / norms[..., None]
        return transformed_points, covariances

    transformed_points = transformed_points / norms[..., None]
    return transformed_points
//  [python]
def linear(
    points: torch.Tensor,
    K_inv: torch.Tensor,
    covariances: Optional[torch.Tensor] = None,
) -> Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
    """Does a linear projection from the image plane onto the unit sphere
    bearing vector, using
    p_3d = K_inv * point
    bv = p_3d / ||p_3d||
    \Sigma_{bv} = J K_inv \Sigma K_inv^\top J^\top

    J = (||p_3d||^2 * I - p_3d * P-3d^\top) / ||p_3d||^3

    Args:
        point (torch.Tensor): image point in homogenous coordinates (...xNx3)
        K_inv (torch.Tensor): inverse camera matrix (...x3x3)
        covariance (Optional[torch.Tensor]): covariances corresponding to the points (...xNx3x3), default = None

    Returns:
        Tuple[torch.Tensor, Optional[torch.Tensor]]: projected points (...xNx3) and covariance matrices (...xNx3x3) (if given)
    """

    transformed_points = torch.einsum("...ij,...nj->...ni", K_inv, points)
    norms = torch.linalg.norm(transformed_points, dim=-1)

    if covariances is not None:
        jacobians = (1 / torch.pow(norms, 3))[..., None, None] * (
            torch.pow(norms, 2)[..., None, None] * torch.eye(3, device=transformed_points.device, dtype=points.dtype)[None, ...]
            - torch.einsum("...i,...j->...ij", transformed_points, transformed_points)
        )
        covariances = torch.einsum(
            "...nij,...jk,...nkl,...ml,...npm->...nip",
            jacobians,
            K_inv,
            covariances,
            K_inv,
            jacobians,
        )
        transformed_points = transformed_points / norms[..., None]
        return transformed_points, covariances

    transformed_points = transformed_points / norms[..., None]
    return transformed_points
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles])
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48)
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles],
        dim=-1)
    return quaternions
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.linalg.norm(axis_angle, ord=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles])
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48)
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles],
        dim=-1)
    return quaternions
//  [python]
def get_weighted_loss(
        self,
        losses,
        shared_parameters,
        **kwargs,
    ):
        grads = {}
        norm_grads = {}

        for i, loss in enumerate(losses):
            g = list(
                torch.autograd.grad(
                    loss,
                    shared_parameters,
                    retain_graph=True,
                )
            )
            grad = torch.cat([torch.flatten(grad) for grad in g])
            norm_term = torch.norm(grad)

            grads[i] = grad
            norm_grads[i] = grad / norm_term

        G = torch.stack(tuple(v for v in grads.values()))
        D = (
            G[
                0,
            ]
            - G[
                1:,
            ]
        )

        U = torch.stack(tuple(v for v in norm_grads.values()))
        U = (
            U[
                0,
            ]
            - U[
                1:,
            ]
        )
        first_element = torch.matmul(
            G[
                0,
            ],
            U.t(),
        )
        try:
            second_element = torch.inverse(torch.matmul(D, U.t()))
        except:
            # workaround for cases where matrix is singular
            second_element = torch.inverse(
                torch.eye(self.n_tasks - 1, device=self.device) * 1e-8
                + torch.matmul(D, U.t())
            )

        alpha_ = torch.matmul(first_element, second_element)
        alpha = torch.cat(
            (torch.tensor(1 - alpha_.sum(), device=self.device).unsqueeze(-1), alpha_)
        )

        loss = torch.sum(losses * alpha)

        return loss, dict(weights=alpha)
//  [python]
def get_weighted_loss(
        self,
        losses,
        shared_parameters,
        **kwargs,
    ):
        grads = {}
        norm_grads = {}

        for i, loss in enumerate(losses):
            g = list(
                torch.autograd.grad(
                    loss,
                    shared_parameters,
                    retain_graph=True,
                )
            )
            grad = torch.cat([torch.flatten(grad) for grad in g])
            norm_term = torch.linalg.norm(grad)

            grads[i] = grad
            norm_grads[i] = grad / norm_term

        G = torch.stack(tuple(v for v in grads.values()))
        D = (
            G[
                0,
            ]
            - G[
                1:,
            ]
        )

        U = torch.stack(tuple(v for v in norm_grads.values()))
        U = (
            U[
                0,
            ]
            - U[
                1:,
            ]
        )
        first_element = torch.matmul(
            G[
                0,
            ],
            U.t(),
        )
        try:
            second_element = torch.inverse(torch.matmul(D, U.t()))
        except:
            # workaround for cases where matrix is singular
            second_element = torch.inverse(
                torch.eye(self.n_tasks - 1, device=self.device) * 1e-8
                + torch.matmul(D, U.t())
            )

        alpha_ = torch.matmul(first_element, second_element)
        alpha = torch.cat(
            (torch.tensor(1 - alpha_.sum(), device=self.device).unsqueeze(-1), alpha_)
        )

        loss = torch.sum(losses * alpha)

        return loss, dict(weights=alpha)
//  [python]
def test_symm_aux_C2x2_LU(env):
    C = env.C[env.keyC]
    T = env.T[env.keyT]
    C2x2= _get_aux_C2x2_LU(C,T)

    tC2x2= C2x2.permute(1,0,3,2)
    print(f"C2x2-C2x2^t {torch.norm(C2x2-tC2x2)}")

    tC2x2= C2x2.permute(1,0,2,3)
    print(f"C2x2-C2x2^t(env) {torch.norm(C2x2-tC2x2)}")

    tC2x2= C2x2.permute(0,1,3,2)
    print(f"C2x2-C2x2^t(aux) {torch.norm(C2x2-tC2x2)}")
//  [python]
def test_symm_aux_C2x2_LU(env):
    C = env.C[env.keyC]
    T = env.T[env.keyT]
    C2x2= _get_aux_C2x2_LU(C,T)

    tC2x2= C2x2.permute(1,0,3,2)
    print(f"C2x2-C2x2^t {torch.linalg.norm(C2x2-tC2x2)}")

    tC2x2= C2x2.permute(1,0,2,3)
    print(f"C2x2-C2x2^t(env) {torch.linalg.norm(C2x2-tC2x2)}")

    tC2x2= C2x2.permute(0,1,3,2)
    print(f"C2x2-C2x2^t(aux) {torch.linalg.norm(C2x2-tC2x2)}")
//  [python]
def test_omp():
    n_samples, n_features, n_atoms, s = 100, 50, 20, 5
    X, D_true, S_true = generate_synthetic_data(n_samples, n_features, n_atoms, s)
    ids, m = sparse.omp(X, D_true, s)
    SM = (m.unsqueeze(1) @ D_true[ids]).squeeze(1)
    error = torch.norm(SM - X)
    assert error < 1e-2, "OMP reconstruction error is too high"
//  [python]
def test_omp():
    n_samples, n_features, n_atoms, s = 100, 50, 20, 5
    X, D_true, S_true = generate_synthetic_data(n_samples, n_features, n_atoms, s)
    ids, m = sparse.omp(X, D_true, s)
    SM = (m.unsqueeze(1) @ D_true[ids]).squeeze(1)
    error = torch.linalg.norm(SM - X)
    assert error < 1e-2, "OMP reconstruction error is too high"
//  [python]
def test_symm_open_C2x2_LU(C, T, a):
    C2x2= _get_open_C2x2_LU_sl(C,T,a)
    tC2x2= C2x2.permute(1,0,2,3)
    print(f"open C2x2-C2x2^t {torch.norm(C2x2-tC2x2)}")

    C2x2= torch.einsum('ijaa->ij',C2x2)
    chi= C.size()[0]
    C2x2= C2x2.reshape(chi,a.size()[3]**2,chi,a.size()[4]**2)
    # |C2x2|--2    |C2x2|--1
    # |____|--3 => |____|--3
    # 0  1         0  2
    C2x2= C2x2.permute(0,2,1,3).contiguous()

    tC2x2= C2x2.permute(1,0,3,2)
    print(f"C2x2-C2x2^t {torch.norm(C2x2-tC2x2)}")

    tC2x2= C2x2.permute(1,0,2,3)
    print(f"C2x2-C2x2^t(env) {torch.norm(C2x2-tC2x2)}")

    tC2x2= C2x2.permute(0,1,3,2)
    print(f"C2x2-C2x2^t(aux) {torch.norm(C2x2-tC2x2)}")
//  [python]
def test_symm_open_C2x2_LU(C, T, a):
    C2x2= _get_open_C2x2_LU_sl(C,T,a)
    tC2x2= C2x2.permute(1,0,2,3)
    print(f"open C2x2-C2x2^t {torch.linalg.norm(C2x2-tC2x2)}")

    C2x2= torch.einsum('ijaa->ij',C2x2)
    chi= C.size()[0]
    C2x2= C2x2.reshape(chi,a.size()[3]**2,chi,a.size()[4]**2)
    # |C2x2|--2    |C2x2|--1
    # |____|--3 => |____|--3
    # 0  1         0  2
    C2x2= C2x2.permute(0,2,1,3).contiguous()

    tC2x2= C2x2.permute(1,0,3,2)
    print(f"C2x2-C2x2^t {torch.linalg.norm(C2x2-tC2x2)}")

    tC2x2= C2x2.permute(1,0,2,3)
    print(f"C2x2-C2x2^t(env) {torch.linalg.norm(C2x2-tC2x2)}")

    tC2x2= C2x2.permute(0,1,3,2)
    print(f"C2x2-C2x2^t(aux) {torch.linalg.norm(C2x2-tC2x2)}")
//  [python]
def forward(self, x) -> torch.Tensor:
        return torch.norm(x, p=self.p, dim=self.dim, keepdim=self.keepdim)
//  [python]
def forward(self, x) -> torch.Tensor:
        return torch.linalg.norm(x, ord=self.p, dim=self.dim, keepdim=self.keepdim)
//  [python]
def is_valid_quaternion(q: torch.tensor) -> bool:
    return allclose(torch.norm(q, dim=-1), 1.0)
//  [python]
def is_valid_quaternion(q: torch.tensor) -> bool:
    return allclose(torch.linalg.norm(q, dim=-1), 1.0)
//  [python]
def normalize_quaternion(q: torch.tensor) -> torch.tensor:
    return q / torch.norm(q, dim=-1, keepdim=True)
//  [python]
def normalize_quaternion(q: torch.tensor) -> torch.tensor:
    return q / torch.linalg.norm(q, dim=-1, keepdim=True)
//  [python]
def fro_norm(mat: torch.Tensor):
    return torch.norm(mat.flatten())
//  [python]
def fro_norm(mat: torch.Tensor):
    return torch.linalg.norm(mat.flatten())
//  [python]
def clip_grad_norm_(params, max_norm, aggregate_norm_fn=None) -> torch.Tensor:
    if isinstance(params, torch.Tensor):
        params = [params]
    params = list(params)
    grads = [p.grad.detach() for p in filter(lambda p: p.grad is not None, params)]
    if len(grads) == 0:
        if len(params) > 0:
            return params[0].new_tensor(0.0)
        else:
            return torch.tensor(0.0)

    if len(grads) == 1:
        total_norm = torch.norm(grads[0], p=2, dtype=torch.float32)
    else:
        if multi_tensor_l2norm_available:
            total_norm = multi_tensor_total_norm(grads)
        else:
            if torch.cuda.is_available():
                warnings.warn(
                    "amp_C fused kernels unavailable, disabling multi_tensor_l2norm; "
                    "you may get better performance by installing NVIDIA's apex library"
                )
                device = torch.cuda.current_device()
            elif grads[0].device.type == "xla":
                device = grads[0].device
            else:
                device = torch.device("cpu")
            total_norm = torch.norm(
                torch.stack(
                    [torch.norm(g, p=2, dtype=torch.float32).to(device) for g in grads]
                )
            )

    if aggregate_norm_fn is not None:
        total_norm = aggregate_norm_fn(total_norm)

    if max_norm > 0:
        max_norm = float(max_norm)
        clip_coef = (max_norm / (total_norm + 1e-6)).clamp_(max=1)
        for g in grads:
            g.mul_(clip_coef)
    return total_norm
//  [python]
def clip_grad_norm_(params, max_norm, aggregate_norm_fn=None) -> torch.Tensor:
    if isinstance(params, torch.Tensor):
        params = [params]
    params = list(params)
    grads = [p.grad.detach() for p in filter(lambda p: p.grad is not None, params)]
    if len(grads) == 0:
        if len(params) > 0:
            return params[0].new_tensor(0.0)
        else:
            return torch.tensor(0.0)

    if len(grads) == 1:
        total_norm = torch.linalg.norm(grads[0], ord=2, dtype=torch.float32)
    else:
        if multi_tensor_l2norm_available:
            total_norm = multi_tensor_total_norm(grads)
        else:
            if torch.cuda.is_available():
                warnings.warn(
                    "amp_C fused kernels unavailable, disabling multi_tensor_l2norm; "
                    "you may get better performance by installing NVIDIA's apex library"
                )
                device = torch.cuda.current_device()
            elif grads[0].device.type == "xla":
                device = grads[0].device
            else:
                device = torch.device("cpu")
            total_norm = torch.linalg.norm(
                torch.stack(
                    [torch.linalg.norm(g, ord=2, dtype=torch.float32).to(device) for g in grads]
                )
            )

    if aggregate_norm_fn is not None:
        total_norm = aggregate_norm_fn(total_norm)

    if max_norm > 0:
        max_norm = float(max_norm)
        clip_coef = (max_norm / (total_norm + 1e-6)).clamp_(max=1)
        for g in grads:
            g.mul_(clip_coef)
    return total_norm
//  [python]
def test_k_svd():
    n_samples, n_features, n_atoms, s = 100, 20, 50, 5
    X, D_true, _ = generate_synthetic_data(n_samples, n_features, n_atoms, s)

    # Initialize with random dictionary
    M_init = torch.randn(n_atoms, n_features)

    # Ensure the dictionary has learned something by checking if the reconstruction error decreases
    ids_initial, m_initial = sparse.omp(X, M_init, s)
    SM_initial = (m_initial.unsqueeze(1) @ M_init[ids_initial]).squeeze(1)
    initial_error = torch.norm(SM_initial - X)

    M_learned, _, _ = sparse.k_svd(X, M_init, s, n_iter=10)
    # M_learned, _ = old_k_svd(X, M_init, s, n_iter=10)

    ids_learned, m_learned = sparse.omp(X, M_learned, s)
    SM_learned = (m_learned.unsqueeze(1) @ M_learned[ids_learned]).squeeze(1)
    learned_error = torch.norm(SM_learned - X)

    print(initial_error)
    print(learned_error)

    assert learned_error < initial_error, "K-SVD didn't reduce the reconstruction error"
//  [python]
def test_k_svd():
    n_samples, n_features, n_atoms, s = 100, 20, 50, 5
    X, D_true, _ = generate_synthetic_data(n_samples, n_features, n_atoms, s)

    # Initialize with random dictionary
    M_init = torch.randn(n_atoms, n_features)

    # Ensure the dictionary has learned something by checking if the reconstruction error decreases
    ids_initial, m_initial = sparse.omp(X, M_init, s)
    SM_initial = (m_initial.unsqueeze(1) @ M_init[ids_initial]).squeeze(1)
    initial_error = torch.linalg.norm(SM_initial - X)

    M_learned, _, _ = sparse.k_svd(X, M_init, s, n_iter=10)
    # M_learned, _ = old_k_svd(X, M_init, s, n_iter=10)

    ids_learned, m_learned = sparse.omp(X, M_learned, s)
    SM_learned = (m_learned.unsqueeze(1) @ M_learned[ids_learned]).squeeze(1)
    learned_error = torch.linalg.norm(SM_learned - X)

    print(initial_error)
    print(learned_error)

    assert learned_error < initial_error, "K-SVD didn't reduce the reconstruction error"
//  [python]
def symsqrt_dist(cov1: torch.Tensor, cov2: torch.Tensor) -> float:
    """Distance between square roots of matrices"""

    cov1 = to_pytorch(cov1)
    cov2 = to_pytorch(cov2)
    cov1 = symsqrt_svd(cov1)
    cov2 = symsqrt_svd(cov2)
    return torch.norm(cov1 - cov2).item()
//  [python]
def symsqrt_dist(cov1: torch.Tensor, cov2: torch.Tensor) -> float:
    """Distance between square roots of matrices"""

    cov1 = to_pytorch(cov1)
    cov2 = to_pytorch(cov2)
    cov1 = symsqrt_svd(cov1)
    cov2 = symsqrt_svd(cov2)
    return torch.linalg.norm(cov1 - cov2).item()
//  [python]
def _reproj_error_HXY(H, X, Y):
    HX = de_homo_py(torch.matmul(H, homo_py(X).t()).t())
    errors = torch.norm(Y - HX, dim=1)
    return torch.mean(errors), errors
//  [python]
def _reproj_error_HXY(H, X, Y):
    HX = de_homo_py(torch.matmul(H, homo_py(X).t()).t())
    errors = torch.linalg.norm(Y - HX, dim=1)
    return torch.mean(errors), errors
//  [python]
def __call__(self, x: torch.Tensor) -> torch.Tensor:
        self.feature_extractor = self.feature_extractor.to(x.device)
        features = self.feature_extractor(x)[self.penultimate_layer_name]

        logits = self._get_logits(features)

        x_p_t = torch.norm(
            torch.matmul(torch.flatten(features, 1) - self.u.to(x.device), self.principal_subspace.to(x.device)), dim=-1
        )
        vlogit = x_p_t * self.alpha
        energy = torch.logsumexp(logits, dim=-1)
        score = -vlogit + energy
        return torch.nan_to_num(score, 1e6)
//  [python]
def __call__(self, x: torch.Tensor) -> torch.Tensor:
        self.feature_extractor = self.feature_extractor.to(x.device)
        features = self.feature_extractor(x)[self.penultimate_layer_name]

        logits = self._get_logits(features)

        x_p_t = torch.linalg.norm(
            torch.matmul(torch.flatten(features, 1) - self.u.to(x.device), self.principal_subspace.to(x.device)), dim=-1
        )
        vlogit = x_p_t * self.alpha
        energy = torch.logsumexp(logits, dim=-1)
        score = -vlogit + energy
        return torch.nan_to_num(score, 1e6)
//  [python]
def end(self):
        self.top_k = 1000 if self.train_features.shape[1] > 1500 else 512

        _logger.info("Train features shape: %s", self.train_features.shape)
        _logger.info("Train logits shape: %s", self.train_logits.shape)

        # calculate eigenvectors of the covariance matrix
        ec = EmpiricalCovariance(assume_centered=True)
        ec.fit(self.train_features.cpu().numpy() - self.u.detach().cpu().numpy())
        eig_vals, eigen_vectors = np.linalg.eig(ec.covariance_)
        determinant = np.linalg.det(ec.covariance_)
        _logger.debug("Determinant: %s", determinant)
        _logger.debug("Eigen values: %s", eig_vals)

        # select largest eigenvectors to get the principal subspace
        largest_eigvals_idx = np.argsort(eig_vals * -1)[self.top_k :]
        self.principal_subspace = torch.from_numpy(
            np.ascontiguousarray((eigen_vectors.T[largest_eigvals_idx]).T)
        ).float()
        _logger.debug("Principal subspace: %s", self.principal_subspace)

        # calculate residual
        x_p_t = torch.matmul(self.train_features.cpu() - self.u.cpu(), self.principal_subspace.cpu())
        vlogits = torch.norm(x_p_t, dim=-1)
        self.alpha = self.train_logits.max(dim=-1)[0].mean() / vlogits.mean()
        _logger.debug("Alpha: %s", self.alpha)

        del self.train_features
//  [python]
def end(self):
        self.top_k = 1000 if self.train_features.shape[1] > 1500 else 512

        _logger.info("Train features shape: %s", self.train_features.shape)
        _logger.info("Train logits shape: %s", self.train_logits.shape)

        # calculate eigenvectors of the covariance matrix
        ec = EmpiricalCovariance(assume_centered=True)
        ec.fit(self.train_features.cpu().numpy() - self.u.detach().cpu().numpy())
        eig_vals, eigen_vectors = np.linalg.eig(ec.covariance_)
        determinant = np.linalg.det(ec.covariance_)
        _logger.debug("Determinant: %s", determinant)
        _logger.debug("Eigen values: %s", eig_vals)

        # select largest eigenvectors to get the principal subspace
        largest_eigvals_idx = np.argsort(eig_vals * -1)[self.top_k :]
        self.principal_subspace = torch.from_numpy(
            np.ascontiguousarray((eigen_vectors.T[largest_eigvals_idx]).T)
        ).float()
        _logger.debug("Principal subspace: %s", self.principal_subspace)

        # calculate residual
        x_p_t = torch.matmul(self.train_features.cpu() - self.u.cpu(), self.principal_subspace.cpu())
        vlogits = torch.linalg.norm(x_p_t, dim=-1)
        self.alpha = self.train_logits.max(dim=-1)[0].mean() / vlogits.mean()
        _logger.debug("Alpha: %s", self.alpha)

        del self.train_features
//  [python]
def _get_M2s_batch(Es_batch):
    from batch_svd import batch_svd # https://github.com/KinglittleQ/torch-batch-svd.git
    # Getting 4 possible poses from E
    Us, Ss, Vs = batch_svd(Es_batch)
    Ws_list = []
    for U, V in zip(Us, Vs):
        W = torch.tensor([[0,-1,0], [1,0,0], [0,0,1]], dtype=Es_batch.dtype, device=Es_batch.device)
        if torch.det(torch.mm(U, torch.mm(W, V.t())))<0:
            W = -W
        Ws_list.append(W)
    Ws = torch.stack(Ws_list)

    t_recover_batch = Us[:, :, 2:3]/torch.norm(Us[:, :, 2:3], p=2, dim=1, keepdim=True)
    R_recover_1_batch = Us @ Ws @ Vs.transpose(1, 2)
    R_recover_2_batch = Us @ Ws.transpose(1, 2) @ Vs.transpose(1, 2) # also an option

    R2s_batch = [R_recover_1_batch, R_recover_2_batch]
    t2s_batch = [t_recover_batch, -t_recover_batch]
    # M2s = [torch.cat((x, y), 1) for x, y in [(x,y) for x in R2s for y in t2s]]
    return R2s_batch, t2s_batch
//  [python]
def _get_M2s_batch(Es_batch):
    from batch_svd import batch_svd # https://github.com/KinglittleQ/torch-batch-svd.git
    # Getting 4 possible poses from E
    Us, Ss, Vs = batch_svd(Es_batch)
    Ws_list = []
    for U, V in zip(Us, Vs):
        W = torch.tensor([[0,-1,0], [1,0,0], [0,0,1]], dtype=Es_batch.dtype, device=Es_batch.device)
        if torch.det(torch.mm(U, torch.mm(W, V.t())))<0:
            W = -W
        Ws_list.append(W)
    Ws = torch.stack(Ws_list)

    t_recover_batch = Us[:, :, 2:3]/torch.linalg.norm(Us[:, :, 2:3], p=2, dim=1, keepdim=True)
    R_recover_1_batch = Us @ Ws @ Vs.transpose(1, 2)
    R_recover_2_batch = Us @ Ws.transpose(1, 2) @ Vs.transpose(1, 2) # also an option

    R2s_batch = [R_recover_1_batch, R_recover_2_batch]
    t2s_batch = [t_recover_batch, -t_recover_batch]
    # M2s = [torch.cat((x, y), 1) for x, y in [(x,y) for x in R2s for y in t2s]]
    return R2s_batch, t2s_batch
//  [python]
def compute_matrix_root_inverse_residuals(
    A: Tensor,
    X_hat: Tensor,
    root: int,
    epsilon: float,
    exponent_multiplier: float,
) -> Tuple[Tensor, Tensor]:
    """Compute residual of matrix root inverse for debugging purposes.

        relative error    = ||X - X_hat||_inf / ||X||_inf
        relative residual = ||A X^r - I||_inf

    Args:
        A (Tensor): Matrix of interest.
        X (Tensor): Computed matrix root inverse.
        root (int): Root of interest.
        epsilon (float): Adds epsilon * I to matrix.
        exponent_multiplier (float): Exponent multiplier to be multiplied to the numerator of the inverse root.

    Returns:
        absolute_error (Tensor): absolute error of matrix root inverse
        relative_error (Tensor): relative error of matrix root inverse
        residual (Tensor): residual of matrix root inverse

    """

    # check shape of matrix
    if len(A.shape) != 2:
        raise ValueError("Matrix is not 2-dimensional!")
    elif A.shape[0] != A.shape[1]:
        raise ValueError("Matrix is not square!")
    elif A.shape != X_hat.shape:
        raise ValueError("Matrix shapes do not match!")

    # compute error by comparing against double precision
    X = matrix_inverse_root(
        A.double(), root, epsilon=epsilon, exponent_multiplier=exponent_multiplier
    )
    relative_error = torch.dist(X, X_hat, p=torch.inf) / torch.norm(X, p=torch.inf)

    # compute residual
    if exponent_multiplier == 1.0:
        X_invr = torch.linalg.matrix_power(X_hat.double(), n=-root)
    else:
        X_invr, _, _ = _matrix_root_eigen(
            X_hat.double(),
            root=1,
            epsilon=0.0,
            inverse=True,
            make_positive_semidefinite=True,
            exponent_multiplier=root / exponent_multiplier,
        )

    A_reg = A.double() + epsilon * torch.eye(
        A.shape[0], dtype=torch.float64, device=A.device
    )
    relative_residual = torch.dist(X_invr, A_reg, p=torch.inf) / torch.norm(
        A_reg, p=torch.inf
    )

    return relative_error, relative_residual
//  [python]
def compute_matrix_root_inverse_residuals(
    A: Tensor,
    X_hat: Tensor,
    root: int,
    epsilon: float,
    exponent_multiplier: float,
) -> Tuple[Tensor, Tensor]:
    """Compute residual of matrix root inverse for debugging purposes.

        relative error    = ||X - X_hat||_inf / ||X||_inf
        relative residual = ||A X^r - I||_inf

    Args:
        A (Tensor): Matrix of interest.
        X (Tensor): Computed matrix root inverse.
        root (int): Root of interest.
        epsilon (float): Adds epsilon * I to matrix.
        exponent_multiplier (float): Exponent multiplier to be multiplied to the numerator of the inverse root.

    Returns:
        absolute_error (Tensor): absolute error of matrix root inverse
        relative_error (Tensor): relative error of matrix root inverse
        residual (Tensor): residual of matrix root inverse

    """

    # check shape of matrix
    if len(A.shape) != 2:
        raise ValueError("Matrix is not 2-dimensional!")
    elif A.shape[0] != A.shape[1]:
        raise ValueError("Matrix is not square!")
    elif A.shape != X_hat.shape:
        raise ValueError("Matrix shapes do not match!")

    # compute error by comparing against double precision
    X = matrix_inverse_root(
        A.double(), root, epsilon=epsilon, exponent_multiplier=exponent_multiplier
    )
    relative_error = torch.dist(X, X_hat, p=torch.inf) / torch.linalg.norm(X, p=torch.inf)

    # compute residual
    if exponent_multiplier == 1.0:
        X_invr = torch.linalg.matrix_power(X_hat.double(), n=-root)
    else:
        X_invr, _, _ = _matrix_root_eigen(
            X_hat.double(),
            root=1,
            epsilon=0.0,
            inverse=True,
            make_positive_semidefinite=True,
            exponent_multiplier=root / exponent_multiplier,
        )

    A_reg = A.double() + epsilon * torch.eye(
        A.shape[0], dtype=torch.float64, device=A.device
    )
    relative_residual = torch.dist(X_invr, A_reg, p=torch.inf) / torch.linalg.norm(
        A_reg, p=torch.inf
    )

    return relative_error, relative_residual
//  [python]
def get_u_hat(self, _u, _w):
        """Enforce w^T u >= -1. When using h(.) = tanh(.), this is a sufficient condition
        for invertibility of the transformation f(z). See Appendix A.1.
        """
        wtu = torch.bmm(_u, torch.transpose(_w, dim0=-2, dim1=-1))

        m_wtu = -1 + torch.log(1 + torch.exp(wtu))
        return _u + (m_wtu - wtu) * _w / torch.norm(_w, p=2, dim=-1, keepdim=True) ** 2
//  [python]
def get_u_hat(self, _u, _w):
        """Enforce w^T u >= -1. When using h(.) = tanh(.), this is a sufficient condition
        for invertibility of the transformation f(z). See Appendix A.1.
        """
        wtu = torch.bmm(_u, torch.transpose(_w, dim0=-2, dim1=-1))

        m_wtu = -1 + torch.log(1 + torch.exp(wtu))
        return _u + (m_wtu - wtu) * _w / torch.linalg.norm(_w, p=2, dim=-1, keepdim=True) ** 2
//  [python]
def _normalize_XY(X, Y):
    """ The Hartley normalization. Following https://github.com/marktao99/python/blob/da2682f8832483650b85b0be295ae7eaf179fcc5/CVP/samples/sfm.py#L157 
    corrected with https://www.mathworks.com/matlabcentral/fileexchange/27541-fundamental-matrix-computation
    and https://en.wikipedia.org/wiki/Eight-point_algorithm#The_normalized_eight-point_algorithm """
    if X.size()[0] != Y.size()[0]:
        raise ValueError("Number of points don't match.")
    X = utils_misc._homo(X)
    mean_1 = torch.mean(X[:, :2], dim=0, keepdim=True)
    S1 = np.sqrt(2) / torch.mean(torch.norm(X[:, :2]-mean_1, 2, dim=1))
    # print(mean_1.size(), S1.size())
    T1 = torch.tensor([[S1,0,-S1*mean_1[0, 0]],[0,S1,-S1*mean_1[0, 1]],[0,0,1]], device=X.device)
    X_normalized = utils_misc._de_homo(torch.mm(T1, X.t()).t()) # ideally zero mean (x, y), and sqrt(2) average norm

    # xxx = X_normalized.numpy()
    # print(np.mean(xxx, axis=0))
    # print(np.mean(np.linalg.norm(xxx, 2, axis=1)))

    Y = utils_misc._homo(Y)
    mean_2 = torch.mean(Y[:, :2], dim=0, keepdim=True)
    S2 = np.sqrt(2) / torch.mean(torch.norm(Y[:, :2]-mean_2, 2, dim=1))
    T2 = torch.tensor([[S2,0,-S2*mean_2[0, 0]],[0,S2,-S2*mean_2[0, 1]],[0,0,1]], device=X.device)
    Y_normalized = utils_misc._de_homo(torch.mm(T2, Y.t()).t())

    return X_normalized, Y_normalized, T1, T2
//  [python]
def _normalize_XY(X, Y):
    """ The Hartley normalization. Following https://github.com/marktao99/python/blob/da2682f8832483650b85b0be295ae7eaf179fcc5/CVP/samples/sfm.py#L157 
    corrected with https://www.mathworks.com/matlabcentral/fileexchange/27541-fundamental-matrix-computation
    and https://en.wikipedia.org/wiki/Eight-point_algorithm#The_normalized_eight-point_algorithm """
    if X.size()[0] != Y.size()[0]:
        raise ValueError("Number of points don't match.")
    X = utils_misc._homo(X)
    mean_1 = torch.mean(X[:, :2], dim=0, keepdim=True)
    S1 = np.sqrt(2) / torch.mean(torch.linalg.norm(X[:, :2]-mean_1, 2, dim=1))
    # print(mean_1.size(), S1.size())
    T1 = torch.tensor([[S1,0,-S1*mean_1[0, 0]],[0,S1,-S1*mean_1[0, 1]],[0,0,1]], device=X.device)
    X_normalized = utils_misc._de_homo(torch.mm(T1, X.t()).t()) # ideally zero mean (x, y), and sqrt(2) average norm

    # xxx = X_normalized.numpy()
    # print(np.mean(xxx, axis=0))
    # print(np.mean(np.linalg.norm(xxx, 2, axis=1)))

    Y = utils_misc._homo(Y)
    mean_2 = torch.mean(Y[:, :2], dim=0, keepdim=True)
    S2 = np.sqrt(2) / torch.mean(torch.linalg.norm(Y[:, :2]-mean_2, 2, dim=1))
    T2 = torch.tensor([[S2,0,-S2*mean_2[0, 0]],[0,S2,-S2*mean_2[0, 1]],[0,0,1]], device=X.device)
    Y_normalized = utils_misc._de_homo(torch.mm(T2, Y.t()).t())

    return X_normalized, Y_normalized, T1, T2
//  [python]
def forward(self, data):
        outputs = super().forward(data, return_entity_embeddings=True)
        kge_outputs = self.kge_model(data, return_entity_embeddings=True)
        blp_scores = outputs['ranking_outputs'].flatten()
        kge_scores = kge_outputs['ranking_outputs'].flatten()

        # compute weighted average of scores instead of adding, if specified
        if self.weighted_average:
            weight = torch.sigmoid(self.alpha)
            scores = weight * blp_scores + (1 - weight) * kge_scores
        else:
            scores = blp_scores + kge_scores
        outputs['ranking_outputs'] = scores

        # regularize BERT and KG entity embeddings to be similar to each other
        if self.regularization:
            head, tail = outputs['head'], outputs['tail']
            kge_head, kge_tail = kge_outputs['head'], kge_outputs['tail']
            emb_norm = torch.norm(kge_head - head) ** 2 + \
                       torch.norm(kge_tail - tail) ** 2
            outputs['regularization'] = self.regularization * emb_norm

        return outputs
//  [python]
def forward(self, data):
        outputs = super().forward(data, return_entity_embeddings=True)
        kge_outputs = self.kge_model(data, return_entity_embeddings=True)
        blp_scores = outputs['ranking_outputs'].flatten()
        kge_scores = kge_outputs['ranking_outputs'].flatten()

        # compute weighted average of scores instead of adding, if specified
        if self.weighted_average:
            weight = torch.sigmoid(self.alpha)
            scores = weight * blp_scores + (1 - weight) * kge_scores
        else:
            scores = blp_scores + kge_scores
        outputs['ranking_outputs'] = scores

        # regularize BERT and KG entity embeddings to be similar to each other
        if self.regularization:
            head, tail = outputs['head'], outputs['tail']
            kge_head, kge_tail = kge_outputs['head'], kge_outputs['tail']
            emb_norm = torch.linalg.norm(kge_head - head) ** 2 + \
                       torch.linalg.norm(kge_tail - tail) ** 2
            outputs['regularization'] = self.regularization * emb_norm

        return outputs
//  [python]
def _approximation_error(matrix: torch.Tensor, s_matrix: torch.Tensor) ->torch.Tensor:
    norm_of_matrix = torch.norm(matrix)
    error = matrix - torch.mm(s_matrix, s_matrix)
    error = torch.norm(error) / norm_of_matrix
    return error
//  [python]
def _approximation_error(matrix: torch.Tensor, s_matrix: torch.Tensor) ->torch.Tensor:
    norm_of_matrix = torch.linalg.norm(matrix)
    error = matrix - torch.mm(s_matrix, s_matrix)
    error = torch.linalg.norm(error) / norm_of_matrix
    return error
//  [python]
def calculate_Kij(iSample, jSample, sigma):
    return torch.exp(-torch.pow(torch.norm(iSample-jSample,p=2),exponent=2)/(2*sigma))
//  [python]
def calculate_Kij(iSample, jSample, sigma):
    return torch.exp(-torch.pow(torch.linalg.norm(iSample-jSample,p=2),exponent=2)/(2*sigma))
//  [python]
def _normalize_XY_batch(X, Y):
    """ The Hartley normalization. Following https://github.com/marktao99/python/blob/da2682f8832483650b85b0be295ae7eaf179fcc5/CVP/samples/sfm.py#L157 
    corrected with https://www.mathworks.com/matlabcentral/fileexchange/27541-fundamental-matrix-computation
    and https://en.wikipedia.org/wiki/Eight-point_algorithm#The_normalized_eight-point_algorithm """
    # X: [batch_size, N, 2]
    if X.size()[1] != Y.size()[1]:
        raise ValueError("Number of points don't match.")
    X = utils_misc._homo(X)
    mean_1s = torch.mean(X[:, :, :2], dim=1, keepdim=True)
    S1s = np.sqrt(2) / torch.mean(torch.norm(X[:, :, :2]-mean_1s, 2, dim=2), dim=1)
    T1_list = []
    for S1, mean_1 in zip(S1s, mean_1s):
        T1_list.append(torch.tensor([[S1,0,-S1*mean_1[0, 0]],[0,S1,-S1*mean_1[0, 1]],[0,0,1]], device=X.device))
    T1s = torch.stack(T1_list)
    X_normalized = utils_misc._de_homo(torch.bmm(T1s, X.transpose(1, 2)).transpose(1, 2)) # ideally zero mean (x, y), and sqrt(2) average norm

    # xxx = X_normalized.numpy()
    # print(np.mean(xxx, axis=0))
    # print(np.mean(np.linalg.norm(xxx, 2, axis=1)))

    Y = utils_misc._homo(Y)
    mean_2s = torch.mean(Y[:, :, :2], dim=1, keepdim=True)
    S2s = np.sqrt(2) / torch.mean(torch.norm(Y[:, :,  :2]-mean_2s, 2, dim=2), dim=1)
    T2_list = []
    for S2, mean_2 in zip(S2s, mean_2s):
        T2_list.append(torch.tensor([[S2,0,-S2*mean_2[0, 0]],[0,S2,-S2*mean_2[0, 1]],[0,0,1]], device=X.device))
    T2s = torch.stack(T2_list)
    Y_normalized = utils_misc._de_homo(torch.bmm(T2s, Y.transpose(1, 2)).transpose(1, 2))

    return X_normalized, Y_normalized, T1s, T2s
//  [python]
def _normalize_XY_batch(X, Y):
    """ The Hartley normalization. Following https://github.com/marktao99/python/blob/da2682f8832483650b85b0be295ae7eaf179fcc5/CVP/samples/sfm.py#L157 
    corrected with https://www.mathworks.com/matlabcentral/fileexchange/27541-fundamental-matrix-computation
    and https://en.wikipedia.org/wiki/Eight-point_algorithm#The_normalized_eight-point_algorithm """
    # X: [batch_size, N, 2]
    if X.size()[1] != Y.size()[1]:
        raise ValueError("Number of points don't match.")
    X = utils_misc._homo(X)
    mean_1s = torch.mean(X[:, :, :2], dim=1, keepdim=True)
    S1s = np.sqrt(2) / torch.mean(torch.linalg.norm(X[:, :, :2]-mean_1s, 2, dim=2), dim=1)
    T1_list = []
    for S1, mean_1 in zip(S1s, mean_1s):
        T1_list.append(torch.tensor([[S1,0,-S1*mean_1[0, 0]],[0,S1,-S1*mean_1[0, 1]],[0,0,1]], device=X.device))
    T1s = torch.stack(T1_list)
    X_normalized = utils_misc._de_homo(torch.bmm(T1s, X.transpose(1, 2)).transpose(1, 2)) # ideally zero mean (x, y), and sqrt(2) average norm

    # xxx = X_normalized.numpy()
    # print(np.mean(xxx, axis=0))
    # print(np.mean(np.linalg.norm(xxx, 2, axis=1)))

    Y = utils_misc._homo(Y)
    mean_2s = torch.mean(Y[:, :, :2], dim=1, keepdim=True)
    S2s = np.sqrt(2) / torch.mean(torch.linalg.norm(Y[:, :,  :2]-mean_2s, 2, dim=2), dim=1)
    T2_list = []
    for S2, mean_2 in zip(S2s, mean_2s):
        T2_list.append(torch.tensor([[S2,0,-S2*mean_2[0, 0]],[0,S2,-S2*mean_2[0, 1]],[0,0,1]], device=X.device))
    T2s = torch.stack(T2_list)
    Y_normalized = utils_misc._de_homo(torch.bmm(T2s, Y.transpose(1, 2)).transpose(1, 2))

    return X_normalized, Y_normalized, T1s, T2s
//  [python]
def normalized_cut_2d(edge_index, pos):
    row, col = edge_index
    edge_attr = torch.norm(pos[row] - pos[col], p=2, dim=1)
    return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0))
//  [python]
def normalized_cut_2d(edge_index, pos):
    row, col = edge_index
    edge_attr = torch.linalg.norm(pos[row] - pos[col], ord=2, dim=1)
    return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0))
//  [python]
def translation_error(t1, t2):
    """
    Torch batch implementation of the rotation error between the estimated and the ground truth rotatiom matrix. 
    Rotation error is defined as r_e = \arccos(\frac{Trace(\mathbf{R}_{ij}^{T}\mathbf{R}_{ij}^{\mathrm{GT}) - 1}{2})
    Args: 
        t1 (torch tensor): Estimated translation vectors [b,3,1]
        t2 (torch tensor): Ground truth translation vectors [b,3,1]
    Returns:
        te (torch tensor): translation error in meters [b,1]
    """
    return torch.norm(t1-t2, dim=(1, 2))
//  [python]
def translation_error(t1, t2):
    """
    Torch batch implementation of the rotation error between the estimated and the ground truth rotatiom matrix. 
    Rotation error is defined as r_e = \arccos(\frac{Trace(\mathbf{R}_{ij}^{T}\mathbf{R}_{ij}^{\mathrm{GT}) - 1}{2})
    Args: 
        t1 (torch tensor): Estimated translation vectors [b,3,1]
        t2 (torch tensor): Ground truth translation vectors [b,3,1]
    Returns:
        te (torch tensor): translation error in meters [b,1]
    """
    return torch.linalg.norm(t1-t2, dim=(1, 2))
//  [python]
def SFR_module(self, output, pc1, feature1, interval):
        
        B = pc1.size()[0]
        N = pc1.size()[2]
        
        ## warped pc1 with scene flow estimation
        pc1_warp = pc1+output
        
        ## estimate rigid transformation using initial scene flow
        ## assume all points static
        mask = torch.ones((pc1.size()[0],pc1.size()[2])).cuda()
        trans = self.rigid_transform_torch(pc1, pc1_warp, mask)
        # from transformation to rigid scene flow
        sf_rg = self.rigid_to_flow(pc1,trans)
        
        # mask for static points approximation by radial projection threshold
        vel_1 = feature1[:,0]
        sf_proj=torch.sum(sf_rg*pc1,dim=1)/(torch.norm(pc1,dim=1))
        residual=(vel_1*interval.unsqueeze(1)-sf_proj)
        mask_s = (abs(residual/vel_1) < self.rigid_thres) 
        
        # when enough points are inliers of a rigid transformation
        # use the rigid transformation to replace individual flow vectors of them
        pre_trans = torch.zeros(trans.size()).cuda()
        sf_agg = torch.zeros(output.size()).cuda()
        for b in range(B):
            if (mask_s[b].sum()/N)>self.rigid_pcs:
                pre_trans[b] = self.rigid_transform_torch(pc1[b].unsqueeze(0), \
                                    pc1_warp[b].unsqueeze(0), mask_s[b].unsqueeze(0))
                # from transformation to rigid scene flow
                sf_agg[b] = self.rigid_to_flow(pc1[b].unsqueeze(0),pre_trans[b].unsqueeze(0))
                sf_agg[b,:,torch.logical_not(mask_s[b])]=output[b,:,torch.logical_not(mask_s[b])]
            else:
                pre_trans[b] = trans[b]
                sf_agg[b] = output[b]
                
        return sf_agg, pre_trans, mask_s
//  [python]
def SFR_module(self, output, pc1, feature1, interval):
        
        B = pc1.size()[0]
        N = pc1.size()[2]
        
        ## warped pc1 with scene flow estimation
        pc1_warp = pc1+output
        
        ## estimate rigid transformation using initial scene flow
        ## assume all points static
        mask = torch.ones((pc1.size()[0],pc1.size()[2])).cuda()
        trans = self.rigid_transform_torch(pc1, pc1_warp, mask)
        # from transformation to rigid scene flow
        sf_rg = self.rigid_to_flow(pc1,trans)
        
        # mask for static points approximation by radial projection threshold
        vel_1 = feature1[:,0]
        sf_proj=torch.sum(sf_rg*pc1,dim=1)/(torch.linalg.norm(pc1,dim=1))
        residual=(vel_1*interval.unsqueeze(1)-sf_proj)
        mask_s = (abs(residual/vel_1) < self.rigid_thres) 
        
        # when enough points are inliers of a rigid transformation
        # use the rigid transformation to replace individual flow vectors of them
        pre_trans = torch.zeros(trans.size()).cuda()
        sf_agg = torch.zeros(output.size()).cuda()
        for b in range(B):
            if (mask_s[b].sum()/N)>self.rigid_pcs:
                pre_trans[b] = self.rigid_transform_torch(pc1[b].unsqueeze(0), \
                                    pc1_warp[b].unsqueeze(0), mask_s[b].unsqueeze(0))
                # from transformation to rigid scene flow
                sf_agg[b] = self.rigid_to_flow(pc1[b].unsqueeze(0),pre_trans[b].unsqueeze(0))
                sf_agg[b,:,torch.logical_not(mask_s[b])]=output[b,:,torch.logical_not(mask_s[b])]
            else:
                pre_trans[b] = trans[b]
                sf_agg[b] = output[b]
                
        return sf_agg, pre_trans, mask_s
//  [python]
def before_update(self, strategy, **kwargs):
        # import pdb
        # pdb.set_trace()
        background_num, baseball_num, bus_num, camera_num, cosplay_num, dress_num, hockey_num, laptop_num, racing_num, soccer_num, sweater_num = torch.unique(strategy.mb_y, return_counts=True)[1]
        beta = (baseball_num + soccer_num + hockey_num + laptop_num + camera_num) / (background_num + cosplay_num + dress_num + racing_num + sweater_num + bus_num) 
        lamda = strategy.i_batch / len(strategy.dataloader) + 1
        alpha = 1.0 * 1 ** lamda
        # x_mean = torch.mean(strategy.mb_x, 0, True)
        if strategy.experience_id != 0:
            for n, w in strategy.model.named_parameters():
                if n == "module.weight":
                    # self.pro_weight(self.Pl, x_mean, w, alpha=alpha_array, stride=2, cnn=False)
                    r = torch.mean(strategy.mb_x, 0, True)
                    k = torch.mm(self.Pl, torch.t(r))
                    # self.Pl = torch.sub(self.Pl, torch.mm(k, torch.t(k)) / (alpha + torch.mm(k, r)))
                    self.Pl = torch.sub(self.Pl, torch.mm(k, torch.t(k)) / (alpha + torch.mm(k, r)))

                    pnorm2 = torch.norm(self.Pl.data,p='fro')
                    self.Pl.data = self.Pl.data / pnorm2

                    mq = torch.mm(self.Pl, torch.linalg.solve(torch.mm(torch.t(self.Pl), self.Pl), torch.t(self.Pl)))
                    e = torch.autograd.Variable(torch.eye(len(mq)).type(torch.cuda.FloatTensor))
                    self.Ql.data = (e - mq).data                                    

                    qnorm2 = torch.norm(self.Ql.data, p='fro')
                    self.Ql.data = self.Ql.data / qnorm2
                    pwq = self.Pl + beta*self.Ql

                    w.grad.data = torch.mm(w.grad.data, torch.t(pwq.data))
//  [python]
def before_update(self, strategy, **kwargs):
        # import pdb
        # pdb.set_trace()
        background_num, baseball_num, bus_num, camera_num, cosplay_num, dress_num, hockey_num, laptop_num, racing_num, soccer_num, sweater_num = torch.unique(strategy.mb_y, return_counts=True)[1]
        beta = (baseball_num + soccer_num + hockey_num + laptop_num + camera_num) / (background_num + cosplay_num + dress_num + racing_num + sweater_num + bus_num) 
        lamda = strategy.i_batch / len(strategy.dataloader) + 1
        alpha = 1.0 * 1 ** lamda
        # x_mean = torch.mean(strategy.mb_x, 0, True)
        if strategy.experience_id != 0:
            for n, w in strategy.model.named_parameters():
                if n == "module.weight":
                    # self.pro_weight(self.Pl, x_mean, w, alpha=alpha_array, stride=2, cnn=False)
                    r = torch.mean(strategy.mb_x, 0, True)
                    k = torch.mm(self.Pl, torch.t(r))
                    # self.Pl = torch.sub(self.Pl, torch.mm(k, torch.t(k)) / (alpha + torch.mm(k, r)))
                    self.Pl = torch.sub(self.Pl, torch.mm(k, torch.t(k)) / (alpha + torch.mm(k, r)))

                    pnorm2 = torch.linalg.norm(self.Pl.data,ord='fro')
                    self.Pl.data = self.Pl.data / pnorm2

                    mq = torch.mm(self.Pl, torch.linalg.solve(torch.mm(torch.t(self.Pl), self.Pl), torch.t(self.Pl)))
                    e = torch.autograd.Variable(torch.eye(len(mq)).type(torch.cuda.FloatTensor))
                    self.Ql.data = (e - mq).data                                    

                    qnorm2 = torch.linalg.norm(self.Ql.data, ord='fro')
                    self.Ql.data = self.Ql.data / qnorm2
                    pwq = self.Pl + beta*self.Ql

                    w.grad.data = torch.mm(w.grad.data, torch.t(pwq.data))
//  [python]
def forward(self, x: Tensor, v: Tensor) -> Tuple[Tensor, Tensor]:
        r"""Applies a gated equivariant operation to node features and vector
        features.

        Args:
            x (torch.Tensor): The scalar features of the nodes.
            v (torch.Tensor): The vector features of the nodes.
        """
        vec1 = torch.norm(self.vec1_proj(v), dim=-2)
        vec2 = self.vec2_proj(v)

        x = torch.cat([x, vec1], dim=-1)
        x, v = torch.split(self.update_net(x), self.out_channels, dim=-1)
        v = v.unsqueeze(1) * vec2

        if self.act is not None:
            x = self.act(x)

        return x, v