jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def solve_reg(x_tr, y_tr, x_te, y_te, lam, y_std=1.0):
    n, d = x_tr.shape
    xtype = x_tr.type()

    if x_tr.shape[0] > x_tr.shape[1]:
        b = x_tr.T @ (y_tr[:, None] if y_tr.dim() == 1 else y_tr)
        y_pred = x_te @ torch.linalg.solve(x_tr.T @ x_tr + lam * torch.eye(d).type(xtype), b)
    else:
        b = y_tr[:, None] if y_tr.dim() == 1 else y_tr
        y_pred = x_te @ (x_tr.T @ torch.linalg.solve(x_tr @ x_tr.T + lam * torch.eye(n).type(xtype), b))

    mse = mean_squared_error((y_pred * y_std).cpu(), (y_te * y_std).cpu())
    acc = ((y_pred.argmax(axis=1) == y_te.argmax(axis=1)) * 1.0).mean().item()
    return mse, acc
//  [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()

        # Repeat w shape properly
        B = repeat(B, '1 t n -> 1 (v t) n', v=self.broadcast)
        P = repeat(P, 'r t n -> r (v t) n', v=self.broadcast)
        Q = repeat(Q, 'r t n -> r (v t) n', v=self.broadcast)
        w = repeat(w, 't n -> (v t) n', v=self.broadcast)

        # 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')
        try:
            R = torch.solve(Q_D, R.to(Q_D))[0] # (H r N)
        except RuntimeError:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).cpu(), Q_D.cpu())).to(Q_D)
        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()

        # Repeat w shape properly
        B = repeat(B, '1 t n -> 1 (v t) n', v=self.broadcast)
        P = repeat(P, 'r t n -> r (v t) n', v=self.broadcast)
        Q = repeat(Q, 'r t n -> r (v t) n', v=self.broadcast)
        w = repeat(w, 't n -> (v t) n', v=self.broadcast)

        # 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')
        try:
            R = torch.linalg.solve(R.to(Q_D), Q_D) # (H r N)
        except torch._C._LinAlgError:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).cpu(), Q_D.cpu())).to(Q_D)
        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, mat_feats):
        W = self.isom_params.data
        
        W = W - W.t() # skew-symmetric A^T = -A
        Id = torch.eye(W.size(0)).cuda()

        W.copy_(torch.solve(Id + W, Id - W)[0].t())  # faster than (Id - W) @ torch.linalg.inv(Id + W)
        W = W.unsqueeze(0)        
        
        mat_feats = W @ mat_feats @ W.transpose(-1, -2)
        return mat_feats
//  [python]
def forward(self, mat_feats):
        W = self.isom_params.data
        
        W = W - W.t() # skew-symmetric A^T = -A
        Id = torch.eye(W.size(0)).cuda()

#       See https://pytorch.org/docs/stable/generated/torch.linalg.inv.html#torch.linalg.inv        
        W.copy_(torch.linalg.solve(Id - W, Id + W).t())  # faster than (Id - W) @ torch.linalg.inv(Id + W)
        W = W.unsqueeze(0)        
        
        mat_feats = W @ mat_feats @ W.transpose(-1, -2)
        return mat_feats
//  [python]
def _solve(A, B):
    ATA = A.transpose(-2, -1) @ A
    ATB = A.transpose(-2, -1) @ B
    return torch.solve(ATB, ATA)[0]
//  [python]
def _solve(A, B):
    ATA = A.transpose(-2, -1) @ A
    ATB = A.transpose(-2, -1) @ B
    return torch.linalg.solve(ATA, ATB)
//  [python]
def get_relative_pose_torch(tsfm_src, tsfm_tgt, dataset):
    """
    We compute relative transformation matrix from source point cloud to target point cloud
    T_src @ X_src = T_tgt @ X_tgt
    X_tgt = T_tgt^-1 @ T_src @ X_src
    T_relative = T_tgt^-1 @ T_src
    """
    if dataset in ['waymo','nuscene']:
        M = torch.solve(tsfm_src, tsfm_tgt)[0]
    else:
        velo2cam = VELO2CAM[dataset]
        velo2cam = torch.from_numpy(velo2cam).float().to(tsfm_src.device)
        M = (velo2cam @ tsfm_src.T @ torch.inverse(tsfm_tgt.T)@ torch.inverse(velo2cam)).T 
    return M
//  [python]
def get_relative_pose_torch(tsfm_src, tsfm_tgt, dataset):
    """
    We compute relative transformation matrix from source point cloud to target point cloud
    T_src @ X_src = T_tgt @ X_tgt
    X_tgt = T_tgt^-1 @ T_src @ X_src
    T_relative = T_tgt^-1 @ T_src
    """
    if dataset in ['waymo','nuscene']:
        M = torch.linalg.solve(tsfm_tgt, tsfm_src)
    else:
        velo2cam = VELO2CAM[dataset]
        velo2cam = torch.from_numpy(velo2cam).float().to(tsfm_src.device)
        M = (velo2cam @ tsfm_src.T @ torch.inverse(tsfm_tgt.T)@ torch.inverse(velo2cam)).T 
    return M
//  [python]
def _lsolve(self, x, b):
        # TODO: A better way to solve it is implemented in `_optimLSolve`
        # Once Pytorch gains support for `LinearOperator`/scipy like `cg`
        # function, it can be used.
        alpha = torch.empty((self._k, self._n))
        beta = torch.empty((self._k, self._m))
        for i in range(self._k):
            A = torch.cat((torch.cat((torch.diag(self._p[i]), x[i]), dim=-1), torch.cat((x[i].T, torch.diag(self._q[i])), dim=-1)), dim=0)
            zeta = torch.solve(b[i].unsqueeze(1), A)[0].squeeze(1)
            alpha[i], beta[i] = zeta[:self._n], zeta[self._n:]
        return alpha, beta
