jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_symeig_A_degenerate(dtype, device, eivaloffset, method):
    # test if the gradient can be stably propagated if the loss function
    # does not depend on which degenerate eigenvectors are used
    # (note: the variable is changed in a certain way so that the degeneracy
    # is kept)

    torch.manual_seed(seed)
    n = 5
    neig = 3
    kwargs = {
        "dtype": dtype,
        "device": device,
    }
    print(dtype)
    # random matrix to be orthogonalized for the eigenvectors
    mat = torch.randn((n, n), **kwargs).requires_grad_()

    # matrix for the loss function
    P2 = torch.randn((n, n), **kwargs).requires_grad_()

    # the degenerate eigenvalues
    a = (torch.tensor([1.0, 2.0, 3.0], **kwargs) + eivaloffset)
    if torch.is_complex(a):
        a = a.real
    a = a.requires_grad_()

    bck_options = {
        "method": "exactsolve",
    }

    def get_loss(a, mat, P2):
        # get the orthogonal vector for the eigenvectors
        P, _ = torch.qr(mat)

        # line up the eigenvalues
        b = torch.cat((a[:2], a[1:2], a[2:], a[2:])).to(dtype)

        # construct the matrix
        diag = torch.diag_embed(b)
        A = torch.matmul(torch.matmul(P.transpose(-2, -1).conj(), diag), P)
        Alinop = LinearOperator.m(A, is_hermitian=True)

        eivals, eivecs = symeig(
            Alinop, neig=neig,
            method=method,
            bck_options=bck_options)
        U = eivecs[:, 1:3]  # the degenerate eigenvectors

        loss = torch.einsum("rc,rc->", torch.matmul(P2, U), U.conj())
        return loss

    gradcheck(get_loss, (a, mat, P2))
    gradgradcheck(get_loss, (a, mat, P2))
//  [python]
def test_symeig_A_degenerate(dtype, device, eivaloffset, method):
    # test if the gradient can be stably propagated if the loss function
    # does not depend on which degenerate eigenvectors are used
    # (note: the variable is changed in a certain way so that the degeneracy
    # is kept)

    torch.manual_seed(seed)
    n = 5
    neig = 3
    kwargs = {
        "dtype": dtype,
        "device": device,
    }
    print(dtype)
    # random matrix to be orthogonalized for the eigenvectors
    mat = torch.randn((n, n), **kwargs).requires_grad_()

    # matrix for the loss function
    P2 = torch.randn((n, n), **kwargs).requires_grad_()

    # the degenerate eigenvalues
    a = (torch.tensor([1.0, 2.0, 3.0], **kwargs) + eivaloffset)
    if torch.is_complex(a):
        a = a.real
    a = a.requires_grad_()

    bck_options = {
        "method": "exactsolve",
    }

    def get_loss(a, mat, P2):
        # get the orthogonal vector for the eigenvectors
        P, _ = torch.linalg.qr(mat)

        # line up the eigenvalues
        b = torch.cat((a[:2], a[1:2], a[2:], a[2:])).to(dtype)

        # construct the matrix
        diag = torch.diag_embed(b)
        A = torch.matmul(torch.matmul(P.transpose(-2, -1).conj(), diag), P)
        Alinop = LinearOperator.m(A, is_hermitian=True)

        eivals, eivecs = symeig(
            Alinop, neig=neig,
            method=method,
            bck_options=bck_options)
        U = eivecs[:, 1:3]  # the degenerate eigenvectors

        loss = torch.einsum("rc,rc->", torch.matmul(P2, U), U.conj())
        return loss

    gradcheck(get_loss, (a, mat, P2))
    gradgradcheck(get_loss, (a, mat, P2))
//  [python]
def __init__(
        self,
        in_channels: int = 3,
        int_channels: int = 256,
        out_channels: int = 128,
        block_out_channels: Tuple[int] = (32, 64, 128, 256),
        latent_channels: int = 4,
        norm_num_groups: int = 32,
        scaling_factor: float = 0.18215,
        pretrained_path: Optional[str] = None,
        encoder: Optional[nn.Module] = None,
        num_mid_blocks: int = 0,
        num_latents: int = 2,
        num_upscalers: int = 1,
        upscale_channels: int = 256,
        parametrization: str = 'gaussian',
        fuse_rgb: bool = False,
        resize_input: bool = False,
        act_fn: str = 'none',
        clamp_output: bool = False,
        freeze_codebook: bool = False,
        skip_encoder: bool = False,
    ) -> None:

        super().__init__()

        self.enable_mid_block = num_mid_blocks > 0
        self.num_mid_blocks = num_mid_blocks
        self.downsample_factor = 2 ** (len(block_out_channels) - 1)
        self.interpolation_factor = self.downsample_factor // (2 ** num_upscalers)
        if "discrete" in parametrization:
            num_embeddings = 128
            if freeze_codebook:
                print('Freezing codebook')
                gen = torch.Generator().manual_seed(42)
                Q = torch.qr(torch.randn(num_embeddings, latent_channels, generator=gen))[0]
                self.codebook = nn.Embedding.from_pretrained(Q, freeze=True)
            else:
                self.codebook = nn.Embedding(num_embeddings, latent_channels, max_norm=None)
            num_latents = num_embeddings // latent_channels
        elif parametrization == 'auto':
            num_latents = 1

        if encoder is None:
            if fuse_rgb:
                in_channels += 3
            self.define_encoder(in_channels, block_out_channels, int_channels, norm_num_groups,
                                latent_channels, num_latents=num_latents, resize_input=resize_input,
                                skip_encoder=skip_encoder)
        else:
            self.encoder = encoder
            self.freeze_encoder()

        self.define_decoder(out_channels, int_channels, norm_num_groups, latent_channels, 
                            num_upscalers=num_upscalers, upscale_channels=upscale_channels)
        self.scaling_factor = scaling_factor
        self.gradient_checkpoint = False
        if pretrained_path is not None:
            self.load_pretrained(pretrained_path)
        self.parametrization = parametrization
        self.interpolation_factor = self.downsample_factor // (2 ** num_upscalers)
        self.num_latents = num_latents
        self.act_fn = act_fn
        self.clamp_output = clamp_output
        print('Interpolation factor: ', self.interpolation_factor)
        print('Parametrization: ', self.parametrization)
        print('Activation function: ', self.act_fn)
        assert self.parametrization in ['gaussian', 'discrete_gumbel_softmax', 'discrete_codebook', 'auto']
        assert self.num_latents in [1, 2, 32]
