jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def iter_sample(self, noiseless=True):
        r"""
        Iteratively constructs a sample from the Gaussian Process posterior.

        Recall that at test input points :math:`X_{new}`, the posterior is
        multivariate Gaussian distributed with mean and covariance matrix
        given by :func:`forward`.

        This method samples lazily from this multivariate Gaussian. The advantage
        of this approach is that later query points can depend upon earlier ones.
        Particularly useful when the querying is to be done by an optimisation
        routine.

        .. note:: The noise parameter ``noise`` (:math:`\epsilon`) together with
            kernel's parameters have been learned from a training procedure (MCMC or
            SVI).

        :param bool noiseless: A flag to decide if we want to add sampling noise
            to the samples beyond the noise inherent in the GP posterior.
        :returns: sampler
        :rtype: function
        """
        noise = self.noise.detach()
        X = self.X.clone().detach()
        y = self.y.clone().detach()
        N = X.size(0)
        Kff = self.kernel(X).contiguous()
        Kff.view(-1)[:: N + 1] += noise  # add noise to the diagonal

        outside_vars = {"X": X, "y": y, "N": N, "Kff": Kff}

        def sample_next(xnew, outside_vars):
            """Repeatedly samples from the Gaussian process posterior,
            conditioning on previously sampled values.
            """
            warn_if_nan(xnew)

            # Variables from outer scope
            X, y, Kff = outside_vars["X"], outside_vars["y"], outside_vars["Kff"]

            # Compute Cholesky decomposition of kernel matrix
            Lff = torch.linalg.cholesky(Kff)
            y_residual = y - self.mean_function(X)

            # Compute conditional mean and variance
            loc, cov = conditional(
                xnew, X, self.kernel, y_residual, None, Lff, False, jitter=self.jitter
            )
            if not noiseless:
                cov = cov + noise

            ynew = torchdist.Normal(
                loc + self.mean_function(xnew), cov.sqrt()
            ).rsample()

            # Update kernel matrix
            N = outside_vars["N"]
            Kffnew = Kff.new_empty(N + 1, N + 1)
            Kffnew[:N, :N] = Kff
            cross = self.kernel(X, xnew).squeeze()
            end = self.kernel(xnew, xnew).squeeze()
            Kffnew[N, :N] = cross
            Kffnew[:N, N] = cross
            # No noise, just jitter for numerical stability
            Kffnew[N, N] = end + self.jitter
            # Heuristic to avoid adding degenerate points
            if Kffnew.logdet() > -15.0:
                outside_vars["Kff"] = Kffnew
                outside_vars["N"] += 1
                outside_vars["X"] = torch.cat((X, xnew))
                outside_vars["y"] = torch.cat((y, ynew))

            return ynew

        return lambda xnew: sample_next(xnew, outside_vars)
//  [python]
def forward(self, Xnew, full_cov=False, noiseless=True):
        r"""
        Computes the mean and covariance matrix (or variance) of Gaussian Process
        posterior on a test input data :math:`X_{new}`:

        .. math:: p(f^* \mid X_{new}, X, y, k, \epsilon) = \mathcal{N}(loc, cov).

        .. note:: The noise parameter ``noise`` (:math:`\epsilon`) together with
            kernel's parameters have been learned from a training procedure (MCMC or
            SVI).

        :param torch.Tensor Xnew: A input data for testing. Note that
            ``Xnew.shape[1:]`` must be the same as ``self.X.shape[1:]``.
        :param bool full_cov: A flag to decide if we want to predict full covariance
            matrix or just variance.
        :param bool noiseless: A flag to decide if we want to include noise in the
            prediction output or not.
        :returns: loc and covariance matrix (or variance) of :math:`p(f^*(X_{new}))`
        :rtype: tuple(torch.Tensor, torch.Tensor)
        """
        self._check_Xnew_shape(Xnew)
        self.set_mode("guide")

        N = self.X.size(0)
        Kff = self.kernel(self.X).contiguous()
        Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to the diagonal
        Lff = torch.cholesky(Kff)

        y_residual = self.y - self.mean_function(self.X)
        loc, cov = conditional(
            Xnew,
            self.X,
            self.kernel,
            y_residual,
            None,
            Lff,
            full_cov,
            jitter=self.jitter,
        )

        if full_cov and not noiseless:
            M = Xnew.size(0)
            cov = cov.contiguous()
            cov.view(-1, M * M)[:, :: M + 1] += self.noise  # add noise to the diagonal
        if not full_cov and not noiseless:
            cov = cov + self.noise.abs()

        return loc + self.mean_function(Xnew), cov
//  [python]
def forward(self, Xnew, full_cov=False, noiseless=True):
        r"""
        Computes the mean and covariance matrix (or variance) of Gaussian Process
        posterior on a test input data :math:`X_{new}`:

        .. math:: p(f^* \mid X_{new}, X, y, k, \epsilon) = \mathcal{N}(loc, cov).

        .. note:: The noise parameter ``noise`` (:math:`\epsilon`) together with
            kernel's parameters have been learned from a training procedure (MCMC or
            SVI).

        :param torch.Tensor Xnew: A input data for testing. Note that
            ``Xnew.shape[1:]`` must be the same as ``self.X.shape[1:]``.
        :param bool full_cov: A flag to decide if we want to predict full covariance
            matrix or just variance.
        :param bool noiseless: A flag to decide if we want to include noise in the
            prediction output or not.
        :returns: loc and covariance matrix (or variance) of :math:`p(f^*(X_{new}))`
        :rtype: tuple(torch.Tensor, torch.Tensor)
        """
        self._check_Xnew_shape(Xnew)
        self.set_mode("guide")

        N = self.X.size(0)
        Kff = self.kernel(self.X).contiguous()
        Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to the diagonal
        Lff = torch.linalg.cholesky(Kff)

        y_residual = self.y - self.mean_function(self.X)
        loc, cov = conditional(
            Xnew,
            self.X,
            self.kernel,
            y_residual,
            None,
            Lff,
            full_cov,
            jitter=self.jitter,
        )

        if full_cov and not noiseless:
            M = Xnew.size(0)
            cov = cov.contiguous()
            cov.view(-1, M * M)[:, :: M + 1] += self.noise  # add noise to the diagonal
        if not full_cov and not noiseless:
            cov = cov + self.noise.abs()

        return loc + self.mean_function(Xnew), cov