//  [python]
def _lsolve(self, x, b):
        # TODO: A better way to solve it is implemented in `_optimLSolve`
        # Once Pytorch gains support for `LinearOperator`/scipy like `cg`
        # function, it can be used.
        alpha = torch.empty((self._k, self._n))
        beta = torch.empty((self._k, self._m))
        for i in range(self._k):
            A = torch.cat((torch.cat((torch.diag(self._p[i]), x[i]), dim=-1), torch.cat((x[i].T, torch.diag(self._q[i])), dim=-1)), dim=0)
            zeta = torch.linalg.solve(A, b[i])
            alpha[i], beta[i] = zeta[:self._n], zeta[self._n:]
        return alpha, beta
//  [python]
def se3_log_map(
        cls, transform: torch.Tensor, eps: float = 1e-4, cos_bound: float = 1e-4
    ):
        """
        SE3 -> se3

        Args:
            transform: shape of (N, 4, 4)
        Returns:
            : shape of (N, 6)
        """
        assert (
            len(transform.shape) == 3
            and transform.shape[1] == 4
            and transform.shape[2] == 4
        ), "Input tensor shape has to be Nx4x4."
        N = transform.shape[0]

        assert torch.allclose(
            transform[:, 3, :3], torch.zeros_like(transform[:, :3, 3])
        ), "All elements of `transform[:, :3, 3]` should be 0."

        R = transform[:, :3, :3]
        log_rot = cls.so3_log_map(R, eps=eps, cos_bound=cos_bound)

        T = transform[:, :3, 3]

        rot_angles = (log_rot**2).sum(-1).clamp(eps).sqrt()
        log_rot_hat = cls.hat(log_rot)
        log_rot_hat_square = torch.bmm(log_rot_hat, log_rot_hat)
        I = torch.eye(3).to(log_rot)[None]

        A = ((1 - torch.cos(rot_angles)) / (rot_angles**2))[:, None, None]
        B = ((rot_angles - torch.sin(rot_angles)) / (rot_angles**3))[:, None, None]

        V = I + A * log_rot_hat + B * log_rot_hat_square
        log_trans = torch.solve(T[:, :, None], V)[0][:, :, 0]

        return torch.cat([log_trans, log_rot], dim=1)
//  [python]
def se3_log_map(
        cls, transform: torch.Tensor, eps: float = 1e-4, cos_bound: float = 1e-4
    ):
        """
        SE3 -> se3

        Args:
            transform: shape of (N, 4, 4)
        Returns:
            : shape of (N, 6)
        """
        assert (
            len(transform.shape) == 3
            and transform.shape[1] == 4
            and transform.shape[2] == 4
        ), "Input tensor shape has to be Nx4x4."
        N = transform.shape[0]

        assert torch.allclose(
            transform[:, 3, :3], torch.zeros_like(transform[:, :3, 3])
        ), "All elements of `transform[:, :3, 3]` should be 0."

        R = transform[:, :3, :3]
        log_rot = cls.so3_log_map(R, eps=eps, cos_bound=cos_bound)

        T = transform[:, :3, 3]

        rot_angles = (log_rot**2).sum(-1).clamp(eps).sqrt()
        log_rot_hat = cls.hat(log_rot)
        log_rot_hat_square = torch.bmm(log_rot_hat, log_rot_hat)
        I = torch.eye(3).to(log_rot)[None]

        A = ((1 - torch.cos(rot_angles)) / (rot_angles**2))[:, None, None]
        B = ((rot_angles - torch.sin(rot_angles)) / (rot_angles**3))[:, None, None]

        V = I + A * log_rot_hat + B * log_rot_hat_square
        log_trans = torch.linalg.solve(V, T[:, :, None])[:, :, 0]

        return torch.cat([log_trans, log_rot], dim=1)
//  [python]
def forward(self, X, y, Xstar):
        
        N = X.shape[0]
        assert Xstar.ndim == 2
        M = Xstar.shape[0]

        Kw11 = self.kerW.matrix(X)
        Kw12 = self.kerW.cross(X, Xstar)
        Kw11InvKw12 = torch.solve(Kw12, Kw11)[1]
        
        Wtr = torch.unsqueeze(torch.permute(self.Wmu, (1,2,0)), -2)
        
        Wstar_mu = torch.permute(torch.tensordot(Wtr, Kw11InvKw12, dims=([3],[0])).squeeze(-2), (2,0,1))
        
        
        Kf11 = self.kerF.matrix(X)
        Kf12 = self.kerF.cross(X, Xstar)
        
        Kf11InvKf12 = torch.solve(Kf12, Kf11)[1]
        
        Ftr = torch.unsqueeze(torch.permute(self.Fmu,(1,0)),-2)
        
        Fstar_mu = torch.permute(torch.tensordot(Ftr, Kf11InvKf12, dims=([2], [0])).squeeze(-2), (1,0))
        
        ystar = torch.bmm(
            Wstar_mu,
            Fstar_mu.unsqueeze(-1)
        ).squeeze(-1)
        
        return ystar
//  [python]
def forward(self, X, y, Xstar):
        
        N = X.shape[0]
        assert Xstar.ndim == 2
        M = Xstar.shape[0]

        Kw11 = self.kerW.matrix(X)
        Kw12 = self.kerW.cross(X, Xstar)
        Kw11InvKw12 = torch.linalg.solve(Kw11, Kw12)
        
        Wtr = torch.unsqueeze(torch.permute(self.Wmu, (1,2,0)), -2)
        
        Wstar_mu = torch.permute(torch.tensordot(Wtr, Kw11InvKw12, dims=([3],[0])).squeeze(-2), (2,0,1))
        
        
        Kf11 = self.kerF.matrix(X)
        Kf12 = self.kerF.cross(X, Xstar)
        
        Kf11InvKf12 = torch.linalg.solve(Kf11, Kf12)
        
        Ftr = torch.unsqueeze(torch.permute(self.Fmu,(1,0)),-2)
        
        Fstar_mu = torch.permute(torch.tensordot(Ftr, Kf11InvKf12, dims=([2], [0])).squeeze(-2), (1,0))
        
        ystar = torch.bmm(
            Wstar_mu,
            Fstar_mu.unsqueeze(-1)
        ).squeeze(-1)
        
        return ystar