//  [python]
def __init__(
        self,
        in_channels: int = 3,
        int_channels: int = 256,
        out_channels: int = 128,
        block_out_channels: Tuple[int] = (32, 64, 128, 256),
        latent_channels: int = 4,
        norm_num_groups: int = 32,
        scaling_factor: float = 0.18215,
        pretrained_path: Optional[str] = None,
        encoder: Optional[nn.Module] = None,
        num_mid_blocks: int = 0,
        num_latents: int = 2,
        num_upscalers: int = 1,
        upscale_channels: int = 256,
        parametrization: str = 'gaussian',
        fuse_rgb: bool = False,
        resize_input: bool = False,
        act_fn: str = 'none',
        clamp_output: bool = False,
        freeze_codebook: bool = False,
        skip_encoder: bool = False,
    ) -> None:

        super().__init__()

        self.enable_mid_block = num_mid_blocks > 0
        self.num_mid_blocks = num_mid_blocks
        self.downsample_factor = 2 ** (len(block_out_channels) - 1)
        self.interpolation_factor = self.downsample_factor // (2 ** num_upscalers)
        if "discrete" in parametrization:
            num_embeddings = 128
            if freeze_codebook:
                print('Freezing codebook')
                gen = torch.Generator().manual_seed(42)
                Q = torch.linalg.qr(torch.randn(num_embeddings, latent_channels, generator=gen))[0]
                self.codebook = nn.Embedding.from_pretrained(Q, freeze=True)
            else:
                self.codebook = nn.Embedding(num_embeddings, latent_channels, max_norm=None)
            num_latents = num_embeddings // latent_channels
        elif parametrization == 'auto':
            num_latents = 1

        if encoder is None:
            if fuse_rgb:
                in_channels += 3
            self.define_encoder(in_channels, block_out_channels, int_channels, norm_num_groups,
                                latent_channels, num_latents=num_latents, resize_input=resize_input,
                                skip_encoder=skip_encoder)
        else:
            self.encoder = encoder
            self.freeze_encoder()

        self.define_decoder(out_channels, int_channels, norm_num_groups, latent_channels, 
                            num_upscalers=num_upscalers, upscale_channels=upscale_channels)
        self.scaling_factor = scaling_factor
        self.gradient_checkpoint = False
        if pretrained_path is not None:
            self.load_pretrained(pretrained_path)
        self.parametrization = parametrization
        self.interpolation_factor = self.downsample_factor // (2 ** num_upscalers)
        self.num_latents = num_latents
        self.act_fn = act_fn
        self.clamp_output = clamp_output
        print('Interpolation factor: ', self.interpolation_factor)
        print('Parametrization: ', self.parametrization)
        print('Activation function: ', self.act_fn)
        assert self.parametrization in ['gaussian', 'discrete_gumbel_softmax', 'discrete_codebook', 'auto']
        assert self.num_latents in [1, 2, 32]
//  [python]
def test_symeig_A_degenerate_req_not_sat(dtype, device):
    # test if the degenerate gradient returns nan if the requirments are not satisfied

    torch.manual_seed(seed)
    n = 5
    neig = 3
    kwargs = {
        "dtype": dtype,
        "device": device,
    }
    # random matrix to be orthogonalized for the eigenvectors
    mat = torch.randn((n, n), **kwargs).requires_grad_()

    # the degenerate eigenvalues
    a = torch.tensor([1.0, 2.0, 3.0], **kwargs).requires_grad_()
    bck_options = {
        "method": "exactsolve",
    }

    def get_loss(a, mat):
        # get the orthogonal vector for the eigenvectors
        P, _ = torch.qr(mat)

        # line up the eigenvalues
        b = torch.cat((a[:2], a[1:2], a[2:], a[2:]))

        # construct the matrix
        diag = torch.diag_embed(b)
        A = torch.matmul(torch.matmul(P.T, diag), P)
        Alinop = LinearOperator.m(A)

        eivals, eivecs = symeig(
            Alinop, neig=neig,
            method="custom_exacteig",
            bck_options=bck_options)
        U = eivecs[:, :3]  # the degenerate eigenvectors are in 1,2
        loss = torch.sum(U ** 4)
        return loss

    with warnings.catch_warnings(record=True) as w, enable_debug():
        loss = get_loss(a, mat)
        loss.backward()
        assert len(w) == 1
        wmsg = str(w[0].message).lower()
        assert "degener" in wmsg
        assert "loss function" in wmsg
        assert "incorrect" in wmsg
        assert w[0].category == MathWarning
//  [python]
def test_symeig_A_degenerate_req_not_sat(dtype, device):
    # test if the degenerate gradient returns nan if the requirments are not satisfied

    torch.manual_seed(seed)
    n = 5
    neig = 3
    kwargs = {
        "dtype": dtype,
        "device": device,
    }
    # random matrix to be orthogonalized for the eigenvectors
    mat = torch.randn((n, n), **kwargs).requires_grad_()

    # the degenerate eigenvalues
    a = torch.tensor([1.0, 2.0, 3.0], **kwargs).requires_grad_()
    bck_options = {
        "method": "exactsolve",
    }

    def get_loss(a, mat):
        # get the orthogonal vector for the eigenvectors
        P, _ = torch.linalg.qr(mat)

        # line up the eigenvalues
        b = torch.cat((a[:2], a[1:2], a[2:], a[2:]))

        # construct the matrix
        diag = torch.diag_embed(b)
        A = torch.matmul(torch.matmul(P.T, diag), P)
        Alinop = LinearOperator.m(A)

        eivals, eivecs = symeig(
            Alinop, neig=neig,
            method="custom_exacteig",
            bck_options=bck_options)
        U = eivecs[:, :3]  # the degenerate eigenvectors are in 1,2
        loss = torch.sum(U ** 4)
        return loss

    with warnings.catch_warnings(record=True) as w, enable_debug():
        loss = get_loss(a, mat)
        loss.backward()
        assert len(w) == 1
        wmsg = str(w[0].message).lower()
        assert "degener" in wmsg
        assert "loss function" in wmsg
        assert "incorrect" in wmsg
        assert w[0].category == MathWarning