//  [python]
def decompose_kernels(gpr,pred_matrix,kernels):
    '''
    A function to calculate different kernels contribution to final prediction

    ------------------Inputs----------------------------
    gpr: an optimized pyro GP regression model
    pred_matrix: a torch tensor of prediction matrix containing the data points for prediction
    kernels: a list of pyro kernels for decomposition

    ------------------Outputs---------------------------
    output: a list of tuples, each tuple contains the mean and covariance of the prediction for each kernel 
    '''
    N = len(gpr.X)
    M = pred_matrix.size(0)
    f_loc = gpr.y - gpr.mean_function(gpr.X)
    latent_shape = f_loc.shape[:-1]
    loc_shape = latent_shape + (M,)
    f_loc = f_loc.permute(-1, *range(len(latent_shape)))
    f_loc_2D = f_loc.reshape(N, -1)
    loc_shape = latent_shape + (M,)
    v_2D = f_loc_2D
    Kff = gpr.kernel(gpr.X).contiguous()
    Kff +=gpr.noise
    Kff.view(-1)[:: N + 1] += gpr.jitter    # add noise to the diagonal
    Lff = torch.cholesky(Kff)
    
    output = []
    for kernel in kernels:
        Kfs = kernel(gpr.X, pred_matrix)
        pack = torch.cat((f_loc_2D, Kfs), dim=1)
        Lffinv_pack = pack.triangular_solve(Lff, upper=False)[0]
        W = Lffinv_pack[:, f_loc_2D.size(1):f_loc_2D.size(1) + M].t()
        Qss = W.matmul(W.t())
        v_2D = Lffinv_pack[:, :f_loc_2D.size(1)]
        loc = W.matmul(v_2D).t().reshape(loc_shape)
        Kss = kernel(pred_matrix)
        cov = Kss - Qss
        output.append((loc, cov))
        
    return output
//  [python]
def decompose_kernels(gpr,pred_matrix,kernels):
    '''
    A function to calculate different kernels contribution to final prediction

    ------------------Inputs----------------------------
    gpr: an optimized pyro GP regression model
    pred_matrix: a torch tensor of prediction matrix containing the data points for prediction
    kernels: a list of pyro kernels for decomposition

    ------------------Outputs---------------------------
    output: a list of tuples, each tuple contains the mean and covariance of the prediction for each kernel 
    '''
    N = len(gpr.X)
    M = pred_matrix.size(0)
    f_loc = gpr.y - gpr.mean_function(gpr.X)
    latent_shape = f_loc.shape[:-1]
    loc_shape = latent_shape + (M,)
    f_loc = f_loc.permute(-1, *range(len(latent_shape)))
    f_loc_2D = f_loc.reshape(N, -1)
    loc_shape = latent_shape + (M,)
    v_2D = f_loc_2D
    Kff = gpr.kernel(gpr.X).contiguous()
    Kff +=gpr.noise
    Kff.view(-1)[:: N + 1] += gpr.jitter    # add noise to the diagonal
    Lff = torch.linalg.cholesky(Kff)
    
    output = []
    for kernel in kernels:
        Kfs = kernel(gpr.X, pred_matrix)
        pack = torch.cat((f_loc_2D, Kfs), dim=1)
        Lffinv_pack = pack.triangular_solve(Lff, upper=False)[0]
        W = Lffinv_pack[:, f_loc_2D.size(1):f_loc_2D.size(1) + M].t()
        Qss = W.matmul(W.t())
        v_2D = Lffinv_pack[:, :f_loc_2D.size(1)]
        loc = W.matmul(v_2D).t().reshape(loc_shape)
        Kss = kernel(pred_matrix)
        cov = Kss - Qss
        output.append((loc, cov))
        
    return output
//  [python]
Xnew = torch.tensor([[2.0, 3.0], [4.0, 6.0]])
X = torch.tensor([[1.0, 5.0], [2.0, 1.0], [3.0, 2.0]])
kernel = Matern52(input_dim=2)
Kff = kernel(X) + torch.eye(3) * 1e-6
Lff = torch.cholesky(Kff)
//  [python]
Xnew = torch.tensor([[2.0, 3.0], [4.0, 6.0]])
X = torch.tensor([[1.0, 5.0], [2.0, 1.0], [3.0, 2.0]])
kernel = Matern52(input_dim=2)
Kff = kernel(X) + torch.eye(3) * 1e-6
Lff = torch.linalg.cholesky(Kff)
//  [python]
def sample_inputs_linalg_cholesky_inverse(op_info, device, dtype, requires_grad=False, **kwargs):
    from torch.testing._internal.common_utils import random_well_conditioned_matrix

    # Cholesky factorization is for positive-definite matrices
    single_well_conditioned_matrix = random_well_conditioned_matrix(S, S, dtype=dtype, device=device)
    batch_well_conditioned_matrices = random_well_conditioned_matrix(2, S, S, dtype=dtype, device=device)
    single_pd = single_well_conditioned_matrix @ single_well_conditioned_matrix.mH
    batch_pd = batch_well_conditioned_matrices @ batch_well_conditioned_matrices.mH

    inputs = (
        torch.zeros(0, 0, dtype=dtype, device=device),  # 0x0 matrix
        torch.zeros(0, 2, 2, dtype=dtype, device=device),  # zero batch of matrices
        single_pd,
        batch_pd
    )
    test_cases = (torch.cholesky(a, upper=False) for a in inputs)
    for l in test_cases:
        # generated lower-triangular samples
        l.requires_grad = requires_grad
        yield SampleInput(l)  # upper=False by default
        yield SampleInput(l.detach().clone().requires_grad_(requires_grad), kwargs=dict(upper=False))

        # generate upper-triangular inputs
        u = l.detach().clone().mT.contiguous().requires_grad_(requires_grad)
        yield SampleInput(u, kwargs=dict(upper=True))
