jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_serialization_parametrization(self):
        r"""Test that it is possible to serialize a parametrized model via state_dict"""
        # A stateful parametrization
        class Orthogonal(nn.Module):
            def __init__(self, n):
                super().__init__()
                self.register_buffer("id", torch.eye(n))
                self.register_buffer("B", torch.empty(n, n))
                init.orthogonal_(self.B)

            def forward(self, X):
                A = X.triu(1)
                A = A - A.T
                return self.B @ torch.solve(self.id - A, self.id + A)[0]

        def get_model():
            model = torch.nn.Sequential(
                torch.nn.Linear(5, 5),
                torch.nn.ReLU(),
                torch.nn.Linear(5, 1),
            )

            parametrize.register_parametrization(model[0], "weight", Orthogonal(5))
            return model

        model = get_model()

        prev_weight = model[0].weight
        prev_B = model[0].parametrizations.weight[0].B

        new_model = get_model()
        with TemporaryFileName() as fname:
            torch.save(model.state_dict(), fname)
            new_model.load_state_dict(torch.load(fname))

        # Integrity tests
        self.assertTrue(parametrize.is_parametrized(new_model[0], "weight"))
        self.assertEqual(prev_weight, new_model[0].weight)
        self.assertEqual(prev_B, new_model[0].parametrizations.weight[0].B)

        # Trying to save the whole parametrized model raises
        with self.assertRaisesRegex(RuntimeError, "state_dict"):
            with TemporaryFileName() as fname:
                torch.save(model, fname)
//  [python]
def test_serialization_parametrization(self):
        r"""Test that it is possible to serialize a parametrized model via state_dict"""
        # A stateful parametrization
        class Orthogonal(nn.Module):
            def __init__(self, n):
                super().__init__()
                self.register_buffer("id", torch.eye(n))
                self.register_buffer("B", torch.empty(n, n))
                init.orthogonal_(self.B)

            def forward(self, X):
                A = X.triu(1)
                A = A - A.T
                return self.B @ torch.linalg.solve(self.id + A, self.id - A)

        def get_model():
            model = torch.nn.Sequential(
                torch.nn.Linear(5, 5),
                torch.nn.ReLU(),
                torch.nn.Linear(5, 1),
            )

            parametrize.register_parametrization(model[0], "weight", Orthogonal(5))
            return model

        model = get_model()

        prev_weight = model[0].weight
        prev_B = model[0].parametrizations.weight[0].B

        new_model = get_model()
        with TemporaryFileName() as fname:
            torch.save(model.state_dict(), fname)
            new_model.load_state_dict(torch.load(fname))

        # Integrity tests
        self.assertTrue(parametrize.is_parametrized(new_model[0], "weight"))
        self.assertEqual(prev_weight, new_model[0].weight)
        self.assertEqual(prev_B, new_model[0].parametrizations.weight[0].B)

        # Trying to save the whole parametrized model raises
        with self.assertRaisesRegex(RuntimeError, "state_dict"):
            with TemporaryFileName() as fname:
                torch.save(model, fname)
//  [python]
def get_weights(self, X: Tensor, Y: Tensor, reg_coeff: float) -> Tensor:
        batch_size, n_samples, n_dim = X.shape
        ones = torch.ones(batch_size, n_samples, 1, device=X.device)
        X = torch.concat([X, ones], dim=-1)

        if n_samples >= n_dim:
            # standard
            A = torch.bmm(X.mT, X)
            A.diagonal(dim1=-2, dim2=-1).add_(reg_coeff)
            B = torch.bmm(X.mT, Y)
            weights = torch.solve(B, A)[0]
        else:
            # Woodbury
            A = torch.bmm(X, X.mT)
            A.diagonal(dim1=-2, dim2=-1).add_(reg_coeff)
            weights = torch.bmm(X.mT, torch.solve(Y, A)[0])

        return weights[:, :-1], weights[:, -1:]
//  [python]
def get_weights(self, X: Tensor, Y: Tensor, reg_coeff: float) -> Tensor:
        batch_size, n_samples, n_dim = X.shape
        ones = torch.ones(batch_size, n_samples, 1, device=X.device)
        X = torch.concat([X, ones], dim=-1)

        if n_samples >= n_dim:
            # standard
            A = torch.bmm(X.mT, X)
            A.diagonal(dim1=-2, dim2=-1).add_(reg_coeff)
            B = torch.bmm(X.mT, Y)
            weights = torch.linalg.solve(A, B)
        else:
            # Woodbury
            A = torch.bmm(X, X.mT)
            A.diagonal(dim1=-2, dim2=-1).add_(reg_coeff)
            weights = torch.bmm(X.mT, torch.linalg.solve(A, Y))

        return weights[:, :-1], weights[:, -1:]
//  [python]
def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
    """Solve the linear equation ax = b."""
    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
    # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
    # mixed input with complex and real tensors.
    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
        if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
            return FC.solve(b, a, return_LU=False)
        else:
            return matmul(inverse(a), b)
    elif torch.is_complex(a) or torch.is_complex(b):
        if torch.is_complex(a) and torch.is_complex(b):
            return torch.solve(b, a)[0]
        else:
            return matmul(inverse(a), b)
    else:
        return torch.solve(b, a)[0]
//  [python]
def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
    """Solve the linear equation ax = b."""
    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
    # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
    # mixed input with complex and real tensors.
    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
        if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
            return FC.solve(b, a, return_LU=False)
        else:
            return matmul(inverse(a), b)
    elif torch.is_complex(a) or torch.is_complex(b):
        if torch.is_complex(a) and torch.is_complex(b):
            return torch.linalg.solve(a, b)
        else:
            return matmul(inverse(a), b)
    else:
        return torch.linalg.solve(a, b)
//  [python]
def solve_lineqn(A, b, eps=1e-5):
    try:
        result = torch.solve(b, A)[0]
    except RuntimeError:
        print(f"WARNING: The matrix was singular")
        result = torch.solve(b, A + eps * torch.eye(A.shape[-1]))[0]
    return result