//  [python]
def __init__(self, dim_inv, dim_spu, n_envs, orthonormal=True):
        super().__init__(dim_inv, dim_spu, n_envs)
        if orthonormal:
            self.scramble, _ = torch.qr(torch.randn(self.dim, self.dim))
        else:
            self.scramble = 10 * torch.randn(self.dim, self.dim)
//  [python]
def __init__(self, dim_inv, dim_spu, n_envs, orthonormal=True):
        super().__init__(dim_inv, dim_spu, n_envs)
        if orthonormal:
            self.scramble, _ = torch.linalg.qr(torch.randn(self.dim, self.dim))
        else:
            self.scramble = 10 * torch.randn(self.dim, self.dim)
//  [python]
def __init__(self, dim_inv, dim_spu, n_envs, orthonormal=True, **kwargs):
        super().__init__(dim_inv, dim_spu, n_envs, **kwargs)

        if orthonormal:
            self.scramble, _ = torch.qr(torch.randn(self.dim, self.dim))
        else:
            self.scramble = 10 * torch.randn(self.dim, self.dim)
//  [python]
def __init__(self, dim_inv, dim_spu, n_envs, orthonormal=True, **kwargs):
        super().__init__(dim_inv, dim_spu, n_envs, **kwargs)

        if orthonormal:
            self.scramble, _ = torch.linalg.qr(torch.randn(self.dim, self.dim))
        else:
            self.scramble = 10 * torch.randn(self.dim, self.dim)
//  [python]
def test_symeig_AM_degenerate(dtype, device, method):
    # same as test_symeig_A_degenerate, but now with the overlap matrix

    torch.manual_seed(seed)
    n = 5
    neig = 3
    kwargs = {
        "dtype": dtype,
        "device": device,
    }
    # random matrix to be orthogonalized for the eigenvectors
    matA = torch.randn((n, n), **kwargs)
    matM = torch.rand((n, n), **kwargs)

    # matrix for the loss function
    P2 = torch.randn((n, n), **kwargs).requires_grad_()

    # the degenerate eigenvalues
    a = torch.tensor([1.0, 2.0, 3.0], **kwargs)
    if torch.is_complex(a):
        a = a.real
    a = a.requires_grad_()

    bck_options = {
        "method": "exactsolve",
    }

    def get_loss(a, matA, matM, P2):
        # get the orthogonal vector for the eigenvectors
        P, _ = torch.qr(matA)
        PM, _ = torch.qr(matM)

        # line up the eigenvalues
        b = torch.cat((a[:2], a[1:2], a[2:], a[2:])).to(dtype)

        # construct the matrix
        diag = torch.diag_embed(b)
        A = torch.matmul(torch.matmul(P.transpose(-2, -1).conj(), diag), P)
        M = torch.matmul(PM.transpose(-2, -1).conj(), PM)
        Alinop = LinearOperator.m(A, is_hermitian=True)
        Mlinop = LinearOperator.m(M, is_hermitian=True)

        eivals, eivecs = symeig(
            Alinop, M=Mlinop, neig=neig,
            method=method,
            bck_options=bck_options)
        U = eivecs[:, 1:3]  # the degenerate eigenvectors

        loss = torch.einsum("rc,rc->", torch.matmul(P2, U), U.conj())
        return loss

    gradcheck(get_loss, (a, matA, matM, P2))
    gradgradcheck(get_loss, (a, matA, matM, P2))
//  [python]
def test_symeig_AM_degenerate(dtype, device, method):
    # same as test_symeig_A_degenerate, but now with the overlap matrix

    torch.manual_seed(seed)
    n = 5
    neig = 3
    kwargs = {
        "dtype": dtype,
        "device": device,
    }
    # random matrix to be orthogonalized for the eigenvectors
    matA = torch.randn((n, n), **kwargs)
    matM = torch.rand((n, n), **kwargs)

    # matrix for the loss function
    P2 = torch.randn((n, n), **kwargs).requires_grad_()

    # the degenerate eigenvalues
    a = torch.tensor([1.0, 2.0, 3.0], **kwargs)
    if torch.is_complex(a):
        a = a.real
    a = a.requires_grad_()

    bck_options = {
        "method": "exactsolve",
    }

    def get_loss(a, matA, matM, P2):
        # get the orthogonal vector for the eigenvectors
        P, _ = torch.linalg.qr(matA)
        PM, _ = torch.linalg.qr(matM)

        # line up the eigenvalues
        b = torch.cat((a[:2], a[1:2], a[2:], a[2:])).to(dtype)

        # construct the matrix
        diag = torch.diag_embed(b)
        A = torch.matmul(torch.matmul(P.transpose(-2, -1).conj(), diag), P)
        M = torch.matmul(PM.transpose(-2, -1).conj(), PM)
        Alinop = LinearOperator.m(A, is_hermitian=True)
        Mlinop = LinearOperator.m(M, is_hermitian=True)

        eivals, eivecs = symeig(
            Alinop, M=Mlinop, neig=neig,
            method=method,
            bck_options=bck_options)
        U = eivecs[:, 1:3]  # the degenerate eigenvectors

        loss = torch.einsum("rc,rc->", torch.matmul(P2, U), U.conj())
        return loss

    gradcheck(get_loss, (a, matA, matM, P2))
    gradgradcheck(get_loss, (a, matA, matM, P2))
//  [python]
def __init__(self, dim_inv, dim_spu, n_envs, orthonormal=True, **kwargs):
        super().__init__(dim_inv, dim_spu, n_envs, **kwargs)
        if orthonormal:
            self.scramble, _ = torch.qr(torch.randn(self.dim, self.dim))
        else:
            self.scramble = 10 * torch.randn(self.dim, self.dim)
//  [python]
def __init__(self, dim_inv, dim_spu, n_envs, orthonormal=True, **kwargs):
        super().__init__(dim_inv, dim_spu, n_envs, **kwargs)
        if orthonormal:
            self.scramble, _ = torch.linalg.qr(torch.randn(self.dim, self.dim))
        else:
            self.scramble = 10 * torch.randn(self.dim, self.dim)