//  [python]
def estimate_filter(self, Q: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
        """Estimate the MIMO prediction filter as
            G(b,f) = Q(b,f) \ R(b,f)
        for each subband in each example in the batch (b, f).

        Args:
            Q: shape (B, F, C, filter_length, C, filter_length)
            R: shape (B, F, C, filter_length, C)

        Returns:
            Complex-valued prediction filter, shape (B, C, F, C, filter_length)
        """
        B, F, C, filter_length, _, _ = Q.shape
        assert (
            filter_length == self.filter_length
        ), f'Shape of Q {Q.shape} is not matching filter length {self.filter_length}'

        # Reshape to analytical dimensions for each (b, f)
        Q = Q.reshape(B, F, C * self.filter_length, C * filter_length)
        R = R.reshape(B, F, C * self.filter_length, C)

        # Diagonal regularization
        if self.diag_reg:
            # Regularization: diag_reg * trace(Q) + eps
            diag_reg = self.diag_reg * torch.diagonal(Q, dim1=-2, dim2=-1).sum(-1).real + self.eps
            # Apply regularization on Q
            Q = Q + torch.diag_embed(diag_reg.unsqueeze(-1) * torch.ones(Q.shape[-1], device=Q.device))

        # Solve for the filter
        G = torch.solve(R, Q)[1]

        # Reshape to desired representation: (B, F, input channels, filter_length, output channels)
        G = G.reshape(B, F, C, filter_length, C)
        # Move output channels to front: (B, output channels, F, input channels, filter_length)
        G = G.permute(0, 4, 1, 2, 3)

        return G
//  [python]
def estimate_filter(self, Q: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
        """Estimate the MIMO prediction filter as
            G(b,f) = Q(b,f) \ R(b,f)
        for each subband in each example in the batch (b, f).

        Args:
            Q: shape (B, F, C, filter_length, C, filter_length)
            R: shape (B, F, C, filter_length, C)

        Returns:
            Complex-valued prediction filter, shape (B, C, F, C, filter_length)
        """
        B, F, C, filter_length, _, _ = Q.shape
        assert (
            filter_length == self.filter_length
        ), f'Shape of Q {Q.shape} is not matching filter length {self.filter_length}'

        # Reshape to analytical dimensions for each (b, f)
        Q = Q.reshape(B, F, C * self.filter_length, C * filter_length)
        R = R.reshape(B, F, C * self.filter_length, C)

        # Diagonal regularization
        if self.diag_reg:
            # Regularization: diag_reg * trace(Q) + eps
            diag_reg = self.diag_reg * torch.diagonal(Q, dim1=-2, dim2=-1).sum(-1).real + self.eps
            # Apply regularization on Q
            Q = Q + torch.diag_embed(diag_reg.unsqueeze(-1) * torch.ones(Q.shape[-1], device=Q.device))

        # Solve for the filter
        G = torch.linalg.solve(Q, R)

        # Reshape to desired representation: (B, F, input channels, filter_length, output channels)
        G = G.reshape(B, F, C, filter_length, C)
        # Move output channels to front: (B, output channels, F, input channels, filter_length)
        G = G.permute(0, 4, 1, 2, 3)

        return G
//  [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 test_cg():
    # simple function with tensor parameters
    def func(x):
        return torch.stack([(x**5).sum(), (x**3).sum()])

    def partial_fvp(v):
        return fvp(func, (x,), (v,), normalize=False)[1]

    x = torch.arange(1.0, 6.0)
    v = torch.ones_like(x)

    jac = torch.func.jacrev(func)(x)
    fisher = jac.T @ jac
    damping = 100

    sol = torch.solve(fisher + damping * torch.eye(fisher.shape[0]), v)[0]
    sol_cg, _ = cg(partial_fvp, v, x0=None, damping=damping, maxiter=10000, tol=1e-10)
    assert torch.allclose(sol, sol_cg, rtol=1e-3)

    # simple complex number example
    A = torch.tensor([[0, -1j], [1j, 0]])

    def mvp(x):
        return A @ x

    b = torch.randn(2, dtype=torch.cfloat)

    sol = torch.solve(A, b)[0]
    sol_cg, _ = cg(mvp, b, x0=None, tol=1e-10)

    assert torch.allclose(sol, sol_cg, rtol=1e-1)

    # function with parameters as a TensorTree
    model = TestModel()

    func_model = model_to_function(model)
    f_per_sample = torch.vmap(func_model, in_dims=(None, 0))

    xs = torch.randn(100, 10)

    def partial_fvp(v):
        return fvp(lambda p: func_model(p, xs), (params,), (v,), normalize=False)[1]

    params = dict(model.named_parameters())
    fisher = empirical_fisher(lambda p: f_per_sample(p, xs), normalize=False)(params)
    damping = 0

    v, _ = tree_ravel(params)
    sol = torch.solve(fisher + damping * torch.eye(fisher.shape[0]), v)[0]
    sol_cg, _ = cg(partial_fvp, params, x0=None, damping=damping, tol=1e-10)

    assert torch.allclose(sol, tree_ravel(sol_cg)[0], rtol=1e-3)
//  [python]
def test_cg():
    # simple function with tensor parameters
    def func(x):
        return torch.stack([(x**5).sum(), (x**3).sum()])

    def partial_fvp(v):
        return fvp(func, (x,), (v,), normalize=False)[1]

    x = torch.arange(1.0, 6.0)
    v = torch.ones_like(x)

    jac = torch.func.jacrev(func)(x)
    fisher = jac.T @ jac
    damping = 100

    sol = torch.linalg.solve(fisher + damping * torch.eye(fisher.shape[0]), v)
    sol_cg, _ = cg(partial_fvp, v, x0=None, damping=damping, maxiter=10000, tol=1e-10)
    assert torch.allclose(sol, sol_cg, rtol=1e-3)

    # simple complex number example
    A = torch.tensor([[0, -1j], [1j, 0]])

    def mvp(x):
        return A @ x

    b = torch.randn(2, dtype=torch.cfloat)

    sol = torch.linalg.solve(A, b)
    sol_cg, _ = cg(mvp, b, x0=None, tol=1e-10)

    assert torch.allclose(sol, sol_cg, rtol=1e-1)

    # function with parameters as a TensorTree
    model = TestModel()

    func_model = model_to_function(model)
    f_per_sample = torch.vmap(func_model, in_dims=(None, 0))

    xs = torch.randn(100, 10)

    def partial_fvp(v):
        return fvp(lambda p: func_model(p, xs), (params,), (v,), normalize=False)[1]

    params = dict(model.named_parameters())
    fisher = empirical_fisher(lambda p: f_per_sample(p, xs), normalize=False)(params)
    damping = 0

    v, _ = tree_ravel(params)
    sol = torch.linalg.solve(fisher + damping * torch.eye(fisher.shape[0]), v)
    sol_cg, _ = cg(partial_fvp, params, x0=None, damping=damping, tol=1e-10)

    assert torch.allclose(sol, tree_ravel(sol_cg)[0], rtol=1e-3)
//  [python]
def mmse_beamformers(h, p):
    k, n = h.shape
    eye_mat = th.eye(n, dtype=h.dtype, device=h.device)
    w = th.solve(eye_mat * k / p + th.conj(th.transpose(h, 0, 1)) @ h, th.conj(th.transpose(h, 0, 1)))[0]
    w = w / th.norm(w, dim=0, keepdim=True) / k ** 0.5
    return w
//  [python]
def mmse_beamformers(h, p):
    k, n = h.shape
    eye_mat = th.eye(n, dtype=h.dtype, device=h.device)
    w = th.linalg.solve(eye_mat * k / p + th.conj(th.transpose(h, 0, 1)) @ h, th.conj(th.transpose(h, 0, 1)))
    w = w / th.norm(w, dim=0, keepdim=True) / k ** 0.5
    return w
//  [python]
def separate(self, mic_index=MIC_INDEX):
        Y_NFTM = torch.einsum("nft, nm -> nftm", self.PSD_NFT, self.G_NM).to(self.TYPE_COMPLEX)
        self.Y_FTM = Y_NFTM.sum(axis=0)
        self.Qx_FTM = torch.einsum("fmi, fti -> ftm", self.Q_FMM, self.X_FTM)
        Qinv_FMM = torch.solve(
            torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device)[None], self.Q_FMM
        )[1]

        self.separated_spec = torch.einsum(
            "fj, ftj, nftj -> nft", Qinv_FMM[:, mic_index], self.Qx_FTM / self.Y_FTM, Y_NFTM
        )
        return self.separated_spec
//  [python]
def separate(self, mic_index=MIC_INDEX):
        Y_NFTM = torch.einsum("nft, nm -> nftm", self.PSD_NFT, self.G_NM).to(self.TYPE_COMPLEX)
        self.Y_FTM = Y_NFTM.sum(axis=0)
        self.Qx_FTM = torch.einsum("fmi, fti -> ftm", self.Q_FMM, self.X_FTM)
        Qinv_FMM = torch.linalg.solve(
            self.Q_FMM, torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device)[None]
        )

        self.separated_spec = torch.einsum(
            "fj, ftj, nftj -> nft", Qinv_FMM[:, mic_index], self.Qx_FTM / self.Y_FTM, Y_NFTM
        )
        return self.separated_spec