//  [python]
def solve_lineqn(A, b, eps=1e-5):
    try:
        result = torch.linalg.solve(A, b)
    except RuntimeError:
        print(f"WARNING: The matrix was singular")
        result = torch.linalg.solve(A + eps * torch.eye(A.shape[-1]), b)
    return result
//  [python]
def calculate_W(S, A):
    '''
    Our new Biharmonic weight formulation
    '''
    # make complementary selector matrix T^T (n x n) w.r.t current S
    pred_idx = torch.argmax(S, dim=-1)
    index = torch.ones((1, A.shape[1]), dtype=torch.float64).to(S.device)
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0)
    # T^T T
    TT = torch.diag_embed(index).to(torch.float64)
    
    # Z = A^{-1} S^T
    Z = torch.solve(torch.transpose(S, -2, -1), A)[0]

    # X = Z(SZ)^{-1}
    X = torch.solve(Z.transpose(-2,-1), torch.bmm(S,Z).transpose(-2,-1))[0].transpose(-2,-1)
    
    # W = S^T + T^T T X
    W_reformulated_full = torch.transpose(S, -2, -1).to(torch.float64) + torch.matmul(TT, X).to(torch.float64)
    return W_reformulated_full
//  [python]
def calculate_W(S, A):
    '''
    Our new Biharmonic weight formulation
    '''
    # make complementary selector matrix T^T (n x n) w.r.t current S
    pred_idx = torch.argmax(S, dim=-1)
    index = torch.ones((1, A.shape[1]), dtype=torch.float64).to(S.device)
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0)
    # T^T T
    TT = torch.diag_embed(index).to(torch.float64)
    
    # Z = A^{-1} S^T
    Z = torch.linalg.solve(A, torch.transpose(S, -2, -1))

    # X = Z(SZ)^{-1}
    X = torch.linalg.solve(torch.bmm(S,Z).transpose(-2,-1), Z.transpose(-2,-1)).transpose(-2,-1)
    
    # W = S^T + T^T T X
    W_reformulated_full = torch.transpose(S, -2, -1).to(torch.float64) + torch.matmul(TT, X).to(torch.float64)
    return W_reformulated_full
//  [python]
def original_calculate_W(S, A):
    '''
    Original Biharmonic weight formulation
    '''
    pred_idx = torch.argmax(S, dim=-1)
    index = torch.ones((1, A.shape[1])).to(S.device)
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0).squeeze().nonzero().squeeze()
    
    T = torch.zeros((1, A.shape[1] - S.shape[1], A.shape[1])).to(S.device)
    T = T.scatter_(index=index.unsqueeze(-1).unsqueeze(0), dim=-1, value=1.0).to(torch.float64)
    
    # (TAT^T)^(-1)
    inv_term = torch.matmul(T, torch.matmul(A, torch.transpose(T, -2, -1)))
    
    # TAS^T
    after_inv = torch.matmul(T, torch.matmul(A, torch.transpose(S, -2, -1)))
    
    # X = (TAT^T)^(-1) TAS^T
    X = torch.solve(after_inv, inv_term)[1]

    # W = S^T - T^T X
    W = torch.transpose(S, -2, -1).to(torch.float64) - torch.matmul(torch.transpose(T, -2, -1), X).to(torch.float64)
    return W
//  [python]
def original_calculate_W(S, A):
    '''
    Original Biharmonic weight formulation
    '''
    pred_idx = torch.argmax(S, dim=-1)
    index = torch.ones((1, A.shape[1])).to(S.device)
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0).squeeze().nonzero().squeeze()
    
    T = torch.zeros((1, A.shape[1] - S.shape[1], A.shape[1])).to(S.device)
    T = T.scatter_(index=index.unsqueeze(-1).unsqueeze(0), dim=-1, value=1.0).to(torch.float64)
    
    # (TAT^T)^(-1)
    inv_term = torch.matmul(T, torch.matmul(A, torch.transpose(T, -2, -1)))
    
    # TAS^T
    after_inv = torch.matmul(T, torch.matmul(A, torch.transpose(S, -2, -1)))
    
    # X = (TAT^T)^(-1) TAS^T
    X = torch.linalg.solve(inv_term, after_inv)

    # W = S^T - T^T X
    W = torch.transpose(S, -2, -1).to(torch.float64) - torch.matmul(torch.transpose(T, -2, -1), X).to(torch.float64)
    return W
//  [python]
def test_initialization_parametrization(self):
        r"""Test that it is possible to initialize a parametrization when it
            implements a `right_inverse` method
        """
        class Skew(nn.Module):
            def forward(self, X):
                A = X.triu(1)
                return A - A.T

            def is_skew(self, A):
                return torch.allclose(A, -A.T, atol=1e-6)

            def right_inverse(self, X):
                if not self.is_skew(X):
                    raise ValueError("The matrix is not skew-symmetric.")
                return X.triu(1)

        # Implements a Cayley map where right_inverse is not quite the inverse of forward
        class Orthogonal(nn.Module):
            def __init__(self, n):
                super().__init__()
                self.register_buffer("B", torch.eye(n))

            def forward(self, X):
                Id = torch.eye(X.size(0))
                return self.B @ torch.solve(Id - X, Id + X)[0]

            def is_orthogonal(self, X):
                Id = torch.eye(X.size(0))
                return torch.allclose(X.T @ X, Id, atol=1e-4)

            def right_inverse(self, X):
                if not self.is_orthogonal(X):
                    raise ValueError("The input is not orthogonal.")
                # cayley(0) == Id, so B @ cayley(0) == B
                self.B = X
                return torch.zeros_like(X)

        N = 5
        model = nn.Linear(N, N)
        # Register the skew-symmetric constraint. The result is now skew-symmetric
        skew = Skew()
        # Make the weight skew-symmetric before registering the parametrization
        with torch.no_grad():
            model.weight.set_(skew(model.weight))
        parametrize.register_parametrization(model, "weight", skew)
        X = torch.rand(N, N)
        # X is not skew-symmetric, so it throws an error
        with self.assertRaises(ValueError):
            model.weight = X
        # Make X skew-symmetric
        X = X - X.T
        model.weight = X
        self.assertEqual(model.parametrizations.weight.original, X.triu(1))
        self.assertEqual(model.weight, X)

        # Having several parametrizations registered should work in the same way
        parametrize.register_parametrization(model, "weight", Orthogonal(N))
        # Register now the Cayley map. The result is now orthogonal
        X = torch.rand(N, N)
        # X is not orthogonal, so it throws an error
        with self.assertRaises(ValueError):
            model.weight = X
        init.orthogonal_(X)
        model.weight = X
        self.assertEqual(model.weight, X)
        self.assertEqual(model.parametrizations.weight.original, torch.zeros_like(X))
