jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def model(self):
        self.set_mode("model")
        N = self.X.size(0)
        x_noise = pyro.sample('xerr',dist.Normal(torch.zeros(N),self.x_noise**0.5).to_event(1))
        if self.X.dim()<=1:
            X_noisy = self.X
            X_noisy = (self.X+x_noise)
        else:
            X_noisy = torch.clone(self.X)
            X_noisy[:,0] += x_noise
       
        Kff = self.kernel(X_noisy).contiguous()
        Kff = Kff + self.noise
        Kff.view(-1)[:: N + 1] += self.jitter
        # Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to diagonal
        Lff = torch.linalg.cholesky(Kff)
        zero_loc = X_noisy.new_zeros(X_noisy.size(0))
        f_loc = zero_loc + self.mean_function(X_noisy)
        if self.y is None:
            f_var = Lff.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:

            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.MultivariateNormal(f_loc, scale_tril=Lff)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )
//  [python]
def plot_gmm(self, ax, data, pred_mean, pred_var, color=None, ymin=None, ymax=None, title=None):
        # ax.cla()
        data = data.cpu()
        if pred_mean is not None and pred_var is not None:
            pred_mean = pred_mean.cpu()
            pred_var = torch.cholesky(pred_var).diagonal(dim1=-2, dim2=-1) if pred_var.dim() == 3 else pred_var
            pred_std = torch.sqrt(pred_var.cpu())
        if data.dim() == 2:
            if color == None:
                color = "lightgrey"
            args = torch.arange(data.size(1))
            theta_args = args * torch.pi / 4
            ax.plot(theta_args.repeat(data.size(0), 1).cpu(), data, "o", color="lightgrey", alpha=0.01, markersize=10, mew=0.0)
        elif data.dim() == 3:
            if color == None:
                color = list(mcolors.TABLEAU_COLORS.keys())
            args = torch.arange(data.size(2))
            theta_args = args * torch.pi / 4
            for i in range(data.size(0)):
                ax.plot(theta_args.repeat(data.size(1), 1).cpu(), data[i], "o", color=color[i], alpha=0.01, markersize=10, mew=0.0)
        if pred_mean is not None and pred_var is not None:
            theta_args_mean = torch.cat([theta_args, theta_args[0].unsqueeze(0)], dim=0)
            pred_mean = torch.cat([pred_mean, pred_mean[:, 0].unsqueeze(1)], dim=1)
            pred_std = torch.cat([pred_std, pred_std[:, 0].unsqueeze(1)], dim=1)
            for i in range(pred_mean.size(0)):
                ax.errorbar(theta_args_mean, pred_mean[i], yerr=pred_std[i], fmt="o-", markersize=4, capsize=5, alpha=0.5, linewidth=3)
        ax.set_xlabel('Channel')
        ax.set_xticks(theta_args)
        ax.set_xticklabels(args.tolist())
        if title != None:
            ax.set_title(title)
        if ymin != None and ymax != None:
            ymin = ymin.cpu()
            ymax = ymax.cpu()
            ax.set_ylim(ymin, ymax)
        ax.yaxis.grid(True)
//  [python]
def plot_gmm(self, ax, data, pred_mean, pred_var, color=None, ymin=None, ymax=None, title=None):
        # ax.cla()
        data = data.cpu()
        if pred_mean is not None and pred_var is not None:
            pred_mean = pred_mean.cpu()
            pred_var = torch.linalg.cholesky(pred_var).diagonal(dim1=-2, dim2=-1) if pred_var.dim() == 3 else pred_var
            pred_std = torch.sqrt(pred_var.cpu())
        if data.dim() == 2:
            if color == None:
                color = "lightgrey"
            args = torch.arange(data.size(1))
            theta_args = args * torch.pi / 4
            ax.plot(theta_args.repeat(data.size(0), 1).cpu(), data, "o", color="lightgrey", alpha=0.01, markersize=10, mew=0.0)
        elif data.dim() == 3:
            if color == None:
                color = list(mcolors.TABLEAU_COLORS.keys())
            args = torch.arange(data.size(2))
            theta_args = args * torch.pi / 4
            for i in range(data.size(0)):
                ax.plot(theta_args.repeat(data.size(1), 1).cpu(), data[i], "o", color=color[i], alpha=0.01, markersize=10, mew=0.0)
        if pred_mean is not None and pred_var is not None:
            theta_args_mean = torch.cat([theta_args, theta_args[0].unsqueeze(0)], dim=0)
            pred_mean = torch.cat([pred_mean, pred_mean[:, 0].unsqueeze(1)], dim=1)
            pred_std = torch.cat([pred_std, pred_std[:, 0].unsqueeze(1)], dim=1)
            for i in range(pred_mean.size(0)):
                ax.errorbar(theta_args_mean, pred_mean[i], yerr=pred_std[i], fmt="o-", markersize=4, capsize=5, alpha=0.5, linewidth=3)
        ax.set_xlabel('Channel')
        ax.set_xticks(theta_args)
        ax.set_xticklabels(args.tolist())
        if title != None:
            ax.set_title(title)
        if ymin != None and ymax != None:
            ymin = ymin.cpu()
            ymax = ymax.cpu()
            ax.set_ylim(ymin, ymax)
        ax.yaxis.grid(True)