//  [python]
def __init__(self,
        sp, # start position; shape = (Batch, 1)
        sv, # start velocity; shape = (Batch, 1)
        sa, # start acceleration; shape = (Batch, 1)
        ep, # end position; shape = (Batch, 1)
        ev, # end velocity; shape = (Batch, 1)
        ea, # end accleration; shape = (Batch, 1)
        et, # end time; shape = (Batch,)
    ):
        self._a0 = sp
        self._a1 = sv
        self._a2 = sa / 2.

        A_r0 = torch.stack([et ** 3, et ** 4, et ** 5], dim=1) # row 0
        A_r1 = torch.stack([3 * et ** 2, 4 * et ** 3, 5 * et ** 4], dim=1) # row 1
        A_r2 = torch.stack([6 * et, 12 * et ** 2, 20 * et ** 3], dim=1) # row 2
        A = torch.stack([A_r0, A_r1, A_r2], dim=1)

        b = torch.stack([(ep - self._a0) - self._a1 * et - self._a2 * et ** 2,
                         (ev - self._a1) - 2 * self._a2 * et,
                         (ea - 2 * self._a2)], dim=1)

        x = torch.solve(b, A)[0]
        self._a3 = x[:, 0]
        self._a4 = x[:, 1]
        self._a5 = x[:, 2]

        self._a0 = self._a0[:, None]
        self._a1 = self._a1[:, None]
        self._a2 = self._a2[:, None]
        self._a3 = self._a3[:, None]
        self._a4 = self._a4[:, None]
        self._a5 = self._a5[:, None]

        self._A = A
        self._b = b
        self._x = x
        self._sp = sp
        self._sv = sv
        self._sa = sa
        self._ep = ep
        self._ev = ev
        self._ea = ea
        self._et = et
//  [python]
def __init__(self,
        sp, # start position; shape = (Batch, 1)
        sv, # start velocity; shape = (Batch, 1)
        sa, # start acceleration; shape = (Batch, 1)
        ep, # end position; shape = (Batch, 1)
        ev, # end velocity; shape = (Batch, 1)
        ea, # end accleration; shape = (Batch, 1)
        et, # end time; shape = (Batch,)
    ):
        self._a0 = sp
        self._a1 = sv
        self._a2 = sa / 2.

        A_r0 = torch.stack([et ** 3, et ** 4, et ** 5], dim=1) # row 0
        A_r1 = torch.stack([3 * et ** 2, 4 * et ** 3, 5 * et ** 4], dim=1) # row 1
        A_r2 = torch.stack([6 * et, 12 * et ** 2, 20 * et ** 3], dim=1) # row 2
        A = torch.stack([A_r0, A_r1, A_r2], dim=1)

        b = torch.stack([(ep - self._a0) - self._a1 * et - self._a2 * et ** 2,
                         (ev - self._a1) - 2 * self._a2 * et,
                         (ea - 2 * self._a2)], dim=1)

        x = torch.linalg.solve(A, b)
        self._a3 = x[:, 0]
        self._a4 = x[:, 1]
        self._a5 = x[:, 2]

        self._a0 = self._a0[:, None]
        self._a1 = self._a1[:, None]
        self._a2 = self._a2[:, None]
        self._a3 = self._a3[:, None]
        self._a4 = self._a4[:, None]
        self._a5 = self._a5[:, None]

        self._A = A
        self._b = b
        self._x = x
        self._sp = sp
        self._sv = sv
        self._sa = sa
        self._ep = ep
        self._ev = ev
        self._ea = ea
        self._et = et