//  [python]
def test_initialization_parametrization(self):
        r"""Test that it is possible to initialize a parametrization when it
            implements a `right_inverse` method
        """
        class Skew(nn.Module):
            def forward(self, X):
                A = X.triu(1)
                return A - A.T

            def is_skew(self, A):
                return torch.allclose(A, -A.T, atol=1e-6)

            def right_inverse(self, X):
                if not self.is_skew(X):
                    raise ValueError("The matrix is not skew-symmetric.")
                return X.triu(1)

        # Implements a Cayley map where right_inverse is not quite the inverse of forward
        class Orthogonal(nn.Module):
            def __init__(self, n):
                super().__init__()
                self.register_buffer("B", torch.eye(n))

            def forward(self, X):
                Id = torch.eye(X.size(0))
                return self.B @ torch.linalg.solve(Id + X, Id - X)

            def is_orthogonal(self, X):
                Id = torch.eye(X.size(0))
                return torch.allclose(X.T @ X, Id, atol=1e-4)

            def right_inverse(self, X):
                if not self.is_orthogonal(X):
                    raise ValueError("The input is not orthogonal.")
                # cayley(0) == Id, so B @ cayley(0) == B
                self.B = X
                return torch.zeros_like(X)

        N = 5
        model = nn.Linear(N, N)
        # Register the skew-symmetric constraint. The result is now skew-symmetric
        skew = Skew()
        # Make the weight skew-symmetric before registering the parametrization
        with torch.no_grad():
            model.weight.set_(skew(model.weight))
        parametrize.register_parametrization(model, "weight", skew)
        X = torch.rand(N, N)
        # X is not skew-symmetric, so it throws an error
        with self.assertRaises(ValueError):
            model.weight = X
        # Make X skew-symmetric
        X = X - X.T
        model.weight = X
        self.assertEqual(model.parametrizations.weight.original, X.triu(1))
        self.assertEqual(model.weight, X)

        # Having several parametrizations registered should work in the same way
        parametrize.register_parametrization(model, "weight", Orthogonal(N))
        # Register now the Cayley map. The result is now orthogonal
        X = torch.rand(N, N)
        # X is not orthogonal, so it throws an error
        with self.assertRaises(ValueError):
            model.weight = X
        init.orthogonal_(X)
        model.weight = X
        self.assertEqual(model.weight, X)
        self.assertEqual(model.parametrizations.weight.original, torch.zeros_like(X))
//  [python]
def test_barycentrics():
    barycentrics = torch.rand((5, 3)) * 0.25

    vertices = torch.randn((5, 4, 3))
    full_barycentrics = torch.cat([1 - barycentrics.sum(dim=-1, keepdim=True), barycentrics], dim=-1)
    points = (vertices * full_barycentrics.unsqueeze(-1)).sum(-2)


    t_mat = (vertices[..., 1:, :] - vertices[..., :1, :]).transpose(-1, -2)
    computed_barycentrics = torch.solve(points - vertices[..., 0, :], t_mat)[0]
    torch.testing.assert_allclose(computed_barycentrics, barycentrics)
//  [python]
def test_barycentrics():
    barycentrics = torch.rand((5, 3)) * 0.25

    vertices = torch.randn((5, 4, 3))
    full_barycentrics = torch.cat([1 - barycentrics.sum(dim=-1, keepdim=True), barycentrics], dim=-1)
    points = (vertices * full_barycentrics.unsqueeze(-1)).sum(-2)


    t_mat = (vertices[..., 1:, :] - vertices[..., :1, :]).transpose(-1, -2)
    computed_barycentrics = torch.linalg.solve(t_mat, points - vertices[..., 0, :])
    torch.testing.assert_allclose(computed_barycentrics, barycentrics)
//  [python]
def calculate_W_pseudo(S, A_inv):
    '''
    Our new Biharmonic weight formulation with torch pseudoinverse
    '''
    # make complementary selector matrix T^T (n x n) w.r.t current S
    pred_idx = torch.argmax(S, dim=-1)
    index = torch.ones((1, A_inv.shape[1]), dtype=torch.float64).to(S.device)
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0)

    # T^T T
    TT = torch.diag_embed(index).to(torch.float64)

    # T
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0).squeeze().nonzero().squeeze()
    T = torch.zeros((1, A_inv.shape[1] - S.shape[1], A_inv.shape[1])).to(S.device)
    T = T.scatter_(index=index.unsqueeze(-1).unsqueeze(0), dim=-1, value=1.0).to(torch.float64)
    
    # Z = A^{-1}S^T
    Z = torch.matmul(A_inv, torch.transpose(S, -2, -1))
    
    # X = Z(SZ)^{-1}
    X = torch.solve(Z.transpose(-2,-1), torch.bmm(S,Z).transpose(-2,-1))[0].transpose(-2,-1)
    
    # W = S^T + T^T T X
    W_pseudo_full = torch.transpose(S, -2, -1).to(torch.float64) + torch.matmul(TT, X).to(torch.float64)
    
    return W_pseudo_full