//  [python]
def sample_inputs_linalg_cholesky_inverse(op_info, device, dtype, requires_grad=False, **kwargs):
    from torch.testing._internal.common_utils import random_well_conditioned_matrix

    # Cholesky factorization is for positive-definite matrices
    single_well_conditioned_matrix = random_well_conditioned_matrix(S, S, dtype=dtype, device=device)
    batch_well_conditioned_matrices = random_well_conditioned_matrix(2, S, S, dtype=dtype, device=device)
    single_pd = single_well_conditioned_matrix @ single_well_conditioned_matrix.mH
    batch_pd = batch_well_conditioned_matrices @ batch_well_conditioned_matrices.mH

    inputs = (
        torch.zeros(0, 0, dtype=dtype, device=device),  # 0x0 matrix
        torch.zeros(0, 2, 2, dtype=dtype, device=device),  # zero batch of matrices
        single_pd,
        batch_pd
    )
    test_cases = (torch.linalg.cholesky(a, upper=False) for a in inputs)
    for l in test_cases:
        # generated lower-triangular samples
        l.requires_grad = requires_grad
        yield SampleInput(l)  # upper=False by default
        yield SampleInput(l.detach().clone().requires_grad_(requires_grad), kwargs=dict(upper=False))

        # generate upper-triangular inputs
        u = l.detach().clone().mT.contiguous().requires_grad_(requires_grad)
        yield SampleInput(u, kwargs=dict(upper=True))
//  [python]
def test_cholesky_autograd(self, device, dtype):
        def func(root):
            x = 0.5 * (root + root.transpose(-1, -2).conj())
            return torch.cholesky(x)

        def run_test(shape):
            root = torch.rand(*shape, dtype=dtype, device=device, requires_grad=True)
            root = root + torch.eye(shape[-1], dtype=dtype, device=device)

            gradcheck(func, root)
            gradgradcheck(func, root)

            root = torch.rand(*shape, dtype=dtype, device=device)
            root = torch.matmul(root, root.transpose(-1, -2).conj())
            root.requires_grad_()
            chol = torch.cholesky(root).sum().backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2).conj())  # Check the gradient is hermitian

        shapes = ((3, 3), (4, 3, 2, 2))
        for shape in shapes:
            run_test(shape)
//  [python]
def test_cholesky_autograd(self, device, dtype):
        def func(root):
            x = 0.5 * (root + root.transpose(-1, -2).conj())
            return torch.linalg.cholesky(x)

        def run_test(shape):
            root = torch.rand(*shape, dtype=dtype, device=device, requires_grad=True)
            root = root + torch.eye(shape[-1], dtype=dtype, device=device)

            gradcheck(func, root)
            gradgradcheck(func, root)

            root = torch.rand(*shape, dtype=dtype, device=device)
            root = torch.matmul(root, root.transpose(-1, -2).conj())
            root.requires_grad_()
            chol = torch.linalg.cholesky(root).sum().backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2).conj())  # Check the gradient is hermitian

        shapes = ((3, 3), (4, 3, 2, 2))
        for shape in shapes:
            run_test(shape)
//  [python]
def reverse_propogator(pastdata=None, futuredata=None, epsilon=1e-6, method='cutoff', device=None):

    """
    Compute the reverse propogator for the time-lagged sequence data with detailed balance contraint.
    This function is used to calculate the loss function for SRVNet.
    cholesky decomposition of reversed instantaneous correlation matrix: LL^{T} = (c_{00} + c_{tt}}) / 2
    propogator = L ((c_{0t} + c_{t0}}) / 2) L^{T}

    Parameters
    ----------
    pastdata: torch.Tensor, default=None
        The past part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    futuredata: torch.Tensor, default=None
        The future part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    epsilon: float, default=1e-6
        The epsilon modification used to calculate the inverse of the instantaneous matrix.
    method: string, default='cutoff'
        The method used to calculate the inverse of the instantaneous matix: : 'regularize', 'cutoff', 'clamp'
    
    Returns
    ----------
    _reverse_propogator: torch.Tensor, shape:(num_features, num_features)
        Reverse propogator matrix for the input time-lagged sequence data.
    """

    c00, c0t, ctt = covariance(pastdata=pastdata, futuredata=futuredata, remove_mean=True, device=device)
    _, ct0, _ = covariance(pastdata=futuredata, futuredata=pastdata, remove_mean=True,  device=device)
    
    c0 = (c00 + ctt) / 2
    ct = (c0t + ct0) / 2

    _lower_triangle = torch.cholesky(c0)
    _lower_triangle_inv = torch.inverse(_lower_triangle)
    _reverse_propogator = torch.matmul(_lower_triangle_inv, torch.matmul(ct, _lower_triangle_inv.t()))

    return _reverse_propogator
//  [python]
def reverse_propogator(pastdata=None, futuredata=None, epsilon=1e-6, method='cutoff', device=None):

    """
    Compute the reverse propogator for the time-lagged sequence data with detailed balance contraint.
    This function is used to calculate the loss function for SRVNet.
    cholesky decomposition of reversed instantaneous correlation matrix: LL^{T} = (c_{00} + c_{tt}}) / 2
    propogator = L ((c_{0t} + c_{t0}}) / 2) L^{T}

    Parameters
    ----------
    pastdata: torch.Tensor, default=None
        The past part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    futuredata: torch.Tensor, default=None
        The future part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    epsilon: float, default=1e-6
        The epsilon modification used to calculate the inverse of the instantaneous matrix.
    method: string, default='cutoff'
        The method used to calculate the inverse of the instantaneous matix: : 'regularize', 'cutoff', 'clamp'
    
    Returns
    ----------
    _reverse_propogator: torch.Tensor, shape:(num_features, num_features)
        Reverse propogator matrix for the input time-lagged sequence data.
    """

    c00, c0t, ctt = covariance(pastdata=pastdata, futuredata=futuredata, remove_mean=True, device=device)
    _, ct0, _ = covariance(pastdata=futuredata, futuredata=pastdata, remove_mean=True,  device=device)
    
    c0 = (c00 + ctt) / 2
    ct = (c0t + ct0) / 2

    _lower_triangle = torch.linalg.cholesky(c0)
    _lower_triangle_inv = torch.inverse(_lower_triangle)
    _reverse_propogator = torch.matmul(_lower_triangle_inv, torch.matmul(ct, _lower_triangle_inv.t()))

    return _reverse_propogator
