jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def solve_A_x_eq_b(self, b, x0=None):
        return torch.solve(self._matrix, b)[0]
//  [python]
def solve_A_x_eq_b(self, b, x0=None):
        return torch.linalg.solve(self._matrix, b)
//  [python]
def solve_system_nonlinear(self):
        """
        Solve the global system of equations to get the nodal displacements.
        """

        # Identify known DOFs from self.node_dof_disp
        known_dof_indices, unknown_dof_indices = self._get_dof_indices()

        # Extract submatrix and subvector
        K_sub = self.global_stiffness[unknown_dof_indices, :][:, unknown_dof_indices]
        R_sub = self.global_load[unknown_dof_indices]
        F_sub = self.global_residual_stress[unknown_dof_indices]
        B_sub = R_sub - F_sub


        # Solve the system
        self.u_incr_sub = torch.solve(B_sub.unsqueeze(1), K_sub)[0]
        
        
        self.global_displacements_temp[unknown_dof_indices] += self.u_incr_sub.squeeze()

        if known_dof_indices.numel() > 0:
            self.global_displacements_temp[known_dof_indices] = self.current_known_disp
//  [python]
def solve_system_nonlinear(self):
        """
        Solve the global system of equations to get the nodal displacements.
        """

        # Identify known DOFs from self.node_dof_disp
        known_dof_indices, unknown_dof_indices = self._get_dof_indices()

        # Extract submatrix and subvector
        K_sub = self.global_stiffness[unknown_dof_indices, :][:, unknown_dof_indices]
        R_sub = self.global_load[unknown_dof_indices]
        F_sub = self.global_residual_stress[unknown_dof_indices]
        B_sub = R_sub - F_sub


        # Solve the system
        self.u_incr_sub = torch.linalg.solve(K_sub, B_sub.unsqueeze(1))
        
        
        self.global_displacements_temp[unknown_dof_indices] += self.u_incr_sub.squeeze()

        if known_dof_indices.numel() > 0:
            self.global_displacements_temp[known_dof_indices] = self.current_known_disp
//  [python]
def __init__(self, sf, ell, X, y, eps=1e-5, kernel="exp"):
        self.sf = sf
        self.ell = ell
        self.X = X
        self.y = y
        self.eps = eps
        self.K = K if kernel == "exp" else Klinear
        self.KXX_inv_y = torch.solve(self.K(X, X, ell, sf, eps), y)[0]
//  [python]
def __init__(self, sf, ell, X, y, eps=1e-5, kernel="exp"):
        self.sf = sf
        self.ell = ell
        self.X = X
        self.y = y
        self.eps = eps
        self.K = K if kernel == "exp" else Klinear
        self.KXX_inv_y = torch.linalg.solve(self.K(X, X, ell, sf, eps), y)[0]
//  [python]
def update_precond_UVd_math_(U, V, d, v, h, step, tiny):
    # type: (Tensor, Tensor, Tensor, Tensor, Tensor, float, float) -> None
    """
    Update preconditioner Q = (I + U*V')*diag(d) with (vector, Hessian-vector product) = (v, h).
    State variables U, V and d are updated inplace.

    U, V, d, v, and h are either matrices or column vectors.
    """
    # balance the numerical dynamic ranges of U and V; optional
    if torch.rand([]) < 0.01:
        normU = torch.linalg.vector_norm(U)
        normV = torch.linalg.vector_norm(V)
        rho = torch.sqrt(normU/normV)
        U.div_(rho)
        V.mul_(rho)

    Qh = IpUVtmatvec(U, V, d*h)
    Ph = d*IpUVtmatvec(V, U, Qh)

    # invQtv = IpUVtsolve(V, U, v/d)
    # invPv = IpUVtsolve(U, V, invQtv)/d
    VtU = V.t().mm(U)
    I = torch.eye(VtU.size(dim=0), dtype=VtU.dtype, device=VtU.device)
    IpVtU = I + VtU
    invQtv = v/d
    # torch's linalg.solve is slow for small matrix
    invQtv = invQtv - V.mm(torch.solve(U.t().mm(invQtv), IpVtU.t())[0])
    invPv  = invQtv - U.mm(torch.solve(V.t().mm(invQtv), IpVtU)[0])
    invPv = invPv/d

    nablaD = Ph*h - v*invPv
    mu = step/(torch.max(torch.abs(nablaD)) + tiny)
    #d = d - mu*d*nablaD
    d.sub_(mu*d*nablaD)

    # update either U or V, not both at the same time
    a, b = Qh, invQtv
    if torch.rand([]) < 0.5:
        # nablaU = Qh.mm(Qh.t().mm(V)) - invQtv.mm(invQtv.t().mm(V))
        # mu = step/(norm_UVt(nablaU, V) + _tiny)
        # U = U - mu*(nablaU + nablaU.mm(V.t().mm(U)))
        atV = a.t().mm(V)
        atVVt = atV.mm(V.t())
        btV = b.t().mm(V)
        btVVt = btV.mm(V.t())
        norm = torch.sqrt(torch.abs( (a.t().mm(a))*(atVVt.mm(atVVt.t())) # abs to avoid sqrt(-0.0)
                                    +(b.t().mm(b))*(btVVt.mm(btVVt.t()))
                                  -2*(a.t().mm(b))*(atVVt.mm(btVVt.t())) ))
        mu = step/(norm + tiny)
        # U = U - mu*( a.mm(atV.mm(IpVtU))
        #             -b.mm(btV.mm(IpVtU)) )
        U.sub_(mu*( a.mm(atV.mm(IpVtU))
                   -b.mm(btV.mm(IpVtU)) ))
    else:
        # nablaV = Qh.mm(Qh.t().mm(U)) - invQtv.mm(invQtv.t().mm(U))
        # mu = step/(norm_UVt(U, nablaV) + _tiny)
        # V = V - mu*(nablaV + V.mm(U.t().mm(nablaV)))
        atU = a.t().mm(U)
        btU = b.t().mm(U)
        UUta = U.mm(atU.t())
        UUtb = U.mm(btU.t())
        norm = torch.sqrt(torch.abs( (UUta.t().mm(UUta))*(a.t().mm(a)) # abs to avoid sqrt(-0.0)
                                    +(UUtb.t().mm(UUtb))*(b.t().mm(b))
                                  -2*(UUta.t().mm(UUtb))*(a.t().mm(b)) ))
        mu = step/(norm + tiny)
        # V = V - mu*( (a + V.mm(atU.t())).mm(atU)
        #             -(b + V.mm(btU.t())).mm(btU) )
        V.sub_(mu*( (a + V.mm(atU.t())).mm(atU)
                   -(b + V.mm(btU.t())).mm(btU) ))