//  [python]
def update_factor(f, A, B):
	t = A.diagonal(dim1=-1, dim2=-2)
	t += 1e-10
	f[:] = torch.solve(B, A)[0].T
//  [python]
def update_factor(f, A, B):
	t = A.diagonal(dim1=-1, dim2=-2)
	t += 1e-10
	f[:] = torch.linalg.solve(A, B).T
//  [python]
def update_u(self, x, u, v):
        # x ≈ u @ t(v) --> u = ?
        *_, M, N = x.shape
        if M >= N:
            u_new = x @ t(torch.linalg.pinv(v))
        else:
            a, b = x @ v, t(v) @ v
            u_new = t(torch.solve(t(a), b)[0])

        return self.project(u_new)
//  [python]
def update_u(self, x, u, v):
        # x ≈ u @ t(v) --> u = ?
        *_, M, N = x.shape
        if M >= N:
            u_new = x @ t(torch.linalg.pinv(v))
        else:
            a, b = x @ v, t(v) @ v
            u_new = torch.linalg.solve(b, t(a))
            u_new = t(u_new)

        return self.project(u_new)
//  [python]
def forward_samples(self, X, y, Xstar, ns=50):
        
        N = X.shape[0]
        assert Xstar.ndim == 2
        M = Xstar.shape[0]
        
        Kw11 = self.kerW.matrix(X)
        Kw12 = self.kerW.cross(X, Xstar)
        Kw11InvKw12 = torch.solve(Kw12, Kw11)[1]
        
        Wtr = torch.unsqueeze(torch.permute(self.Wmu, (1,2,0)), -2)
        
        Wstar_mu = torch.permute(torch.tensordot(Wtr, Kw11InvKw12, dims=([3],[0])).squeeze(-2), (2,0,1))
        
        Kw22 = self.kerW.matrix(Xstar)

        Lw11 = torch.linalg.cholesky(Kw11)
        vw = torch.linalg.solve(Lw11, Kw12)
        
        Wstar_V = Kw22-torch.mm(vw.T, vw)
        Wstar_std = torch.diag(torch.sqrt(torch.diag(Wstar_V)))
        
        Kf11 = self.kerF.matrix(X)
        Kf12 = self.kerF.cross(X, Xstar)
        Kf22 = self.kerF.matrix(Xstar)
        
        Kf11InvKf12 = torch.solve(Kf12, Kf11)[1]
        
        Ftr = torch.unsqueeze(torch.permute(self.Fmu,(1,0)),-2)
        
        Fstar_mu = torch.permute(torch.tensordot(Ftr, Kf11InvKf12, dims=([2], [0])).squeeze(-2), (1,0))
        
        Lf11 = torch.linalg.cholesky(Kf11)
        vf = torch.linalg.solve(Lf11, Kf12)
        
        Fstar_V = Kf22-torch.mm(vf.T, vf)
        
        Fstar_std = torch.diag(torch.sqrt(torch.diag(Fstar_V)))
        
        epsiW = torch.randn(size=[ns, M, self.dimd1, self.dimd2, self.dimK]).to(self.dummy.device)
        epsiF = torch.randn(size=[ns, M, self.dimK]).to(self.dummy.device)
        
        LM = Wstar_std
        LK = torch.tril(self.chol_W_K)
        Ld1 = torch.tril(self.chol_W_d1)
        Ld2 = torch.tril(self.chol_W_d2)
        
        AM = Fstar_std
        AK = torch.tril(self.chol_F_K)
        
        Wposts_fact =  tl.tenalg.multi_mode_dot(
            tensor=epsiW, 
            matrix_or_vec_list=[
                LM, 
                Ld1,
                Ld2, 
                LK
            ],
            modes=[1,2,3,4]
        ) + Wstar_mu.reshape([M, self.dimd1, self.dimd2, self.dimK])
        
        Wposts = Wposts_fact.reshape([ns, M, self.dimd1*self.dimd2, self.dimK])
        
        
        Fposts =  tl.tenalg.multi_mode_dot(
            tensor=epsiF, 
            matrix_or_vec_list=[
                AM,
                AK
            ],
            modes=[1,2]
        ) + Fstar_mu
        
        Yposts = torch.einsum(
            'sbij, sbjk->sbik',
            Wposts,
            Fposts.unsqueeze(-1)
        ).squeeze(-1) # samples batch matrix multiplication
        
        
        return Yposts