//  [python]
def forward(self, Xnew, full_cov=False):
        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 notest_cov[i:i+1]t.
        :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()
        if self.noise.dim() <=1:
            Kff.view(-1)[:: N + 1] += self.jitter + self.noise
        elif self.noise.dim() ==2:
            Kff = Kff + self.noise
            Kff.view(-1)[:: N + 1] += self.jitter
        # 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,
        )

        return loc + self.mean_function(Xnew), cov
//  [python]
def forward(self, Xnew, full_cov=False):
        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 notest_cov[i:i+1]t.
        :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()
        if self.noise.dim() <=1:
            Kff.view(-1)[:: N + 1] += self.jitter + self.noise
        elif self.noise.dim() ==2:
            Kff = Kff + self.noise
            Kff.view(-1)[:: N + 1] += self.jitter
        # 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,
        )

        return loc + self.mean_function(Xnew), cov
//  [python]
def negative_log_likelihood_res(self, x, y, l1, h1, l2, h2):
        # y = [ytr_h, ytr_l]
        y_num, y_dimension = y[0].shape
        Sigma = self.kernel_res(x, x, l1, h1, l2, h2) + self.log_beta.exp().pow(-1) * torch.eye(
        x.size(0)) + JITTER * torch.eye(x.size(0))

        L = torch.cholesky(Sigma)
        #option 1 (use this if torch supports)
        y_res = y[0] - (-self.b * (h1 - l1)).exp() * y[1]
        Gamma,_ = torch.triangular_solve(y_res, L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension

        return nll
//  [python]
def negative_log_likelihood_res(self, x, y, l1, h1, l2, h2):
        # y = [ytr_h, ytr_l]
        y_num, y_dimension = y[0].shape
        Sigma = self.kernel_res(x, x, l1, h1, l2, h2) + self.log_beta.exp().pow(-1) * torch.eye(
        x.size(0)) + JITTER * torch.eye(x.size(0))

        L = torch.linalg.cholesky(Sigma)
        #option 1 (use this if torch supports)
        y_res = y[0] - (-self.b * (h1 - l1)).exp() * y[1]
        Gamma,_ = torch.triangular_solve(y_res, L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension

        return nll
//  [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()
        if self.noise.dim() <=1:
            Kff.view(-1)[:: N + 1] += self.jitter + self.noise
        elif self.noise.dim() ==2:
            Kff = Kff + self.noise
            Kff.view(-1)[:: N + 1] += self.jitter

        # 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()
        if self.noise.dim() <=1:
            Kff.view(-1)[:: N + 1] += self.jitter + self.noise
        elif self.noise.dim() ==2:
            Kff = Kff + self.noise
            Kff.view(-1)[:: N + 1] += self.jitter

        # 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 inverse_torch(A, approach='svd'):
    """
    Custom inverse function for inverting large covariance matrices with Pytorch

    see function header of inverse() for different approaches
    """
    if approach == 'gaussian':  # not recommended
        return torch.linalg.inv(A)

    elif approach == 'cholesky':  # if on CUDA will use CPU therefore slow
        L = torch.cholesky(A, upper=False)
        Ia = torch.eye(L.shape[0], dtype=L.dtype, device=L.device)
        L_inv = torch.linalg.solve_triangular(L, Ia, upper=False)
        return torch.matmul(L_inv.T, L_inv)

    elif approach == 'qr':
        Q, R = torch.linalg.qr(A, mode='complete')
        Ir = torch.eye(R.shape[0], dtype=R.dtype, device=R.device)
        R_inv = torch.linalg.solve_triangular(R, Ir, upper=True)
        return torch.matmul(R_inv, Q.T)

    elif approach == 'svd':  # if on CUDA will use GPU unless A is complex
        U, S, Vh = torch.linalg.svd(A, full_matrices=True)
        S = _reciprocal_pos_vals(S)
        return torch.matmul(Vh.T * S, U.T)

    else:
        raise ValueError('Specified matrix inversion approach is not recognized.')
//  [python]
def inverse_torch(A, approach='svd'):
    """
    Custom inverse function for inverting large covariance matrices with Pytorch

    see function header of inverse() for different approaches
    """
    if approach == 'gaussian':  # not recommended
        return torch.linalg.inv(A)

    elif approach == 'cholesky':  # if on CUDA will use CPU therefore slow
        L = torch.linalg.cholesky(A, upper=False)
        Ia = torch.eye(L.shape[0], dtype=L.dtype, device=L.device)
        L_inv = torch.linalg.solve_triangular(L, Ia, upper=False)
        return torch.matmul(L_inv.T, L_inv)

    elif approach == 'qr':
        Q, R = torch.linalg.qr(A, mode='complete')
        Ir = torch.eye(R.shape[0], dtype=R.dtype, device=R.device)
        R_inv = torch.linalg.solve_triangular(R, Ir, upper=True)
        return torch.matmul(R_inv, Q.T)

    elif approach == 'svd':  # if on CUDA will use GPU unless A is complex
        U, S, Vh = torch.linalg.svd(A, full_matrices=True)
        S = _reciprocal_pos_vals(S)
        return torch.matmul(Vh.T * S, U.T)

    else:
        raise ValueError('Specified matrix inversion approach is not recognized.')
//  [python]
def negative_log_likelihood(self, ytr_m):
        ytr_m = (ytr_m - self.Ymean.expand_as(ytr_m)) / (self.Ystd.expand_as(ytr_m) + EPS)
        
        y_num, y_dimension = self.Y.shape
        Sigma = self.kernel(self.X, self.X) + self.log_beta.exp().pow(-1) * torch.eye(
            self.X.size(0)) + JITTER * torch.eye(self.X.size(0))

        L = torch.cholesky(Sigma)
        # option 1 (use this if torch supports)
        # Gamma, _ = torch.triangular_solve(self.Y - ytr_m @ self.rho, L, upper=False)
       
        # option 2
        gamma = L.inverse() @ (self.Y - ytr_m * self.rho)     # we can use this as an alternative because L is a lower triangular matrix.

        nll = 0.5 * (gamma ** 2).sum() + L.diag().log().sum() * y_dimension \
              + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return nll
//  [python]
def negative_log_likelihood(self, ytr_m):
        ytr_m = (ytr_m - self.Ymean.expand_as(ytr_m)) / (self.Ystd.expand_as(ytr_m) + EPS)
        
        y_num, y_dimension = self.Y.shape
        Sigma = self.kernel(self.X, self.X) + self.log_beta.exp().pow(-1) * torch.eye(
            self.X.size(0)) + JITTER * torch.eye(self.X.size(0))

        L = torch.linalg.cholesky(Sigma)
        # option 1 (use this if torch supports)
        # Gamma, _ = torch.triangular_solve(self.Y - ytr_m @ self.rho, L, upper=False)
       
        # option 2
        gamma = L.inverse() @ (self.Y - ytr_m * self.rho)     # we can use this as an alternative because L is a lower triangular matrix.

        nll = 0.5 * (gamma ** 2).sum() + L.diag().log().sum() * y_dimension \
              + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return nll
//  [python]
def _cal_log_det(self, var):
        """
        cal log_det in log space, which can prevent overflow
        var: torch.Tensor (1,k,d,)
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)
        for k in range(self.n_components):
            log_det[k] = 2 * torch.log(torch.diagonal(torch.cholesky(var[0,k]))).sum()
            # evals, evecs = torch.linalg.eig(var[0,k])
            # log_det[k] = torch.log(evals.real).sum()
        return log_det.unsqueeze(-1)
//  [python]
def _cal_log_det(self, var):
        """
        cal log_det in log space, which can prevent overflow
        var: torch.Tensor (1,k,d,)
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)
        for k in range(self.n_components):
            log_det[k] = 2 * torch.log(torch.diagonal(torch.linalg.cholesky(var[0,k]))).sum()
            # evals, evecs = torch.linalg.eig(var[0,k])
            # log_det[k] = torch.log(evals.real).sum()
        return log_det.unsqueeze(-1)
//  [python]
def fit(self, dataset: torch.utils.data.DataLoader) -> None:
        errors = []
        mean = 0
        total = 0

        # Compute mean and covariance over the entire validation dataset
        for i, (b_inputs, b_targets) in enumerate(dataset):
            b_inputs = tuple(b_inp.to(torch_utils.get_device(self.model)) for b_inp in b_inputs)
            b_targets = tuple(b_tar.to(torch_utils.get_device(self.model)) for b_tar in b_targets)
            with torch.no_grad():
                pred = self.model(b_inputs)

            target, = b_targets

            if i == 0:
                # Use all datapoints in the first window
                error = target - pred
            else:
                # In all subsequent windows, only the last datapoint will be new
                error = target[-1:] - pred[-1:]
            error.abs_()
            errors.append(error.reshape(-1, error.shape[-1]))

            mean += torch.sum(error, dim=(0, 1))
            total += error.shape[0] * error.shape[1]

        mean /= total

        errors = torch.cat(errors, dim=0)
        errors -= mean
        cov = torch.matmul(errors.T, errors)
        cov /= total - 1

        # Add a small epsilon to the diagonal of the matrix to make it non-singular
        cov.diagonal().add_(1e-5)

        # This construction ensures that the resulting precision matrix is pos. semi-definite, even if the condition
        # number of the cov matrix is large
        cholesky = torch.cholesky(cov)
        precision = cov
        torch.cholesky_inverse(cholesky, out=precision)

        self.register_buffer('mean', mean)
        self.register_buffer('precision', precision)
//  [python]
def fit(self, dataset: torch.utils.data.DataLoader) -> None:
        errors = []
        mean = 0
        total = 0

        # Compute mean and covariance over the entire validation dataset
        for i, (b_inputs, b_targets) in enumerate(dataset):
            b_inputs = tuple(b_inp.to(torch_utils.get_device(self.model)) for b_inp in b_inputs)
            b_targets = tuple(b_tar.to(torch_utils.get_device(self.model)) for b_tar in b_targets)
            with torch.no_grad():
                pred = self.model(b_inputs)

            target, = b_targets

            if i == 0:
                # Use all datapoints in the first window
                error = target - pred
            else:
                # In all subsequent windows, only the last datapoint will be new
                error = target[-1:] - pred[-1:]
            error.abs_()
            errors.append(error.reshape(-1, error.shape[-1]))

            mean += torch.sum(error, dim=(0, 1))
            total += error.shape[0] * error.shape[1]

        mean /= total

        errors = torch.cat(errors, dim=0)
        errors -= mean
        cov = torch.matmul(errors.T, errors)
        cov /= total - 1

        # Add a small epsilon to the diagonal of the matrix to make it non-singular
        cov.diagonal().add_(1e-5)

        # This construction ensures that the resulting precision matrix is pos. semi-definite, even if the condition
        # number of the cov matrix is large
        cholesky = torch.linalg.cholesky(cov)
        precision = cov
        torch.cholesky_inverse(cholesky, out=precision)

        self.register_buffer('mean', mean)
        self.register_buffer('precision', precision)
//  [python]
def negative_log_likelihood(self):
        
        X = self.dkl_nn(self.X)
        
        y_num, y_dimension = self.Y.shape
        Sigma = self.kernel(X, X) + self.log_beta.exp().pow(-1) * torch.eye(
            X.size(0)) + JITTER * torch.eye(X.size(0))

        L = torch.cholesky(Sigma)
        #option 1 (use this if torch supports)
        Gamma,_ = torch.triangular_solve(self.Y, L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return nll
//  [python]
def negative_log_likelihood(self):
        
        X = self.dkl_nn(self.X)
        
        y_num, y_dimension = self.Y.shape
        Sigma = self.kernel(X, X) + self.log_beta.exp().pow(-1) * torch.eye(
            X.size(0)) + JITTER * torch.eye(X.size(0))

        L = torch.linalg.cholesky(Sigma)
        #option 1 (use this if torch supports)
        Gamma,_ = torch.triangular_solve(self.Y, L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return nll
//  [python]
def forward(
        self,
        x: TensorType[..., 'dim'],
        i: TensorType[..., 1],
        _return_all: Optional[bool] = False, # For internal use only
        **kwargs,
    ) -> Tuple[TensorType[..., 'dim'], TensorType[..., 'dim']]:

        noise_gaussian = torch.randn_like(x)

        if self.is_time_series:
            cov = self.noise.covariance(**kwargs)
            L = torch.cholesky(cov)
            noise = L @ noise_gaussian
        else:
            noise = noise_gaussian

        beta_int = self.beta_fn.integral(i)

        mean = x * torch.exp(-beta_int / 2)
        std = (1 - torch.exp(-beta_int)).clamp(1e-5).sqrt()

        y = mean + std * noise

        if _return_all:
            return y, noise, mean, std, cov if self.is_time_series else None

        if self.predict_gaussian_noise:
            return y, noise_gaussian
        else:
            return y, noise
//  [python]
def forward(
        self,
        x: TensorType[..., 'dim'],
        i: TensorType[..., 1],
        _return_all: Optional[bool] = False, # For internal use only
        **kwargs,
    ) -> Tuple[TensorType[..., 'dim'], TensorType[..., 'dim']]:

        noise_gaussian = torch.randn_like(x)

        if self.is_time_series:
            cov = self.noise.covariance(**kwargs)
            L = torch.linalg.cholesky(cov)
            noise = L @ noise_gaussian
        else:
            noise = noise_gaussian

        beta_int = self.beta_fn.integral(i)

        mean = x * torch.exp(-beta_int / 2)
        std = (1 - torch.exp(-beta_int)).clamp(1e-5).sqrt()

        y = mean + std * noise

        if _return_all:
            return y, noise, mean, std, cov if self.is_time_series else None

        if self.predict_gaussian_noise:
            return y, noise_gaussian
        else:
            return y, noise
//  [python]
def prepare(self, percdamp: float = 0.01, actorder: bool = False) -> None:
        """
        Prepares the Hessian matrix inverse for further computations.

        Args:
            percdamp: The percentage of damping to be applied.
            actorder: Whether to reorder the Hessian matrix based on activation order.
        """
        with torch.inference_mode():
            hessian = self.hessian
            del self.hessian

            # Activation order
            if actorder:
                self.perm = torch.argsort(torch.diag(hessian), descending=True)
                hessian = hessian[self.perm][:, self.perm]
            else:
                self.perm = torch.arange(self.rows, device=self.dev)

            # Damping
            damp = percdamp * torch.mean(torch.diag(hessian))
            diag = torch.arange(self.rows, device=self.dev)

            # Inverse of H
            attempts = 0
            while True:
                try:
                    hessian[diag, diag] += damp
                    hessian_inv = torch.cholesky(hessian)
                    hessian_inv = torch.cholesky_inverse(hessian_inv)

                    # The Cholesky inverse will sometimes fail to compute due to accumulated rounding errors when H
                    # is very large (e.g. 70B MLP down proj) and a lot of calibration data is used (e.g. 100 rows of
                    # 4096 tokens). This won't always throw an exception and sometimes just results in a NaN tensor.
                    if torch.any(torch.isnan(hessian_inv)): raise RuntimeError
                    hessian_inv = torch.cholesky(hessian_inv,
                                                        upper=True)
                    hessian_inv = hessian_inv.contiguous()
                    break
                except RuntimeError:
                    # If inverting failed, assume there were non-positive eigenvalues, so apply more damping to shift
                    # the eigenvalues in a positive direction.
                    logger.warning(" !! Warning: Applied additional damping")
                    attempts += 1
                    if attempts == 10:
                        raise ValueError("Hessian is not invertible")

            self.hessian_inv = hessian_inv
//  [python]
def prepare(self, percdamp: float = 0.01, actorder: bool = False) -> None:
        """
        Prepares the Hessian matrix inverse for further computations.

        Args:
            percdamp: The percentage of damping to be applied.
            actorder: Whether to reorder the Hessian matrix based on activation order.
        """
        with torch.inference_mode():
            hessian = self.hessian
            del self.hessian

            # Activation order
            if actorder:
                self.perm = torch.argsort(torch.diag(hessian), descending=True)
                hessian = hessian[self.perm][:, self.perm]
            else:
                self.perm = torch.arange(self.rows, device=self.dev)

            # Damping
            damp = percdamp * torch.mean(torch.diag(hessian))
            diag = torch.arange(self.rows, device=self.dev)

            # Inverse of H
            attempts = 0
            while True:
                try:
                    hessian[diag, diag] += damp
                    hessian_inv = torch.linalg.cholesky(hessian)
                    hessian_inv = torch.cholesky_inverse(hessian_inv)

                    # The Cholesky inverse will sometimes fail to compute due to accumulated rounding errors when H
                    # is very large (e.g. 70B MLP down proj) and a lot of calibration data is used (e.g. 100 rows of
                    # 4096 tokens). This won't always throw an exception and sometimes just results in a NaN tensor.
                    if torch.any(torch.isnan(hessian_inv)): raise RuntimeError
                    hessian_inv = torch.linalg.cholesky(hessian_inv,
                                                        upper=True)
                    hessian_inv = hessian_inv.contiguous()
                    break
                except RuntimeError:
                    # If inverting failed, assume there were non-positive eigenvalues, so apply more damping to shift
                    # the eigenvalues in a positive direction.
                    logger.warning(" !! Warning: Applied additional damping")
                    attempts += 1
                    if attempts == 10:
                        raise ValueError("Hessian is not invertible")

            self.hessian_inv = hessian_inv
//  [python]
def ode_sample(
        self,
        model: Callable,
        num_samples: int,
        device: str = None,
        **kwargs,
    ) -> TensorType['num_samples', 'dim']:
        if self.is_time_series:
            cov = self.noise.covariance(**kwargs)
            L = torch.cholesky(cov)
        else:
            L = None

        def drift(i, y):
            score = self._get_score(model, y, i=i, L=L, **kwargs)
            if self.is_time_series:
                # Have to include `cov` since g(t) = "scalar" * L @ dW
                score = cov @ score
            return -0.5 * self.beta_fn(i) * (y + score)

        x = self.noise(*num_samples, **kwargs).to(device)
        t = torch.Tensor([self.t1, 0]).to(device)
        y = odeint(drift, x, t, method='rk4', options={'step_size': .01})[1]
        # y = odeint(drift, x, t, rtol=1e-6, atol=1e-5)[1]

        return y
//  [python]
def ode_sample(
        self,
        model: Callable,
        num_samples: int,
        device: str = None,
        **kwargs,
    ) -> TensorType['num_samples', 'dim']:
        if self.is_time_series:
            cov = self.noise.covariance(**kwargs)
            L = torch.linalg.cholesky(cov)
        else:
            L = None

        def drift(i, y):
            score = self._get_score(model, y, i=i, L=L, **kwargs)
            if self.is_time_series:
                # Have to include `cov` since g(t) = "scalar" * L @ dW
                score = cov @ score
            return -0.5 * self.beta_fn(i) * (y + score)

        x = self.noise(*num_samples, **kwargs).to(device)
        t = torch.Tensor([self.t1, 0]).to(device)
        y = odeint(drift, x, t, method='rk4', options={'step_size': .01})[1]
        # y = odeint(drift, x, t, rtol=1e-6, atol=1e-5)[1]

        return y
//  [python]
def _eval_kl_div(self, t_list):
        
        tvec = torch.vstack(t_list)
        Kt = self.KerT.matrix(tvec)
        L_Kt = torch.cholesky(Kt)
        
        d, K, T = self.A.d, self.A.K, self.A.T
        
        # compute logdet pA
        log_det_pA = 2*d*K*torch.sum(torch.log(torch.diag(L_Kt)))
        
        L_d = torch.tril(self.A.chol_d)
        L_K = torch.tril(self.A.chol_K)
        L_T = torch.tril(self.A.chol_T)
        
        Sd = torch.mm(L_d, L_d.T)
        SK = torch.mm(L_K, L_K.T)
        ST = torch.mm(L_T, L_T.T)
        
        # compute logdet qA
        log_det_qA = 2*K*T*torch.sum(torch.log(torch.diag(L_d))) +\
            2*d*T*torch.sum(torch.log(torch.diag(L_K))) +\
            2*d*K*torch.sum(torch.log(torch.diag(L_T)))
        
        
        # trace term
        tr_term = torch.trace(Sd)*torch.trace(SK)*torch.trace(torch.inverse(Kt)@ST)

        # quadratic term
        L_Kt_inv = torch.inverse(L_Kt)
        mu = self.A.mu
        #print(mu.shape)
        
        Id = torch.eye(d).to(self.dummy.device)
        IK = torch.eye(K).to(self.dummy.device)
        
        quad_tensor = tl.tenalg.mode_dot(mu, Id, mode=0)
        quad_tensor = tl.tenalg.mode_dot(quad_tensor, IK, mode=1)
        quad_tensor = tl.tenalg.mode_dot(quad_tensor, L_Kt_inv, mode=2)
        
        quad_term = torch.sum(torch.square(quad_tensor))
        
        
        kld = 0.5*(log_det_pA - log_det_qA - d*K*T + tr_term + quad_term)
        
        return kld
//  [python]
def _eval_kl_div(self, t_list):
        
        tvec = torch.vstack(t_list)
        Kt = self.KerT.matrix(tvec)
        L_Kt = torch.linalg.cholesky(Kt)
        
        d, K, T = self.A.d, self.A.K, self.A.T
        
        # compute logdet pA
        log_det_pA = 2*d*K*torch.sum(torch.log(torch.diag(L_Kt)))
        
        L_d = torch.tril(self.A.chol_d)
        L_K = torch.tril(self.A.chol_K)
        L_T = torch.tril(self.A.chol_T)
        
        Sd = torch.mm(L_d, L_d.T)
        SK = torch.mm(L_K, L_K.T)
        ST = torch.mm(L_T, L_T.T)
        
        # compute logdet qA
        log_det_qA = 2*K*T*torch.sum(torch.log(torch.diag(L_d))) +\
            2*d*T*torch.sum(torch.log(torch.diag(L_K))) +\
            2*d*K*torch.sum(torch.log(torch.diag(L_T)))
        
        
        # trace term
        tr_term = torch.trace(Sd)*torch.trace(SK)*torch.trace(torch.linalg.inv(Kt)@ST)

        # quadratic term
        L_Kt_inv = torch.linalg.inv(L_Kt)
        mu = self.A.mu
        #print(mu.shape)
        
        Id = torch.eye(d).to(self.dummy.device)
        IK = torch.eye(K).to(self.dummy.device)
        
        quad_tensor = tl.tenalg.mode_dot(mu, Id, mode=0)
        quad_tensor = tl.tenalg.mode_dot(quad_tensor, IK, mode=1)
        quad_tensor = tl.tenalg.mode_dot(quad_tensor, L_Kt_inv, mode=2)
        
        quad_term = torch.sum(torch.square(quad_tensor))
        
        
        kld = 0.5*(log_det_pA - log_det_qA - d*K*T + tr_term + quad_term)
        
        return kld
//  [python]
def inv_logdet(v, Y_ct, R, R_inv, logdet_R):
    # Returns matrix-vector product (Y Y^T + R)^{-1} times v for matrix Y=Y_ct and any choice of vector/matrix v. Also returns the log-determinant of (Y Y^T + R).
    # Supports batch operation
    # Y_ct: (*bs, N_ensem, y_dim), R: (y_dim, y_dim), v: (*bs, bs2, y_dim)
    # out (invv): (*bs, bs2, y_dim)
    device = Y_ct.device
    N_ensem = Y_ct.shape[-2]
    y_dim = Y_ct.shape[-1]
    if N_ensem >= y_dim:
        YYT_R = Y_ct.transpose(-1, -2) @ Y_ct + R
        YYT_R_chol = torch.cholesky(YYT_R)
        logdet = 2 * YYT_R_chol.diagonal(dim1=-2, dim2=-1).log().sum(-1)
        invv = torch.cholesky_solve(v.transpose(-1,-2), YYT_R_chol).transpose(-1,-2)
    else:
        YTRinv = Y_ct @ R_inv  # (*bs, N_ensem, y_dim)
        YTRinvv = YTRinv @ v.transpose(-1, -2)  # (*bs, N_ensem, bs2)
        I_YTRinvY = torch.eye(N_ensem, device=device) + YTRinv @ Y_ct.transpose(-1, -2)  # (*bs, N_ensem, N_ensem)
        sc = power_iter(I_YTRinvY,n_iter=1)
        I_YTRinvY_sc = I_YTRinvY/sc
        I_YTRinvY_chol_sc = torch.cholesky(I_YTRinvY_sc)  # (*bs, N_ensem, N_ensem)
        I_YTRinvY_inv_YTRinvv = 1/sc * torch.cholesky_solve(YTRinvv, I_YTRinvY_chol_sc)   # (*bs, N_ensem, bs2)
        invv = v @ R_inv - I_YTRinvY_inv_YTRinvv.transpose(-1,-2) @ YTRinv  # (*bs, bs2, y_dim)
        logdet = y_dim * torch.log(sc).squeeze(-1).squeeze(-1) + 2 * I_YTRinvY_chol_sc.diagonal(dim1=-2, dim2=-1).log().sum(-1) + logdet_R
    return invv, logdet
//  [python]
def inv_logdet(v, Y_ct, R, R_inv, logdet_R):
    # Returns matrix-vector product (Y Y^T + R)^{-1} times v for matrix Y=Y_ct and any choice of vector/matrix v. Also returns the log-determinant of (Y Y^T + R).
    # Supports batch operation
    # Y_ct: (*bs, N_ensem, y_dim), R: (y_dim, y_dim), v: (*bs, bs2, y_dim)
    # out (invv): (*bs, bs2, y_dim)
    device = Y_ct.device
    N_ensem = Y_ct.shape[-2]
    y_dim = Y_ct.shape[-1]
    if N_ensem >= y_dim:
        YYT_R = Y_ct.transpose(-1, -2) @ Y_ct + R
        YYT_R_chol = torch.linalg.cholesky(YYT_R)
        logdet = 2 * YYT_R_chol.diagonal(dim1=-2, dim2=-1).log().sum(-1)
        invv = torch.cholesky_solve(v.transpose(-1,-2), YYT_R_chol).transpose(-1,-2)
    else:
        YTRinv = Y_ct @ R_inv  # (*bs, N_ensem, y_dim)
        YTRinvv = YTRinv @ v.transpose(-1, -2)  # (*bs, N_ensem, bs2)
        I_YTRinvY = torch.eye(N_ensem, device=device) + YTRinv @ Y_ct.transpose(-1, -2)  # (*bs, N_ensem, N_ensem)
        sc = power_iter(I_YTRinvY,n_iter=1)
        I_YTRinvY_sc = I_YTRinvY/sc
        I_YTRinvY_chol_sc = torch.linalg.cholesky(I_YTRinvY_sc)  # (*bs, N_ensem, N_ensem)
        I_YTRinvY_inv_YTRinvv = 1/sc * torch.cholesky_solve(YTRinvv, I_YTRinvY_chol_sc)   # (*bs, N_ensem, bs2)
        invv = v @ R_inv - I_YTRinvY_inv_YTRinvv.transpose(-1,-2) @ YTRinv  # (*bs, bs2, y_dim)
        logdet = y_dim * torch.log(sc).squeeze(-1).squeeze(-1) + 2 * I_YTRinvY_chol_sc.diagonal(dim1=-2, dim2=-1).log().sum(-1) + logdet_R
    return invv, logdet
//  [python]
def log_prob(
        self,
        model: Callable,
        x: Union[TensorType[..., 'dim'], TensorType[..., 'seq_len', 'dim']],
        num_samples: int = 1,
        **kwargs,
    ) -> TensorType[..., 1]:
        model.train() # Allows backprop through RNN
        self._e = torch.randn(num_samples, *x.shape).to(x)

        if self.is_time_series:
            cov = self.noise.covariance(**kwargs)
            L = torch.cholesky(cov)
        else:
            L = None

        def drift(i, state):
            y, _ = state
            with torch.set_grad_enabled(True):
                y = y.requires_grad_(True)
                score = self._get_score(model, y, i=i, L=L, **kwargs)
                if self.is_time_series:
                    # Have to include `cov` since g(t) = "scalar" * L @ dW
                    score = cov @ score
                dy = -0.5 * self.beta_fn(i) * (y + score)
                divergence = divergence_approx(dy, y, self._e, num_samples=num_samples)
            return dy, -divergence

        interval = torch.Tensor([0, self.t1]).to(x)

        # states = odeint(drift, (x, torch.zeros_like(x).to(x)), interval, rtol=1e-6, atol=1e-5)
        states = odeint(drift, (x, torch.zeros_like(x).to(x)), interval,
            method='rk4', options={'step_size': .01})
        y, div = states[0][-1], states[1][-1]

        if self.is_time_series:
            p0 = td.Independent(torch.distributions.MultivariateNormal(
                torch.zeros_like(y).transpose(-1, -2),
                cov.unsqueeze(-3).repeat_interleave(self.dim, dim=-3),
            ), 1)
            log_prob = p0.log_prob(y.transpose(-1, -2)) - div.sum([-1, -2])
            log_prob = log_prob / x.shape[-2]
        else:
            p0 = td.Independent(td.Normal(torch.zeros_like(y), torch.ones_like(y)), 1)
            log_prob = p0.log_prob(y) - div.sum(-1)

        return log_prob.unsqueeze(-1)
//  [python]
def log_prob(
        self,
        model: Callable,
        x: Union[TensorType[..., 'dim'], TensorType[..., 'seq_len', 'dim']],
        num_samples: int = 1,
        **kwargs,
    ) -> TensorType[..., 1]:
        model.train() # Allows backprop through RNN
        self._e = torch.randn(num_samples, *x.shape).to(x)

        if self.is_time_series:
            cov = self.noise.covariance(**kwargs)
            L = torch.linalg.cholesky(cov)
        else:
            L = None

        def drift(i, state):
            y, _ = state
            with torch.set_grad_enabled(True):
                y = y.requires_grad_(True)
                score = self._get_score(model, y, i=i, L=L, **kwargs)
                if self.is_time_series:
                    # Have to include `cov` since g(t) = "scalar" * L @ dW
                    score = cov @ score
                dy = -0.5 * self.beta_fn(i) * (y + score)
                divergence = divergence_approx(dy, y, self._e, num_samples=num_samples)
            return dy, -divergence

        interval = torch.Tensor([0, self.t1]).to(x)

        # states = odeint(drift, (x, torch.zeros_like(x).to(x)), interval, rtol=1e-6, atol=1e-5)
        states = odeint(drift, (x, torch.zeros_like(x).to(x)), interval,
            method='rk4', options={'step_size': .01})
        y, div = states[0][-1], states[1][-1]

        if self.is_time_series:
            p0 = td.Independent(torch.distributions.MultivariateNormal(
                torch.zeros_like(y).transpose(-1, -2),
                cov.unsqueeze(-3).repeat_interleave(self.dim, dim=-3),
            ), 1)
            log_prob = p0.log_prob(y.transpose(-1, -2)) - div.sum([-1, -2])
            log_prob = log_prob / x.shape[-2]
        else:
            p0 = td.Independent(td.Normal(torch.zeros_like(y), torch.ones_like(y)), 1)
            log_prob = p0.log_prob(y) - div.sum(-1)

        return log_prob.unsqueeze(-1)
//  [python]
def tallqr(V, MV=None):
    # faster QR for tall and skinny matrix
    # V: (*BV, na, nguess)
    # MV: (*BM, na, nguess) where M is the basis to make Q M-orthogonal
    # if MV is None, then MV=V
    if MV is None:
        MV = V
    VTV = torch.matmul(V.transpose(-2, -1), MV)  # (*BMV, nguess, nguess)
    R = torch.cholesky(VTV.transpose(-2, -1).conj()).transpose(-2, -1).conj()  # (*BMV, nguess, nguess)
    Rinv = torch.inverse(R)  # (*BMV, nguess, nguess)
    Q = torch.matmul(V, Rinv)
    return Q, R
//  [python]
def tallqr(V, MV=None):
    # faster QR for tall and skinny matrix
    # V: (*BV, na, nguess)
    # MV: (*BM, na, nguess) where M is the basis to make Q M-orthogonal
    # if MV is None, then MV=V
    if MV is None:
        MV = V
    VTV = torch.matmul(V.transpose(-2, -1), MV)  # (*BMV, nguess, nguess)
    R = torch.linalg.cholesky(VTV.transpose(-2, -1).conj()).transpose(-2, -1).conj()  # (*BMV, nguess, nguess)
    Rinv = torch.inverse(R)  # (*BMV, nguess, nguess)
    Q = torch.matmul(V, Rinv)
    return Q, R
//  [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 test_olse_shrinkage(p: int, n: int, dtype: torch.dtype):
    torch.manual_seed(42)

    # Number of matrices
    N = 1000

    # Generate a random covariance matrix
    A = torch.randn(N, p, p, dtype=dtype)
    S_true = A @ A.mH / p
    torch.linalg.diagonal(S_true).add_(1e-3)

    # Generate random Gaussian vectors with this covariance matrix
    scale_tril = torch.cholesky(S_true)
    X = torch.randn(N, n, p, dtype=dtype) @ scale_tril.mH
    assert X.shape == (N, n, p)

    # Compute the sample covariance matrix
    X_centered = X - X.mean(dim=0, keepdim=True)
    S_hat = (X_centered.mH @ X_centered) / n

    # Apply shrinkage
    S_olse = optimal_linear_shrinkage(S_hat, n)

    # Check that the Frobenius norm of the difference has decreased
    norm_naive = torch.norm(S_hat - S_true, dim=(-1, -2)).mean()
    norm_olse = torch.norm(S_olse - S_true, dim=(-1, -2)).mean()

    assert norm_olse <= norm_naive
//  [python]
def test_olse_shrinkage(p: int, n: int, dtype: torch.dtype):
    torch.manual_seed(42)

    # Number of matrices
    N = 1000

    # Generate a random covariance matrix
    A = torch.randn(N, p, p, dtype=dtype)
    S_true = A @ A.mH / p
    torch.linalg.diagonal(S_true).add_(1e-3)

    # Generate random Gaussian vectors with this covariance matrix
    scale_tril = torch.linalg.cholesky(S_true)
    X = torch.randn(N, n, p, dtype=dtype) @ scale_tril.mH
    assert X.shape == (N, n, p)

    # Compute the sample covariance matrix
    X_centered = X - X.mean(dim=0, keepdim=True)
    S_hat = (X_centered.mH @ X_centered) / n

    # Apply shrinkage
    S_olse = optimal_linear_shrinkage(S_hat, n)

    # Check that the Frobenius norm of the difference has decreased
    norm_naive = torch.norm(S_hat - S_true, dim=(-1, -2)).mean()
    norm_olse = torch.norm(S_olse - S_true, dim=(-1, -2)).mean()

    assert norm_olse <= norm_naive