//  [python]
def update_precond_UVd_math_(U, V, d, v, h, step, tiny):
    # type: (Tensor, Tensor, Tensor, Tensor, Tensor, float, float) -> None
    """
    Update preconditioner Q = (I + U*V')*diag(d) with (vector, Hessian-vector product) = (v, h).
    State variables U, V and d are updated inplace.

    U, V, d, v, and h are either matrices or column vectors.
    """
    # balance the numerical dynamic ranges of U and V; optional
    if torch.rand([]) < 0.01:
        normU = torch.linalg.vector_norm(U)
        normV = torch.linalg.vector_norm(V)
        rho = torch.sqrt(normU/normV)
        U.div_(rho)
        V.mul_(rho)

    Qh = IpUVtmatvec(U, V, d*h)
    Ph = d*IpUVtmatvec(V, U, Qh)

    # invQtv = IpUVtsolve(V, U, v/d)
    # invPv = IpUVtsolve(U, V, invQtv)/d
    VtU = V.t().mm(U)
    I = torch.eye(VtU.size(dim=0), dtype=VtU.dtype, device=VtU.device)
    IpVtU = I + VtU
    invQtv = v/d
    # torch's linalg.solve is slow for small matrix
    invQtv = invQtv - V.mm(torch.linalg.solve(IpVtU.t(), U.t().mm(invQtv)))
    invPv  = invQtv - U.mm(torch.linalg.solve(IpVtU,     V.t().mm(invQtv)))
    invPv = invPv/d

    nablaD = Ph*h - v*invPv
    mu = step/(torch.max(torch.abs(nablaD)) + tiny)
    #d = d - mu*d*nablaD
    d.sub_(mu*d*nablaD)

    # update either U or V, not both at the same time
    a, b = Qh, invQtv
    if torch.rand([]) < 0.5:
        # nablaU = Qh.mm(Qh.t().mm(V)) - invQtv.mm(invQtv.t().mm(V))
        # mu = step/(norm_UVt(nablaU, V) + _tiny)
        # U = U - mu*(nablaU + nablaU.mm(V.t().mm(U)))
        atV = a.t().mm(V)
        atVVt = atV.mm(V.t())
        btV = b.t().mm(V)
        btVVt = btV.mm(V.t())
        norm = torch.sqrt(torch.abs( (a.t().mm(a))*(atVVt.mm(atVVt.t())) # abs to avoid sqrt(-0.0)
                                    +(b.t().mm(b))*(btVVt.mm(btVVt.t()))
                                  -2*(a.t().mm(b))*(atVVt.mm(btVVt.t())) ))
        mu = step/(norm + tiny)
        # U = U - mu*( a.mm(atV.mm(IpVtU))
        #             -b.mm(btV.mm(IpVtU)) )
        U.sub_(mu*( a.mm(atV.mm(IpVtU))
                   -b.mm(btV.mm(IpVtU)) ))
    else:
        # nablaV = Qh.mm(Qh.t().mm(U)) - invQtv.mm(invQtv.t().mm(U))
        # mu = step/(norm_UVt(U, nablaV) + _tiny)
        # V = V - mu*(nablaV + V.mm(U.t().mm(nablaV)))
        atU = a.t().mm(U)
        btU = b.t().mm(U)
        UUta = U.mm(atU.t())
        UUtb = U.mm(btU.t())
        norm = torch.sqrt(torch.abs( (UUta.t().mm(UUta))*(a.t().mm(a)) # abs to avoid sqrt(-0.0)
                                    +(UUtb.t().mm(UUtb))*(b.t().mm(b))
                                  -2*(UUta.t().mm(UUtb))*(a.t().mm(b)) ))
        mu = step/(norm + tiny)
        # V = V - mu*( (a + V.mm(atU.t())).mm(atU)
        #             -(b + V.mm(btU.t())).mm(btU) )
        V.sub_(mu*( (a + V.mm(atU.t())).mm(atU)
                   -(b + V.mm(btU.t())).mm(btU) ))
