jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.solve(S, dz.unsqueeze(1))).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.linalg.solve(S, H))
        # *Joseph form* of covariance update for numerical stability.
        S_inv_R = torch.linalg.solve(S, R)
        P = ImKH.mm(self.cov).mm(ImKH.transpose(-1, -2)) + K_prefix.mm(
            torch.linalg.solve(S, K_prefix.mm(S_inv_R).transpose(-1, -2))
        )

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

        return state, (dz, S)
//  [python]
def cayley_map(X):
    # compute (I+X/2)(I-X/2)^{-1}
    n = X.size(-1)
    Id = torch.eye(n, dtype=X.dtype, device=X.device)
    return torch.solve(Id.add(X, alpha=0.5), Id.add(X, alpha=-0.5))[1]
//  [python]
def cayley_map(X):
    # compute (I+X/2)(I-X/2)^{-1}
    n = X.size(-1)
    Id = torch.eye(n, dtype=X.dtype, device=X.device)
    return torch.linalg.solve(Id.add(X, alpha=-0.5), Id.add(X, alpha=0.5))
//  [python]
def get_wout(s, sd,sdd, y0s, t,a1s,a0s,fs):

    a0_batch = torch.cat([var_(t) for var_ in a0s], 1)
    a1_batch = torch.cat([var_(t) for var_ in a1s], 1)
    f_batch = torch.cat([var_(t) for var_ in fs], 1)
    WS = []
    for i in range(f_batch.shape[1]):
        y0 = y0s[i,:].reshape(1,-1)
        a0 = a0_batch[:,i].reshape(-1,1)
        a1 = a1_batch[:,i].reshape(-1,1)
        f = f_batch[:,i].reshape(-1,1)

        # print(a0,a1,f)
        DH = (sdd + a1 * sd + a0 * s)
        D0 = -f

        h0m = s[0].reshape(-1, 1)
        h0d = sd[0].reshape(-1, 1)
        W0 = torch.solve(torch.cat((DH.t() @ DH + h0m @ h0m.t() + h0d @ h0d.t(),-DH.t() @ D0 + h0m @ (y0[:,0].reshape(1, -1)) + h0d @ (y0[:,1].reshape(1, -1))),1)[0]
        WS.append(W0)
    nWS = (torch.cat(WS)).reshape(f_batch.shape[1],-1)
    return nWS.t()
//  [python]
def get_wout(s, sd,sdd, y0s, t,a1s,a0s,fs):

    a0_batch = torch.cat([var_(t) for var_ in a0s], 1)
    a1_batch = torch.cat([var_(t) for var_ in a1s], 1)
    f_batch = torch.cat([var_(t) for var_ in fs], 1)
    WS = []
    for i in range(f_batch.shape[1]):
        y0 = y0s[i,:].reshape(1,-1)
        a0 = a0_batch[:,i].reshape(-1,1)
        a1 = a1_batch[:,i].reshape(-1,1)
        f = f_batch[:,i].reshape(-1,1)

        # print(a0,a1,f)
        DH = (sdd + a1 * sd + a0 * s)
        D0 = -f

        h0m = s[0].reshape(-1, 1)
        h0d = sd[0].reshape(-1, 1)
        W0 = torch.linalg.solve(DH.t() @ DH + h0m @ h0m.t() + h0d @ h0d.t(),
                                -DH.t() @ D0 + h0m @ (y0[:,0].reshape(1, -1)) + h0d @ (y0[:,1].reshape(1, -1)))
        WS.append(W0)
    nWS = (torch.cat(WS)).reshape(f_batch.shape[1],-1)
    return nWS.t()
//  [python]
def _lambda_fnc(self, A, x, y):
        # Note this would be -eigenvalues due to the -lambda*h in the Lagrangian form in DDNs
        # while the commonly-used one is +lambda*h, both are correct due to the equality constraint.
        # Just need to pay attention that when calculating B and H, change the sign of eigenvalues
        # for lambda.
        batch, m, n = y.shape
        lambds = torch.solve(
            torch.einsum('bnm,bkm->bnk', A, A),
            torch.einsum('bnm,bm->bn', A, self.dfdy_fnc(x, y).view(batch, m * n)))[1]

        return lambds
//  [python]
def _lambda_fnc(self, A, x, y):
        # Note this would be -eigenvalues due to the -lambda*h in the Lagrangian form in DDNs
        # while the commonly-used one is +lambda*h, both are correct due to the equality constraint.
        # Just need to pay attention that when calculating B and H, change the sign of eigenvalues
        # for lambda.
        batch, m, n = y.shape
        lambds = torch.linalg.solve(
            torch.einsum('bnm,bkm->bnk', A, A),
            torch.einsum('bnm,bm->bn', A, self.dfdy_fnc(x, y).view(batch, m * n)))

        return lambds
//  [python]
def _get_steering_vector_power(
            self,
            psd_s: torch.Tensor,
            psd_n: torch.Tensor,
            reference_vector: torch.Tensor
    ) -> torch.Tensor:
        r"""Estimate the steering vector by the power method.

        Args:
            psd_s (torch.tensor): covariance matrix of speech
                Tensor of dimension `(..., freq, channel, channel)`
            psd_n (torch.Tensor): covariance matrix of noise
                Tensor of dimension `(..., freq, channel, channel)`
            reference_vector (torch.Tensor): one-hot reference channel matrix

        Returns:
            torch.Tensor: the enhanced STFT
                Tensor of dimension `(..., freq, channel, 1)`
        """
        phi = torch.solve(psd_s, psd_n)[0]  # psd_n.inv() @ psd_s
        stv = torch.einsum("...fec,...c->...fe", [phi, reference_vector])
        stv = stv.unsqueeze(-1)
        stv = torch.matmul(phi, stv)
        stv = torch.matmul(psd_s, stv)
        return stv
//  [python]
def _get_steering_vector_power(
            self,
            psd_s: torch.Tensor,
            psd_n: torch.Tensor,
            reference_vector: torch.Tensor
    ) -> torch.Tensor:
        r"""Estimate the steering vector by the power method.

        Args:
            psd_s (torch.tensor): covariance matrix of speech
                Tensor of dimension `(..., freq, channel, channel)`
            psd_n (torch.Tensor): covariance matrix of noise
                Tensor of dimension `(..., freq, channel, channel)`
            reference_vector (torch.Tensor): one-hot reference channel matrix

        Returns:
            torch.Tensor: the enhanced STFT
                Tensor of dimension `(..., freq, channel, 1)`
        """
        phi = torch.linalg.solve(psd_n, psd_s)  # psd_n.inv() @ psd_s
        stv = torch.einsum("...fec,...c->...fe", [phi, reference_vector])
        stv = stv.unsqueeze(-1)
        stv = torch.matmul(phi, stv)
        stv = torch.matmul(psd_s, stv)
        return stv
//  [python]
def test_exact_tree(Guide):
    is_exact = Guide not in (
        AutoNormal,
        AutoDiagonalNormal,
        AutoNormalMessenger,
        AutoHierarchicalNormalMessenger,
        AutoRegressiveMessenger,
    )

    def model(data):
        x = pyro.sample("x", dist.Normal(0, 1))
        with pyro.plate("data", len(data)):
            y = pyro.sample("y", dist.Normal(x, 1))
            pyro.sample("obs", dist.Normal(y, 1), obs=data)
        return {"x": x, "y": y}

    data = torch.randn(2)
    g = Gaussian(
        log_normalizer=torch.zeros(()),
        info_vec=torch.zeros(5),
        precision=torch.tensor(
            [
                [3, -1, -1, 0, 0],  # x
                [-1, 2, 0, -1, 0],  # y[0]
                [-1, 0, 2, 0, -1],  # y[1]
                [0, -1, 0, 1, 0],  # obs[0]
                [0, 0, -1, 0, 1],  # obs[1]
            ],
            dtype=data.dtype,
        ),
    )
    g_cond = g.condition(data)
    mean = torch.solve(g_cond.info_vec.unsqueeze(1), g_cond.precision).solution.squeeze()
    std = torch.inverse(g_cond.precision).diag().sqrt()
    expected_mean = {"x": mean[0], "y": mean[1:]}
    expected_std = {"x": std[0], "y": std[1:]}
    expected_loss = float(g.event_logsumexp() - g_cond.event_logsumexp())

    guide = Guide(model)
    Elbo = JitTrace_ELBO
    if Guide is AutoRegressiveMessenger:
        Elbo = Trace_ELBO  # currently fails with jit
    elbo = Elbo(num_particles=100, vectorize_particles=True, ignore_jit_warnings=True)
    num_steps = 500
    optim = ClippedAdam({"lr": 0.05, "lrd": 0.1 ** (1 / num_steps)})
    svi = SVI(model, guide, optim, elbo)
    for step in range(num_steps):
        svi.step(data)

    guide.train(False)
    guide.requires_grad_(False)
    with torch.no_grad():
        # Check moments.
        vectorize = pyro.plate("particles", 50000, dim=-2)
        guide_trace = poutine.trace(vectorize(guide)).get_trace(data)
        samples = poutine.replay(vectorize(model), guide_trace)(data)
        for name in ["x", "y"]:
            actual_mean = samples[name].mean(0).squeeze()
            actual_std = samples[name].std(0).squeeze()
            assert_close(actual_mean, expected_mean[name], atol=0.05)
            if is_exact:
                assert_close(actual_std, expected_std[name], rtol=0.05)

        if is_exact:
            # Check ELBO loss.
            actual_loss = elbo.loss(model, guide, data)
            assert_close(actual_loss, expected_loss, atol=0.01)
//  [python]
def test_exact_tree(Guide):
    is_exact = Guide not in (
        AutoNormal,
        AutoDiagonalNormal,
        AutoNormalMessenger,
        AutoHierarchicalNormalMessenger,
        AutoRegressiveMessenger,
    )

    def model(data):
        x = pyro.sample("x", dist.Normal(0, 1))
        with pyro.plate("data", len(data)):
            y = pyro.sample("y", dist.Normal(x, 1))
            pyro.sample("obs", dist.Normal(y, 1), obs=data)
        return {"x": x, "y": y}

    data = torch.randn(2)
    g = Gaussian(
        log_normalizer=torch.zeros(()),
        info_vec=torch.zeros(5),
        precision=torch.tensor(
            [
                [3, -1, -1, 0, 0],  # x
                [-1, 2, 0, -1, 0],  # y[0]
                [-1, 0, 2, 0, -1],  # y[1]
                [0, -1, 0, 1, 0],  # obs[0]
                [0, 0, -1, 0, 1],  # obs[1]
            ],
            dtype=data.dtype,
        ),
    )
    g_cond = g.condition(data)
    mean = torch.linalg.solve(g_cond.precision, g_cond.info_vec)
    std = torch.inverse(g_cond.precision).diag().sqrt()
    expected_mean = {"x": mean[0], "y": mean[1:]}
    expected_std = {"x": std[0], "y": std[1:]}
    expected_loss = float(g.event_logsumexp() - g_cond.event_logsumexp())

    guide = Guide(model)
    Elbo = JitTrace_ELBO
    if Guide is AutoRegressiveMessenger:
        Elbo = Trace_ELBO  # currently fails with jit
    elbo = Elbo(num_particles=100, vectorize_particles=True, ignore_jit_warnings=True)
    num_steps = 500
    optim = ClippedAdam({"lr": 0.05, "lrd": 0.1 ** (1 / num_steps)})
    svi = SVI(model, guide, optim, elbo)
    for step in range(num_steps):
        svi.step(data)

    guide.train(False)
    guide.requires_grad_(False)
    with torch.no_grad():
        # Check moments.
        vectorize = pyro.plate("particles", 50000, dim=-2)
        guide_trace = poutine.trace(vectorize(guide)).get_trace(data)
        samples = poutine.replay(vectorize(model), guide_trace)(data)
        for name in ["x", "y"]:
            actual_mean = samples[name].mean(0).squeeze()
            actual_std = samples[name].std(0).squeeze()
            assert_close(actual_mean, expected_mean[name], atol=0.05)
            if is_exact:
                assert_close(actual_std, expected_std[name], rtol=0.05)

        if is_exact:
            # Check ELBO loss.
            actual_loss = elbo.loss(model, guide, data)
            assert_close(actual_loss, expected_loss, atol=0.01)
//  [python]
def torch_detrending_filter(signal, device, lam=50):
    T = len(signal)
    I = sparse.eye(T).toarray()
    D2 = sparse.diags([1, -2, 1], [0, 1, 2], shape=(T - 2, T)).toarray()

    # torch Tensor
    sr = signal

    # numpy
    op = np.multiply(np.transpose(D2),lam**2)
    op2 = np.matmul(op, D2) + I

    # torch Tensor
    op2 = torch.FloatTensor(op2).to(device)

    op3 = torch.solve(sr, op2)[0] # np.matmul(np.linalg.inv(op2), sr), ref: An Advanced Detrending Method With Application to HRV Analysis (2002, IEEE TOBE)

    detrended_signal = torch.matmul(torch.FloatTensor(I).to(device), sr) - op3

    return detrended_signal
//  [python]
def torch_detrending_filter(signal, device, lam=50):
    T = len(signal)
    I = sparse.eye(T).toarray()
    D2 = sparse.diags([1, -2, 1], [0, 1, 2], shape=(T - 2, T)).toarray()

    # torch Tensor
    sr = signal

    # numpy
    op = np.multiply(np.transpose(D2),lam**2)
    op2 = np.matmul(op, D2) + I

    # torch Tensor
    op2 = torch.FloatTensor(op2).to(device)

    op3 = torch.linalg.solve(op2, sr) # np.matmul(np.linalg.inv(op2), sr), ref: An Advanced Detrending Method With Application to HRV Analysis (2002, IEEE TOBE)

    detrended_signal = torch.matmul(torch.FloatTensor(I).to(device), sr) - op3

    return detrended_signal
//  [python]
def get_dLdx(self, batch, m, n, B, H, dLdy):
        Q = dLdy.view(batch, 1, m * n)

        if self.backprop_inverse_mode == 'solve':
            HHT = torch.einsum('bmn,bkn->bmk', H, H)
            HQT = torch.einsum('bmn,bkn->bmk', H, Q)
            QHinv = torch.solve(HQT, HHT)[0].permute(0, 2, 1)
        else:
            QHinv = torch.einsum('bmn,bnk->bmk', Q, torch.pinverse(H))

        dLdx = -torch.einsum('bmn,bnk->bmk', QHinv, B).view(batch, m, m)
        if self.enable_symmetric: dLdx = 0.5 * (dLdx + dLdx.permute(0, 2, 1))

        return dLdx
//  [python]
def get_dLdx(self, batch, m, n, B, H, dLdy):
        Q = dLdy.view(batch, 1, m * n)

        if self.backprop_inverse_mode == 'solve':
            HHT = torch.einsum('bmn,bkn->bmk', H, H)
            HQT = torch.einsum('bmn,bkn->bmk', H, Q)
            QHinv = torch.linalg.solve(HHT, HQT).permute(0, 2, 1)
        else:
            QHinv = torch.einsum('bmn,bnk->bmk', Q, torch.linalg.pinv(H))

        dLdx = -torch.einsum('bmn,bnk->bmk', QHinv, B).view(batch, m, m)
        if self.enable_symmetric: dLdx = 0.5 * (dLdx + dLdx.permute(0, 2, 1))

        return dLdx
//  [python]
def _get_mvdr_vector(
        self,
        psd_s: torch.Tensor,
        psd_n: torch.Tensor,
        reference_vector: torch.Tensor,
        solution: str = 'ref_channel',
        diagonal_loading: bool = True,
        diag_eps: float = 1e-7,
        eps: float = 1e-8,
    ) -> torch.Tensor:
        r"""Compute beamforming vector by the reference channel selection method.

        Args:
            psd_s (torch.Tensor): psd matrix of target speech
            psd_n (torch.Tensor): psd matrix of noise
            reference_vector (torch.Tensor): one-hot reference channel matrix
            solution (str, optional): the solution to estimate the beamforming weight
                (Default: ``ref_channel``)
            diagonal_loading (bool, optional): whether to apply diagonal loading to psd_n
                (Default: ``True``)
            diag_eps (float, optional): The coefficient multipied to the identity matrix for diagonal loading
                (Default: 1e-7)
            eps (float, optional): a value added to the denominator in mask normalization. Default: 1e-8

        Returns:
            torch.Tensor: the mvdr beamforming weight matrix
        """
        if diagonal_loading:
            psd_n = self._tik_reg(psd_n, reg=diag_eps, eps=eps)
        if solution == "ref_channel":
            numerator = torch.solve(psd_s, psd_n)[0]  # psd_n.inv() @ psd_s
            # ws: (..., C, C) / (...,) -> (..., C, C)
            ws = numerator / (_get_mat_trace(numerator)[..., None, None] + eps)
            # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
            beamform_vector = torch.einsum("...fec,...c->...fe", [ws, reference_vector])
        else:
            if solution == "stv_evd":
                stv = self._get_steering_vector_evd(psd_s)
            else:
                stv = self._get_steering_vector_power(psd_s, psd_n, reference_vector)
            # numerator = psd_n.inv() @ stv
            numerator = torch.solve(stv, psd_n)[0].squeeze(-1)  # (..., freq, channel)
            # denominator = stv^H @ psd_n.inv() @ stv
            denominator = torch.einsum("...d,...d->...", [stv.conj().squeeze(-1), numerator])
            # normalzie the numerator
            scale = stv.squeeze(-1)[..., self.ref_channel, None].conj()
            beamform_vector = numerator * scale / (denominator.real.unsqueeze(-1) + eps)

        return beamform_vector
//  [python]
def _get_mvdr_vector(
        self,
        psd_s: torch.Tensor,
        psd_n: torch.Tensor,
        reference_vector: torch.Tensor,
        solution: str = 'ref_channel',
        diagonal_loading: bool = True,
        diag_eps: float = 1e-7,
        eps: float = 1e-8,
    ) -> torch.Tensor:
        r"""Compute beamforming vector by the reference channel selection method.

        Args:
            psd_s (torch.Tensor): psd matrix of target speech
            psd_n (torch.Tensor): psd matrix of noise
            reference_vector (torch.Tensor): one-hot reference channel matrix
            solution (str, optional): the solution to estimate the beamforming weight
                (Default: ``ref_channel``)
            diagonal_loading (bool, optional): whether to apply diagonal loading to psd_n
                (Default: ``True``)
            diag_eps (float, optional): The coefficient multipied to the identity matrix for diagonal loading
                (Default: 1e-7)
            eps (float, optional): a value added to the denominator in mask normalization. Default: 1e-8

        Returns:
            torch.Tensor: the mvdr beamforming weight matrix
        """
        if diagonal_loading:
            psd_n = self._tik_reg(psd_n, reg=diag_eps, eps=eps)
        if solution == "ref_channel":
            numerator = torch.linalg.solve(psd_n, psd_s)  # psd_n.inv() @ psd_s
            # ws: (..., C, C) / (...,) -> (..., C, C)
            ws = numerator / (_get_mat_trace(numerator)[..., None, None] + eps)
            # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
            beamform_vector = torch.einsum("...fec,...c->...fe", [ws, reference_vector])
        else:
            if solution == "stv_evd":
                stv = self._get_steering_vector_evd(psd_s)
            else:
                stv = self._get_steering_vector_power(psd_s, psd_n, reference_vector)
            # numerator = psd_n.inv() @ stv
            numerator = torch.linalg.solve(psd_n, stv).squeeze(-1)  # (..., freq, channel)
            # denominator = stv^H @ psd_n.inv() @ stv
            denominator = torch.einsum("...d,...d->...", [stv.conj().squeeze(-1), numerator])
            # normalzie the numerator
            scale = stv.squeeze(-1)[..., self.ref_channel, None].conj()
            beamform_vector = numerator * scale / (denominator.real.unsqueeze(-1) + eps)

        return beamform_vector
//  [python]
def backward(ctx, grad_barycentrics):
        barycentrics, vertices, = ctx.saved_tensors
        grad_vertices = None
        grad_points = None
        if ctx.needs_input_grad[1] or ctx.needs_input_grad[2]:
            t_mat = (vertices[..., 1:, :] - vertices[..., :1, :])
            m_vec = torch.solve(grad_barycentrics, t_mat)[0]
            full_barycentrics = torch.cat([1.0-barycentrics.sum(-1, keepdim=True), barycentrics], -1)
        if ctx.needs_input_grad[1]:
            grad_vertices = (full_barycentrics.unsqueeze(-1) * m_vec.unsqueeze(-2)).mul_(-1.0)
        if ctx.needs_input_grad[2]:
            grad_points = m_vec
        return grad_barycentrics, grad_vertices, grad_points
//  [python]
def backward(ctx, grad_barycentrics):
        barycentrics, vertices, = ctx.saved_tensors
        grad_vertices = None
        grad_points = None
        if ctx.needs_input_grad[1] or ctx.needs_input_grad[2]:
            t_mat = (vertices[..., 1:, :] - vertices[..., :1, :])
            m_vec = torch.linalg.solve(t_mat, grad_barycentrics)
            full_barycentrics = torch.cat([1.0-barycentrics.sum(-1, keepdim=True), barycentrics], -1)
        if ctx.needs_input_grad[1]:
            grad_vertices = (full_barycentrics.unsqueeze(-1) * m_vec.unsqueeze(-2)).mul_(-1.0)
        if ctx.needs_input_grad[2]:
            grad_points = m_vec
        return grad_barycentrics, grad_vertices, grad_points
//  [python]
def proju(self, x: Tensor, u: Tensor) -> Tensor:
        """Arbitrary matrix C projects to skew-Hermitian

                    B := (C - C^†) / 2

        then 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 proju(self, x: Tensor, u: Tensor) -> Tensor:
        """Arbitrary matrix C projects to skew-Hermitian

                    B := (C - C^†) / 2

        then 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 _solve_interpolation(
        train_points,
        train_values,
        regularization_weight
):

    k = train_values.shape[2]
    b, n, d = train_points.shape
    c = train_points.double()
    f = train_values.double()

    A = _phi(_pairwise_squared_distance_matrix(c))

    if regularization_weight > 0:
        batch_identity_matrix = torch.eye(n).unsqueeze(0)
        A = A + regularization_weight * batch_identity_matrix  # broadcast
    ones = torch.ones([b, n, 1]).to(train_points.device)

    B = torch.cat([c, ones], dim=2)
    left_block = torch.cat([A, B.transpose(2, 1)], dim=1)
    num_b_cols = B.shape[2]
    lhs_zeros = torch.zeros([b, num_b_cols, num_b_cols]).to(train_points.device)
    right_block = torch.cat([B, lhs_zeros], dim=1)
    lhs = torch.cat([left_block, right_block], dim=2)
    rhs_zores = torch.zeros([b, d + 1, k]).to(train_points.device)
    rhs = torch.cat([f, rhs_zores], dim=1)

    w_v = torch.matmul(torch.inverse(lhs), rhs)
    w = w_v[:, :n, :]
    v = w_v[:, n:, :]

    return w, v
//  [python]
def _solve_interpolation(
        train_points,
        train_values,
        regularization_weight
):

    k = train_values.shape[2]
    b, n, d = train_points.shape
    c = train_points.double()
    f = train_values.double()

    A = _phi(_pairwise_squared_distance_matrix(c))

    if regularization_weight > 0:
        batch_identity_matrix = torch.eye(n).unsqueeze(0)
        A = A + regularization_weight * batch_identity_matrix  # broadcast
    ones = torch.ones([b, n, 1]).to(train_points.device)

    B = torch.cat([c, ones], dim=2)
    left_block = torch.cat([A, B.transpose(2, 1)], dim=1)
    num_b_cols = B.shape[2]
    lhs_zeros = torch.zeros([b, num_b_cols, num_b_cols]).to(train_points.device)
    right_block = torch.cat([B, lhs_zeros], dim=1)
    lhs = torch.cat([left_block, right_block], dim=2)
    rhs_zores = torch.zeros([b, d + 1, k]).to(train_points.device)
    rhs = torch.cat([f, rhs_zores], dim=1)

    w_v = torch.linalg.solve(lhs, rhs)
    # w_v = torch.matmul(lhs.inverse() , rhs)
    w = w_v[:, :n, :]
    v = w_v[:, n:, :]

    return w, v
//  [python]
def forward(self,batch_x):
        batch_x = torch.mean(batch_x, dim=(3, 4))
        # batch_x = torch.permute(batch_x, (0, 2, 1))  # (B, N, C)
        batch_size, num_features, N  = batch_x.shape
        # batch_x = batch_x.view(batch_size,-1,num_features,N)

        sig_mean = torch.mean(batch_x, dim=2)

        signal_norm_r = batch_x[:, 0, :] / sig_mean[:, 0].unsqueeze(1)
        signal_norm_g = batch_x[:, 1, :] / sig_mean[:, 1].unsqueeze(1)
        signal_norm_b = batch_x[:, 2, :] / sig_mean[:, 2].unsqueeze(1)

        pbv_n = torch.stack([torch.std(signal_norm_r, dim=1),
                             torch.std(signal_norm_g, dim=1),
                             torch.std(signal_norm_b, dim=1)], dim=0)

        pbv_d = torch.sqrt(torch.var(signal_norm_r, dim=1) +
                           torch.var(signal_norm_g, dim=1) +
                           torch.var(signal_norm_b, dim=1))

        pbv = pbv_n / pbv_d

        C = torch.stack([signal_norm_r, signal_norm_g, signal_norm_b], dim=1)
        # C = torch.permute(C,(0,2,1))
        Ct = torch.transpose(C, 1, 2)
        Q = torch.matmul(C, Ct)
        W = torch.solve(pbv.permute(1, 0).unsqueeze(2), Q)[0].squeeze(2)

        A = torch.matmul(Ct, W.unsqueeze(2))
        B = torch.matmul(pbv.T.unsqueeze(1), W.unsqueeze(2))
        bvp = A / B
        return bvp.squeeze(dim=2)
//  [python]
def forward(self,batch_x):
        batch_x = torch.mean(batch_x, dim=(3, 4))
        # batch_x = torch.permute(batch_x, (0, 2, 1))  # (B, N, C)
        batch_size, num_features, N  = batch_x.shape
        # batch_x = batch_x.view(batch_size,-1,num_features,N)

        sig_mean = torch.mean(batch_x, dim=2)

        signal_norm_r = batch_x[:, 0, :] / sig_mean[:, 0].unsqueeze(1)
        signal_norm_g = batch_x[:, 1, :] / sig_mean[:, 1].unsqueeze(1)
        signal_norm_b = batch_x[:, 2, :] / sig_mean[:, 2].unsqueeze(1)

        pbv_n = torch.stack([torch.std(signal_norm_r, dim=1),
                             torch.std(signal_norm_g, dim=1),
                             torch.std(signal_norm_b, dim=1)], dim=0)

        pbv_d = torch.sqrt(torch.var(signal_norm_r, dim=1) +
                           torch.var(signal_norm_g, dim=1) +
                           torch.var(signal_norm_b, dim=1))

        pbv = pbv_n / pbv_d

        C = torch.stack([signal_norm_r, signal_norm_g, signal_norm_b], dim=1)
        # C = torch.permute(C,(0,2,1))
        Ct = torch.transpose(C, 1, 2)
        Q = torch.matmul(C, Ct)
        W = torch.linalg.solve(Q, pbv.permute(1, 0))

        A = torch.matmul(Ct, W.unsqueeze(2))
        B = torch.matmul(pbv.T.unsqueeze(1), W.unsqueeze(2))
        bvp = A / B
        return bvp.squeeze(dim=2)
//  [python]
def _lm_step(f, X, Y, Cinv, L, Lup, Ldn, epsilon, L_min, L_max):
    # Forward
    fY = f(X)
    dY = Y - fY

    # Jacobian
    J = jacfwd(f)(X)
    J = J.to(dtype=X.dtype)
    if Cinv.ndim == 1:
        chi2 = (dY**2 * Cinv).sum(-1)
    else:
        chi2 = (dY @ Cinv @ dY).sum(-1)

    # Gradient
    if Cinv.ndim == 1:
        grad = J.T @ (dY * Cinv)
    else:
        grad = J.T @ Cinv @ dY

    # Hessian
    if Cinv.ndim == 1:
        hess = J.T @ (J * Cinv.reshape(-1, 1))
    else:
        hess = J.T @ Cinv @ J
    hess_perturb = L * torch.eye(hess.shape[0], device=hess.device)
    hess = hess + hess_perturb

    # Step
    h = torch.solve(grad, hess)[0]

    # New chi^2
    fYnew = f(X + h)
    dYnew = Y - fYnew
    if Cinv.ndim == 1:
        chi2_new = (dYnew**2 * Cinv).sum(-1)
    else:
        chi2_new = (dYnew @ Cinv @ dYnew).sum(-1)

    # Test
    rho = (chi2 - chi2_new) / torch.abs(h @ (L * torch.dot(torch.diag(hess), h) + grad))  # fmt: skip

    # Update
    X = torch.where(rho >= epsilon, X + h, X)
    chi2 = torch.where(rho > epsilon, chi2_new, chi2)
    L = torch.clamp(torch.where(rho >= epsilon, L / Ldn, L * Lup), L_min, L_max)

    return X, L, chi2
//  [python]
def _lm_step(f, X, Y, Cinv, L, Lup, Ldn, epsilon, L_min, L_max):
    # Forward
    fY = f(X)
    dY = Y - fY

    # Jacobian
    J = jacfwd(f)(X)
    J = J.to(dtype=X.dtype)
    if Cinv.ndim == 1:
        chi2 = (dY**2 * Cinv).sum(-1)
    else:
        chi2 = (dY @ Cinv @ dY).sum(-1)

    # Gradient
    if Cinv.ndim == 1:
        grad = J.T @ (dY * Cinv)
    else:
        grad = J.T @ Cinv @ dY

    # Hessian
    if Cinv.ndim == 1:
        hess = J.T @ (J * Cinv.reshape(-1, 1))
    else:
        hess = J.T @ Cinv @ J
    hess_perturb = L * torch.eye(hess.shape[0], device=hess.device)
    hess = hess + hess_perturb

    # Step
    h = torch.linalg.solve(hess, grad)

    # New chi^2
    fYnew = f(X + h)
    dYnew = Y - fYnew
    if Cinv.ndim == 1:
        chi2_new = (dYnew**2 * Cinv).sum(-1)
    else:
        chi2_new = (dYnew @ Cinv @ dYnew).sum(-1)

    # Test
    rho = (chi2 - chi2_new) / torch.abs(h @ (L * torch.dot(torch.diag(hess), h) + grad))  # fmt: skip

    # Update
    X = torch.where(rho >= epsilon, X + h, X)
    chi2 = torch.where(rho > epsilon, chi2_new, chi2)
    L = torch.clamp(torch.where(rho >= epsilon, L / Ldn, L * Lup), L_min, L_max)

    return X, L, chi2
//  [python]
def get_enclosed_indices(self):
        start_mm, x_edge_mm, y_edge_mm, z_edge_mm = self.params
        start_mm = torch.tensor(start_mm, dtype=torch.float, device=self.torch_device)
        x_edge_mm = torch.tensor(x_edge_mm, dtype=torch.float, device=self.torch_device)
        y_edge_mm = torch.tensor(y_edge_mm, dtype=torch.float, device=self.torch_device)
        z_edge_mm = torch.tensor(z_edge_mm, dtype=torch.float, device=self.torch_device)

        start_voxels = start_mm / self.voxel_spacing
        x_edge_voxels = x_edge_mm / self.voxel_spacing
        y_edge_voxels = y_edge_mm / self.voxel_spacing
        z_edge_voxels = z_edge_mm / self.voxel_spacing

        target_vector = torch.stack(torch.meshgrid(torch.arange(self.volume_dimensions_voxels[0], dtype=torch.float, device=self.torch_device),
                                                   torch.arange(
                                                       self.volume_dimensions_voxels[1], dtype=torch.float, device=self.torch_device),
                                                   torch.arange(
                                                       self.volume_dimensions_voxels[2], dtype=torch.float, device=self.torch_device),
                                                   indexing='ij'), dim=-1)
        target_vector -= start_voxels

        matrix = torch.stack((x_edge_voxels, y_edge_voxels, z_edge_voxels))

        result = torch.solve(target_vector, matrix.T.expand((target_vector.shape[:-1]+matrix.shape)))[0]
        del target_vector

        norm_vector = torch.tensor([1/torch.linalg.norm(x_edge_voxels),
                                    1/torch.linalg.norm(y_edge_voxels),
                                    1/torch.linalg.norm(z_edge_voxels)], device=self.torch_device)

        filled_mask_bool = (0 <= result) & (result <= 1 - norm_vector)

        volume_fractions = torch.zeros(tuple(self.volume_dimensions_voxels),
                                       dtype=torch.float, device=self.torch_device)
        filled_mask = torch.all(filled_mask_bool, dim=-1)

        volume_fractions[filled_mask] = 1

        return filled_mask.cpu().numpy(), volume_fractions[filled_mask].cpu().numpy()
//  [python]
def get_enclosed_indices(self):
        start_mm, x_edge_mm, y_edge_mm, z_edge_mm = self.params
        start_mm = torch.tensor(start_mm, dtype=torch.float, device=self.torch_device)
        x_edge_mm = torch.tensor(x_edge_mm, dtype=torch.float, device=self.torch_device)
        y_edge_mm = torch.tensor(y_edge_mm, dtype=torch.float, device=self.torch_device)
        z_edge_mm = torch.tensor(z_edge_mm, dtype=torch.float, device=self.torch_device)

        start_voxels = start_mm / self.voxel_spacing
        x_edge_voxels = x_edge_mm / self.voxel_spacing
        y_edge_voxels = y_edge_mm / self.voxel_spacing
        z_edge_voxels = z_edge_mm / self.voxel_spacing

        target_vector = torch.stack(torch.meshgrid(torch.arange(self.volume_dimensions_voxels[0], dtype=torch.float, device=self.torch_device),
                                                   torch.arange(
                                                       self.volume_dimensions_voxels[1], dtype=torch.float, device=self.torch_device),
                                                   torch.arange(
                                                       self.volume_dimensions_voxels[2], dtype=torch.float, device=self.torch_device),
                                                   indexing='ij'), dim=-1)
        target_vector -= start_voxels

        matrix = torch.stack((x_edge_voxels, y_edge_voxels, z_edge_voxels))

        result = torch.linalg.solve(matrix.T.expand((target_vector.shape[:-1]+matrix.shape)), target_vector)
        del target_vector

        norm_vector = torch.tensor([1/torch.linalg.norm(x_edge_voxels),
                                    1/torch.linalg.norm(y_edge_voxels),
                                    1/torch.linalg.norm(z_edge_voxels)], device=self.torch_device)

        filled_mask_bool = (0 <= result) & (result <= 1 - norm_vector)

        volume_fractions = torch.zeros(tuple(self.volume_dimensions_voxels),
                                       dtype=torch.float, device=self.torch_device)
        filled_mask = torch.all(filled_mask_bool, dim=-1)

        volume_fractions[filled_mask] = 1

        return filled_mask.cpu().numpy(), volume_fractions[filled_mask].cpu().numpy()
//  [python]
def __init__(self, L, p, A_init, mu, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self.W = nn.Linear(L, p, bias=False)
        self.B = nn.Linear(p, p, bias=False)

        # init
        M = A_init.T @ A_init + mu * torch.eye(p)

        self.W.weight.data = torch.solve(A_init.T, M)[0]
        self.B.weight.data = torch.solve(mu * torch.eye(p), M)[0]
//  [python]
def __init__(self, L, p, A_init, mu, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self.W = nn.Linear(L, p, bias=False)
        self.B = nn.Linear(p, p, bias=False)

        # init
        M = A_init.T @ A_init + mu * torch.eye(p)

        self.W.weight.data = torch.linalg.solve(M, A_init.T)
        self.B.weight.data = torch.linalg.solve(M, mu * torch.eye(p))
//  [python]
def setUp(self) -> None:
        # The device can be specialised by a daughter class
        if not hasattr(self, "device"):
            self.device = torch.device("cpu")

        self.RTOL = 1e-2

        self.A_shape = (4, 4)
        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device)
        self.A = self.A + self.A.t()
        self.A_csr = self.A.to_sparse_csr()
        self.B_shape = (4, 2)
        self.B = torch.randn(self.B_shape, dtype=torch.float64, device=self.device)

        self.x_ref = torch.solve(self.B, self.A)[0]
//  [python]
def setUp(self) -> None:
        # The device can be specialised by a daughter class
        if not hasattr(self, "device"):
            self.device = torch.device("cpu")

        self.RTOL = 1e-2

        self.A_shape = (4, 4)
        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device)
        self.A = self.A + self.A.t()
        self.A_csr = self.A.to_sparse_csr()
        self.B_shape = (4, 2)
        self.B = torch.randn(self.B_shape, dtype=torch.float64, device=self.device)

        self.x_ref = torch.linalg.solve(self.A, self.B)
//  [python]
def predict_f(self, Xnew, noisy_targets=False):
        """
        Compute the predictive mean vector and covariance matrix for the specified inputs.

        :param Xnew: torch.Tensor, test inputs at which to perform predictions
        :param noisy_targets: bool, specify if predictive covariances are computed for noisy targets
        :return: tuple, predictive mean, predictive covariance matrix
        """
        # Throw exception if data not assigned
        if not self.data_assigned:
            raise Exception('Assign data first')

        # Ensure devices match
        Xnew = Xnew.to(self.X.device)

        # Compute covariance matrices for test/training inputs (s for test, t for train)
        K_st = self.kern.K(Xnew, self.X)
        K_ts = K_st.T
        K_ss = self.kern.K(Xnew)
        K_tt = self.kern.K(self.X) + torch.eye(self.X.shape[0], dtype=self.X.dtype, device=self.X.device) * self.sn2

        # Compute predictive mean
        L = self.cholesky_factor(K_tt, jitter_level=self.jitter).to(dtype=self.Y.dtype)
        A0 = torch.solve(self.Y, L)[0]
        A1 = torch.solve(A0, L.T)[0].to(dtype=K_st.dtype)
        fmean = torch.mm(K_st, A1)

        # Compute predictive variance
        L = L.to(dtype=K_ts.dtype)
        V = torch.solve(K_ts, L)[0]
        if noisy_targets:
            K_ss += torch.eye(Xnew.shape[0], dtype=self.X.dtype, device=self.X.device) * self.sn2
        fvar = K_ss - torch.mm(V.T, V)

        return fmean, fvar
//  [python]
def predict_f(self, Xnew, noisy_targets=False):
        """
        Compute the predictive mean vector and covariance matrix for the specified inputs.

        :param Xnew: torch.Tensor, test inputs at which to perform predictions
        :param noisy_targets: bool, specify if predictive covariances are computed for noisy targets
        :return: tuple, predictive mean, predictive covariance matrix
        """
        # Throw exception if data not assigned
        if not self.data_assigned:
            raise Exception('Assign data first')

        # Ensure devices match
        Xnew = Xnew.to(self.X.device)

        # Compute covariance matrices for test/training inputs (s for test, t for train)
        K_st = self.kern.K(Xnew, self.X)
        K_ts = K_st.T
        K_ss = self.kern.K(Xnew)
        K_tt = self.kern.K(self.X) + torch.eye(self.X.shape[0], dtype=self.X.dtype, device=self.X.device) * self.sn2

        # Compute predictive mean
        L = self.cholesky_factor(K_tt, jitter_level=self.jitter).to(dtype=self.Y.dtype)
        A0 = torch.linalg.solve(L, self.Y)
        A1 = torch.linalg.solve(L.T, A0).to(dtype=K_st.dtype)
        fmean = torch.mm(K_st, A1)

        # Compute predictive variance
        L = L.to(dtype=K_ts.dtype)
        V = torch.linalg.solve(L, K_ts)
        if noisy_targets:
            K_ss += torch.eye(Xnew.shape[0], dtype=self.X.dtype, device=self.X.device) * self.sn2
        fvar = K_ss - torch.mm(V.T, V)

        return fmean, fvar
//  [python]
def ridge_regression(X, Y, lamda=0, bias=True):
    if bias:
        X = torch.cat([torch.ones(*X.shape[:-1], 1, device=X.device), X], dim=-1)
    X_t = X.transpose(-2, -1)

    n_samples, n_dim = X.shape[-2:]
    try:
        if n_samples >= n_dim:
            # standard
            A = X_t @ X
            A.diagonal(dim1=-2, dim2=-1).add_(lamda)
            B = X_t @ Y
            weights = torch.solve(B, A)[0]
        else:
            # Woodbury
            A = X @ X_t
            A.diagonal(dim1=-2, dim2=-1).add_(lamda)
            weights = X_t @ torch.solve(Y, A)[0]
    except RuntimeError:
        print("The matrix is singular. Using pseudoinverse to solve.")
        assert lamda == 0
        weights = torch.lstsq(Y, X)[0]
    return weights
//  [python]
def ridge_regression(X, Y, lamda=0, bias=True):
    if bias:
        X = torch.cat([torch.ones(*X.shape[:-1], 1, device=X.device), X], dim=-1)
    X_t = X.transpose(-2, -1)

    n_samples, n_dim = X.shape[-2:]
    try:
        if n_samples >= n_dim:
            # standard
            A = X_t @ X
            A.diagonal(dim1=-2, dim2=-1).add_(lamda)
            B = X_t @ Y
            weights = torch.linalg.solve(A, B)
        else:
            # Woodbury
            A = X @ X_t
            A.diagonal(dim1=-2, dim2=-1).add_(lamda)
            weights = X_t @ torch.linalg.solve(A, Y)
    except torch.linalg.LinAlgError:
        print("The matrix is singular. Using pseudoinverse to solve.")
        assert lamda == 0
        weights = torch.linalg.lstsq(X, Y).solution
    return weights
//  [python]
def test_solver_gradient_j4t(self):
        # Sparse solver:
        As1 = self.A_csr.detach().to(torch.float32).clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().to(torch.float32).clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = tsgujax.sparse_solve_j4t(As1, Bd1)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.A.detach().to(torch.float32).clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().to(torch.float32).clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.solve(Bd2, Ad2)[0]
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_solver_gradient_j4t(self):
        # Sparse solver:
        As1 = self.A_csr.detach().to(torch.float32).clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().to(torch.float32).clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = tsgujax.sparse_solve_j4t(As1, Bd1)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.A.detach().to(torch.float32).clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().to(torch.float32).clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.solve(Ad2, Bd2)
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def auto_diff(seed):
    torch.manual_seed(seed)
    theta = torch.randn((32,32), requires_grad=True)
    K = theta * 2
    x = torch.randn((32))
    b = K @ x
    b = b.clone().detach().requires_grad_(True)

    xhat = torch.solve(b, K)[0]

    loss = xhat.mean()
    loss.backward()

    return theta.grad, b.grad
//  [python]
def auto_diff(seed):
    torch.manual_seed(seed)
    theta = torch.randn((32,32), requires_grad=True)
    K = theta * 2
    x = torch.randn((32))
    b = K @ x
    b = b.clone().detach().requires_grad_(True)

    xhat = torch.linalg.solve(K, b)

    loss = xhat.mean()
    loss.backward()

    return theta.grad, b.grad
//  [python]
def _compute_acv_estimator_variance(self, covariance, ratios, N):
        big_C = covariance[:, 1:, 1:]
        c_bar = covariance[:, 0, 1:] \
            / torch.sqrt(covariance[:, 0, 0]).unsqueeze(1)

        F, F0 = self._compute_acv_F_and_F0(ratios)
        a = (F0.unsqueeze(0) * c_bar)

        alpha = torch.solve(a.unsqueeze(2), big_C * F)[0]
        R_squared = (a.unsqueeze(2) * alpha).sum(1).sum(1)
        variance = covariance[:, 0, 0] / N * (1 - R_squared)

        return variance
//  [python]
def _compute_acv_estimator_variance(self, covariance, ratios, N):
        big_C = covariance[:, 1:, 1:]
        c_bar = covariance[:, 0, 1:] \
            / torch.sqrt(covariance[:, 0, 0]).unsqueeze(1)

        F, F0 = self._compute_acv_F_and_F0(ratios)
        a = (F0.unsqueeze(0) * c_bar)

        alpha = torch.linalg.solve(big_C * F, a.unsqueeze(2))
        R_squared = (a.unsqueeze(2) * alpha).sum(1).sum(1)
        variance = covariance[:, 0, 0] / N * (1 - R_squared)

        return variance
//  [python]
def background_update(W, H, C_XX, C_XbarX, eps=1e-5):
    """
    Recomputes J based on W, H, and C_XX = E[X X^H] and C_XbarX = E[ X_bar X^H ]
    """
    n_src, n_freq, n_chan = W.shape[-3:]

    A1 = torch.einsum("...sfc,...cfd->...fsd", W, C_XX)
    A2 = torch.einsum("...sfdt,...dtfc->...fsc", H, C_XbarX)
    A = A1 + A2  # (..., n_freq, n_src, n_chan)

    mat = A[..., :n_src]
    rhs = A[..., n_src:]

    with torch.no_grad():
        # normalize the rows of A without changing the solution
        # for numerical stability
        norm = torch.linalg.norm(mat.detach(), dim=-1, keepdim=True)
        weights = 1.0 / torch.clamp(norm, min=eps)

    load = eps * torch.eye(n_src).type_as(mat)

    # make eigenvalues positive and scale
    mat2 = mat * weights
    rhs = rhs * weights
    mat = torch.einsum("...km,...kn->...mn", mat2.conj(), mat2)
    rhs = torch.einsum("...km,...kn->...mn", mat2.conj(), rhs)

    J_H = torch.solve(rhs, mat + load)[0]
    J = J_H.conj().moveaxis(-1, -3)  # (..., n_chan - n_src, n_freq, n_src)

    return J
//  [python]
def background_update(W, H, C_XX, C_XbarX, eps=1e-5):
    """
    Recomputes J based on W, H, and C_XX = E[X X^H] and C_XbarX = E[ X_bar X^H ]
    """
    n_src, n_freq, n_chan = W.shape[-3:]

    A1 = torch.einsum("...sfc,...cfd->...fsd", W, C_XX)
    A2 = torch.einsum("...sfdt,...dtfc->...fsc", H, C_XbarX)
    A = A1 + A2  # (..., n_freq, n_src, n_chan)

    mat = A[..., :n_src]
    rhs = A[..., n_src:]

    with torch.no_grad():
        # normalize the rows of A without changing the solution
        # for numerical stability
        norm = torch.linalg.norm(mat.detach(), dim=-1, keepdim=True)
        weights = 1.0 / torch.clamp(norm, min=eps)

    load = eps * torch.eye(n_src).type_as(mat)

    # make eigenvalues positive and scale
    mat2 = mat * weights
    rhs = rhs * weights
    mat = torch.einsum("...km,...kn->...mn", mat2.conj(), mat2)
    rhs = torch.einsum("...km,...kn->...mn", mat2.conj(), rhs)

    J_H = torch.linalg.solve(mat + load, rhs)
    J = J_H.conj().moveaxis(-1, -3)  # (..., n_chan - n_src, n_freq, n_src)

    return J
//  [python]
def calculate_composition_weights(
    datasets: Union[Dataset, List[Dataset]], property: str
) -> Tuple[torch.Tensor, List[int]]:
    """Calculate the composition weights for a dataset.

    For now, it assumes per-system properties.

    :param dataset: Dataset to calculate the composition weights for.
    :returns: Composition weights for the dataset, as well as the
        list of species that the weights correspond to.
    """
    if not isinstance(datasets, list):
        datasets = [datasets]

    species = get_all_species(datasets)
    # note that this is sorted, and the composition weights are sorted
    # as well, because the species are sorted in the composition features

    targets = torch.stack(
        [
            sample._asdict()[property].block().values
            for dataset in datasets
            for sample in dataset
        ]
    )
    targets = targets.squeeze(dim=(1, 2))  # remove component and property dimensions

    structure_list = [
        sample._asdict()["system"] for dataset in datasets for sample in dataset
    ]

    dtype = structure_list[0].positions.dtype
    composition_features = torch.empty((len(structure_list), len(species)), dtype=dtype)
    for i, structure in enumerate(structure_list):
        for j, s in enumerate(species):
            composition_features[i, j] = torch.sum(structure.types == s)

    regularizer = 1e-20
    while regularizer:
        if regularizer > 1e5:
            raise RuntimeError(
                "Failed to solve the linear system to calculate the "
                "composition weights. The dataset is probably too small "
                "or ill-conditioned."
            )
        try:
            solution = torch.solve(
                composition_features.T @ targets,
                composition_features.T @ composition_features
                + regularizer
                * torch.eye(
                    composition_features.shape[1],
                    dtype=composition_features.dtype,
                    device=composition_features.device,
                ),
            )[0]
            break
        except RuntimeError:
            regularizer *= 10.0

    return solution, species