//  [python]
def test_conditional_whiten(Xnew, X, kernel, f_loc, f_scale_tril, loc, cov):
    if f_scale_tril is None:
        return

    loc0, cov0 = conditional(
        Xnew, X, kernel, f_loc, f_scale_tril, full_cov=True, whiten=False
    )
    Kff = kernel(X) + torch.eye(3) * 1e-6
    Lff = torch.cholesky(Kff)
    whiten_f_loc = torch.cholesky_solve(f_loc, Lff)
    whiten_f_scale_tril = torch.cholesky_solve(f_scale_tril, Lff)
    loc1, cov1 = conditional(
        Xnew, X, kernel, whiten_f_loc, whiten_f_scale_tril, full_cov=True, whiten=True
    )

    assert_equal(loc0, loc1)
    assert_equal(cov0, cov1)
//  [python]
def test_conditional_whiten(Xnew, X, kernel, f_loc, f_scale_tril, loc, cov):
    if f_scale_tril is None:
        return

    loc0, cov0 = conditional(
        Xnew, X, kernel, f_loc, f_scale_tril, full_cov=True, whiten=False
    )
    Kff = kernel(X) + torch.eye(3) * 1e-6
    Lff = torch.linalg.cholesky(Kff)
    whiten_f_loc = Lff.inverse().matmul(f_loc)
    whiten_f_scale_tril = Lff.inverse().matmul(f_scale_tril)
    loc1, cov1 = conditional(
        Xnew, X, kernel, whiten_f_loc, whiten_f_scale_tril, full_cov=True, whiten=True
    )

    assert_equal(loc0, loc1)
    assert_equal(cov0, cov1)
//  [python]
def fit(self, dataset: Tensor) -> None:
        """Fit a KDE model to the input dataset.
        Args:
          dataset (Tensor): Input dataset.
        Returns:
            None
        """

        num_samples, dimension = dataset.shape

        # compute scott's bandwidth factor
        factor = num_samples ** (-1 / (dimension + 4))

        cov_mat = self.cov(dataset.T)
        inv_cov_mat = torch.linalg.inv(cov_mat)
        inv_cov = inv_cov_mat / factor**2

        # transform data to account for bandwidth
        bw_transform = torch.cholesky(inv_cov)
        dataset = torch.matmul(dataset, bw_transform)

        #
        norm = torch.prod(torch.diag(bw_transform))
        norm *= math.pow((2 * math.pi), (-dimension / 2))

        self.bw_transform = bw_transform
        self.dataset = dataset
        self.norm = norm
//  [python]
def fit(self, dataset: Tensor) -> None:
        """Fit a KDE model to the input dataset.
        Args:
          dataset (Tensor): Input dataset.
        Returns:
            None
        """

        num_samples, dimension = dataset.shape

        # compute scott's bandwidth factor
        factor = num_samples ** (-1 / (dimension + 4))

        cov_mat = self.cov(dataset.T)
        inv_cov_mat = torch.linalg.inv(cov_mat)
        inv_cov = inv_cov_mat / factor**2

        # transform data to account for bandwidth
        bw_transform = torch.linalg.cholesky(inv_cov)
        dataset = torch.matmul(dataset, bw_transform)

        #
        norm = torch.prod(torch.diag(bw_transform))
        norm *= math.pow((2 * math.pi), (-dimension / 2))

        self.bw_transform = bw_transform
        self.dataset = dataset
        self.norm = norm
//  [python]
def normal_samples(mean, var, n_samples, generator=None):
    """Produce samples from a batch of Normal distributions either parameterized
    by a diagonal or full covariance given by `var`.

    Parameters
    ----------
    mean : torch.Tensor
        `(batch_size, output_dim)`
    var : torch.Tensor
        (co)variance of the Normal distribution
        `(batch_size, output_dim, output_dim)` or `(batch_size, output_dim)`
    generator : torch.Generator
        random number generator
    """
    assert mean.ndim == 2, 'Invalid input shape of mean, should be 2-dimensional.'
    _, output_dim = mean.shape
    randn_samples = torch.randn((output_dim, n_samples), device=mean.device, 
                                dtype=mean.dtype, generator=generator)
    
    if mean.shape == var.shape:
        # diagonal covariance
        scaled_samples = var.sqrt().unsqueeze(-1) * randn_samples.unsqueeze(0)
        return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
    elif mean.shape == var.shape[:2] and var.shape[-1] == mean.shape[1]:
        # full covariance
        scale = torch.cholesky(var)
        scaled_samples = torch.matmul(scale, randn_samples.unsqueeze(0))  # expand batch dim
        return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
    else:
        raise ValueError('Invalid input shapes.')
//  [python]
def normal_samples(mean, var, n_samples, generator=None):
    """Produce samples from a batch of Normal distributions either parameterized
    by a diagonal or full covariance given by `var`.

    Parameters
    ----------
    mean : torch.Tensor
        `(batch_size, output_dim)`
    var : torch.Tensor
        (co)variance of the Normal distribution
        `(batch_size, output_dim, output_dim)` or `(batch_size, output_dim)`
    generator : torch.Generator
        random number generator
    """
    assert mean.ndim == 2, 'Invalid input shape of mean, should be 2-dimensional.'
    _, output_dim = mean.shape
    randn_samples = torch.randn((output_dim, n_samples), device=mean.device, 
                                dtype=mean.dtype, generator=generator)
    
    if mean.shape == var.shape:
        # diagonal covariance
        scaled_samples = var.sqrt().unsqueeze(-1) * randn_samples.unsqueeze(0)
        return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
    elif mean.shape == var.shape[:2] and var.shape[-1] == mean.shape[1]:
        # full covariance
        scale = torch.linalg.cholesky(var)
        scaled_samples = torch.matmul(scale, randn_samples.unsqueeze(0))  # expand batch dim
        return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
    else:
        raise ValueError('Invalid input shapes.')