//  [python]
def forward_samples(self, X, y, Xstar, ns=50):
        
        N = X.shape[0]
        assert Xstar.ndim == 2
        M = Xstar.shape[0]
        
        Kw11 = self.kerW.matrix(X)
        Kw12 = self.kerW.cross(X, Xstar)
        Kw11InvKw12 = torch.linalg.solve(Kw11, Kw12)
        
        Wtr = torch.unsqueeze(torch.permute(self.Wmu, (1,2,0)), -2)
        
        Wstar_mu = torch.permute(torch.tensordot(Wtr, Kw11InvKw12, dims=([3],[0])).squeeze(-2), (2,0,1))
        
        Kw22 = self.kerW.matrix(Xstar)

        Lw11 = torch.linalg.cholesky(Kw11)
        vw = torch.linalg.solve(Lw11, Kw12)
        
        Wstar_V = Kw22-torch.mm(vw.T, vw)
        Wstar_std = torch.diag(torch.sqrt(torch.diag(Wstar_V)))
        
        Kf11 = self.kerF.matrix(X)
        Kf12 = self.kerF.cross(X, Xstar)
        Kf22 = self.kerF.matrix(Xstar)
        
        Kf11InvKf12 = torch.linalg.solve(Kf11, Kf12)
        
        Ftr = torch.unsqueeze(torch.permute(self.Fmu,(1,0)),-2)
        
        Fstar_mu = torch.permute(torch.tensordot(Ftr, Kf11InvKf12, dims=([2], [0])).squeeze(-2), (1,0))
        
        Lf11 = torch.linalg.cholesky(Kf11)
        vf = torch.linalg.solve(Lf11, Kf12)
        
        Fstar_V = Kf22-torch.mm(vf.T, vf)
        
        Fstar_std = torch.diag(torch.sqrt(torch.diag(Fstar_V)))
        
        epsiW = torch.randn(size=[ns, M, self.dimd1, self.dimd2, self.dimK]).to(self.dummy.device)
        epsiF = torch.randn(size=[ns, M, self.dimK]).to(self.dummy.device)
        
        LM = Wstar_std
        LK = torch.tril(self.chol_W_K)
        Ld1 = torch.tril(self.chol_W_d1)
        Ld2 = torch.tril(self.chol_W_d2)
        
        AM = Fstar_std
        AK = torch.tril(self.chol_F_K)
        
        Wposts_fact =  tl.tenalg.multi_mode_dot(
            tensor=epsiW, 
            matrix_or_vec_list=[
                LM, 
                Ld1,
                Ld2, 
                LK
            ],
            modes=[1,2,3,4]
        ) + Wstar_mu.reshape([M, self.dimd1, self.dimd2, self.dimK])
        
        Wposts = Wposts_fact.reshape([ns, M, self.dimd1*self.dimd2, self.dimK])
        
        
        Fposts =  tl.tenalg.multi_mode_dot(
            tensor=epsiF, 
            matrix_or_vec_list=[
                AM,
                AK
            ],
            modes=[1,2]
        ) + Fstar_mu
        
        Yposts = torch.einsum(
            'sbij, sbjk->sbik',
            Wposts,
            Fposts.unsqueeze(-1)
        ).squeeze(-1) # samples batch matrix multiplication
        
        
        return Yposts
//  [python]
def non_orth_basis_proj(directions, point):
    lam = torch.solve(directions.T@directions, directions.T@point)[0]
    return lam
//  [python]
def non_orth_basis_proj(directions, point):
    lam = torch.linalg.solve(directions.T@directions, directions.T@point)
    return lam
//  [python]
def proportional_to_joint_diagnostics(
    potential_function: Callable,
    q: Distribution,
    proposal: Optional[Distribution] = None,
    N: int = int(5e4),
) -> float:
    r"""This will evaluate the posteriors quality by investingating its importance
    weights. If q is a perfect posterior approximation then $q(\theta) \propto
    p(\theta, x_o)$. Thus we should be able to fit a line to $(q(\theta),
    p(\theta, x_o))$, whereas the slope will be proportional to the normalizing
    constant. The quality of a linear fit is hence a direct metric for the quality of q.
    We use R2 statistic.

    NOTE: In our experience this metric does distinguish "good" from "ok", but
    becomes less sensitive to distinguish "very bad" from "ok".

    Args:
        potential_function: Potential function of target.
        q: Variational distribution, should be proportional to the potential_function
        proposal: Proposal for samples. Typically this is q.
        N: Number of samples involved in the test.

    Returns:
        float: Quality metric

    """

    with torch.no_grad():
        if proposal is None:
            samples = q.sample(Size((N,)))
        else:
            samples = proposal.sample(Size((N,)))
        log_q = q.log_prob(samples)
        log_potential = potential_function(samples)

        X = log_q.exp().unsqueeze(-1)
        Y = log_potential.exp().unsqueeze(-1)
        w = torch.solve(X.T @ X, X.T @ Y)[0]  # Linear regression

        residuals = Y - w * X
        var_res = torch.sum(residuals ** 2)
        var_tot = torch.sum((Y - Y.mean()) ** 2)
        r2 = 1 - var_res / var_tot  # R2 statistic to evaluate fit
    return r2.item()
//  [python]
def proportional_to_joint_diagnostics(
    potential_function: Callable,
    q: Distribution,
    proposal: Optional[Distribution] = None,
    N: int = int(5e4),
) -> float:
    r"""This will evaluate the posteriors quality by investingating its importance
    weights. If q is a perfect posterior approximation then $q(\theta) \propto
    p(\theta, x_o)$. Thus we should be able to fit a line to $(q(\theta),
    p(\theta, x_o))$, whereas the slope will be proportional to the normalizing
    constant. The quality of a linear fit is hence a direct metric for the quality of q.
    We use R2 statistic.

    NOTE: In our experience this metric does distinguish "good" from "ok", but
    becomes less sensitive to distinguish "very bad" from "ok".

    Args:
        potential_function: Potential function of target.
        q: Variational distribution, should be proportional to the potential_function
        proposal: Proposal for samples. Typically this is q.
        N: Number of samples involved in the test.

    Returns:
        float: Quality metric

    """

    with torch.no_grad():
        if proposal is None:
            samples = q.sample(Size((N,)))
        else:
            samples = proposal.sample(Size((N,)))
        log_q = q.log_prob(samples)
        log_potential = potential_function(samples)

        X = log_q.exp().unsqueeze(-1)
        Y = log_potential.exp().unsqueeze(-1)
        w = torch.linalg.solve(X.T @ X, X.T @ Y)  # Linear regression

        residuals = Y - w * X
        var_res = torch.sum(residuals ** 2)
        var_tot = torch.sum((Y - Y.mean()) ** 2)
        r2 = 1 - var_res / var_tot  # R2 statistic to evaluate fit
    return r2.item()
//  [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()

        # Repeat w shape properly
        B = repeat(B, '1 t n -> 1 (v t) n', v=self.broadcast)
        P = repeat(P, 'r t n -> r (v t) n', v=self.broadcast)
        Q = repeat(Q, 'r t n -> r (v t) n', v=self.broadcast)
        w = repeat(w, 't n -> (v t) n', v=self.broadcast)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt - 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')
        try:
            R = torch.solve(Q_D, R)[0] # (H R N)
        except:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).contiguous().detach().cpu(), Q_D.contiguous().detach().cpu())).to(Q_D)
        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 + 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()

        # Repeat w shape properly
        B = repeat(B, '1 t n -> 1 (v t) n', v=self.broadcast)
        P = repeat(P, 'r t n -> r (v t) n', v=self.broadcast)
        Q = repeat(Q, 'r t n -> r (v t) n', v=self.broadcast)
        w = repeat(w, 't n -> (v t) n', v=self.broadcast)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt - 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')
        try:
            R = torch.linalg.solve(R, Q_D) # (H R N)
        except:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).contiguous().detach().cpu(), Q_D.contiguous().detach().cpu())).to(Q_D)
        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 + w, # (H N)
        }