//  [python]
def init_rot(batch_size: int = 10):
        """
        Randomly generate a batch of `batch_size` 3x3 rotation matrices.
        """
        device = torch.device("cuda:0")

        # TODO(dnovotny): replace with random_rotation from random_rotation.py
        rot = []
        for _ in range(batch_size):
            r = torch.qr(torch.randn((3, 3), device=device))[0]
            f = torch.randint(2, (3,), device=device, dtype=torch.float32)
            if f.sum() % 2 == 0:
                f = 1 - f
            rot.append(r * (2 * f - 1).float())
        rot = torch.stack(rot)

        return rot
//  [python]
def init_rot(batch_size: int = 10):
        """
        Randomly generate a batch of `batch_size` 3x3 rotation matrices.
        """
        device = torch.device("cuda:0")

        # TODO(dnovotny): replace with random_rotation from random_rotation.py
        rot = []
        for _ in range(batch_size):
            r = torch.linalg.qr(torch.randn((3, 3), device=device))[0]
            f = torch.randint(2, (3,), device=device, dtype=torch.float32)
            if f.sum() % 2 == 0:
                f = 1 - f
            rot.append(r * (2 * f - 1).float())
        rot = torch.stack(rot)

        return rot
//  [python]
def __init__(self, channel):
        """ Invertible MLP """
        super().__init__()

        weight = torch.randn(channel, channel)
        bias = torch.randn(1, channel, 1, 1)
        q, _ = torch.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
        self.bias = nn.Parameter(bias)
//  [python]
def __init__(self, channel):
        """ Invertible MLP """
        super().__init__()

        weight = torch.randn(channel, channel)
        bias = torch.randn(1, channel, 1, 1)
        q, _ = torch.linalg.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
        self.bias = nn.Parameter(bias)
//  [python]
def orthogonalize(matrix: torch.Tensor, eps=torch.tensor(1e-16)):
    if matrix.shape[-1] == 1:
        matrix.div_(torch.maximum(matrix.norm(), eps))
    else:
        q, r = torch.qr(matrix)
        matrix.copy_(q)
//  [python]
def orthogonalize(matrix: torch.Tensor, eps=torch.tensor(1e-16)):
    if matrix.shape[-1] == 1:
        matrix.div_(torch.maximum(matrix.norm(), eps))
    else:
        matrix.copy_(torch.linalg.qr(matrix).Q)
//  [python]
def rand_phase(self, n=1):
        qr = torch.randn((n, self.n, self.n),device=device, dtype=dtype)
        q, r = torch.qr(qr)
        r_diag_sign = torch.sign(torch.diagonal(r, dim1=-2, dim2=-1))
        q *= r_diag_sign[:, None]
        q_det_sign = torch.sign(torch.det(q))
        q[:, :, 0] *= q_det_sign[:, None]
        return q
//  [python]
def rand_phase(self, n=1):
        qr = torch.randn((n, self.n, self.n),device=device, dtype=dtype)
        q, r = torch.linalg.qr(qr)
        r_diag_sign = torch.sign(torch.diagonal(r, dim1=-2, dim2=-1))
        q *= r_diag_sign[:, None]
        q_det_sign = torch.sign(torch.det(q))
        q[:, :, 0] *= q_det_sign[:, None]
        return q
//  [python]
def qr(x: B.TorchNumeric, mode="reduced"):
    """
    Return a QR decomposition of a matrix x.
    """
    Q, R = torch.qr(x, mode=mode)
    return Q, R
//  [python]
def qr(x: B.TorchNumeric, mode="reduced"):
    """
    Return a QR decomposition of a matrix x.
    """
    Q, R = torch.linalg.qr(x, mode=mode)
    return Q, R
//  [python]
def iwasawa_decomposition(self, x):
            h, an = torch.qr(x, some=True)

            diag_sign = torch.diag_embed(torch.diagonal(torch.sign(an), dim1=-2, dim2=-1))
            h = torch.bmm(h, diag_sign)
            an = torch.bmm(diag_sign, an)

            a = torch.diagonal(an, dim1=-2, dim2=-1)

            a_inv = torch.div(torch.ones_like(a), a)
            a_inv = torch.diag_embed(a_inv)

            n = torch.bmm(a_inv, an)

            return h, a, n
//  [python]
def iwasawa_decomposition(self, x):
            h, an = torch.linalg.qr(x, mode='complete')

            diag_sign = torch.diag_embed(torch.diagonal(torch.sign(an), dim1=-2, dim2=-1))
            h = torch.bmm(h, diag_sign)
            an = torch.bmm(diag_sign, an)

            a = torch.diagonal(an, dim1=-2, dim2=-1)

            a_inv = torch.div(torch.ones_like(a), a)
            a_inv = torch.diag_embed(a_inv)

            n = torch.bmm(a_inv, an)

            return h, a, n
//  [python]
def _random_ortho(n: int, m: int) -> torch.Tensor:
    q, _ = torch.qr(torch.randn(n, m))
    return q
//  [python]
def _random_ortho(n: int, m: int) -> torch.Tensor:
    q, _ = torch.linalg.qr(torch.randn(n, m))
    return q