//  [python]
def test_cholesky_errors_and_warnings(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        # cholesky requires the input to be a square matrix or batch of square matrices
        A = torch.randn(2, 3, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, r'must be batches of square matrices'):
            torch.cholesky(A)
        A = torch.randn(2, 2, 3, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, r'must be batches of square matrices'):
            torch.cholesky(A)
        with self.assertRaisesRegex(np.linalg.LinAlgError, r'Last 2 dimensions of the array must be square'):
            np.linalg.cholesky(A.cpu().numpy())

        # cholesky requires the input to be at least 2 dimensional tensor
        A = torch.randn(2, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, r'must have at least 2 dimensions'):
            torch.cholesky(A)
        with self.assertRaisesRegex(np.linalg.LinAlgError,
                                    r'1-dimensional array given\. Array must be at least two-dimensional'):
            np.linalg.cholesky(A.cpu().numpy())

        # if the input matrix is singular, an error should be raised
        A = torch.eye(3, 3, dtype=dtype, device=device)
        A[-1, -1] = 0  # Now A is singular
        with self.assertRaisesRegex(RuntimeError, r'U\(3,3\) is zero, singular U\.'):
            torch.cholesky(A)
        with self.assertRaisesRegex(np.linalg.LinAlgError, r'Matrix is not positive definite'):
            np.linalg.cholesky(A.cpu().numpy())

        # if at least one matrix in the batch is singular, an error should be raised
        A = torch.eye(3, 3, dtype=dtype, device=device)
        A = A.reshape((1, 3, 3))
        A = A.repeat(5, 1, 1)
        A[4, -1, -1] = 0  # Now A[4] is singular
        with self.assertRaisesRegex(RuntimeError, r'For batch 4: U\(3,3\) is zero, singular U\.'):
            torch.cholesky(A)

        # if out tensor with wrong shape is passed a warning is given
        A = random_hermitian_pd_matrix(3, dtype=dtype, device=device)
        out = torch.empty(2, 3, dtype=dtype, device=device)
        with warnings.catch_warnings(record=True) as w:
            # Trigger warning
            torch.cholesky(A, out=out)
            # Check warning occurs
            self.assertEqual(len(w), 1)
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))

        # dtypes should match
        out = torch.empty_like(A).to(torch.int)
        with self.assertRaisesRegex(RuntimeError, "result dtype Int does not match self dtype"):
            torch.cholesky(A, out=out)
//  [python]
def test_cholesky_errors_and_warnings(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        # cholesky requires the input to be a square matrix or batch of square matrices
        A = torch.randn(2, 3, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, r'must be batches of square matrices'):
            torch.linalg.cholesky(A)
        A = torch.randn(2, 2, 3, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, r'must be batches of square matrices'):
            torch.linalg.cholesky(A)
        with self.assertRaisesRegex(np.linalg.LinAlgError, r'Last 2 dimensions of the array must be square'):
            np.linalg.cholesky(A.cpu().numpy())

        # cholesky requires the input to be at least 2 dimensional tensor
        A = torch.randn(2, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, r'must have at least 2 dimensions'):
            torch.linalg.cholesky(A)
        with self.assertRaisesRegex(np.linalg.LinAlgError,
                                    r'1-dimensional array given\. Array must be at least two-dimensional'):
            np.linalg.cholesky(A.cpu().numpy())

        # if the input matrix is singular, an error should be raised
        A = torch.eye(3, 3, dtype=dtype, device=device)
        A[-1, -1] = 0  # Now A is singular
        with self.assertRaisesRegex(RuntimeError, r'U\(3,3\) is zero, singular U\.'):
            torch.linalg.cholesky(A)
        with self.assertRaisesRegex(np.linalg.LinAlgError, r'Matrix is not positive definite'):
            np.linalg.cholesky(A.cpu().numpy())

        # if at least one matrix in the batch is singular, an error should be raised
        A = torch.eye(3, 3, dtype=dtype, device=device)
        A = A.reshape((1, 3, 3))
        A = A.repeat(5, 1, 1)
        A[4, -1, -1] = 0  # Now A[4] is singular
        with self.assertRaisesRegex(RuntimeError, r'For batch 4: U\(3,3\) is zero, singular U\.'):
            torch.linalg.cholesky(A)

        # if out tensor with wrong shape is passed a warning is given
        A = random_hermitian_pd_matrix(3, dtype=dtype, device=device)
        out = torch.empty(2, 3, dtype=dtype, device=device)
        with warnings.catch_warnings(record=True) as w:
            # Trigger warning
            torch.linalg.cholesky(A, out=out)
            # Check warning occurs
            self.assertEqual(len(w), 1)
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))

        # dtypes should match
        out = torch.empty_like(A).to(torch.int)
        with self.assertRaisesRegex(RuntimeError, "result dtype Int does not match self dtype"):
            torch.linalg.cholesky(A, out=out)
//  [python]
def torch_reduction_matrix(sigma: Tensor, reduction_method="pseudo"):
    import torch

    if reduction_method == "cholesky":
        C = torch.cholesky(sigma)
        return torch.inverse(C.t())
    elif reduction_method == "svd":
        u, s, _ = torch.linalg.svd(sigma)
        return u @ torch.diag(torch.sqrt(1 / s))
    elif reduction_method == "pseudo" or reduction_method == "pinv":
        return torch.linalg.pinv(sigma)
    elif reduction_method == "inverse" or reduction_method == "inv":
        return torch.linalg.inv(sigma)
    else:
        raise ValueError(f"Unknown reduction method {reduction_method}")
//  [python]
def torch_reduction_matrix(sigma: Tensor, reduction_method="pseudo"):
    import torch

    if reduction_method == "cholesky":
        C = torch.linalg.cholesky(sigma)
        return torch.linalg.inv(C.T)
    elif reduction_method == "svd":
        u, s, _ = torch.linalg.svd(sigma)
        return u @ torch.diag(torch.sqrt(1 / s))
    elif reduction_method == "pseudo" or reduction_method == "pinv":
        return torch.linalg.pinv(sigma)
    elif reduction_method == "inverse" or reduction_method == "inv":
        return torch.linalg.inv(sigma)
    else:
        raise ValueError(f"Unknown reduction method {reduction_method}")