//  [python]
def forward(self, input: torch.Tensor, mask_s: torch.Tensor, mask_n: torch.Tensor) -> torch.Tensor:
        """Return processed signal.
        The output has either one channel (M=1) if a ref_channel is selected,
        or the same number of channels as the input (M=C) if ref_channel is None.

        Args:
            input: Input signal, complex tensor with shape (B, C, F, T)
            mask_s: Mask for the desired signal, shape (B, F, T)
            mask_n: Mask for the undesired noise, shape (B, F, T)

        Returns:
            Processed signal, shape (B, M, F, T)
        """
        iodtype = input.dtype

        with torch.cuda.amp.autocast(enabled=False):
            # Convert to double
            input = input.cdouble()
            mask_s = mask_s.double()
            mask_n = mask_n.double()

            # Calculate signal statistics
            psd_s = self.psd(input, mask_s)
            psd_n = self.psd(input, mask_n)

            if self.rank == 'one':
                # Calculate filter W using (18) in [1]
                # Diagonal regularization
                if self.diag_reg:
                    psd_n = self.apply_diag_reg(psd_n)

                # MIMO filter
                # (B, F, C, C)
                W = torch.solve(psd_s, psd_n)[0]
                lam = self.trace(W, keepdim=True).real
                W = W / (self.beta + lam + self.eps)
            elif self.rank == 'full':
                # Calculate filter W using (15) in [1]
                psd_sn = psd_s + self.beta * psd_n

                if self.diag_reg:
                    psd_sn = self.apply_diag_reg(psd_sn)

                # MIMO filter
                # (B, F, C, C)
                W = torch.solve(psd_s, psd_sn)[0]
            else:
                raise RuntimeError(f'Unexpected rank {self.rank}')

            if torch.jit.isinstance(self.ref_channel, int):
                # Fixed ref channel
                # (B, F, C, 1)
                W = W[..., self.ref_channel].unsqueeze(-1)
            elif self.ref_estimator is not None:
                # Estimate ref channel tensor (one-hot or soft across C)
                # (B, C)
                ref_channel_tensor = self.ref_estimator(W=W, psd_s=psd_s, psd_n=psd_n).to(W.dtype)
                # Weighting across channels
                # (B, F, C, 1)
                W = torch.sum(W * ref_channel_tensor[:, None, None, :], dim=-1, keepdim=True)

            output = self.apply_filter(input=input, filter=W)

            # Optional: postfilter
            if self.postfilter == 'ban':
                output = self.apply_ban(input=output, filter=W, psd_n=psd_n)

        return output.to(iodtype)
//  [python]
def forward(self, input: torch.Tensor, mask_s: torch.Tensor, mask_n: torch.Tensor) -> torch.Tensor:
        """Return processed signal.
        The output has either one channel (M=1) if a ref_channel is selected,
        or the same number of channels as the input (M=C) if ref_channel is None.

        Args:
            input: Input signal, complex tensor with shape (B, C, F, T)
            mask_s: Mask for the desired signal, shape (B, F, T)
            mask_n: Mask for the undesired noise, shape (B, F, T)

        Returns:
            Processed signal, shape (B, M, F, T)
        """
        iodtype = input.dtype

        with torch.cuda.amp.autocast(enabled=False):
            # Convert to double
            input = input.cdouble()
            mask_s = mask_s.double()
            mask_n = mask_n.double()

            # Calculate signal statistics
            psd_s = self.psd(input, mask_s)
            psd_n = self.psd(input, mask_n)

            if self.rank == 'one':
                # Calculate filter W using (18) in [1]
                # Diagonal regularization
                if self.diag_reg:
                    psd_n = self.apply_diag_reg(psd_n)

                # MIMO filter
                # (B, F, C, C)
                W = torch.linalg.solve(psd_n, psd_s)
                lam = self.trace(W, keepdim=True).real
                W = W / (self.beta + lam + self.eps)
            elif self.rank == 'full':
                # Calculate filter W using (15) in [1]
                psd_sn = psd_s + self.beta * psd_n

                if self.diag_reg:
                    psd_sn = self.apply_diag_reg(psd_sn)

                # MIMO filter
                # (B, F, C, C)
                W = torch.linalg.solve(psd_sn, psd_s)
            else:
                raise RuntimeError(f'Unexpected rank {self.rank}')

            if torch.jit.isinstance(self.ref_channel, int):
                # Fixed ref channel
                # (B, F, C, 1)
                W = W[..., self.ref_channel].unsqueeze(-1)
            elif self.ref_estimator is not None:
                # Estimate ref channel tensor (one-hot or soft across C)
                # (B, C)
                ref_channel_tensor = self.ref_estimator(W=W, psd_s=psd_s, psd_n=psd_n).to(W.dtype)
                # Weighting across channels
                # (B, F, C, 1)
                W = torch.sum(W * ref_channel_tensor[:, None, None, :], dim=-1, keepdim=True)

            output = self.apply_filter(input=input, filter=W)

            # Optional: postfilter
            if self.postfilter == 'ban':
                output = self.apply_ban(input=output, filter=W, psd_n=psd_n)

        return output.to(iodtype)
//  [python]
def cayley_batch(self, data):
        """
        Perform the Cayley parametrization on a batch of skew-symmetric matrices.

        Args:
            data: A batch of skew-symmetric matrices of shape (b, r, c).
        """
        b, r, c = data.shape
        # Ensure the input matrix is skew-symmetric
        skew_mat = 0.5 * (data - data.transpose(1, 2))
        id_mat = torch.eye(r, device=data.device).unsqueeze(0).expand(b, r, c)

        # Perform the Cayley parametrization
        Q = torch.solve(id_mat - skew_mat, id_mat + skew_mat)[1]

        return Q