//  [python]
def test_so3_log_singularity(self, batch_size: int = 100):
        """
        Tests whether the `so3_log_map` is robust to the input matrices
        who's rotation angles are close to the numerically unstable region
        (i.e. matrices with low rotation angles).
        """
        # generate random rotations with a tiny angle
        device = torch.device("cuda:0")
        identity = torch.eye(3, device=device)
        rot180 = identity * torch.tensor([[1.0, -1.0, -1.0]], device=device)
        r = [identity, rot180]
        # add random rotations and random almost orthonormal matrices
        r.extend(
            [
                torch.qr(identity + torch.randn_like(identity) * 1e-4)[0]
                + float(i > batch_size // 2) * (0.5 - torch.rand_like(identity)) * 1e-3
                # this adds random noise to the second half
                # of the random orthogonal matrices to generate
                # near-orthogonal matrices
                for i in range(batch_size - 2)
            ]
        )
        r = torch.stack(r)
        r.requires_grad = True
        # the log of the rotation matrix r
        r_log = so3_log_map(r, cos_bound=1e-4, eps=1e-2)
        # tests whether all outputs are finite
        self.assertTrue(torch.isfinite(r_log).all())
        # tests whether the gradient is not None and all finite
        loss = r.sum()
        loss.backward()
        self.assertIsNotNone(r.grad)
        self.assertTrue(torch.isfinite(r.grad).all())
//  [python]
def test_so3_log_singularity(self, batch_size: int = 100):
        """
        Tests whether the `so3_log_map` is robust to the input matrices
        who's rotation angles are close to the numerically unstable region
        (i.e. matrices with low rotation angles).
        """
        # generate random rotations with a tiny angle
        device = torch.device("cuda:0")
        identity = torch.eye(3, device=device)
        rot180 = identity * torch.tensor([[1.0, -1.0, -1.0]], device=device)
        r = [identity, rot180]
        # add random rotations and random almost orthonormal matrices
        r.extend(
            [
                torch.linalg.qr(identity + torch.randn_like(identity) * 1e-4)[0]
                + float(i > batch_size // 2) * (0.5 - torch.rand_like(identity)) * 1e-3
                # this adds random noise to the second half
                # of the random orthogonal matrices to generate
                # near-orthogonal matrices
                for i in range(batch_size - 2)
            ]
        )
        r = torch.stack(r)
        r.requires_grad = True
        # the log of the rotation matrix r
        r_log = so3_log_map(r, cos_bound=1e-4, eps=1e-2)
        # tests whether all outputs are finite
        self.assertTrue(torch.isfinite(r_log).all())
        # tests whether the gradient is not None and all finite
        loss = r.sum()
        loss.backward()
        self.assertIsNotNone(r.grad)
        self.assertTrue(torch.isfinite(r.grad).all())
//  [python]
def _alternative_compress(self, name, tensor, P, Q, E):
        with torch.no_grad():
            M = tensor.view(tensor.shape[0], -1)
            if self._compute_p: 
                Q.copy_(torch.qr(Q)[0])
                P.copy_((M + E) @ Q)
                E.add_(M - P @ Q.T)
            else:
                P.copy_(torch.qr(P)[0])
                Q.copy_((M + E).T @ P)
                E.add_(M - P @ Q.T)
//  [python]
def _alternative_compress(self, name, tensor, P, Q, E):
        with torch.no_grad():
            M = tensor.view(tensor.shape[0], -1)
            if self._compute_p: 
                Q.copy_(torch.linalg.qr(Q).Q)
                P.copy_((M + E) @ Q)
                E.add_(M - P @ Q.T)
            else:
                P.copy_(torch.linalg.qr(P).Q)
                Q.copy_((M + E).T @ P)
                E.add_(M - P @ Q.T)
//  [python]
def test_so3_cos_bound(self, batch_size: int = 100):
        """
        Checks that for an identity rotation `R=I`, the so3_rotation_angle returns
        non-finite gradients when `cos_bound=None` and finite gradients
        for `cos_bound > 0.0`.
        """
        # generate random rotations with a tiny angle to generate cases
        # with the gradient singularity
        device = torch.device("cuda:0")
        identity = torch.eye(3, device=device)
        rot180 = identity * torch.tensor([[1.0, -1.0, -1.0]], device=device)
        r = [identity, rot180]
        r.extend(
            [
                torch.qr(identity + torch.randn_like(identity) * 1e-4)[0]
                for _ in range(batch_size - 2)
            ]
        )
        r = torch.stack(r)
        r.requires_grad = True
        for is_grad_finite in (True, False):
            # clear the gradients and decide the cos_bound:
            #     for is_grad_finite we run so3_rotation_angle with cos_bound
            #     set to a small float, otherwise we set to 0.0
            r.grad = None
            cos_bound = 1e-4 if is_grad_finite else 0.0
            # compute the angles of r
            angles = so3_rotation_angle(r, cos_bound=cos_bound)
            # tests whether all outputs are finite in both cases
            self.assertTrue(torch.isfinite(angles).all())
            # compute the gradients
            loss = angles.sum()
            loss.backward()
            # tests whether the gradient is not None for both cases
            self.assertIsNotNone(r.grad)
            if is_grad_finite:
                # all grad values have to be finite
                self.assertTrue(torch.isfinite(r.grad).all())
//  [python]
def test_so3_cos_bound(self, batch_size: int = 100):
        """
        Checks that for an identity rotation `R=I`, the so3_rotation_angle returns
        non-finite gradients when `cos_bound=None` and finite gradients
        for `cos_bound > 0.0`.
        """
        # generate random rotations with a tiny angle to generate cases
        # with the gradient singularity
        device = torch.device("cuda:0")
        identity = torch.eye(3, device=device)
        rot180 = identity * torch.tensor([[1.0, -1.0, -1.0]], device=device)
        r = [identity, rot180]
        r.extend(
            [
                torch.linalg.qr(identity + torch.randn_like(identity) * 1e-4)[0]
                for _ in range(batch_size - 2)
            ]
        )
        r = torch.stack(r)
        r.requires_grad = True
        for is_grad_finite in (True, False):
            # clear the gradients and decide the cos_bound:
            #     for is_grad_finite we run so3_rotation_angle with cos_bound
            #     set to a small float, otherwise we set to 0.0
            r.grad = None
            cos_bound = 1e-4 if is_grad_finite else 0.0
            # compute the angles of r
            angles = so3_rotation_angle(r, cos_bound=cos_bound)
            # tests whether all outputs are finite in both cases
            self.assertTrue(torch.isfinite(angles).all())
            # compute the gradients
            loss = angles.sum()
            loss.backward()
            # tests whether the gradient is not None for both cases
            self.assertIsNotNone(r.grad)
            if is_grad_finite:
                # all grad values have to be finite
                self.assertTrue(torch.isfinite(r.grad).all())
//  [python]
def rsample(self, nsamples: int) -> Tensor:
        """Produces n uniform samples over SU(self.dim)."""
        final = torch.zeros(nsamples, self.dim, self.dim).to(torch.complex64)
        for k in range(nsamples):
            z = torch.complex(
                torch.randn(self.dim, self.dim),
                torch.randn(self.dim, self.dim),
            )
            q, r = torch.qr(z / np.sqrt(2.0))
            diag = r.diag()
            ph = diag / diag.abs()
            q = q * ph * q
            q = q / (q.det() ** (1 / self.dim))
            final[k] = q
        return final
//  [python]
def rsample(self, nsamples: int) -> Tensor:
        """Produces n uniform samples over SU(self.dim)."""
        final = torch.zeros(nsamples, self.dim, self.dim).to(torch.complex64)
        for k in range(nsamples):
            z = torch.complex(
                torch.randn(self.dim, self.dim),
                torch.randn(self.dim, self.dim),
            )
            q, r = torch.linalg.qr(z / np.sqrt(2.0))
            diag = r.diag()
            ph = diag / diag.abs()
            q = q * ph * q
            q = q / (q.det() ** (1 / self.dim))
            final[k] = q
        return final
//  [python]
def orthogonalize(w):
    """
    Runs QR decomposition of weights and returns orthogonalized weights.
    :param w: 4D or 2D weights
    :return: orthogonalized weights of the same shape as input
    """
    flattened, is_transposed = weights2mat(w)
    rows, cols = flattened.shape
    assert rows >= cols, flattened.shape
    q, r = torch.qr(flattened)
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph
    if is_transposed:
        q.t_()
    return q.view_as(w)
//  [python]
def orthogonalize(w):
    """
    Runs QR decomposition of weights and returns orthogonalized weights.
    :param w: 4D or 2D weights
    :return: orthogonalized weights of the same shape as input
    """
    flattened, is_transposed = weights2mat(w)
    rows, cols = flattened.shape
    assert rows >= cols, flattened.shape
    q, r = torch.linalg.qr(flattened)
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph
    if is_transposed:
        q.t_()
    return q.view_as(w)
//  [python]
def project_jac_to_manifold(jac, x_in, x_out, manifold):
    bases = []
    # Compute a basis each for x, z, and x1
    for pos in [x_in, x_out]:
        bs, dim = pos.shape[0], prod(pos.shape[1:])
        # This is a (bs, dim, manifold_dim) tensor
        tangents = torch.stack([
            random_tangent_vec(manifold, pos, n_samples=bs).reshape(bs, dim)
            for _ in range(manifold.dim)
        ], -1)
        q, r = torch.qr(tangents)
        basis = q
        bases.append(basis)
    x_in_basis, x_out_basis = bases

    # Project the Jacobian after reshaping to bs x out_dim x in_dim
    x_in_dim = prod(x_in.shape[1:])
    x_out_dim = prod(x_out.shape[1:])
    jac_vec = jac.reshape(jac.shape[0], x_out_dim, x_in_dim)
    return torch.bmm(
        torch.bmm(x_out_basis.transpose(-1, -2), jac_vec),
        x_in_basis
    )
//  [python]
def project_jac_to_manifold(jac, x_in, x_out, manifold):
    bases = []
    # Compute a basis each for x, z, and x1
    for pos in [x_in, x_out]:
        bs, dim = pos.shape[0], prod(pos.shape[1:])
        # This is a (bs, dim, manifold_dim) tensor
        tangents = torch.stack([
            random_tangent_vec(manifold, pos, n_samples=bs).reshape(bs, dim)
            for _ in range(manifold.dim)
        ], -1)
        basis, _ = torch.linalg.qr(tangents)
        bases.append(basis)
    x_in_basis, x_out_basis = bases

    # Project the Jacobian after reshaping to bs x out_dim x in_dim
    x_in_dim = prod(x_in.shape[1:])
    x_out_dim = prod(x_out.shape[1:])
    jac_vec = jac.reshape(jac.shape[0], x_out_dim, x_in_dim)
    return torch.bmm(
        torch.bmm(x_out_basis.transpose(-1, -2), jac_vec),
        x_in_basis
    )
//  [python]
def subspace(A):
    """ Return the matrix of the subspace
        Arg:
            A: Bases of a linear subspace (batchsize, num_bases, emb_dim)
        Return:
            S: Orthonormalized bases of a linear subspace (batchsize, num_bases, emb_dim)
        Example:
            >>> A = torch.randn(5, 4, 300)
            >>> subspace(A)
    """ 
    # orthonormalize
    S, _ = torch.qr(torch.transpose(A, 1, 2))
    return torch.transpose(S, 1, 2)
//  [python]
def subspace(A):
    """ Return the matrix of the subspace
        Arg:
            A: Bases of a linear subspace (batchsize, num_bases, emb_dim)
        Return:
            S: Orthonormalized bases of a linear subspace (batchsize, num_bases, emb_dim)
        Example:
            >>> A = torch.randn(5, 4, 300)
            >>> subspace(A)
    """ 
    # orthonormalize
    S, _ = torch.linalg.qr(torch.transpose(A, 1, 2))
    return torch.transpose(S, 1, 2)
//  [python]
def linalgqr(input, mode):
        q, r = torch.qr(input, mode)
        out = [q, r]
        return out
//  [python]
def linalgqr(input, mode):
        q, r = torch.linalg.qr(input, mode)
        out = [q, r]
        return out
//  [python]
def _solve(self):

        H = torch.diag(self.potential) + self.K
        _, psi = torch.symeig(H, eigenvectors=True) 

        return psi[:, 0]
//  [python]
def _solve(self):

        H = torch.diag(self.potential) + self.K
        _, psi = torch.linalg.eigh(H)

        return psi[:, 0]
//  [python]
def orthogonal_(tensor: torch.Tensor, gain: float = 1) -> torch.Tensor:
    """
    **Overview**:
        Fills the input ``Tensor`` with a (semi) orthogonal matrix, as described in this paper <link https://arxiv.org/pdf/1312.6120.pdf link>.
        The input tensor must have at least 2 dimensions, and for tensors with more than 2 dimensions the trailing dimensions are flattened.
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    # Initialize a new tensor with normal distribution. The shape is the same as the input tensor.
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    # If rows < cols, transpose the original tensor for computational efficiency.
    if rows < cols:
        flattened.t_()

    # Compute the QR factorization, Q is an orthogonal matrix and R is an upper triangular matrix.
    # <link https://en.wikipedia.org/wiki/QR_decomposition link>
    q, r = torch.qr(flattened)
    # Although Q is orthogonal, each value of Q is not uniformly distributed. To make Q uniform, we can use the equation below: $$Q^* = Q sign(diag(R))$$. Proof for this equation can be viewed in this paper: <link https://arxiv.org/pdf/math-ph/0609050.pdf link>.
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    # If rows < cols, transpose the output tensor to match the shape of original tensor.
    if rows < cols:
        q.t_()

    # Using ``torch.no_grad()`` here can make sure that these operations won't be added to the computational graph used by PyTorch's autograd system, thus improving efficiency.
    with torch.no_grad():
        # Reshape the result and copy the weight from q.
        tensor.view_as(q).copy_(q)
        # Multiply an optional scaling factor.
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal_(tensor: torch.Tensor, gain: float = 1) -> torch.Tensor:
    """
    **Overview**:
        Fills the input ``Tensor`` with a (semi) orthogonal matrix, as described in this paper <link https://arxiv.org/pdf/1312.6120.pdf link>.
        The input tensor must have at least 2 dimensions, and for tensors with more than 2 dimensions the trailing dimensions are flattened.
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    # Initialize a new tensor with normal distribution. The shape is the same as the input tensor.
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    # If rows < cols, transpose the original tensor for computational efficiency.
    if rows < cols:
        flattened.t_()

    # Compute the QR factorization, Q is an orthogonal matrix and R is an upper triangular matrix.
    # <link https://en.wikipedia.org/wiki/QR_decomposition link>
    q, r = torch.linalg.qr(flattened)
    # Although Q is orthogonal, each value of Q is not uniformly distributed. To make Q uniform, we can use the equation below: $$Q^* = Q sign(diag(R))$$. Proof for this equation can be viewed in this paper: <link https://arxiv.org/pdf/math-ph/0609050.pdf link>.
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    # If rows < cols, transpose the output tensor to match the shape of original tensor.
    if rows < cols:
        q.t_()

    # Using ``torch.no_grad()`` here can make sure that these operations won't be added to the computational graph used by PyTorch's autograd system, thus improving efficiency.
    with torch.no_grad():
        # Reshape the result and copy the weight from q.
        tensor.view_as(q).copy_(q)
        # Multiply an optional scaling factor.
        tensor.mul_(gain)
    return tensor
//  [python]
def _compute_kfe(self, mod):

        state=self.state[mod]

        # ==== Query xxt and ggt ====
        keys = ['xxt', 'ggt', 'num_locations']
        run_avg_x, run_avg_g, run_avg_nl = self._get_store_keys(mod, keys)

        xxt, nfe = run_avg_x.avg, run_avg_x.N
        ggt, nbe = run_avg_g.avg, run_avg_g.N

        num_locations = run_avg_nl.sum

        # ==== compute kfe_x, kfe_gy, m2 ====
        Ex, state['kfe_x'] = torch.symeig(xxt, eigenvectors=True)
        Eg, state['kfe_gy'] = torch.symeig(ggt, eigenvectors=True)

        state['m2'] = Eg.unsqueeze(1) * Ex.unsqueeze(0) * num_locations
//  [python]
def _compute_kfe(self, mod):

        state=self.state[mod]

        # ==== Query xxt and ggt ====
        keys = ['xxt', 'ggt', 'num_locations']
        run_avg_x, run_avg_g, run_avg_nl = self._get_store_keys(mod, keys)

        xxt, nfe = run_avg_x.avg, run_avg_x.N
        ggt, nbe = run_avg_g.avg, run_avg_g.N

        num_locations = run_avg_nl.sum

        # ==== compute kfe_x, kfe_gy, m2 ====
        Ex, state['kfe_x'] = torch.linalg.eigh(xxt)
        Eg, state['kfe_gy'] = torch.linalg.eigh(ggt)

        state['m2'] = Eg.unsqueeze(1) * Ex.unsqueeze(0) * num_locations
//  [python]
def backward(self, losses, **kwargs):

        grads = self._get_grads(losses, mode='backward')
        if self.rep_grad:
            per_grads, grads = grads[0], grads[1]
        
        M = torch.matmul(grads, grads.t()) # [num_tasks, num_tasks]
        lmbda, V = torch.symeig(M, eigenvectors=True)
        tol = (
            torch.max(lmbda)
            * max(M.shape[-2:])
            * torch.finfo().eps
        )
        rank = sum(lmbda > tol)

        order = torch.argsort(lmbda, dim=-1, descending=True)
        lmbda, V = lmbda[order][:rank], V[:, order][:, :rank]

        sigma = torch.diag(1 / lmbda.sqrt())
        B = lmbda[-1].sqrt() * ((V @ sigma) @ V.t())
        alpha = B.sum(0)

        if self.rep_grad:
            self._backward_new_grads(alpha, per_grads=per_grads)
        else:
            self._backward_new_grads(alpha, grads=grads)
        return alpha.detach().cpu().numpy()
//  [python]
def backward(self, losses, **kwargs):

        grads = self._get_grads(losses, mode='backward')
        if self.rep_grad:
            per_grads, grads = grads[0], grads[1]
        
        M = torch.matmul(grads, grads.t()) # [num_tasks, num_tasks]
        lmbda, V = torch.linalg.eigh(M)
        tol = (
            torch.max(lmbda)
            * max(M.shape[-2:])
            * torch.finfo().eps
        )
        rank = sum(lmbda > tol)

        order = torch.argsort(lmbda, dim=-1, descending=True)
        lmbda, V = lmbda[order][:rank], V[:, order][:, :rank]

        sigma = torch.diag(1 / lmbda.sqrt())
        B = lmbda[-1].sqrt() * ((V @ sigma) @ V.t())
        alpha = B.sum(0)

        if self.rep_grad:
            self._backward_new_grads(alpha, per_grads=per_grads)
        else:
            self._backward_new_grads(alpha, grads=grads)
        return alpha.detach().cpu().numpy()
//  [python]
def test_se3_log_singularity(self, batch_size: int = 100):
        """
        Tests whether the `se3_log_map` is robust to the input matrices
        whose rotation angles and translations are close to the numerically
        unstable region (i.e. matrices with low rotation angles
        and 0 translation).
        """
        # generate random rotations with a tiny angle
        device = torch.device("cuda:0")
        identity = torch.eye(3, device=device)
        rot180 = identity * torch.tensor([[1.0, -1.0, -1.0]], device=device)
        r = [identity, rot180]
        r.extend(
            [
                torch.qr(identity + torch.randn_like(identity) * 1e-6)[0]
                + float(i > batch_size // 2) * (0.5 - torch.rand_like(identity)) * 1e-8
                # this adds random noise to the second half
                # of the random orthogonal matrices to generate
                # near-orthogonal matrices
                for i in range(batch_size - 2)
            ]
        )
        r = torch.stack(r)
        # tiny translations
        t = torch.randn(batch_size, 3, dtype=r.dtype, device=device) * 1e-6
        # create the transform matrix
        transform = torch.zeros(batch_size, 4, 4, dtype=torch.float32, device=device)
        transform[:, :3, :3] = r
        transform[:, 3, :3] = t
        transform[:, 3, 3] = 1.0
        transform.requires_grad = True
        # the log of the transform
        log_transform = se3_log_map(transform, eps=1e-4, cos_bound=1e-4)
        # tests whether all outputs are finite
        self.assertTrue(torch.isfinite(log_transform).all())
        # tests whether all gradients are finite and not None
        loss = log_transform.sum()
        loss.backward()
        self.assertIsNotNone(transform.grad)
        self.assertTrue(torch.isfinite(transform.grad).all())
//  [python]
def test_se3_log_singularity(self, batch_size: int = 100):
        """
        Tests whether the `se3_log_map` is robust to the input matrices
        whose rotation angles and translations are close to the numerically
        unstable region (i.e. matrices with low rotation angles
        and 0 translation).
        """
        # generate random rotations with a tiny angle
        device = torch.device("cuda:0")
        identity = torch.eye(3, device=device)
        rot180 = identity * torch.tensor([[1.0, -1.0, -1.0]], device=device)
        r = [identity, rot180]
        r.extend(
            [
                torch.linalg.qr(identity + torch.randn_like(identity) * 1e-6)[0]
                + float(i > batch_size // 2) * (0.5 - torch.rand_like(identity)) * 1e-8
                # this adds random noise to the second half
                # of the random orthogonal matrices to generate
                # near-orthogonal matrices
                for i in range(batch_size - 2)
            ]
        )
        r = torch.stack(r)
        # tiny translations
        t = torch.randn(batch_size, 3, dtype=r.dtype, device=device) * 1e-6
        # create the transform matrix
        transform = torch.zeros(batch_size, 4, 4, dtype=torch.float32, device=device)
        transform[:, :3, :3] = r
        transform[:, 3, :3] = t
        transform[:, 3, 3] = 1.0
        transform.requires_grad = True
        # the log of the transform
        log_transform = se3_log_map(transform, eps=1e-4, cos_bound=1e-4)
        # tests whether all outputs are finite
        self.assertTrue(torch.isfinite(log_transform).all())
        # tests whether all gradients are finite and not None
        loss = log_transform.sum()
        loss.backward()
        self.assertIsNotNone(transform.grad)
        self.assertTrue(torch.isfinite(transform.grad).all())
//  [python]
def cal_leading_eigenvector(self, M, method='power'):
        """
        Calculate the leading eigenvector using power iteration algorithm or torch.symeig
        Input: 
            - M:      [bs, num_corr, num_corr] the compatibility matrix 
            - method: select different method for calculating the learding eigenvector.
        Output: 
            - solution: [bs, num_corr] leading eigenvector
        """
        if method == 'power':
            leading_eig = torch.ones_like(M[:, :, 0:1])
            leading_eig_last = leading_eig
            for i in range(self.num_iterations):
                leading_eig = torch.bmm(M, leading_eig)
                leading_eig = leading_eig / (torch.norm(leading_eig, dim=1, keepdim=True) + 1e-06)
                if torch.allclose(leading_eig, leading_eig_last):
                    break
                leading_eig_last = leading_eig
            leading_eig = leading_eig.squeeze(-1)
            return leading_eig
        elif method == 'eig':
            e, v = torch.symeig(M, eigenvectors=True)
            leading_eig = v[:, :, -1]
            return leading_eig
        else:
            exit(-1)
//  [python]
def cal_leading_eigenvector(self, M, method='power'):
        """
        Calculate the leading eigenvector using power iteration algorithm or torch.symeig
        Input: 
            - M:      [bs, num_corr, num_corr] the compatibility matrix 
            - method: select different method for calculating the learding eigenvector.
        Output: 
            - solution: [bs, num_corr] leading eigenvector
        """
        if method == 'power':
            leading_eig = torch.ones_like(M[:, :, 0:1])
            leading_eig_last = leading_eig
            for i in range(self.num_iterations):
                leading_eig = torch.bmm(M, leading_eig)
                leading_eig = leading_eig / (torch.norm(leading_eig, dim=1, keepdim=True) + 1e-06)
                if torch.allclose(leading_eig, leading_eig_last):
                    break
                leading_eig_last = leading_eig
            leading_eig = leading_eig.squeeze(-1)
            return leading_eig
        elif method == 'eig':
            e, v = torch.linalg.eigh(M)
            leading_eig = v[:, :, -1]
            return leading_eig
        else:
            exit(-1)
//  [python]
def __call__(self, data):
        pos_centered = data.pos - data.pos.mean(axis=0)
        cov_matrix = pos_centered.T.mm(pos_centered) / len(pos_centered)
        eig, v = torch.symeig(cov_matrix, eigenvectors=True)
        data.eigenvalues = eig
        data.eigenvectors = v
        return data
//  [python]
def __call__(self, data):
        pos_centered = data.pos - data.pos.mean(axis=0)
        cov_matrix = pos_centered.T.mm(pos_centered) / len(pos_centered)
        eig, v = torch.linalg.eigh(cov_matrix)
        data.eigenvalues = eig
        data.eigenvectors = v
        return data