//  [python]
def set_variances(self, var):
        self.var[:] = var
        if self.covariance_type == "diag":
            self.var_chol[:] = var.sqrt()
        elif self.covariance_type == "full":
            self.var_chol[:] = torch.cholesky(var)
        else:
            raise Exception("[GaussianMixture] __init__ got invalid covariance_type: {}".format(self.covariance_type))
        self.gmm.covariances_ = var.detach().cpu().numpy()
        self.gmm.precisions_cholesky_ = torch.cholesky(torch.inverse(var)).detach().cpu().numpy()
//  [python]
def set_variances(self, var):
        self.var[:] = var
        if self.covariance_type == "diag":
            self.var_chol[:] = var.sqrt()
        elif self.covariance_type == "full":
            self.var_chol[:] = torch.linalg.cholesky(var)
        else:
            raise Exception("[GaussianMixture] __init__ got invalid covariance_type: {}".format(self.covariance_type))
        self.gmm.covariances_ = var.detach().cpu().numpy()
        self.gmm.precisions_cholesky_ = torch.linalg.cholesky(torch.linalg.inv(var)).detach().cpu().numpy()
//  [python]
def load_state_dict(self, state_dict):
        super().load_state_dict(state_dict)
        self.gmm.means_ = self.mu.detach().cpu().numpy()
        self.gmm.covariances_ = self.var.detach().cpu().numpy()
        self.gmm.weights_ = self.pi.detach().cpu().numpy()
        if self.covariance_type == "diag":
            self.gmm.precisions_cholesky_ = torch.linalg.inv(self.var).detach().cpu().numpy()
        elif self.covariance_type == "full":
            self.gmm.precisions_cholesky_ = torch.cholesky(torch.linalg.inv(self.var)).detach().cpu().numpy()
        else:
            raise Exception("[GaussianMixture] __init__ got invalid covariance_type: {}".format(self.covariance_type))
//  [python]
def load_state_dict(self, state_dict):
        super().load_state_dict(state_dict)
        self.gmm.means_ = self.mu.detach().cpu().numpy()
        self.gmm.covariances_ = self.var.detach().cpu().numpy()
        self.gmm.weights_ = self.pi.detach().cpu().numpy()
        if self.covariance_type == "diag":
            self.gmm.precisions_cholesky_ = torch.linalg.inv(self.var).detach().cpu().numpy()
        elif self.covariance_type == "full":
            self.gmm.precisions_cholesky_ = torch.linalg.cholesky(torch.linalg.inv(self.var)).detach().cpu().numpy()
        else:
            raise Exception("[GaussianMixture] __init__ got invalid covariance_type: {}".format(self.covariance_type))
//  [python]
def predictive_dist_exact(train_data, test_data, train_labels, noise_vars, kernel_fun):
    """
    Heteroskedastic GP predictive distribution with the exact kernel.
    This function is used to compute the ground truth for KL divergence computation.

    train_data: Real-valued training inputs
    test_data: Real-valued test inputs
    train_labels: Real-valued training labels
    noise_vars: Heteroskedastic noise variances
    kernel_fun: callable k(X, Y)
    """

    with torch.no_grad():
        k_xx = kernel_fun(train_data, train_data)
        k_xt = kernel_fun(train_data, test_data)
        k_tt = kernel_fun(test_data, test_data)

        f_test_mean_full = torch.zeros(len(test_data), train_labels.shape[1], device=train_labels.device)
        f_test_stds_full = torch.zeros(len(test_data), train_labels.shape[1], device=train_labels.device)

        for out_dim in tqdm(range(train_labels.shape[1])):
            train_labels_tmp = train_labels[:, out_dim].unsqueeze(1)
            c_xx = k_xx + noise_vars[:, out_dim].diag()
            # c_xx = (c_xx + c_xx.t()) / 2.
            L_xx = torch.cholesky(c_xx)
            L_X_t = torch.triangular_solve(k_xt, L_xx, upper=False)[0]

            f_test_mean_full[:, out_dim] = (k_xt.t() @ cholesky_solve(train_labels_tmp, L_xx)).view(-1)
            f_test_stds_full[:, out_dim] = k_tt.diagonal() - (L_X_t**2).sum(dim=0)
        f_test_stds_full[f_test_stds_full < 0] = 0
        f_test_stds_full = f_test_stds_full.sqrt()

    return f_test_mean_full, f_test_stds_full
//  [python]
def predictive_dist_exact(train_data, test_data, train_labels, noise_vars, kernel_fun):
    """
    Heteroskedastic GP predictive distribution with the exact kernel.
    This function is used to compute the ground truth for KL divergence computation.

    train_data: Real-valued training inputs
    test_data: Real-valued test inputs
    train_labels: Real-valued training labels
    noise_vars: Heteroskedastic noise variances
    kernel_fun: callable k(X, Y)
    """

    with torch.no_grad():
        k_xx = kernel_fun(train_data, train_data)
        k_xt = kernel_fun(train_data, test_data)
        k_tt = kernel_fun(test_data, test_data)

        f_test_mean_full = torch.zeros(len(test_data), train_labels.shape[1], device=train_labels.device)
        f_test_stds_full = torch.zeros(len(test_data), train_labels.shape[1], device=train_labels.device)

        for out_dim in tqdm(range(train_labels.shape[1])):
            train_labels_tmp = train_labels[:, out_dim].unsqueeze(1)
            c_xx = k_xx + noise_vars[:, out_dim].diag()
            # c_xx = (c_xx + c_xx.t()) / 2.
            L_xx = torch.linalg.cholesky(c_xx)
            L_X_t = torch.triangular_solve(k_xt, L_xx, upper=False)[0]

            f_test_mean_full[:, out_dim] = (k_xt.t() @ cholesky_solve(train_labels_tmp, L_xx)).view(-1)
            f_test_stds_full[:, out_dim] = k_tt.diagonal() - (L_X_t**2).sum(dim=0)
        f_test_stds_full[f_test_stds_full < 0] = 0
        f_test_stds_full = f_test_stds_full.sqrt()

    return f_test_mean_full, f_test_stds_full