//  [python]
def setup_step_linear(self):
        w = torch.view_as_complex(self.w)
        V = torch.view_as_complex(self.V)
        BV = torch.view_as_complex(self.BV)
        CV = torch.view_as_complex(self.CV)
        w = torch.cat([w, w.conj()], dim=-1)
        V = torch.cat([V, V.conj()], dim=-1)
        BV = torch.cat([BV, BV.conj()], dim=-1)
        CV = torch.cat([CV, CV.conj()], dim=-1)
        pV = BV[..., -1, :]
        qV = CV[..., -1, :]
        BV = BV[..., 0, :]
        CV = CV[..., 0, :]
        dt = torch.exp(self.log_dt)
        d = (2./dt.unsqueeze(-1) - w).reciprocal() # (H, N)
        r = 1 + contract('... n, ... n, ... n -> ...', qV, d, pV.conj())
        r = r.reciprocal()
        A_f = torch.diag_embed(2./dt[:, None] + w) - contract('... n, ... m -> ... n m', pV.conj(), qV)
        A_b = torch.diag_embed(d) - contract('... p, ... p, ..., ... q, ... q -> ... p q', d, pV.conj(), r, qV, d)
        dA = A_b @ A_f
        dA_L = power(self.L, dA)
        I = torch.eye(dA.size(-1)).to(dA)
        dCV = torch.solve(CV.unsqueeze(-1), I - dA_L.transpose(-1, -2))[0].squeeze(-1)

        self.step_params = {
            'd': d,
            'r': r.unsqueeze(-1)*d*qV,
            # 'r': r,
            'dCV': dCV,
            'pV': pV,
            'qV': qV,
            'BV': BV,
            'd1': 2./dt.unsqueeze(-1) + w,
        }
        N = w.shape[-1]
        self.step_params = {k: v[...,:N//2] for k, v in self.step_params.items()}
//  [python]
def setup_step_linear(self):
        w = torch.view_as_complex(self.w)
        V = torch.view_as_complex(self.V)
        BV = torch.view_as_complex(self.BV)
        CV = torch.view_as_complex(self.CV)
        w = torch.cat([w, w.conj()], dim=-1)
        V = torch.cat([V, V.conj()], dim=-1)
        BV = torch.cat([BV, BV.conj()], dim=-1)
        CV = torch.cat([CV, CV.conj()], dim=-1)
        pV = BV[..., -1, :]
        qV = CV[..., -1, :]
        BV = BV[..., 0, :]
        CV = CV[..., 0, :]
        dt = torch.exp(self.log_dt)
        d = (2./dt.unsqueeze(-1) - w).reciprocal() # (H, N)
        r = 1 + contract('... n, ... n, ... n -> ...', qV, d, pV.conj())
        r = r.reciprocal()
        A_f = torch.diag_embed(2./dt[:, None] + w) - contract('... n, ... m -> ... n m', pV.conj(), qV)
        A_b = torch.diag_embed(d) - contract('... p, ... p, ..., ... q, ... q -> ... p q', d, pV.conj(), r, qV, d)
        dA = A_b @ A_f
        dA_L = power(self.L, dA)
        I = torch.eye(dA.size(-1)).to(dA)
        dCV = torch.linalg.solve(I - dA_L.transpose(-1, -2), CV.unsqueeze(-1)).squeeze(-1)

        self.step_params = {
            'd': d,
            'r': r.unsqueeze(-1)*d*qV,
            # 'r': r,
            'dCV': dCV,
            'pV': pV,
            'qV': qV,
            'BV': BV,
            'd1': 2./dt.unsqueeze(-1) + w,
        }
        N = w.shape[-1]
        self.step_params = {k: v[...,:N//2] for k, v in self.step_params.items()}
//  [python]
def log3x3(x: Tensor):
    eigs = su3_to_eigs(x)
    q, _ = torch.solve(
        torch.log(eigs).unsqueeze(-1), (
            1e-6 * torch.eye(3).unsqueeze(0)
            + eigs.unsqueeze(-1) ** (
                torch.tensor([0, 1, 2]).unsqueeze(0).unsqueeze(0)
            )
        )
    )
    q = q.unsqueeze(-1)
    eye = torch.eye(x.shape[-1]).reshape(1, x.shape[-1], x.shape[-1])
    eye = eye.repeat(x.shape[0], 1, 1)

    return q[:, 0] * eye + q[:, 1] * x + q[:, 2] * x @ x
//  [python]
def log3x3(x: Tensor):
    eigs = su3_to_eigs(x)
    q, _ = torch.linalg.solve(
        torch.log(eigs).unsqueeze(-1), (
            1e-6 * torch.eye(3).unsqueeze(0)
            + eigs.unsqueeze(-1) ** (
                torch.tensor([0, 1, 2]).unsqueeze(0).unsqueeze(0)
            )
        )
    )
    q = q.unsqueeze(-1)
    eye = torch.eye(x.shape[-1]).reshape(1, x.shape[-1], x.shape[-1])
    eye = eye.repeat(x.shape[0], 1, 1)

    return q[:, 0] * eye + q[:, 1] * x + q[:, 2] * x @ x
//  [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 test_root_inv_decomposition_no_cholesky(self):
        with settings.max_cholesky_size(0):
            linear_op = self.create_linear_op()
            test_mat = torch.randn(*linear_op.batch_shape, linear_op.size(-1), 5)
            # Check that cholesky is not called
            with mock.patch.object(linear_op, "cholesky") as chol_mock:
                root_approx = linear_op.root_inv_decomposition()
                res = root_approx.matmul(test_mat)
                actual = torch.solve(test_mat, linear_op.to_dense())[0]
                self.assertAllClose(res, actual, rtol=0.05, atol=0.02)
                chol_mock.assert_not_called()
//  [python]
def test_root_inv_decomposition_no_cholesky(self):
        with settings.max_cholesky_size(0):
            linear_op = self.create_linear_op()
            test_mat = torch.randn(*linear_op.batch_shape, linear_op.size(-1), 5)
            # Check that cholesky is not called
            with mock.patch.object(linear_op, "cholesky") as chol_mock:
                root_approx = linear_op.root_inv_decomposition()
                res = root_approx.matmul(test_mat)
                actual = torch.linalg.solve(linear_op.to_dense(), test_mat)
                self.assertAllClose(res, actual, rtol=0.05, atol=0.02)
                chol_mock.assert_not_called()
//  [python]
def compat_proju(self, u: Tensor, x: Tensor) -> Tensor:
        """Arbitrary matrix C projects to skew-hermitian B := (C - C^H) / 2

        Make traceless with tr(B - (tr(B) / N) * I) = tr(B) - tr(B) = 0
        """
        _, n, _ = x.shape
        algebra_elem = torch.solve(x, u)[0]  # X^{-1} u
        # do projection in lie algebra
        B = (algebra_elem - algebra_elem.conj().transpose(-2, -1)) / 2.
        trace = torch.einsum('bii->b', B)
        B = B - (
            (1 / n) * trace.unsqueeze(-1).unsqueeze(-1)
            * torch.eye(n).repeat(x.shape[0], 1, 1)
        )
        assert torch.abs(torch.mean(torch.einsum('bii->b', B))) < 1e-6

        return B
//  [python]
def compat_proju(self, u: Tensor, x: Tensor) -> Tensor:
        """Arbitrary matrix C projects to skew-hermitian B := (C - C^H) / 2

        Make traceless with tr(B - (tr(B) / N) * I) = tr(B) - tr(B) = 0
        """
        _, n, _ = x.shape
        algebra_elem = torch.linalg.solve(u, x)[0]  # X^{-1} u
        # do projection in lie algebra
        B = (algebra_elem - algebra_elem.conj().transpose(-2, -1)) / 2.
        trace = torch.einsum('bii->b', B)
        B = B - (
            (1 / n) * trace.unsqueeze(-1).unsqueeze(-1)
            * torch.eye(n).repeat(x.shape[0], 1, 1)
        )
        assert torch.abs(torch.mean(torch.einsum('bii->b', B))) < 1e-6

        return B
//  [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 forward(self, x):
        B, D, H, W = x.shape
        N = H * W
        identity = torch.eye(N).cuda()
        # reshape x, s.t. we can use the gmp formulation as a global pooling operation
        x = x.view(B, D, N)
        x = x.permute(0, 2, 1)
        # compute the linear kernel
        K = torch.bmm(x, x.permute(0, 2, 1))
        # solve the linear system (K + lambda * I) * alpha = ones
        A = K + self.lamb * identity
        o = torch.ones(B, N, 1).cuda()
        alphas, _ = torch.gesv(o, A)
        alphas = alphas.view(B, 1, -1)        
        xi = torch.bmm(alphas, x)
        xi = xi.view(B, -1)
        return xi
//  [python]
def forward(self, x):
        B, D, H, W = x.shape
        N = H * W
        identity = torch.eye(N).cuda()
        # reshape x, s.t. we can use the gmp formulation as a global pooling operation
        x = x.view(B, D, N)
        x = x.permute(0, 2, 1)
        # compute the linear kernel
        K = torch.bmm(x, x.permute(0, 2, 1))
        # solve the linear system (K + lambda * I) * alpha = ones
        A = K + self.lamb * identity
        o = torch.ones(B, N, 1).cuda()
        #alphas, _ = torch.gesv(o, A) # tested using pytorch 1.0.1
        alphas = torch.linalg.solve(A,o) # TODO check it again
        alphas = alphas.view(B, 1, -1)        
        xi = torch.bmm(alphas, x)
        xi = xi.view(B, -1)
        return xi
//  [python]
def inverse(self, inputs):
        """Cost:
            output = O(D^3 + D^2N)
            logabsdet = O(D^3)
        where:
            D = num of features
            N = num of inputs
        """
        outputs = inputs - self.bias
        outputs = torch.solve(outputs.t(), self._weight)[0]  # Linear-system solver.
        outputs = outputs.t()
        return outputs
//  [python]
def inverse(self, inputs):
        """Cost:
            output = O(D^3 + D^2N)
            logabsdet = O(D^3)
        where:
            D = num of features
            N = num of inputs
        """
        outputs = inputs - self.bias
        outputs = torch.linalg.solve(self._weight, outputs.t())  # Linear-system solver.
        outputs = outputs.t()
        return outputs
//  [python]
def get_tta_loss(self,out):
        cfg = self.cfg
        num_classes = self.num_classes
        outputs = out / cfg.OPTIM.TEMP
        softmax_prob = F.softmax(outputs, dim=1)
        eps = cfg.MODEL.EPS
        if cfg.OPTIM.ADAPT == "ent":
            tta_loss = -(outputs.softmax(1) * outputs.log_softmax(1)).sum(1)
        elif cfg.OPTIM.ADAPT == "conjugate":
            smax_inp = softmax_prob
            eye = torch.eye(num_classes).to(outputs.device)
            eye = eye.reshape((1, num_classes, num_classes))
            eye = eye.repeat(outputs.shape[0], 1, 1)
            t2 = eps * torch.diag_embed(smax_inp)
            smax_inp = torch.unsqueeze(smax_inp, 2)
            t3 = eps*torch.bmm(smax_inp, torch.transpose(smax_inp, 1, 2))
            matrix = eye + t2 - t3
            y_star = torch.solve(smax_inp, matrix)[0]
            y_star = torch.squeeze(y_star)
            pseudo_prob = y_star
            tta_loss = torch.logsumexp(outputs, dim=1) - (pseudo_prob * outputs - eps * pseudo_prob *(1-softmax_prob)).sum(dim=1)
        elif cfg.OPTIM.ADAPT == "softmax_pl":
            tta_loss = torch.logsumexp(outputs, dim=1) - (softmax_prob * outputs - eps * softmax_prob * (1-softmax_prob)).sum(dim=1)
        elif cfg.OPTIM.ADAPT == "hard_pl":
            yp = outputs.max(1)[1]
            y_star = 1 * F.one_hot(yp, num_classes=num_classes)
            thresh_idxs = torch.where(outputs.softmax(1).max(1)[0] > 0.75)
            tta_loss = torch.logsumexp(outputs[thresh_idxs], dim=1) - torch.sum(y_star[thresh_idxs]*outputs[thresh_idxs], dim=1) + torch.sum(eps*y_star[thresh_idxs]*(1 - F.softmax(outputs[thresh_idxs], dim=1)), dim=1)
        elif cfg.OPTIM.ADAPT == "rpl":
            p = F.softmax(outputs, dim=1)
            y_pl = outputs.max(1)[1]
            Yg = torch.gather(p, 1, torch.unsqueeze(y_pl, 1))
            tta_loss = (1- (Yg**0.8))/0.8
        else:
            pass 
        return tta_loss
//  [python]
def get_tta_loss(self,out):
        cfg = self.cfg
        num_classes = self.num_classes
        outputs = out / cfg.OPTIM.TEMP
        softmax_prob = F.softmax(outputs, dim=1)
        eps = cfg.MODEL.EPS
        if cfg.OPTIM.ADAPT == "ent":
            tta_loss = -(outputs.softmax(1) * outputs.log_softmax(1)).sum(1)
        elif cfg.OPTIM.ADAPT == "conjugate":
            smax_inp = softmax_prob
            eye = torch.eye(num_classes).to(outputs.device)
            eye = eye.reshape((1, num_classes, num_classes))
            eye = eye.repeat(outputs.shape[0], 1, 1)
            t2 = eps * torch.diag_embed(smax_inp)
            smax_inp = torch.unsqueeze(smax_inp, 2)
            t3 = eps*torch.bmm(smax_inp, torch.transpose(smax_inp, 1, 2))
            matrix = eye + t2 - t3
            y_star = torch.linalg.solve(matrix, smax_inp)
            y_star = torch.squeeze(y_star)
            pseudo_prob = y_star
            tta_loss = torch.logsumexp(outputs, dim=1) - (pseudo_prob * outputs - eps * pseudo_prob *(1-softmax_prob)).sum(dim=1)
        elif cfg.OPTIM.ADAPT == "softmax_pl":
            tta_loss = torch.logsumexp(outputs, dim=1) - (softmax_prob * outputs - eps * softmax_prob * (1-softmax_prob)).sum(dim=1)
        elif cfg.OPTIM.ADAPT == "hard_pl":
            yp = outputs.max(1)[1]
            y_star = 1 * F.one_hot(yp, num_classes=num_classes)
            thresh_idxs = torch.where(outputs.softmax(1).max(1)[0] > 0.75)
            tta_loss = torch.logsumexp(outputs[thresh_idxs], dim=1) - torch.sum(y_star[thresh_idxs]*outputs[thresh_idxs], dim=1) + torch.sum(eps*y_star[thresh_idxs]*(1 - F.softmax(outputs[thresh_idxs], dim=1)), dim=1)
        elif cfg.OPTIM.ADAPT == "rpl":
            p = F.softmax(outputs, dim=1)
            y_pl = outputs.max(1)[1]
            Yg = torch.gather(p, 1, torch.unsqueeze(y_pl, 1))
            tta_loss = (1- (Yg**0.8))/0.8
        else:
            pass 
        return tta_loss
//  [python]
def _transp_follow_one(
        self, x: torch.Tensor, v: torch.Tensor, *, u: torch.Tensor
    ) -> torch.Tensor:
        a = self._amat(x, u)
        rhs = v + 1 / 2 * a @ v
        lhs = -1 / 2 * a
        lhs[..., torch.arange(a.shape[-2]), torch.arange(x.shape[-2])] += 1
        qv = torch.solve(rhs, lhs)[1]
        return qv
//  [python]
def _transp_follow_one(
        self, x: torch.Tensor, v: torch.Tensor, *, u: torch.Tensor
    ) -> torch.Tensor:
        a = self._amat(x, u)
        rhs = v + 1 / 2 * a @ v
        lhs = -1 / 2 * a
        lhs[..., torch.arange(a.shape[-2]), torch.arange(x.shape[-2])] += 1
        qv = torch.linalg.solve(lhs, rhs)
        return qv
//  [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 _kalman_gain(cov: Tensor, Ht: Tensor, system_covariance: Tensor) -> Tensor:
        covs_measured = cov @ Ht
        A = system_covariance.permute(0, 2, 1)
        B = covs_measured.permute(0, 2, 1)
        Kt = torch.solve(B, A)[0]
        K = Kt.permute(0, 2, 1)
        return K
//  [python]
def _kalman_gain(cov: Tensor, Ht: Tensor, system_covariance: Tensor) -> Tensor:
        covs_measured = cov @ Ht
        A = system_covariance.permute(0, 2, 1)
        B = covs_measured.permute(0, 2, 1)
        Kt = torch.linalg.solve(A, B)
        K = Kt.permute(0, 2, 1)
        return K
//  [python]
def matrix_pade_approximant(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.mm(p_app)
    #There are 4 options to compute the MPA: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    #It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    #Please check which one is faster before running the code;
    return torch.solve(p_sqrt, q_sqrt)[0]
//  [python]
def matrix_pade_approximant(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.mm(p_app)
    #There are 4 options to compute the MPA: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    #It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    #Please check which one is faster before running the code;
    return torch.linalg.solve(q_sqrt, p_sqrt)
//  [python]
def matrix_pade_approximant_inverse(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.bmm(p_app)
    #There are 4 options to compute the MPA_inverse: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    #It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    #Please check which one is faster before running the code;
    return torch.solve(q_sqrt, p_sqrt)[0]
//  [python]
def matrix_pade_approximant_inverse(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.bmm(p_app)
    #There are 4 options to compute the MPA_inverse: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    #It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    #Please check which one is faster before running the code;
    return torch.linalg.solve(p_sqrt, q_sqrt)
//  [python]
def condition_on_pt(m, P, H, R, y):
    S = R + H @ P @ H.T
    K = torch.solve(H @ P, S + 1e-6)[0].T
    Sigma_cond = P - K @ S @ K.T
    mu_cond = m + K @ (y - H @ m)
    return mu_cond, Sigma_cond
//  [python]
def condition_on_pt(m, P, H, R, y):
    S = R + H @ P @ H.T
    K = torch.linalg.solve(S + 1e-6, H @ P).T
    Sigma_cond = P - K @ S @ K.T
    mu_cond = m + K @ (y - H @ m)
    return mu_cond, Sigma_cond
//  [python]
def forward(self, observations: torch.Tensor) -> torch.Tensor:
        # observations: (T, L, B, D)
        T, B, D = observations.shape

        assert D == self.observation_dim

        s_cov = F.softplus(self.state_cov)
        obs_cov = F.softplus(self.obs_cov)

        log_likelihood = 0
        state_mean = torch.zeros(B, self.state_dim, dtype=observations.dtype, device=observations.device)
        state_cov = torch.eye(self.state_dim, dtype=observations.dtype, device=observations.device)

        for t in range(T):
            observation = observations[t]
            state_upd_mean = state_mean @ self.F.T
            state_update_cov = self.F @ state_cov @ self.F.T + s_cov.diag_embed()

            innovation = observation - state_upd_mean @ self.H.T
            innovation_cov = self.H @ state_update_cov @ self.H.T + obs_cov.diag_embed()

            expanded_innovation_cov = innovation_cov.view(1, self.observation_dim, self.observation_dim)
            expanded_innovation_cov = expanded_innovation_cov.expand(B, self.observation_dim, self.observation_dim)
            solve_mean = torch.solve(innovation.unsqueeze(2), expanded_innovation_cov)[0].squeeze(2)
            solve_cov = torch.solve(self.H, innovation_cov)[0]

            PH = state_update_cov @ self.H.T

            state_mean = state_upd_mean + solve_mean @ PH.T
            state_cov = state_update_cov - PH @ solve_cov @ state_update_cov

            log_likelihood = log_likelihood - 0.5 * (torch.inner(innovation, solve_mean) + torch.logdet(innovation_cov)
                                                     + self.observation_dim * math.log(2 * math.pi))

        return log_likelihood
//  [python]
def forward(self, observations: torch.Tensor) -> torch.Tensor:
        # observations: (T, L, B, D)
        T, B, D = observations.shape

        assert D == self.observation_dim

        s_cov = F.softplus(self.state_cov)
        obs_cov = F.softplus(self.obs_cov)

        log_likelihood = 0
        state_mean = torch.zeros(B, self.state_dim, dtype=observations.dtype, device=observations.device)
        state_cov = torch.eye(self.state_dim, dtype=observations.dtype, device=observations.device)

        for t in range(T):
            observation = observations[t]
            state_upd_mean = state_mean @ self.F.T
            state_update_cov = self.F @ state_cov @ self.F.T + s_cov.diag_embed()

            innovation = observation - state_upd_mean @ self.H.T
            innovation_cov = self.H @ state_update_cov @ self.H.T + obs_cov.diag_embed()

            expanded_innovation_cov = innovation_cov.view(1, self.observation_dim, self.observation_dim)
            expanded_innovation_cov = expanded_innovation_cov.expand(B, self.observation_dim, self.observation_dim)
            solve_mean = torch.linalg.solve(expanded_innovation_cov, innovation)
            solve_cov = torch.linalg.solve(innovation_cov, self.H)

            PH = state_update_cov @ self.H.T

            state_mean = state_upd_mean + solve_mean @ PH.T
            state_cov = state_update_cov - PH @ solve_cov @ state_update_cov

            log_likelihood = log_likelihood - 0.5 * (torch.inner(innovation, solve_mean) + torch.logdet(innovation_cov)
                                                     + self.observation_dim * math.log(2 * math.pi))

        return log_likelihood
//  [python]
def anderson_normal_equation(X, relaxation=1.0, regularization = 0.0):
    # Anderson Acceleration
    # Take a matrix X of iterates such that X[:,i] = g(X[:,i-1])
    # Return acceleration for X[:,-1]

    assert X.ndim==2, "X must be a matrix"
    assert regularization >= 0.0, "regularization for least-squares must be >=0.0"

    # Compute residuals
    DX =  X[:,1:] -  X[:,:-1] # DX[:,i] =  X[:,i+1] -  X[:,i]
    DR = DX[:,1:] - DX[:,:-1] # DR[:,i] = DX[:,i+1] - DX[:,i] = X[:,i+2] - 2*X[:,i+1] + X[:,i]

    # # use QR factorization
    # q, r = torch.qr(DR)
    # gamma, _ = torch.triangular_solve( (q.t()@DX[:,-1]).unsqueeze(1), r )
    # gamma = gamma.squeeze(1)

    # solve unconstrained least-squares problem
    if regularization != 0.0: 
       RR = DR.t()@DR + regularization * torch.eye(DR.size(1))
    else: 
       RR = DR.t()@DR

    projected_residual = DR.t()@DX[:,-1].unsqueeze(1)
    
    gamma = torch.solve( projected_residual, RR)[0]
    gamma = gamma.view(-1)

    # compute acceleration
    extr = X[:,-2] + DX[:,-1] - (DX[:,:-1]+DR)@gamma

    if relaxation!=1:
        assert relaxation>0, "relaxation must be positive"
        # compute solution of the contraint optimization problem s.t. gamma = X[:,1:]@alpha
        alpha = torch.zeros(gamma.numel()+1).to(DX.device)
        alpha[0]    = gamma[0]
        alpha[1:-1] = gamma[1:] - gamma[:-1]
        alpha[-1]   = 1 - gamma[-1]
        extr = relaxation*extr + (1-relaxation)*X[:,:-1]@alpha

    return extr
//  [python]
def anderson_normal_equation(X, relaxation=1.0, regularization = 0.0):
    # Anderson Acceleration
    # Take a matrix X of iterates such that X[:,i] = g(X[:,i-1])
    # Return acceleration for X[:,-1]

    assert X.ndim==2, "X must be a matrix"
    assert regularization >= 0.0, "regularization for least-squares must be >=0.0"

    # Compute residuals
    DX =  X[:,1:] -  X[:,:-1] # DX[:,i] =  X[:,i+1] -  X[:,i]
    DR = DX[:,1:] - DX[:,:-1] # DR[:,i] = DX[:,i+1] - DX[:,i] = X[:,i+2] - 2*X[:,i+1] + X[:,i]

    # # use QR factorization
    # q, r = torch.qr(DR)
    # gamma, _ = torch.triangular_solve( (q.t()@DX[:,-1]).unsqueeze(1), r )
    # gamma = gamma.squeeze(1)

    # solve unconstrained least-squares problem
    if regularization != 0.0: 
       RR = DR.t()@DR + regularization * torch.eye(DR.size(1))
    else: 
       RR = DR.t()@DR

    projected_residual = DR.t()@DX[:,-1].unsqueeze(1)
    
    gamma = torch.linalg.solve( RR, projected_residual )
    gamma = gamma.view(-1)

    # compute acceleration
    extr = X[:,-2] + DX[:,-1] - (DX[:,:-1]+DR)@gamma

    if relaxation!=1:
        assert relaxation>0, "relaxation must be positive"
        # compute solution of the contraint optimization problem s.t. gamma = X[:,1:]@alpha
        alpha = torch.zeros(gamma.numel()+1).to(DX.device)
        alpha[0]    = gamma[0]
        alpha[1:-1] = gamma[1:] - gamma[:-1]
        alpha[-1]   = 1 - gamma[-1]
        extr = relaxation*extr + (1-relaxation)*X[:,:-1]@alpha

    return extr
//  [python]
def _solve_ABE(A: torch.Tensor, B: torch.Tensor, E: torch.Tensor):
    # A: (*BA, na, na) matrix
    # B: (*BB, na, ncols) matrix
    # E: (*BE, ncols) matrix
    na = A.shape[-1]
    BA, BB, BE = normalize_bcast_dims(A.shape[:-2], B.shape[:-2], E.shape[:-1])
    E = E.reshape(1, *BE, E.shape[-1]).transpose(0, -1)  # (ncols, *BE, 1)
    B = B.reshape(1, *BB, *B.shape[-2:]).transpose(0, -1)  # (ncols, *BB, na, 1)

    # NOTE: The line below is very inefficient for large na and ncols
    AE = A - torch.diag_embed(E.repeat_interleave(repeats=na, dim=-1), dim1=-2, dim2=-1)  # (ncols, *BAE, na, na)
    r = torch.solve(B, AE)[0]  # (ncols, *BAEM, na, 1)
    r = r.transpose(0, -1).squeeze(0)  # (*BAEM, na, ncols)
    return r
//  [python]
def _solve_ABE(A: torch.Tensor, B: torch.Tensor, E: torch.Tensor):
    # A: (*BA, na, na) matrix
    # B: (*BB, na, ncols) matrix
    # E: (*BE, ncols) matrix
    na = A.shape[-1]
    BA, BB, BE = normalize_bcast_dims(A.shape[:-2], B.shape[:-2], E.shape[:-1])
    E = E.reshape(1, *BE, E.shape[-1]).transpose(0, -1)  # (ncols, *BE, 1)
    B = B.reshape(1, *BB, *B.shape[-2:]).transpose(0, -1)  # (ncols, *BB, na, 1)

    # NOTE: The line below is very inefficient for large na and ncols
    AE = A - torch.diag_embed(E.repeat_interleave(repeats=na, dim=-1), dim1=-2, dim2=-1)  # (ncols, *BAE, na, na)
    r = torch.linalg.solve(AE, B)  # (ncols, *BAEM, na, 1)
    r = r.transpose(0, -1).squeeze(0)  # (*BAEM, na, ncols)
    return r
//  [python]
def exactsolve(A: LinearOperator, B: torch.Tensor,
               E: Union[torch.Tensor, None],
               M: Union[LinearOperator, None]):
    """
    Solve the linear equation by contructing the full matrix of LinearOperators.

    Warnings
    --------
    * As this method construct the linear operators explicitly, it might requires
      a large memory.
    """
    # A: (*BA, na, na)
    # B: (*BB, na, ncols)
    # E: (*BE, ncols)
    # M: (*BM, na, na)
    if E is None:
        Amatrix = A.fullmatrix()  # (*BA, na, na)
        x = torch.solve(B, Amatrix)[1]  # (*BAB, na, ncols)
    elif M is None:
        Amatrix = A.fullmatrix()
        x = _solve_ABE(Amatrix, B, E)
    else:
        Mmatrix = M.fullmatrix()  # (*BM, na, na)
        L = torch.linalg.cholesky(Mmatrix)  # (*BM, na, na)
        Linv = torch.inverse(L)  # (*BM, na, na)
        LinvT = Linv.transpose(-2, -1).conj()  # (*BM, na, na)
        A2 = torch.matmul(Linv, A.mm(LinvT))  # (*BAM, na, na)
        B2 = torch.matmul(Linv, B)  # (*BBM, na, ncols)

        X2 = _solve_ABE(A2, B2, E)  # (*BABEM, na, ncols)
        x = torch.matmul(LinvT, X2)  # (*BABEM, na, ncols)
    return x
//  [python]
def exactsolve(A: LinearOperator, B: torch.Tensor,
               E: Union[torch.Tensor, None],
               M: Union[LinearOperator, None]):
    """
    Solve the linear equation by contructing the full matrix of LinearOperators.

    Warnings
    --------
    * As this method construct the linear operators explicitly, it might requires
      a large memory.
    """
    # A: (*BA, na, na)
    # B: (*BB, na, ncols)
    # E: (*BE, ncols)
    # M: (*BM, na, na)
    if E is None:
        Amatrix = A.fullmatrix()  # (*BA, na, na)
        x = torch.linalg.solve(Amatrix, B)  # (*BAB, na, ncols)
    elif M is None:
        Amatrix = A.fullmatrix()
        x = _solve_ABE(Amatrix, B, E)
    else:
        Mmatrix = M.fullmatrix()  # (*BM, na, na)
        L = torch.linalg.cholesky(Mmatrix)  # (*BM, na, na)
        Linv = torch.inverse(L)  # (*BM, na, na)
        LinvT = Linv.transpose(-2, -1).conj()  # (*BM, na, na)
        A2 = torch.matmul(Linv, A.mm(LinvT))  # (*BAM, na, na)
        B2 = torch.matmul(Linv, B)  # (*BBM, na, ncols)

        X2 = _solve_ABE(A2, B2, E)  # (*BABEM, na, ncols)
        x = torch.matmul(LinvT, X2)  # (*BABEM, na, ncols)
    return x
//  [python]
def bilinear(dt, A, B=None):
        """
        dt: (H 1) timescales (or H N)
        A: (H N N)
        B: (H N)
        """
        N = A.shape[-1]
        I = torch.eye(N).to(A)
        A_backwards = I - dt[:, None] / 2 * A # Doesn't quite make sense if dt has shape (H N)
        A_forwards = I + dt[:, None] / 2 * A

        if B is None:
            dB = None
        else:
            dB = dt * 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: (H 1) timescales (or H N)
        A: (H N N)
        B: (H N)
        """
        N = A.shape[-1]
        I = torch.eye(N).to(A)
        A_backwards = I - dt[:, None] / 2 * A # Doesn't quite make sense if dt has shape (H N)
        A_forwards = I + dt[:, None] / 2 * A

        if B is None:
            dB = None
        else:
            dB = dt * 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 transp(self, x: torch.Tensor, y: torch.Tensor, v: torch.Tensor) -> torch.Tensor:

        xinvy = torch.solve(y.double(), x.double())[0]
        s, U = torch.linalg.eig(xinvy.transpose(-2,-1))
        s = s.real
        U = U.real

        Ut = U.transpose(-2,-1)
        Esqm = torch.solve(Ut, torch.diag_embed(s.sqrt()) @ Ut)[0].transpose(-2,-1).to(y.dtype)

        return Esqm @ v @ Esqm.transpose(-1,-2)
//  [python]
def transp(self, x: torch.Tensor, y: torch.Tensor, v: torch.Tensor) -> torch.Tensor:

        xinvy = torch.linalg.solve(x.double(),y.double())
        s, U = torch.linalg.eig(xinvy.transpose(-2,-1))
        s = s.real
        U = U.real

        Ut = U.transpose(-2,-1)
        Esqm = torch.linalg.solve(Ut, torch.diag_embed(s.sqrt()) @ Ut).transpose(-2,-1).to(y.dtype)

        return Esqm @ v @ Esqm.transpose(-1,-2)
//  [python]
def _torch_solve_cast(input: Tensor, A: Tensor) -> tuple[Tensor, Tensor]:
	"""Helper function to make torch.solve work with other than fp32/64.

	The function torch.solve is only implemented for fp32/64 which makes
	impossible to be used by fp16 or others. What this function does, is cast
	input data type to fp32, apply torch.svd, and cast back to the input dtype.
	"""
	if not isinstance(input, Tensor):
		raise AssertionError(f"Input must be Tensor. Got: {type(input)}.")
	dtype = input.dtype
	if dtype not in (torch.float32, torch.float64):
		dtype = torch.float32

	out = torch.solve(A.to(dtype), input.to(dtype))

	return out.to(input.dtype), out
//  [python]
def _torch_solve_cast(input: Tensor, A: Tensor) -> tuple[Tensor, Tensor]:
	"""Helper function to make torch.solve work with other than fp32/64.

	The function torch.solve is only implemented for fp32/64 which makes
	impossible to be used by fp16 or others. What this function does, is cast
	input data type to fp32, apply torch.svd, and cast back to the input dtype.
	"""
	if not isinstance(input, Tensor):
		raise AssertionError(f"Input must be Tensor. Got: {type(input)}.")
	dtype = input.dtype
	if dtype not in (torch.float32, torch.float64):
		dtype = torch.float32

	out = solve(A.to(dtype), input.to(dtype))

	return out.to(input.dtype), out
//  [python]
def broydens_method(f, J, x0, tol=1e-5, max_iter=100):
    """
    Broyden's method for finding roots of a function.

    Usage:
        f = lambda x: x**2 - 1
        J = lambda x: 2*x
        broydens_method(f, J, 0.5)
        tensor(1.0000)

    """
    x = x0
    B = J(x)
    for _ in range(max_iter):
        s = torch.solve(f(x), B)[0]
        x_new = x + s
        if torch.norm(x_new - x) < tol:
            return x_new
        y = f(x_new) - f(x)
        B += (y - B @ s) @ s.T / (s.T @ s)
        x = x_new
    return x
//  [python]
def broydens_method(f, J, x0, tol=1e-5, max_iter=100):
    """
    Broyden's method for finding roots of a function.

    Usage:
        f = lambda x: x**2 - 1
        J = lambda x: 2*x
        broydens_method(f, J, 0.5)
        tensor(1.0000)

    """
    x = x0
    B = J(x)
    for _ in range(max_iter):
        s = torch.linalg.solve(-f(x), B)[0]
        x_new = x + s
        if torch.norm(x_new - x) < tol:
            return x_new
        y = f(x_new) - f(x)
        B += (y - B @ s) @ s.T / (s.T @ s)
        x = x_new
    return x
//  [python]
def world_to_cam(self, cam: str, frame: int) -> torch.Tensor:
        """
        Get the world space to camera space transformation matrix.
        [batch_size, 4, 4]
        """
        return torch.solve(self.T[cam], self.world_to_car(frame))[1]
//  [python]
def world_to_cam(self, cam: str, frame: int) -> torch.Tensor:
        """
        Get the world space to camera space transformation matrix.
        [batch_size, 4, 4]
        """
        return torch.linalg.solve(self.T[cam], self.world_to_car(frame))
//  [python]
def _solve_sytem(self, Atb: torch.Tensor, AtA: torch.Tensor) -> torch.Tensor:
        return torch.solve(Atb, AtA)[0].squeeze(2)
//  [python]
def _solve_sytem(self, Atb: torch.Tensor, AtA: torch.Tensor) -> torch.Tensor:
        return torch.linalg.solve(AtA, Atb).squeeze(2)
//  [python]
def forward(self, 
                shared_parameters,
                losses: list[torch.Tensor, ...],
               ) -> tuple[torch.Tensor, torch.Tensor]:
        # >>> Loss Balance
        L_t, g, u = [], [], []
        
        for i, l in enumerate(losses):
            if self.method in ['hybrid','loss']:
                # If 'hybrid' or 'loss' method selected, then learnable scale s_t is used
                L_t.append(l * self.e.pow(self.s_t[i]) - self.s_t[i])
            else:
                # Else, the losses are unscaled
                L_t.append(l)

            g.append(torch.flatten(torch.autograd.grad(L_t[-1], shared_parameters, retain_graph=True)))
            u.append(g[-1] / torch.linalg.norm(g[-1]))
            
        g_t = torch.stack([_g for _g in g])
        u_t = torch.stack([_u for _u in u])
       
        # >>> Gradient Balance
        D = g_t[0,...] - g_t[1:,...]
        U = u_t[0,...] - u_t[1:,...]

        arg0 = D.matmul(U.t())
        arg1 = torch.eye(arg0.shape[-1].data)
        alpha_2T = g_t[0,...].matmul(U.t()).matmul(torch.solve(arg1,arg0)[0])
        
        alpha = torch.cat([torch.ones(1, device=self.device) - alpha_2T.sum(dim=0, keepdim=True), 
                           alpha_2T], dim=0).squeeze()
        
        # All methods return two outputs: gradient loss and individual losses
        if self.method=='hybrid':
            return torch.sum(L_t * alpha), L_t
        elif self.method=='gradient':
            return torch.sum(L_t * alpha), None
        elif self.method=='loss':
            return None, L_t