//  [python]
def cayley_batch(self, data):
        """
        Perform the Cayley parametrization on a batch of skew-symmetric matrices.

        Args:
            data: A batch of skew-symmetric matrices of shape (b, r, c).
        """
        b, r, c = data.shape
        # Ensure the input matrix is skew-symmetric
        skew_mat = 0.5 * (data - data.transpose(1, 2))
        id_mat = torch.eye(r, device=data.device).unsqueeze(0).expand(b, r, c)

        # Perform the Cayley parametrization
        Q = torch.linalg.solve(id_mat + skew_mat, id_mat - skew_mat, left=False)

        return Q
//  [python]
def kl_divergence_mvn(self, g):
        term1 = torch.log(torch.linalg.det(g.cov)) - torch.log(torch.linalg.det(self.cov))
        term2 = torch.trace(torch.solve(self.cov, g.cov)[0]) - self.manifold.dim_T
        v = g.manifold.log_map(self.mean, base=g.mean)
        term3 = v.transpose(-2, -1) @ g.cov_inv @ v
        return (term1 + term2 + term3) / 2
//  [python]
def kl_divergence_mvn(self, g):
        term1 = torch.log(torch.linalg.det(g.cov)) - torch.log(torch.linalg.det(self.cov))
        term2 = torch.trace(torch.linalg.solve(g.cov, self.cov)) - self.manifold.dim_T
        v = g.manifold.log_map(self.mean, base=g.mean)
        term3 = v.transpose(-2, -1) @ g.cov_inv @ v
        return (term1 + term2 + term3) / 2
//  [python]
def _test_minres(self, rhs_shape, shifts=None, matrix_batch_shape=torch.Size([])):
        size = rhs_shape[-2] if len(rhs_shape) > 1 else rhs_shape[-1]
        rhs = torch.randn(rhs_shape, dtype=torch.float64)

        matrix = torch.randn(*matrix_batch_shape, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(size, dtype=torch.float64).mul_(1e-1))

        # Compute solves with minres
        if shifts is not None:
            shifts = shifts.type_as(rhs)

        settings = torchsparsegradutils.utils.MINRESSettings(minres_tolerance=1e-6)
        solves = minres(matrix, rhs=rhs, value=-1, shifts=shifts, settings=settings)

        # Make sure that we're not getting weird batch dim effects
        while matrix.dim() < len(rhs_shape):
            matrix = matrix.unsqueeze(0)

        # Maybe add shifts
        if shifts is not None:
            matrix = matrix - torch.mul(
                torch.eye(size, dtype=torch.float64), shifts.view(*shifts.shape, *[1 for _ in matrix.shape])
            )

        # Compute solves exactly
        actual = torch.solve(-matrix, rhs.unsqueeze(-1) if rhs.dim() == 1 else rhs)[0]
        if rhs.dim() == 1:
            actual = actual.squeeze(-1)

        self.assertAllClose(solves, actual, atol=1e-3, rtol=1e-4)
//  [python]
def _test_minres(self, rhs_shape, shifts=None, matrix_batch_shape=torch.Size([])):
        size = rhs_shape[-2] if len(rhs_shape) > 1 else rhs_shape[-1]
        rhs = torch.randn(rhs_shape, dtype=torch.float64)

        matrix = torch.randn(*matrix_batch_shape, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(size, dtype=torch.float64).mul_(1e-1))

        # Compute solves with minres
        if shifts is not None:
            shifts = shifts.type_as(rhs)

        settings = torchsparsegradutils.utils.MINRESSettings(minres_tolerance=1e-6)
        solves = minres(matrix, rhs=rhs, value=-1, shifts=shifts, settings=settings)

        # Make sure that we're not getting weird batch dim effects
        while matrix.dim() < len(rhs_shape):
            matrix = matrix.unsqueeze(0)

        # Maybe add shifts
        if shifts is not None:
            matrix = matrix - torch.mul(
                torch.eye(size, dtype=torch.float64), shifts.view(*shifts.shape, *[1 for _ in matrix.shape])
            )

        # Compute solves exactly
        actual = torch.linalg.solve(-matrix, rhs.unsqueeze(-1) if rhs.dim() == 1 else rhs)
        if rhs.dim() == 1:
            actual = actual.squeeze(-1)

        self.assertAllClose(solves, actual, atol=1e-3, rtol=1e-4)
//  [python]
def update(self, measurement):
        """
        Use measurement to update state estimate in-place and return
        innovation. The innovation is useful, e.g., for evaluating filter
        consistency or updating model likelihoods when the ``EKFState`` is part
        of an ``IMMFState``.

        :param: measurement.
        :returns: EKF State, Innovation mean and covariance.
        """
        if self._time is not None:
            assert (
                self._time == measurement.time
            ), "State time and measurement time must be aligned!"
        if self._frame_num is not None:
            assert (
                self._frame_num == measurement.frame_num
            ), "State time and measurement time must be aligned!"

        x = self._mean
        x_pv = self._dynamic_model.mean2pv(x)
        P = self.cov
        H = measurement.jacobian(x_pv)[:, : self.dimension]
        R = measurement.cov
        z = measurement.mean
        z_predicted = measurement(x_pv)
        dz = measurement.geodesic_difference(z, z_predicted)
        S = H.mm(P).mm(H.transpose(-1, -2)) + R  # innovation cov

        K_prefix = self._cov.mm(H.transpose(-1, -2))
        dx = K_prefix.mm(torch.solve(dz.unsqueeze(1), S)[0]).squeeze(1)  # K*dz
        x = self._dynamic_model.geodesic_difference(x, -dx)

        I = eye_like(x, self._dynamic_model.dimension)  # noqa: E741
        ImKH = I - K_prefix.mm(torch.solve(H, S)[0])
        # *Joseph form* of covariance update for numerical stability.
        S_inv_R = torch.solve(R, S)[0]
        P = ImKH.mm(self.cov).mm(ImKH.transpose(-1, -2)) + K_prefix.mm(
            torch.solve(K_prefix.mm(S_inv_R).transpose(-1, -2), S)[0]
        )

        pred_mean = x
        pred_cov = P
        state = EKFState(
            self._dynamic_model, pred_mean, pred_cov, self._time, self._frame_num
        )

        return state, (dz, S)