//  [python]
def log_marginal_likelihood(self, train_data, train_labels, log_noise_vars, heteroskedastic=False):
        """
        Computes the log marginal likelihood on train_data.
        We assume 0 covariance between the individual heteroskedastic GPs.

        This function is only implemented for real-valued random features for now!

        train_data: training inputs
        train_labels: training labels
        log_noise_vars: (heteroskedastic) log noise variances
        heteroskedastic: whether the noise variances are heteroskedastic
        """

        # step 1: feature extraction (to store gradients for kernel parameters)
        if self.feature_encoder is not None:
            train_features = self.feature_encoder.forward(train_data)
        else:
            train_features = train_data

        # step 2: compute marginal log likelihood for every gp
        lml = 0
        for out_dim in range(train_labels.shape[1]):
            y_tmp = train_labels[:, out_dim].unsqueeze(1)
            if heteroskedastic:
                noise_var_tmp = log_noise_vars[:, out_dim].exp().unsqueeze(1)
            else:
                noise_var_tmp = log_noise_vars.exp() * torch.ones_like(y_tmp)

            sigma_inv = train_features.t() @ (train_features / noise_var_tmp)
            sigma_inv = sigma_inv + torch.eye(len(sigma_inv), device=sigma_inv.device)
            L_sigma_inv = torch.cholesky(sigma_inv)
            xTy = train_features.t() @ (y_tmp / noise_var_tmp)
            q_inv_xTy = torch.triangular_solve(xTy, L_sigma_inv, upper=False)[0]

            lml += -0.5 * noise_var_tmp.log().sum() - L_sigma_inv.diagonal().log().sum()
            lml += -0.5 * (y_tmp * (y_tmp / noise_var_tmp)).sum()
            lml += 0.5 * q_inv_xTy.pow(2).sum()
            lml += -len(y_tmp) / 2 * np.log(2. * np.pi)

        return lml
//  [python]
def log_marginal_likelihood(self, train_data, train_labels, log_noise_vars, heteroskedastic=False):
        """
        Computes the log marginal likelihood on train_data.
        We assume 0 covariance between the individual heteroskedastic GPs.

        This function is only implemented for real-valued random features for now!

        train_data: training inputs
        train_labels: training labels
        log_noise_vars: (heteroskedastic) log noise variances
        heteroskedastic: whether the noise variances are heteroskedastic
        """

        # step 1: feature extraction (to store gradients for kernel parameters)
        if self.feature_encoder is not None:
            train_features = self.feature_encoder.forward(train_data)
        else:
            train_features = train_data

        # step 2: compute marginal log likelihood for every gp
        lml = 0
        for out_dim in range(train_labels.shape[1]):
            y_tmp = train_labels[:, out_dim].unsqueeze(1)
            if heteroskedastic:
                noise_var_tmp = log_noise_vars[:, out_dim].exp().unsqueeze(1)
            else:
                noise_var_tmp = log_noise_vars.exp() * torch.ones_like(y_tmp)

            sigma_inv = train_features.t() @ (train_features / noise_var_tmp)
            sigma_inv = sigma_inv + torch.eye(len(sigma_inv), device=sigma_inv.device)
            L_sigma_inv = torch.linalg.cholesky(sigma_inv)
            xTy = train_features.t() @ (y_tmp / noise_var_tmp)
            q_inv_xTy = torch.triangular_solve(xTy, L_sigma_inv, upper=False)[0]

            lml += -0.5 * noise_var_tmp.log().sum() - L_sigma_inv.diagonal().log().sum()
            lml += -0.5 * (y_tmp * (y_tmp / noise_var_tmp)).sum()
            lml += 0.5 * q_inv_xTy.pow(2).sum()
            lml += -len(y_tmp) / 2 * np.log(2. * np.pi)

        return lml
//  [python]
def is_pd(Q):
        """checks whether Q (respectively, its entry-matrices in each batch,
        are positive definite.
        """
        try:
            t.cholesky(Q)
            return True
        except RuntimeError as e:
            raise RuntimeError("Q is not PD") from e
//  [python]
def is_pd(Q):
        """checks whether Q (respectively, its entry-matrices in each batch,
        are positive definite.
        """
        try:
            torch.linalg.cholesky(Q)
            return True
        except RuntimeError as e:
            raise RuntimeError("Q is not PD") from e
//  [python]
def fit(self, x):
        x = x.cpu().numpy()
        self.gmm.fit(x)
        self.mu[:] = torch.tensor(self.gmm.means_, device=self.device, dtype=torch.float, requires_grad=False)
        self.var[:] = torch.tensor(self.gmm.covariances_, device=self.device, dtype=torch.float, requires_grad=False)
        self.pi[:] = torch.tensor(self.gmm.weights_, device=self.device, dtype=torch.float, requires_grad=False)
        if self.covariance_type == "diag":
            self.var_chol[:] = self.var.sqrt()
        elif self.covariance_type == "full":
            self.var_chol[:] = torch.cholesky(self.var)
        else:
            raise Exception("[GaussianMixture] __init__ got invalid covariance_type: {}".format(self.covariance_type))
//  [python]
def fit(self, x):
        x = x.cpu().numpy()
        self.gmm.fit(x)
        self.mu[:] = torch.tensor(self.gmm.means_, device=self.device, dtype=torch.float, requires_grad=False)
        self.var[:] = torch.tensor(self.gmm.covariances_, device=self.device, dtype=torch.float, requires_grad=False)
        self.pi[:] = torch.tensor(self.gmm.weights_, device=self.device, dtype=torch.float, requires_grad=False)
        if self.covariance_type == "diag":
            self.var_chol[:] = self.var.sqrt()
        elif self.covariance_type == "full":
            self.var_chol[:] = torch.linalg.cholesky(self.var)
        else:
            raise Exception("[GaussianMixture] __init__ got invalid covariance_type: {}".format(self.covariance_type))