//  [python]
def calculate_W_pseudo(S, A_inv):
    '''
    Our new Biharmonic weight formulation with torch pseudoinverse
    '''
    # make complementary selector matrix T^T (n x n) w.r.t current S
    pred_idx = torch.argmax(S, dim=-1)
    index = torch.ones((1, A_inv.shape[1]), dtype=torch.float64).to(S.device)
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0)

    # T^T T
    TT = torch.diag_embed(index).to(torch.float64)

    # T
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0).squeeze().nonzero().squeeze()
    T = torch.zeros((1, A_inv.shape[1] - S.shape[1], A_inv.shape[1])).to(S.device)
    T = T.scatter_(index=index.unsqueeze(-1).unsqueeze(0), dim=-1, value=1.0).to(torch.float64)
    
    # Z = A^{-1}S^T
    Z = torch.matmul(A_inv, torch.transpose(S, -2, -1))
    
    # X = Z(SZ)^{-1}
    X = torch.linalg.solve(torch.bmm(S,Z).transpose(-2,-1), Z.transpose(-2,-1)).transpose(-2,-1)
    
    # W = S^T + T^T T X
    W_pseudo_full = torch.transpose(S, -2, -1).to(torch.float64) + torch.matmul(TT, X).to(torch.float64)
    
    return W_pseudo_full
//  [python]
def nullspace_project(self, x: Tensor) -> Tensor:
        """Project the given data onto the nullspace of the classifier."""

        # https://en.wikipedia.org/wiki/Projection_(linear_algebra)
        A = self.linear.weight.data.T
        P = A @ torch.solve(A.mT @ A, A.mT)[1]
        return x - x @ P
//  [python]
def nullspace_project(self, x: Tensor) -> Tensor:
        """Project the given data onto the nullspace of the classifier."""

        # https://en.wikipedia.org/wiki/Projection_(linear_algebra)
        A = self.linear.weight.data.T
        P = A @ torch.linalg.solve(A.mT @ A, A.mT)
        return x - x @ P
//  [python]
def test_barycentrics_grad():
    gt_barycentrics = torch.rand((5, 3)) * 0.25
    vertices = torch.randn((5, 4, 3))
    vertices.detach_().requires_grad_(True)
    full_barycentrics = torch.cat([1 - gt_barycentrics.sum(dim=-1, keepdim=True), gt_barycentrics], dim=-1)
    points = (vertices * full_barycentrics.unsqueeze(-1)).sum(-2)
    points.detach_().requires_grad_(True)

    t_mat = (vertices[..., 1:, :] - vertices[..., :1, :])
    barycentrics = torch.solve(points - vertices[..., 0, :], t_mat.transpose(-1, -2))[0]
    torch.testing.assert_allclose(barycentrics, gt_barycentrics)

    # Compute grads automatically
    barycentrics.retain_grad()
    comb = torch.randn((5, 3))
    (barycentrics * comb).sum().backward()

    # Compute grads manually
    t_mat = (vertices[..., 1:, :] - vertices[..., :1, :])
    grad_barycentrics = barycentrics.grad
    m_vec = torch.solve(grad_barycentrics, t_mat)[0]
    points_grad = m_vec
    torch.testing.assert_allclose(points_grad, points.grad)

    grad_vertices = -(full_barycentrics.unsqueeze(-1) * m_vec.unsqueeze(-2))
    torch.testing.assert_allclose(grad_vertices, vertices.grad)
//  [python]
def test_barycentrics_grad():
    gt_barycentrics = torch.rand((5, 3)) * 0.25
    vertices = torch.randn((5, 4, 3))
    vertices.detach_().requires_grad_(True)
    full_barycentrics = torch.cat([1 - gt_barycentrics.sum(dim=-1, keepdim=True), gt_barycentrics], dim=-1)
    points = (vertices * full_barycentrics.unsqueeze(-1)).sum(-2)
    points.detach_().requires_grad_(True)

    t_mat = (vertices[..., 1:, :] - vertices[..., :1, :])
    barycentrics = torch.linalg.solve(t_mat.transpose(-1, -2), points - vertices[..., 0, :])
    torch.testing.assert_allclose(barycentrics, gt_barycentrics)

    # Compute grads automatically
    barycentrics.retain_grad()
    comb = torch.randn((5, 3))
    (barycentrics * comb).sum().backward()

    # Compute grads manually
    t_mat = (vertices[..., 1:, :] - vertices[..., :1, :])
    grad_barycentrics = barycentrics.grad
    m_vec = torch.linalg.solve(t_mat, grad_barycentrics)
    points_grad = m_vec
    torch.testing.assert_allclose(points_grad, points.grad)

    grad_vertices = -(full_barycentrics.unsqueeze(-1) * m_vec.unsqueeze(-2))
    torch.testing.assert_allclose(grad_vertices, vertices.grad)