//  [python]
def predictive_dist(self, train_data, test_data, train_labels, noise_vars):
        with torch.no_grad():
            # step 1: extract random features (optional)
            if self.feature_encoder is not None:
                train_features = self.feature_encoder.forward(train_data)
                test_features = self.feature_encoder.forward(test_data)
            else:
                train_features = train_data
                test_features = test_data

            if train_features.dtype in [torch.complex64, torch.complex128]:
                train_labels = train_labels.type(train_features.dtype)
                noise_vars = noise_vars.type(train_features.dtype)
            
            f_test_mean_full = torch.zeros(len(test_data), train_labels.shape[1], device=train_labels.device)
            f_test_stds_full = torch.zeros(len(test_data), train_labels.shape[1], device=train_labels.device)

            for out_dim in range(train_labels.shape[1]):
                y_tmp = train_labels[:, out_dim].unsqueeze(1)
                noise_var_tmp = noise_vars[:, out_dim].unsqueeze(1)
                # bottleneck (cannot be accelerated without using a lot of memory)
                sigma_inv = train_features.conj().t() @ (train_features / noise_var_tmp)
                sigma_inv = sigma_inv + torch.eye(len(sigma_inv), device=train_labels.device)
                xTy = train_features.conj().t() @ (y_tmp / noise_var_tmp)
                L_sigma_inv = torch.cholesky(sigma_inv)
                alpha = cholesky_solve(xTy, L_sigma_inv)

                f_test_mean = test_features @ alpha
                if f_test_mean.dtype in [torch.complex64, torch.complex128]:
                    f_test_mean = f_test_mean.real
                # bottleneck (cannot be accelerated without using a lot of memory)
                L_sigma_X_t = torch.triangular_solve(test_features.conj().t(), L_sigma_inv, upper=False)[0]
                f_test_stds = L_sigma_X_t.t().norm(dim=-1)

                del sigma_inv, L_sigma_inv, xTy, L_sigma_X_t

                f_test_mean_full[:, out_dim] = f_test_mean.view(-1)
                f_test_stds_full[:, out_dim] = f_test_stds

            return f_test_mean_full, f_test_stds_full
//  [python]
def predictive_dist(self, train_data, test_data, train_labels, noise_vars):
        with torch.no_grad():
            # step 1: extract random features (optional)
            if self.feature_encoder is not None:
                train_features = self.feature_encoder.forward(train_data)
                test_features = self.feature_encoder.forward(test_data)
            else:
                train_features = train_data
                test_features = test_data

            if train_features.dtype in [torch.complex64, torch.complex128]:
                train_labels = train_labels.type(train_features.dtype)
                noise_vars = noise_vars.type(train_features.dtype)
            
            f_test_mean_full = torch.zeros(len(test_data), train_labels.shape[1], device=train_labels.device)
            f_test_stds_full = torch.zeros(len(test_data), train_labels.shape[1], device=train_labels.device)

            for out_dim in range(train_labels.shape[1]):
                y_tmp = train_labels[:, out_dim].unsqueeze(1)
                noise_var_tmp = noise_vars[:, out_dim].unsqueeze(1)
                # bottleneck (cannot be accelerated without using a lot of memory)
                sigma_inv = train_features.conj().t() @ (train_features / noise_var_tmp)
                sigma_inv = sigma_inv + torch.eye(len(sigma_inv), device=train_labels.device)
                xTy = train_features.conj().t() @ (y_tmp / noise_var_tmp)
                L_sigma_inv = torch.linalg.cholesky(sigma_inv)
                alpha = cholesky_solve(xTy, L_sigma_inv)

                f_test_mean = test_features @ alpha
                if f_test_mean.dtype in [torch.complex64, torch.complex128]:
                    f_test_mean = f_test_mean.real
                # bottleneck (cannot be accelerated without using a lot of memory)
                L_sigma_X_t = torch.triangular_solve(test_features.conj().t(), L_sigma_inv, upper=False)[0]
                f_test_stds = L_sigma_X_t.t().norm(dim=-1)

                del sigma_inv, L_sigma_inv, xTy, L_sigma_X_t

                f_test_mean_full[:, out_dim] = f_test_mean.view(-1)
                f_test_stds_full[:, out_dim] = f_test_stds

            return f_test_mean_full, f_test_stds_full
//  [python]
def test_cholesky_hermitian_grad(self, device, dtype):
        # Check that the gradient is Hermitian (or symmetric)
        def run_test(shape):
            root = torch.rand(*shape, dtype=dtype, device=device)
            root = torch.matmul(root, root.transpose(-1, -2).conj())
            root.requires_grad_()
            chol = torch.cholesky(root).sum().backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2).conj())

        shapes = ((3, 3), (1, 1, 3, 3))
        for shape in shapes:
            run_test(shape)
//  [python]
def test_cholesky_hermitian_grad(self, device, dtype):
        # Check that the gradient is Hermitian (or symmetric)
        def run_test(shape):
            root = torch.rand(*shape, dtype=dtype, device=device)
            root = torch.matmul(root, root.transpose(-1, -2).conj())
            root.requires_grad_()
            chol = torch.linalg.cholesky(root).sum().backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2).conj())

        shapes = ((3, 3), (1, 1, 3, 3))
        for shape in shapes:
            run_test(shape)
//  [python]
def _apply_non_linear_constraint_at_t(  # noqa: WPS210, WPS231
    loc_t: Tensor, cov_t: Tensor, max_iter: int, tol: float, mode: InequalityConstraint
):
    nx = loc_t.shape[0]
    eye = torch.eye(nx).to(loc_t.device)
    chol_t: Tensor = torch.cholesky(cov_t)
    chinv: Tensor = torch.solve(eye, chol_t)[0]
    chinv_loc_t = chinv.matmul(loc_t)
    if mode == "nonneg":
        a_matrix = chinv.numpy()
        b = chinv_loc_t.squeeze(-1).numpy()
        loc_t_np, _ = nnls(a_matrix, b)
        return torch.from_numpy(loc_t_np).unsqueeze(-1)
    if mode == "monotone":
        diff_matrix = _construct_diff_matrix(nx)
        diff_matrix_inv = np.linalg.inv(diff_matrix)
        a_matrix = np.matmul(chinv.numpy(), diff_matrix_inv)
        b = chinv_loc_t.squeeze(-1).numpy()
        loc_t_np_diff, _ = nnls(a_matrix, b)
        loc_t_np = np.matmul(diff_matrix_inv, loc_t_np_diff)
        return torch.from_numpy(loc_t_np).unsqueeze(-1)