//  [python]
def forward(self, X: torch.Tensor) -> torch.Tensor:
        n, k = X.size(-2), X.size(-1)
        transposed = n < k
        if transposed:
            X = X.transpose(-2, -1)
            n, k = k, n
        # Here n > k and X is a tall matrix
        if self.orthogonal_map == _OrthMaps.matrix_exp or self.orthogonal_map == _OrthMaps.cayley:
            # We just need n x k - k(k-1)/2 parameters
            X = X.tril()
            if n != k:
                # Embed into a square matrix
                X = torch.cat([X, X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1)
            A = X - X.transpose(-2, -1).conj()
            # A is skew-symmetric (or skew-hermitian)
            if self.orthogonal_map == _OrthMaps.matrix_exp:
                Q = torch.matrix_exp(A)
            elif self.orthogonal_map == _OrthMaps.cayley:
                # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1}
                Id = torch.eye(n, dtype=A.dtype, device=A.device)
                Q = torch.solve(torch.add(Id, A, alpha=0.5), torch.add(Id, A, alpha=-0.5))[1]
            # Q is now orthogonal (or unitary) of size (..., n, n)
            if n != k:
                Q = Q[..., :k]
            # Q is now the size of the X (albeit perhaps transposed)
        else:
            # X is real here, as we do not support householder with complex numbers
            A = X.tril(diagonal=-1)
            tau = 2. / (1. + (A * A).sum(dim=-2))
            Q = torch.linalg.householder_product(A, tau)
            # The diagonal of X is 1's and -1's
            # We do not want to differentiate through this or update the diagonal of X hence the casting
            Q = Q * X.diagonal(dim1=-2, dim2=-1).int().unsqueeze(-2)

        if hasattr(self, "base"):
            Q = self.base @ Q
        if transposed:
            Q = Q.transpose(-2, -1)
        return Q
//  [python]
def forward(self, X: torch.Tensor) -> torch.Tensor:
        n, k = X.size(-2), X.size(-1)
        transposed = n < k
        if transposed:
            X = X.transpose(-2, -1)
            n, k = k, n
        # Here n > k and X is a tall matrix
        if self.orthogonal_map == _OrthMaps.matrix_exp or self.orthogonal_map == _OrthMaps.cayley:
            # We just need n x k - k(k-1)/2 parameters
            X = X.tril()
            if n != k:
                # Embed into a square matrix
                X = torch.cat([X, X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1)
            A = X - X.transpose(-2, -1).conj()
            # A is skew-symmetric (or skew-hermitian)
            if self.orthogonal_map == _OrthMaps.matrix_exp:
                Q = torch.matrix_exp(A)
            elif self.orthogonal_map == _OrthMaps.cayley:
                # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1}
                Id = torch.eye(n, dtype=A.dtype, device=A.device)
                Q = torch.linalg.solve(torch.add(Id, A, alpha=-0.5), torch.add(Id, A, alpha=0.5))
            # Q is now orthogonal (or unitary) of size (..., n, n)
            if n != k:
                Q = Q[..., :k]
            # Q is now the size of the X (albeit perhaps transposed)
        else:
            # X is real here, as we do not support householder with complex numbers
            A = X.tril(diagonal=-1)
            tau = 2. / (1. + (A * A).sum(dim=-2))
            Q = torch.linalg.householder_product(A, tau)
            # The diagonal of X is 1's and -1's
            # We do not want to differentiate through this or update the diagonal of X hence the casting
            Q = Q * X.diagonal(dim1=-2, dim2=-1).int().unsqueeze(-2)

        if hasattr(self, "base"):
            Q = self.base @ Q
        if transposed:
            Q = Q.transpose(-2, -1)
        return Q
//  [python]
def test_barycentrics_util():
    from tetranerf.utils.extension import add_barycentrics_grad
    
    gt_barycentrics = torch.rand((5, 3)) * 0.25
    vertices = torch.randn((5, 4, 3))
    vertices.detach_().requires_grad_(True)
    full_barycentrics = torch.cat([1 - gt_barycentrics.sum(dim=-1, keepdim=True), gt_barycentrics], dim=-1)
    points = (vertices * full_barycentrics.unsqueeze(-1)).sum(-2)
    points.detach_().requires_grad_(True)

    t_mat = (vertices[..., 1:, :] - vertices[..., :1, :])
    barycentrics = torch.solve(points - vertices[..., 0, :], t_mat.transpose(-1, -2))[0]
    torch.testing.assert_allclose(barycentrics, gt_barycentrics)

    # Compute grads automatically
    barycentrics.retain_grad()
    comb = torch.randn((5, 3))
    (barycentrics * comb).sum().backward()

    # Compute grads manually
    barycentrics2 = barycentrics.detach().requires_grad_(True)
    vertices2 = vertices.detach().requires_grad_(True)
    points2 = points.detach().requires_grad_(True)
    barycentrics2 = add_barycentrics_grad(barycentrics2, vertices2, points2)
    (barycentrics2 * comb).sum().backward()
    torch.testing.assert_allclose(points2.grad, points.grad)
    torch.testing.assert_allclose(vertices2.grad, vertices.grad)
//  [python]
def test_barycentrics_util():
    from tetranerf.utils.extension import add_barycentrics_grad
    
    gt_barycentrics = torch.rand((5, 3)) * 0.25
    vertices = torch.randn((5, 4, 3))
    vertices.detach_().requires_grad_(True)
    full_barycentrics = torch.cat([1 - gt_barycentrics.sum(dim=-1, keepdim=True), gt_barycentrics], dim=-1)
    points = (vertices * full_barycentrics.unsqueeze(-1)).sum(-2)
    points.detach_().requires_grad_(True)

    t_mat = (vertices[..., 1:, :] - vertices[..., :1, :])
    barycentrics = torch.linalg.solve(t_mat.transpose(-1, -2), points - vertices[..., 0, :])
    torch.testing.assert_allclose(barycentrics, gt_barycentrics)

    # Compute grads automatically
    barycentrics.retain_grad()
    comb = torch.randn((5, 3))
    (barycentrics * comb).sum().backward()

    # Compute grads manually
    barycentrics2 = barycentrics.detach().requires_grad_(True)
    vertices2 = vertices.detach().requires_grad_(True)
    points2 = points.detach().requires_grad_(True)
    barycentrics2 = add_barycentrics_grad(barycentrics2, vertices2, points2)
    (barycentrics2 * comb).sum().backward()
    torch.testing.assert_allclose(points2.grad, points.grad)
    torch.testing.assert_allclose(vertices2.grad, vertices.grad)
//  [python]
def fit(c, f, lambd=0.):
        device = c.device
        
        n = c.shape[0]
        f_dim = f.shape[1]

        U = TPS.u(TPS.d(c, c))
        K = U + torch.eye(n, device=device) * lambd

        P = torch.ones((n, 4), device=device)
        P[:, 1:] = c

        v = torch.zeros((n+4, f_dim), device=device)
        v[:n, :] = f

        A = torch.zeros((n+4, n+4), device=device)
        A[:n, :n] = K
        A[:n, -4:] = P
        A[-4:, :n] = P.t()

        theta = torch.solve(v, A)[0]
        return theta
//  [python]
def fit(c, f, lambd=0.):
        device = c.device
        
        n = c.shape[0]
        f_dim = f.shape[1]

        U = TPS.u(TPS.d(c, c))
        K = U + torch.eye(n, device=device) * lambd

        P = torch.ones((n, 4), device=device)
        P[:, 1:] = c

        v = torch.zeros((n+4, f_dim), device=device)
        v[:n, :] = f

        A = torch.zeros((n+4, n+4), device=device)
        A[:n, :n] = K
        A[:n, -4:] = P
        A[-4:, :n] = P.t()

        theta = torch.linalg.solve(A,v)
        #theta = torch.solve(v, A)[0]
        return theta
//  [python]
def get_geometry(
        self,
        fH,
        fW,
        lidar2img,
        img_metas,
    ):
        B, N, _, _ = lidar2img.shape
        device = lidar2img.device
        # import pdb;pdb.set_trace()
        if self.frustum == None:
            self.frustum = self.create_frustum(fH,fW,img_metas)
            self.frustum = self.frustum.to(device)
            # self.D = self.frustum.shape[0]
        
        points = self.frustum.view(1,1,self.D, fH, fW, 3) \
                 .repeat(B,N,1,1,1,1)
        lidar2img = lidar2img.view(B,N,1,1,1,4,4)
        # img2lidar = torch.inverse(lidar2img)
        points = torch.cat(
            (points, torch.ones_like(points[..., :1])), -1)
        points = torch.solve(points.unsqueeze(-1).to(torch.float32), 
                                    lidar2img.to(torch.float32))[0].squeeze(-1)
        # points = torch.matmul(img2lidar.to(torch.float32),
        #                       points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        # import pdb;pdb.set_trace()
        eps = 1e-5
        points = points[..., 0:3] / torch.maximum(
            points[..., 3:4], torch.ones_like(points[..., 3:4]) * eps)

        return points
//  [python]
def get_geometry(
        self,
        fH,
        fW,
        lidar2img,
        img_metas,
    ):
        B, N, _, _ = lidar2img.shape
        device = lidar2img.device
        # import pdb;pdb.set_trace()
        if self.frustum == None:
            self.frustum = self.create_frustum(fH,fW,img_metas)
            self.frustum = self.frustum.to(device)
            # self.D = self.frustum.shape[0]
        
        points = self.frustum.view(1,1,self.D, fH, fW, 3) \
                 .repeat(B,N,1,1,1,1)
        lidar2img = lidar2img.view(B,N,1,1,1,4,4)
        # img2lidar = torch.inverse(lidar2img)
        points = torch.cat(
            (points, torch.ones_like(points[..., :1])), -1)
        points = torch.linalg.solve(lidar2img.to(torch.float32), 
                                    points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        # points = torch.matmul(img2lidar.to(torch.float32),
        #                       points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        # import pdb;pdb.set_trace()
        eps = 1e-5
        points = points[..., 0:3] / torch.maximum(
            points[..., 3:4], torch.ones_like(points[..., 3:4]) * eps)

        return points
//  [python]
def func_slnr_max_vec(h_vec, eta_vec: float):
    batch_size, user_k, bases_n = h_vec.shape

    w_mmse_vec = torch.empty(
        (batch_size, bases_n, user_k), dtype=h_vec.dtype, device=h_vec.device
    )
    eye_div_eta = torch.eye(bases_n, dtype=h_vec.dtype, device=h_vec.device) / eta_vec

    effective_channel = h_vec.conj().transpose(2, 1)
    projected_channel_vec = eye_div_eta.unsqueeze(0) + torch.matmul(
        effective_channel, h_vec
    )
    # assert projected_channel_vec.shape(batch_size, user_k, user_k)
    for bs in range(batch_size):
        for k in range(user_k):
            projected_channel = torch.solve(
                effective_channel[bs, :, k].unsqueeze(1), projected_channel_vec[bs]
            )[1]
            w_mmse_vec[bs, :, k] = projected_channel

    w_mmse_vec = w_mmse_vec / (torch.norm(w_mmse_vec, dim=1, keepdim=True) * user_k)
    return w_mmse_vec
//  [python]
def func_slnr_max_vec(h_vec, eta_vec: float):
    batch_size, user_k, bases_n = h_vec.shape

    w_mmse_vec = torch.empty(
        (batch_size, bases_n, user_k), dtype=h_vec.dtype, device=h_vec.device
    )
    eye_div_eta = torch.eye(bases_n, dtype=h_vec.dtype, device=h_vec.device) / eta_vec

    effective_channel = h_vec.conj().transpose(2, 1)
    projected_channel_vec = eye_div_eta.unsqueeze(0) + torch.matmul(
        effective_channel, h_vec
    )
    # assert projected_channel_vec.shape(batch_size, user_k, user_k)
    for bs in range(batch_size):
        for k in range(user_k):
            projected_channel = torch.linalg.solve(
                projected_channel_vec[bs], effective_channel[bs, :, k]
            )
            w_mmse_vec[bs, :, k] = projected_channel

    w_mmse_vec = w_mmse_vec / (torch.norm(w_mmse_vec, dim=1, keepdim=True) * user_k)
    return w_mmse_vec
//  [python]
def _inverse_lattice(self) -> Tuple[Tensor, Tensor]:
        """Get inverse lattice vectors."""
        # Build a mask for zero vectors in 1D/ 2D lattice vectors
        mask_zero = self.latvec.eq(0).all(-1)
        _latvec = self.latvec + torch.diag_embed(mask_zero.type(
            self.latvec.dtype))

        eye = torch.eye(_latvec.shape[-1],
                        device=self._device, dtype=self._dtype)

        if self._n_batch is not None:
            eye = eye.repeat(self._n_batch, 1, 1)

        # Inverse lattice vectors
        invlat = torch.solve(_latvec, eye)[1].transpose(-1, -2)

        invlat[mask_zero] = 0

        return invlat, mask_zero
//  [python]
def _inverse_lattice(self) -> Tuple[Tensor, Tensor]:
        """Get inverse lattice vectors."""
        # Build a mask for zero vectors in 1D/ 2D lattice vectors
        mask_zero = self.latvec.eq(0).all(-1)
        _latvec = self.latvec + torch.diag_embed(mask_zero.type(
            self.latvec.dtype))

        eye = torch.eye(_latvec.shape[-1],
                        device=self._device, dtype=self._dtype)

        if self._n_batch is not None:
            eye = eye.repeat(self._n_batch, 1, 1)

        # Inverse lattice vectors
        invlat = torch.linalg.solve(_latvec, eye).transpose(-1, -2)

        invlat[mask_zero] = 0

        return invlat, mask_zero
//  [python]
def solve(self, v, regul=1e-8, impl="solve"):
        """
        solves v = Ax in x
        """
        # TODO: test
        if impl == "solve":
            # TODO: reuse LU decomposition once it is computed
            inv_v = torch.solve(
                v.get_flat_representation().view(-1, 1),
                self.data + regul * torch.eye(self.size(0), device=self.data.device),
            )[0]
            return PVector(v.layer_collection, vector_repr=inv_v[:, 0])
        elif impl == "eigendecomposition":
            v_eigenbasis = self.project_to_diag(v)
            inv_v_eigenbasis = v_eigenbasis / (self.evals + regul)
            return self.project_from_diag(inv_v_eigenbasis)
        else:
            raise NotImplementedError
//  [python]
def solve(self, v, regul=1e-8, impl="solve"):
        """
        solves v = Ax in x
        """
        # TODO: test
        if impl == "solve":
            # TODO: reuse LU decomposition once it is computed
            inv_v = torch.linalg.solve(
                self.data + regul * torch.eye(self.size(0), device=self.data.device),
                v.get_flat_representation().view(-1, 1),
            )
            return PVector(v.layer_collection, vector_repr=inv_v[:, 0])
        elif impl == "eigendecomposition":
            v_eigenbasis = self.project_to_diag(v)
            inv_v_eigenbasis = v_eigenbasis / (self.evals + regul)
            return self.project_from_diag(inv_v_eigenbasis)
        else:
            raise NotImplementedError
//  [python]
def bilinear(dt, A, B=None):
    """
    dt: (...) timescales
    A: (... N N)
    B: (... N)
    """
    N = A.shape[-1]
    I = torch.eye(N).to(A)
    A_backwards = I - dt[:, None, None] / 2 * A
    A_forwards = I + dt[:, None, None] / 2 * A

    if B is None:
        dB = None
    else:
        dB = dt[..., None] * torch.solve(
            B.unsqueeze(-1), A_backwards
        )[0].squeeze(-1) # (... N)

    dA = torch.solve(A_forwards, A_backwards)[0]  # (... N N)
    return dA, dB
//  [python]
def bilinear(dt, A, B=None):
    """
    dt: (...) timescales
    A: (... N N)
    B: (... N)
    """
    N = A.shape[-1]
    I = torch.eye(N).to(A)
    A_backwards = I - dt[:, None, None] / 2 * A
    A_forwards = I + dt[:, None, None] / 2 * A

    if B is None:
        dB = None
    else:
        dB = dt[..., None] * torch.linalg.solve(
            A_backwards, B.unsqueeze(-1)
        ).squeeze(-1) # (... N)

    dA = torch.linalg.solve(A_backwards, A_forwards)  # (... N N)
    return dA, dB
//  [python]
def solve(self, vs, regul=1e-8):
        vs_dict = vs.get_dict_representation()
        out_dict = dict()
        for layer_id, layer in self.generator.layer_collection.layers.items():
            v = vs_dict[layer_id][0].view(-1)
            if layer.bias is not None:
                v = torch.cat([v, vs_dict[layer_id][1].view(-1)])
            block = self.data[layer_id]

            inv_v = torch.solve(
                v.view(-1, 1),
                block + regul * torch.eye(block.size(0), device=block.device),
            )[1]
            inv_v_tuple = (inv_v[: layer.weight.numel()].view(*layer.weight.size),)
            if layer.bias is not None:
                inv_v_tuple = (
                    inv_v_tuple[0],
                    inv_v[layer.weight.numel() :].view(*layer.bias.size),
                )

            out_dict[layer_id] = inv_v_tuple
        return PVector(layer_collection=vs.layer_collection, dict_repr=out_dict)
//  [python]
def solve(self, vs, regul=1e-8):
        vs_dict = vs.get_dict_representation()
        out_dict = dict()
        for layer_id, layer in self.generator.layer_collection.layers.items():
            v = vs_dict[layer_id][0].view(-1)
            if layer.bias is not None:
                v = torch.cat([v, vs_dict[layer_id][1].view(-1)])
            block = self.data[layer_id]

            inv_v = torch.linalg.solve(
                block + regul * torch.eye(block.size(0), device=block.device),
                v.view(-1, 1),
            )
            inv_v_tuple = (inv_v[: layer.weight.numel()].view(*layer.weight.size),)
            if layer.bias is not None:
                inv_v_tuple = (
                    inv_v_tuple[0],
                    inv_v[layer.weight.numel() :].view(*layer.bias.size),
                )

            out_dict[layer_id] = inv_v_tuple
        return PVector(layer_collection=vs.layer_collection, dict_repr=out_dict)
//  [python]
def _setup_linear(self):
        """ Create parameters that allow fast linear stepping of state """
        w = self._w()
        B = _r2c(self.B) # (H N)
        P = _r2c(self.P)
        Q = P.conj() if self.Q is None else _r2c(self.Q)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (torch.eye(self.rank, dtype=w.dtype, device=w.device) + 2*contract('r h n, h n, s h n -> h r s', Q, D, P).real) # (H r r)
        Q_D = rearrange(Q*D, 'r h n -> h r n')
        R = torch.solve(Q_D, R.to(Q_D))[0] # (H r N)
        R = rearrange(R, 'h r n -> r h n')

        self.step_params = {
            "D": D, # (H N)
            "R": R, # (r H N)
            "P": P, # (r H N)
            "Q": Q, # (r H N)
            "B": B, # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w, # (H N)
        }
//  [python]
def _setup_linear(self):
        """ Create parameters that allow fast linear stepping of state """
        w = self._w()
        B = _r2c(self.B) # (H N)
        P = _r2c(self.P)
        Q = P.conj() if self.Q is None else _r2c(self.Q)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (torch.eye(self.rank, dtype=w.dtype, device=w.device) + 2*contract('r h n, h n, s h n -> h r s', Q, D, P).real) # (H r r)
        Q_D = rearrange(Q*D, 'r h n -> h r n')
        R = torch.linalg.solve(R.to(Q_D), Q_D) # (H r N)
        R = rearrange(R, 'h r n -> r h n')

        self.step_params = {
            "D": D, # (H N)
            "R": R, # (r H N)
            "P": P, # (r H N)
            "Q": Q, # (r H N)
            "B": B, # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w, # (H N)
        }
//  [python]
def score_parameters(self) -> List[Tensor]:
        """
        :return: List of Tensors the same shapes as the given Parameters where
            each Parameter's elements are scored based on the blockwise OBS
        """
        scores = [None] * len(self._params)
        block_finv_w = [None] * len(self._params)

        if self._is_main_proc:
            for i, finv in enumerate(self._finvs):
                if self._mask_type == "unstructured":
                    scores[i] = (
                        (self._params[i].data.view(-1) ** 2).to(self._devices[i])
                        / (2.0 * finv.diag() + self._eps)
                    ).view(self._params[i].shape)
                else:  # self._mask_type == "block8":
                    block_w = self._params[i].data.view(-1, 8).to(finv.dev)  # (d/Q, Q)
                    block_finv = (
                        torch.cat(
                            [
                                finv.f_inv[:, i : i + 8, i : i + 8]
                                for i in range(0, finv.B, 8)
                            ],
                            dim=1,
                        )
                        .reshape((finv.d // finv.B, finv.B // 8, 8, 8))
                        .reshape((finv.d // 8, 8, 8))
                    )  # (Q, d/Q, Q) -> (d/Q, Q, Q)
                    block_finv_w[i] = torch.solve(
                        block_w,
                        block_finv,
                    )[0]  # (d/Q, Q)
                    score = 0.5 * torch.einsum(
                        "bi,bi->b", block_w, block_finv_w[i]
                    )  # d/Q
                    scores[i] = (
                        score.unsqueeze(1)
                        .expand(-1, 8)
                        .reshape(self._params[i].data.shape)
                    )

            # make sure pruned ones will stay pruned
            for i, score in enumerate(scores):
                score[self._masks[i] == 0] = float("-inf")

        self._broadcast_list_from_main(scores)
        if self._mask_type == "block8":
            self._broadcast_list_from_main(block_finv_w)
            self._block_finv_w = block_finv_w  # cache for OBS weight update

        return scores
//  [python]
def score_parameters(self) -> List[Tensor]:
        """
        :return: List of Tensors the same shapes as the given Parameters where
            each Parameter's elements are scored based on the blockwise OBS
        """
        scores = [None] * len(self._params)
        block_finv_w = [None] * len(self._params)

        if self._is_main_proc:
            for i, finv in enumerate(self._finvs):
                if self._mask_type == "unstructured":
                    scores[i] = (
                        (self._params[i].data.view(-1) ** 2).to(self._devices[i])
                        / (2.0 * finv.diag() + self._eps)
                    ).view(self._params[i].shape)
                else:  # self._mask_type == "block8":
                    block_w = self._params[i].data.view(-1, 8).to(finv.dev)  # (d/Q, Q)
                    block_finv = (
                        torch.cat(
                            [
                                finv.f_inv[:, i : i + 8, i : i + 8]
                                for i in range(0, finv.B, 8)
                            ],
                            dim=1,
                        )
                        .reshape((finv.d // finv.B, finv.B // 8, 8, 8))
                        .reshape((finv.d // 8, 8, 8))
                    )  # (Q, d/Q, Q) -> (d/Q, Q, Q)
                    block_finv_w[i] = torch.linalg.solve(
                        block_finv,
                        block_w,
                    )  # (d/Q, Q)
                    score = 0.5 * torch.einsum(
                        "bi,bi->b", block_w, block_finv_w[i]
                    )  # d/Q
                    scores[i] = (
                        score.unsqueeze(1)
                        .expand(-1, 8)
                        .reshape(self._params[i].data.shape)
                    )

            # make sure pruned ones will stay pruned
            for i, score in enumerate(scores):
                score[self._masks[i] == 0] = float("-inf")

        self._broadcast_list_from_main(scores)
        if self._mask_type == "block8":
            self._broadcast_list_from_main(block_finv_w)
            self._block_finv_w = block_finv_w  # cache for OBS weight update

        return scores
//  [python]
def setup_step(self, mode='dense'):
        """ Set up dA, dB, dC discretized parameters for stepping """
        self._setup_state()

        # Calculate original C
        dA_L = power(self.L, self.dA)
        I = torch.eye(self.dA.size(-1)).to(dA_L)
        C = _conj(_r2c(self.C)) # (H C N)

        dC = torch.solve(
            C.unsqueeze(-1),
            I - dA_L.transpose(-1, -2),
        ).solution.squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == 'linear':
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2*self.dC[:, :, :self.N]
        elif mode == 'diagonal':
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print("Diagonalization error:", torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA))

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract('h n m, h m -> h n', V_inv, self.dB)
            self.dC = contract('h n m, c h n -> c h m', V, self.dC)

        elif mode == 'dense':
            pass
        else: raise NotImplementedError("NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}")
//  [python]
def setup_step(self, mode='dense'):
        """ Set up dA, dB, dC discretized parameters for stepping """
        self._setup_state()

        # Calculate original C
        dA_L = power(self.L, self.dA)
        I = torch.eye(self.dA.size(-1)).to(dA_L)
        C = _conj(_r2c(self.C)) # (H C N)

        dC = torch.linalg.solve(
            I - dA_L.transpose(-1, -2),
            C.unsqueeze(-1),
        ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == 'linear':
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2*self.dC[:, :, :self.N]
        elif mode == 'diagonal':
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print("Diagonalization error:", torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA))

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract('h n m, h m -> h n', V_inv, self.dB)
            self.dC = contract('h n m, c h n -> c h m', V, self.dC)

        elif mode == 'dense':
            pass
        else: raise NotImplementedError("NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}")
//  [python]
def precision_weighted_inner_prod(self, b, reduce_dim=True):
        assert b.shape[-1] == 1
        prod = (torch.solve(b, self.scale_tril)[0]**2).sum(-2)
        return prod.squeeze(-1) if reduce_dim else prod