jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def sde_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

        is_time_series = self.is_time_series

        x = self.noise(*num_samples, **kwargs).to(device)
        shape = x.shape
        x = x.transpose(-2, -1).flatten(0, -2)

        class SDE(nn.Module):
            noise_type = 'general' if is_time_series else 'diagonal'
            sde_type = 'ito'

            def __init__(self, beta_fn, _get_score):
                super().__init__()
                self.beta_fn = beta_fn
                self._get_score = _get_score

            def f(self, i, inp):
                i = -i
                inp = inp.view(*shape) # Reshape back to original

                score = self._get_score(model, inp, i=i, L=L, **kwargs)
                if is_time_series:
                    score = cov @ score

                dx = self.beta_fn(i) * (0.5 * inp + score)

                if is_time_series:
                    return dx.transpose(-1, -2).flatten(0, -2)
                return dx.view(-1, shape[-1])

            def g(self, i, inp):
                i = -i
                beta = -self.beta_fn(i).sqrt()

                if is_time_series:
                    return (beta * L).repeat_interleave(shape[-1], dim=0)
                return beta.view(1, 1).repeat(np.prod(shape[:-1]), shape[-1]).to(device)

        sde = SDE(self.beta_fn, self._get_score)
        interval = torch.Tensor([-self.t1, 0]).to(device) # Time from -t1 to 0

        step_size = self.t1 / 100
        if not is_time_series:
            x = x.view(-1, shape[-1])
        else:
            x = x.view(-1, shape[-2])
        y = sdeint(sde, x, interval, dt=step_size)[-1]
        y = y.view(*shape)

        return y
//  [python]
def sde_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

        is_time_series = self.is_time_series

        x = self.noise(*num_samples, **kwargs).to(device)
        shape = x.shape
        x = x.transpose(-2, -1).flatten(0, -2)

        class SDE(nn.Module):
            noise_type = 'general' if is_time_series else 'diagonal'
            sde_type = 'ito'

            def __init__(self, beta_fn, _get_score):
                super().__init__()
                self.beta_fn = beta_fn
                self._get_score = _get_score

            def f(self, i, inp):
                i = -i
                inp = inp.view(*shape) # Reshape back to original

                score = self._get_score(model, inp, i=i, L=L, **kwargs)
                if is_time_series:
                    score = cov @ score

                dx = self.beta_fn(i) * (0.5 * inp + score)

                if is_time_series:
                    return dx.transpose(-1, -2).flatten(0, -2)
                return dx.view(-1, shape[-1])

            def g(self, i, inp):
                i = -i
                beta = -self.beta_fn(i).sqrt()

                if is_time_series:
                    return (beta * L).repeat_interleave(shape[-1], dim=0)
                return beta.view(1, 1).repeat(np.prod(shape[:-1]), shape[-1]).to(device)

        sde = SDE(self.beta_fn, self._get_score)
        interval = torch.Tensor([-self.t1, 0]).to(device) # Time from -t1 to 0

        step_size = self.t1 / 100
        if not is_time_series:
            x = x.view(-1, shape[-1])
        else:
            x = x.view(-1, shape[-2])
        y = sdeint(sde, x, interval, dt=step_size)[-1]
        y = y.view(*shape)

        return y
//  [python]
def _estimate_shock_interval_for_one_f(mean_i: Tensor, cov_i: Tensor, n_samples: int):
    nt, nx, _ = mean_i.shape
    device = mean_i.device
    idx = rearrange(torch.arange(0, nx, device=device), "nx -> 1 1 nx")
    try:
        chol: Tensor = torch.cholesky(cov_i).unsqueeze(0)
    except:
        chol: Tensor = torch.cholesky(
            cov_i + torch.eye(201).unsqueeze(0).cuda() * 1e-8
        ).unsqueeze(0)
    z = torch.randn(n_samples, *mean_i.shape, device=device)
    y = mean_i.unsqueeze(0) + chol.matmul(z)
    y = rearrange(y, "ns nt nx 1 -> ns nt nx")
    less_than_zero = y <= 0
    objective = less_than_zero * idx + (~less_than_zero) * nx
    return torch.argmin(objective, dim=2)
//  [python]
def _estimate_shock_interval_for_one_f(mean_i: Tensor, cov_i: Tensor, n_samples: int):
    nt, nx, _ = mean_i.shape
    device = mean_i.device
    idx = rearrange(torch.arange(0, nx, device=device), "nx -> 1 1 nx")
    try:
        chol: Tensor = torch.linalg.cholesky(cov_i).unsqueeze(0)
    except:
        chol: Tensor = torch.linalg.cholesky(
            cov_i + torch.eye(201).unsqueeze(0).cuda() * 1e-8
        ).unsqueeze(0)
    z = torch.randn(n_samples, *mean_i.shape, device=device)
    y = mean_i.unsqueeze(0) + chol.matmul(z)
    y = rearrange(y, "ns nt nx 1 -> ns nt nx")
    less_than_zero = y <= 0
    objective = less_than_zero * idx + (~less_than_zero) * nx
    return torch.argmin(objective, dim=2)
//  [python]
def finalize(self, loss, target_labels):
        """
        Compute the Hessian of the parameters wrt ``loss``

        :param torch.Tensor loss: the output of evaluating a loss function such as
                                  `pyro.infer.Trace_ELBO().differentiable_loss` on the model, guide and design.
        :param list target_labels: list indicating the sample sites that are targets, i.e. for which information gain
                                   should be measured.
        """
        # set self.training = False
        self.eval()
        for l, mu_l in self.means.items():
            if l not in target_labels:
                continue
            hess_l = self._hessian_diag(loss, mu_l, event_shape=(self.w_sizes[l],))
            cov_l = rinverse(hess_l)
            self.scale_trils[l] = torch.cholesky(
                cov_l.transpose(-2, -1)
            ).transpose(-2, -1)
//  [python]
def finalize(self, loss, target_labels):
        """
        Compute the Hessian of the parameters wrt ``loss``

        :param torch.Tensor loss: the output of evaluating a loss function such as
                                  `pyro.infer.Trace_ELBO().differentiable_loss` on the model, guide and design.
        :param list target_labels: list indicating the sample sites that are targets, i.e. for which information gain
                                   should be measured.
        """
        # set self.training = False
        self.eval()
        for l, mu_l in self.means.items():
            if l not in target_labels:
                continue
            hess_l = self._hessian_diag(loss, mu_l, event_shape=(self.w_sizes[l],))
            cov_l = rinverse(hess_l)
            self.scale_trils[l] = torch.linalg.cholesky(
                cov_l.transpose(-2, -1)
            ).transpose(-2, -1)
//  [python]
def test_lower_cholesky_transform(batch_shape, dim):
    t = torch.distributions.transform_to(torch.distributions.constraints.lower_cholesky)
    x = torch.randn(batch_shape + (dim, dim))
    y = t(x)
    assert y.shape == x.shape
    actual = t.inv(y.matmul(y.transpose(-1, -2)))
    assert_close(actual, y)
    x2 = t.inv(y)
    assert x2.shape == x.shape
    y2 = t(x2)
    assert_close(y2, y)
//  [python]
def test_lower_cholesky_transform(batch_shape, dim):
    t = torch.distributions.transform_to(torch.distributions.constraints.lower_cholesky)
    x = torch.randn(batch_shape + (dim, dim))
    y = t(x)
    assert y.shape == x.shape
    actual = torch.linalg.cholesky(y.matmul(y.transpose(-1, -2)))
    assert_close(actual, y)
    x2 = t.inv(y)
    assert x2.shape == x.shape
    y2 = t(x2)
    assert_close(y2, y)
//  [python]
def exp08_properties_in_linear_systems(b, n, rhs):
    B = torch.randn((n, rhs), dtype=torch.float64)
    C = torch.zeros((n, rhs), dtype=torch.float64)

    # SPD
    A = torch.randn((n, n), dtype=torch.float64)
    A = A + torch.transpose(A, 0, 1) + torch.eye(n, dtype=torch.float64) * n
    logger.info('A is symmetric: {}'.format(torch.allclose(A, torch.transpose(A, 0, 1))))
    try:
        torch.cholesky(A)
        logger.info('A is SPD: True')
    except RuntimeError:
        logger.info('A is SPD: False')
    b.benchmark('solve_spd', solve, A, B, C)

    # Symmetric
    A = torch.randn((n, n), dtype=torch.float64)
    A = A + torch.transpose(A, 0, 1)
    A[1, 1] = -1.0
    logger.info('A is symmetric: {}'.format(torch.allclose(A, torch.transpose(A, 0, 1))))
    try:
        torch.cholesky(A)
        logger.info('A is SPD: True')
    except RuntimeError:
        logger.info('A is SPD: False')
    b.benchmark('solve_sym', solve, A, B, C)

    # Triangular
    A = torch.randn((n, n), dtype=torch.float64)
    A = torch.tril(A)
    logger.info('A is triangular: {}'.format(torch.allclose(A, torch.tril(A))))
    b.benchmark('solve_tri', solve, A, B, C)

    # Diagonal
    A = torch.randn((n, n), dtype=torch.float64)
    A = torch.diag(torch.diag(A))
    logger.info('A is diagonal: {}'.format(torch.allclose(A, torch.diag(torch.diag(A)))))
    b.benchmark('solve_dia', solve, A, B, C)
//  [python]
def exp08_properties_in_linear_systems(b, n, rhs):
    B = torch.randn((n, rhs), dtype=torch.float64)
    C = torch.zeros((n, rhs), dtype=torch.float64)

    # SPD
    A = torch.randn((n, n), dtype=torch.float64)
    A = A + torch.transpose(A, 0, 1) + torch.eye(n, dtype=torch.float64) * n
    logger.info('A is symmetric: {}'.format(torch.allclose(A, torch.transpose(A, 0, 1))))
    try:
        torch.linalg.cholesky(A)
        logger.info('A is SPD: True')
    except RuntimeError:
        logger.info('A is SPD: False')
    b.benchmark('solve_spd', solve, A, B, C)

    # Symmetric
    A = torch.randn((n, n), dtype=torch.float64)
    A = A + torch.transpose(A, 0, 1)
    A[1, 1] = -1.0
    logger.info('A is symmetric: {}'.format(torch.allclose(A, torch.transpose(A, 0, 1))))
    try:
        torch.linalg.cholesky(A)
        logger.info('A is SPD: True')
    except RuntimeError:
        logger.info('A is SPD: False')
    b.benchmark('solve_sym', solve, A, B, C)

    # Triangular
    A = torch.randn((n, n), dtype=torch.float64)
    A = torch.tril(A)
    logger.info('A is triangular: {}'.format(torch.allclose(A, torch.tril(A))))
    b.benchmark('solve_tri', solve, A, B, C)

    # Diagonal
    A = torch.randn((n, n), dtype=torch.float64)
    A = torch.diag(torch.diag(A))
    logger.info('A is diagonal: {}'.format(torch.allclose(A, torch.diag(torch.diag(A)))))
    b.benchmark('solve_dia', solve, A, B, C)
//  [python]
def fasterprune(self, sparsity, prune_n=0, prune_m=0, blocksize=128, percdamp=0.01):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        tick = time.time()

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros(self.rows, device=self.dev)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H

        mask = None

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            if prune_n == 0:
                if mask is not None:
                    mask1 = mask[:, i1:i2]
                else:
                    tmp = W1**2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
                    thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
                    mask1 = tmp <= thresh
            else:
                mask1 = torch.zeros_like(W1) == 1

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                if prune_n != 0 and i % prune_m == 0:
                    tmp = (
                        W1[:, i : (i + prune_m)] ** 2
                        / (torch.diag(Hinv1)[i : (i + prune_m)].reshape((1, -1))) ** 2
                    )
                    mask1.scatter_(
                        1, i + torch.topk(tmp, prune_n, dim=1, largest=False)[1], True
                    )

                q = w.clone()
                q[mask1[:, i]] = 0

                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d**2

                err1 = (w - q) / d
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            W[:, i1:i2] = Q1
            Losses += torch.sum(Losses1, 1) / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

        torch.cuda.synchronize()
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        self.layer.weight.data = W.reshape(self.layer.weight.shape).to(
            self.layer.weight.data.dtype
        )
//  [python]
def fasterprune(self, sparsity, prune_n=0, prune_m=0, blocksize=128, percdamp=0.01):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        tick = time.time()

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros(self.rows, device=self.dev)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.linalg.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H, upper=True)
        Hinv = H

        mask = None

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            if prune_n == 0:
                if mask is not None:
                    mask1 = mask[:, i1:i2]
                else:
                    tmp = W1**2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
                    thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
                    mask1 = tmp <= thresh
            else:
                mask1 = torch.zeros_like(W1) == 1

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                if prune_n != 0 and i % prune_m == 0:
                    tmp = (
                        W1[:, i : (i + prune_m)] ** 2
                        / (torch.diag(Hinv1)[i : (i + prune_m)].reshape((1, -1))) ** 2
                    )
                    mask1.scatter_(
                        1, i + torch.topk(tmp, prune_n, dim=1, largest=False)[1], True
                    )

                q = w.clone()
                q[mask1[:, i]] = 0

                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d**2

                err1 = (w - q) / d
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            W[:, i1:i2] = Q1
            Losses += torch.sum(Losses1, 1) / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

        torch.cuda.synchronize()
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        self.layer.weight.data = W.reshape(self.layer.weight.shape).to(
            self.layer.weight.data.dtype
        )
//  [python]
def get_w_metric(self):
        # NOTE: in lora pruning, the importance metric will be: lora_B @ grad(lora_A) + grad(lora_B) @ lora_A - grad(lora_B) @ grad(lora_A)
        if self.metric_type == 'Weight':
            self.W_metric = torch.abs(self.weight)
        elif self.metric_type == 'Wanda':
            self.W_metric = torch.abs(self.weight) * torch.sqrt(self.scaler_row.reshape((1,-1)))
            self.scaler_row = None
        elif self.metric_type == 'SparseGPT' or self.metric_type == 'SparseGPT-Git':
            percdamp = 0.01 # Percent of the average Hessian diagonal used for dampening
            dead = torch.diag(self.Hessian) == 0
            self.Hessian[dead, dead] = 1
            # NOTE: SparseGPT updates the weight in cols of zero diag
            self.weight.data[:, dead] = 0
            damp = percdamp * torch.mean(torch.diag(self.Hessian))
            diag = torch.arange(self.columns, device=self.device)
            self.Hessian[diag, diag] += damp
            Hinv = torch.cholesky(self.Hessian)
            Hinv = torch.cholesky_inverse(Hinv)
            # NOTE: use cholesky_ex as Hinv many not be complex Hermitian matrix
            Hinv, _ = torch.linalg.cholesky_ex(Hinv, upper=True)
            Hinv = torch.diag(Hinv).reshape((1, -1))
            if self.metric_type == 'SparseGPT-Git':
                self.W_metric = self.weight ** 2 / Hinv ** 2
            else:
                self.W_metric = self.weight ** 2 / Hinv

        if self.wise_dim == 'row':
            self.sort_indices = torch.sort(self.W_metric, dim=-1, stable=True)[1]
        elif self.wise_dim == 'column':
            self.sort_indices = torch.sort(self.W_metric, dim=0, stable=True)[1]
        else:
            raise NotImplementedError(f"Invalid wise dim: {self.wise_dim}")
//  [python]
def get_w_metric(self):
        # NOTE: in lora pruning, the importance metric will be: lora_B @ grad(lora_A) + grad(lora_B) @ lora_A - grad(lora_B) @ grad(lora_A)
        if self.metric_type == 'Weight':
            self.W_metric = torch.abs(self.weight)
        elif self.metric_type == 'Wanda':
            self.W_metric = torch.abs(self.weight) * torch.sqrt(self.scaler_row.reshape((1,-1)))
            self.scaler_row = None
        elif self.metric_type == 'SparseGPT' or self.metric_type == 'SparseGPT-Git':
            percdamp = 0.01 # Percent of the average Hessian diagonal used for dampening
            dead = torch.diag(self.Hessian) == 0
            self.Hessian[dead, dead] = 1
            # NOTE: SparseGPT updates the weight in cols of zero diag
            self.weight.data[:, dead] = 0
            damp = percdamp * torch.mean(torch.diag(self.Hessian))
            diag = torch.arange(self.columns, device=self.device)
            self.Hessian[diag, diag] += damp
            Hinv = torch.linalg.cholesky(self.Hessian)
            Hinv = torch.cholesky_inverse(Hinv)
            # NOTE: use cholesky_ex as Hinv many not be complex Hermitian matrix
            Hinv, _ = torch.linalg.cholesky_ex(Hinv, upper=True)
            Hinv = torch.diag(Hinv).reshape((1, -1))
            if self.metric_type == 'SparseGPT-Git':
                self.W_metric = self.weight ** 2 / Hinv ** 2
            else:
                self.W_metric = self.weight ** 2 / Hinv

        if self.wise_dim == 'row':
            self.sort_indices = torch.sort(self.W_metric, dim=-1, stable=True)[1]
        elif self.wise_dim == 'column':
            self.sort_indices = torch.sort(self.W_metric, dim=0, stable=True)[1]
        else:
            raise NotImplementedError(f"Invalid wise dim: {self.wise_dim}")
//  [python]
def _compute_precision_cholesky(covariances, covariance_type):
    estimate_precision_error_message = (
        "Fitting the mixture model failed because some components have "
        "ill-defined empirical covariance (for instance caused by singleton "
        "or collapsed samples). Try to decrease the number of components, "
        "or increase reg_covar."
    )

    if covariance_type == "full":
        n_components, n_features, _ = covariances.shape
        precisions_chol = torch.empty((n_components, n_features, n_features))
        for k, covariance in enumerate(covariances):
            try:
                cov_chol = torch.cholesky(torch.FloatTensor(covariance), upper=False)
            except RuntimeError:
                raise ValueError(estimate_precision_error_message)
            precisions_chol[k] = torch.triangular_solve(torch.eye(n_features), cov_chol, upper=False)[0].T
        return precisions_chol
    else:
        raise NotImplementedError(
            "Hummingbird does not currently support {} covariance type for BayesianGaussianMixture. The supported covariance type is full.".format(
                covariance_type
            )
        )
//  [python]
def _compute_precision_cholesky(covariances, covariance_type):
    estimate_precision_error_message = (
        "Fitting the mixture model failed because some components have "
        "ill-defined empirical covariance (for instance caused by singleton "
        "or collapsed samples). Try to decrease the number of components, "
        "or increase reg_covar."
    )

    if covariance_type == "full":
        n_components, n_features, _ = covariances.shape
        precisions_chol = torch.empty((n_components, n_features, n_features))
        for k, covariance in enumerate(covariances):
            try:
                cov_chol = torch.linalg.cholesky(torch.FloatTensor(covariance), upper=False)
            except torch.linalg.LinAlgError:
                raise ValueError(estimate_precision_error_message)
            precisions_chol[k] = torch.linalg.solve_triangular(cov_chol, torch.eye(n_features), upper=False).T
        return precisions_chol
    else:
        raise NotImplementedError(
            "Hummingbird does not currently support {} covariance type for BayesianGaussianMixture. The supported covariance type is full.".format(
                covariance_type
            )
        )
//  [python]
def warm_start(self, sym_idx, y, learning_fraction=0.3, single_circular=False):
        """
        Train covariance matrix with known symbols.

        :params sym_idx: Sent symbol index
        :param y: Received complex symbols
        :param learning_fraction: Learning rate
        :param single_circular: Toggle if a a single sigma is learned for all symbols
        """
        with torch.no_grad():
            meanfree_vals = torch.empty(0).to(y.device)
            for i in np.arange(self.M.item()):
                idx = (sym_idx == i).nonzero()
                idx = idx[
                    torch.randperm(len(idx))[
                        0 : int(np.ceil(learning_fraction * len(idx)))
                    ]
                ]
                rx_sym = y[idx].squeeze()
                if y.dtype == torch.complex64:
                    rx_sym = torch.stack(
                        (torch.real(rx_sym), torch.imag(rx_sym))
                    ).squeeze()
                else:
                    rx_sym = torch.transpose(rx_sym.unsqueeze(1), 0, 1)

                rx_mean = torch.mean(rx_sym, 1).unsqueeze(1)
                self.noise_mean[i] = rx_mean

                rx_meanfree = torch.sub(rx_sym, rx_mean)
                if not single_circular:
                    CovMatrix = torch.matmul(
                        rx_meanfree, torch.transpose(rx_meanfree, 0, 1)
                    ) / len(idx)
                    self.cov_gen_mat[i] = torch.cholesky(CovMatrix)
                else:
                    meanfree_vals = torch.cat(
                        [
                            meanfree_vals,
                            torch.flatten(rx_meanfree),
                        ]
                    )
            if single_circular:
                sigma = torch.sqrt(torch.sum(meanfree_vals**2) / len(meanfree_vals))
                for i in np.arange(self.M.item()):
                    self.cov_gen_mat[i] = sigma * torch.eye(self.dim, device=y.device)
//  [python]
def warm_start(self, sym_idx, y, learning_fraction=0.3, single_circular=False):
        """
        Train covariance matrix with known symbols.

        :params sym_idx: Sent symbol index
        :param y: Received complex symbols
        :param learning_fraction: Learning rate
        :param single_circular: Toggle if a a single sigma is learned for all symbols
        """
        with torch.no_grad():
            meanfree_vals = torch.empty(0).to(y.device)
            for i in np.arange(self.M.item()):
                idx = (sym_idx == i).nonzero()
                idx = idx[
                    torch.randperm(len(idx))[
                        0 : int(np.ceil(learning_fraction * len(idx)))
                    ]
                ]
                rx_sym = y[idx].squeeze()
                if y.dtype == torch.complex64:
                    rx_sym = torch.stack(
                        (torch.real(rx_sym), torch.imag(rx_sym))
                    ).squeeze()
                else:
                    rx_sym = torch.transpose(rx_sym.unsqueeze(1), 0, 1)

                rx_mean = torch.mean(rx_sym, 1).unsqueeze(1)
                self.noise_mean[i] = rx_mean

                rx_meanfree = torch.sub(rx_sym, rx_mean)
                if not single_circular:
                    CovMatrix = torch.matmul(
                        rx_meanfree, torch.transpose(rx_meanfree, 0, 1)
                    ) / len(idx)
                    self.cov_gen_mat[i] = torch.linalg.cholesky(CovMatrix)
                else:
                    meanfree_vals = torch.cat(
                        [
                            meanfree_vals,
                            torch.flatten(rx_meanfree),
                        ]
                    )
            if single_circular:
                sigma = torch.sqrt(torch.sum(meanfree_vals**2) / len(meanfree_vals))
                for i in np.arange(self.M.item()):
                    self.cov_gen_mat[i] = sigma * torch.eye(self.dim, device=y.device)
//  [python]
def hist_match(target, source, mode="cdf", eps=1e-2):
    target = target.permute(0, 3, 1, 2)  # -> b, c, h, w
    source = source.permute(0, 3, 1, 2)

    if mode == "cdf":
        b, c, h, w = target.shape
        matched = cdf_match(target, source).clamp(target.min(), target.max()).to(target)

    else:
        # based on https://github.com/ProGamerGov/Neural-Tools/blob/master/linear-color-transfer.py#L36

        mu_t = target.mean((2, 3), keepdim=True)
        hist_t = (target - mu_t).view(target.size(1), -1)  # [c, b * h * w]
        cov_t = hist_t @ hist_t.T / hist_t.shape[1] + eps * torch.eye(hist_t.shape[0]).to(hist_t)

        mu_s = source.mean((2, 3), keepdim=True)
        hist_s = (source - mu_s).view(source.size(1), -1)
        cov_s = hist_s @ hist_s.T / hist_s.shape[1] + eps * torch.eye(hist_s.shape[0]).to(hist_s)

        if mode == "chol":
            chol_t = torch.cholesky(cov_t)
            chol_s = torch.cholesky(cov_s)
            matched = chol_s @ torch.inverse(chol_t) @ hist_t

        elif mode == "pca":
            eva_t, eve_t = torch.symeig(cov_t, eigenvectors=True, upper=True)
            Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
            eva_s, eve_s = torch.symeig(cov_s, eigenvectors=True, upper=True)
            Qs = eve_s @ torch.sqrt(torch.diag(eva_s)) @ eve_s.T
            matched = Qs @ torch.inverse(Qt) @ hist_t

        elif mode == "sym":
            eva_t, eve_t = torch.symeig(cov_t, eigenvectors=True, upper=True)
            Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
            Qt_Cs_Qt = Qt @ cov_s @ Qt
            eva_QtCsQt, eve_QtCsQt = torch.symeig(Qt_Cs_Qt, eigenvectors=True, upper=True)
            QtCsQt = eve_QtCsQt @ torch.sqrt(torch.diag(eva_QtCsQt)) @ eve_QtCsQt.T
            matched = torch.inverse(Qt) @ QtCsQt @ torch.inverse(Qt) @ hist_t

        matched = matched.view(*target.shape) + mu_s

    return matched.permute(0, 2, 3, 1)
//  [python]
def hist_match(target, source, mode="cdf", eps=1e-2):
    target = target.permute(0, 3, 1, 2)  # -> b, c, h, w
    source = source.permute(0, 3, 1, 2)

    if mode == "cdf":
        b, c, h, w = target.shape
        matched = cdf_match(target, source).clamp(target.min(), target.max()).to(target)

    else:
        # based on https://github.com/ProGamerGov/Neural-Tools/blob/master/linear-color-transfer.py#L36

        mu_t = target.mean((2, 3), keepdim=True)
        hist_t = (target - mu_t).view(target.size(1), -1)  # [c, b * h * w]
        cov_t = hist_t @ hist_t.T / hist_t.shape[1] + eps * torch.eye(hist_t.shape[0]).to(hist_t)

        mu_s = source.mean((2, 3), keepdim=True)
        hist_s = (source - mu_s).view(source.size(1), -1)
        cov_s = hist_s @ hist_s.T / hist_s.shape[1] + eps * torch.eye(hist_s.shape[0]).to(hist_s)

        if mode == "chol":
            chol_t = torch.linalg.cholesky(cov_t)
            chol_s = torch.linalg.cholesky(cov_s)
            matched = chol_s @ torch.inverse(chol_t) @ hist_t

        elif mode == "pca":
            eva_t, eve_t = torch.symeig(cov_t, eigenvectors=True, upper=True)
            Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
            eva_s, eve_s = torch.symeig(cov_s, eigenvectors=True, upper=True)
            Qs = eve_s @ torch.sqrt(torch.diag(eva_s)) @ eve_s.T
            matched = Qs @ torch.inverse(Qt) @ hist_t

        elif mode == "sym":
            eva_t, eve_t = torch.symeig(cov_t, eigenvectors=True, upper=True)
            Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
            Qt_Cs_Qt = Qt @ cov_s @ Qt
            eva_QtCsQt, eve_QtCsQt = torch.symeig(Qt_Cs_Qt, eigenvectors=True, upper=True)
            QtCsQt = eve_QtCsQt @ torch.sqrt(torch.diag(eva_QtCsQt)) @ eve_QtCsQt.T
            matched = torch.inverse(Qt) @ QtCsQt @ torch.inverse(Qt) @ hist_t

        matched = matched.view(*target.shape) + mu_s

    return matched.permute(0, 2, 3, 1)
//  [python]
def _cholesky(self, K, add_jitter=False):
        if add_jitter:
            K = K + (self.jitter * K.diagonal().mean()).repeat(K.shape[0]).diagflat()
        try:
            return torch.cholesky(K)
        except RuntimeError as e:
            print("ERROR:", e.args[0], file=sys.__stdout__)
            if K.isnan().any():
                print("ERROR: kernel matrix has NaNs!", file=sys.__stdout__)
            if K.isinf().any():
                print("ERROR: kernel matrix has infinities!", file=sys.__stdout__)
            self.print_parameters()
            plot_gram(K)
            raise CholeskyException(e.args[0], K, self)
//  [python]
def _cholesky(self, K, add_jitter=False):
        if add_jitter:
            K = K + (self.jitter * K.diagonal().mean()).repeat(K.shape[0]).diagflat()
        try:
            return torch.linalg.cholesky(K)
        except RuntimeError as e:
            print("ERROR:", e.args[0], file=sys.__stdout__)
            if K.isnan().any():
                print("ERROR: kernel matrix has NaNs!", file=sys.__stdout__)
            if K.isinf().any():
                print("ERROR: kernel matrix has infinities!", file=sys.__stdout__)
            self.print_parameters()
            plot_gram(K)
            raise CholeskyException(e.args[0], K, self)
//  [python]
def inverse(self, inputs, context=None):
        inputs_jitter = inputs + self.eye * self.eps
        if self.checkargs:
            self.check_pos_def(inputs_jitter)

        outputs = torch.cholesky(inputs_jitter, upper=False)
        diagonal = torch.diagonal(outputs, dim1=-2, dim2=-1)
        logabsdet = self.N * np.log(2.) + (self.powers * diagonal.log()).sum(1)

        return outputs, -logabsdet
//  [python]
def inverse(self, inputs, context=None):
        inputs_jitter = inputs + self.eye * self.eps
        if self.checkargs:
            self.check_pos_def(inputs_jitter)

        outputs = torch.linalg.cholesky(inputs_jitter, upper=False)
        diagonal = torch.diagonal(outputs, dim1=-2, dim2=-1)
        logabsdet = self.N * np.log(2.) + (self.powers * diagonal.log()).sum(1)

        return outputs, -logabsdet
//  [python]
def test_cholesky_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(A_dims, b_dims, upper):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_hermitian_pd_matrix(A_matrix_size, *A_batch_dims,
                                           dtype=dtype, device='cpu')
            b = torch.randn(*b_dims, dtype=dtype, device='cpu')
            x_exp = torch.tensor(solve(A.numpy(), b.numpy()), dtype=dtype, device=device)
            A, b = A.to(dtype=dtype, device=device), b.to(dtype=dtype, device=device)
            L = torch.cholesky(A, upper=upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)
            # https://github.com/pytorch/pytorch/issues/42695
            x = torch.cholesky_solve(b, L, upper=upper, out=x)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        for upper in [True, False]:
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), upper)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), upper)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), upper)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), upper)
//  [python]
def test_cholesky_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(A_dims, b_dims, upper):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_hermitian_pd_matrix(A_matrix_size, *A_batch_dims,
                                           dtype=dtype, device='cpu')
            b = torch.randn(*b_dims, dtype=dtype, device='cpu')
            x_exp = torch.tensor(solve(A.numpy(), b.numpy()), dtype=dtype, device=device)
            A, b = A.to(dtype=dtype, device=device), b.to(dtype=dtype, device=device)
            L = torch.linalg.cholesky(A, upper=upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)
            # https://github.com/pytorch/pytorch/issues/42695
            x = torch.cholesky_solve(b, L, upper=upper, out=x)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        for upper in [True, False]:
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), upper)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), upper)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), upper)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), upper)
//  [python]
def test_cholesky_inverse(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(shape, batch, upper, contiguous):
            A = random_hermitian_pd_matrix(shape, *batch, dtype=dtype, device=device)
            if A.numel() > 0 and not contiguous:
                A = A.mT
                self.assertFalse(A.is_contiguous())
            L = torch.cholesky(A)
            expected_inverse = torch.inverse(A)
            L = L.mH if upper else L
            actual_inverse = torch.cholesky_inverse(L, upper)
            self.assertEqual(actual_inverse, expected_inverse)

        shapes = (0, 3, 5)
        batches = ((), (0,), (3, ), (2, 2))
        for shape, batch, upper, contiguous in list(itertools.product(shapes, batches, (True, False), (True, False))):
            run_test(shape, batch, upper, contiguous)

        # check the out= variant
        A = random_hermitian_pd_matrix(3, 2, dtype=dtype, device=device)
        L = torch.cholesky(A)

        # There are two code paths currently for the out= variant
        # 1. When 'out' tensor is in Fortran (column-major) memory format
        # then the fast route is taken and the storage is reused directly in the computations
        # 2. When 'out' tensor is not in Fortran format then a temporary tensor is allocated internally
        # and the result is copied from the temporary tensor to 'out' tensor

        # This test checks the first code path
        out = torch.empty_like(A)
        out_t = out.mT.clone(memory_format=torch.contiguous_format)
        out = out_t.mT
        ans = torch.cholesky_inverse(L, out=out)
        self.assertEqual(ans, out)
        expected = torch.inverse(A)
        self.assertEqual(expected, out)

        # This test checks the second code path
        out = torch.empty_like(A)
        ans = torch.cholesky_inverse(L, out=out)
        self.assertEqual(ans, out)
        expected = torch.inverse(A)
        self.assertEqual(expected, out)
//  [python]
def test_cholesky_inverse(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(shape, batch, upper, contiguous):
            A = random_hermitian_pd_matrix(shape, *batch, dtype=dtype, device=device)
            if A.numel() > 0 and not contiguous:
                A = A.mT
                self.assertFalse(A.is_contiguous())
            L = torch.linalg.cholesky(A)
            expected_inverse = torch.inverse(A)
            L = L.mH if upper else L
            actual_inverse = torch.cholesky_inverse(L, upper)
            self.assertEqual(actual_inverse, expected_inverse)

        shapes = (0, 3, 5)
        batches = ((), (0,), (3, ), (2, 2))
        for shape, batch, upper, contiguous in list(itertools.product(shapes, batches, (True, False), (True, False))):
            run_test(shape, batch, upper, contiguous)

        # check the out= variant
        A = random_hermitian_pd_matrix(3, 2, dtype=dtype, device=device)
        L = torch.linalg.cholesky(A)

        # There are two code paths currently for the out= variant
        # 1. When 'out' tensor is in Fortran (column-major) memory format
        # then the fast route is taken and the storage is reused directly in the computations
        # 2. When 'out' tensor is not in Fortran format then a temporary tensor is allocated internally
        # and the result is copied from the temporary tensor to 'out' tensor

        # This test checks the first code path
        out = torch.empty_like(A)
        out_t = out.mT.clone(memory_format=torch.contiguous_format)
        out = out_t.mT
        ans = torch.cholesky_inverse(L, out=out)
        self.assertEqual(ans, out)
        expected = torch.inverse(A)
        self.assertEqual(expected, out)

        # This test checks the second code path
        out = torch.empty_like(A)
        ans = torch.cholesky_inverse(L, out=out)
        self.assertEqual(ans, out)
        expected = torch.inverse(A)
        self.assertEqual(expected, out)
//  [python]
def exp_lower(self, mat):
        return torch.cholesky(mat)
//  [python]
def exp_lower(self, mat):
        return torch.linalg.cholesky(mat)
//  [python]
def robust_cholesky(matrix: torch.Tensor) -> torch.Tensor:
    """
    Implements a Cholesky decomposition.
    If the Cholesky decomposition fails, it is retried with increasing added jitter on the diagonal.
    :param matrix: Symmetric positive semi-definite matrix to factorize.
    :return: Approximate cholesky factor L such that (approximately) LL^T = matrix
    """
    eps = 1e-5 * matrix.trace() / matrix.shape[-1]
    L = None
    for i in range(10):
        try:
            L = torch.cholesky(matrix)
            break
        except RuntimeError:
            print('Increasing jitter for Cholesky decomposition', flush=True)
            matrix += eps * torch.eye(matrix.shape[-1], device=matrix.device, dtype=matrix.dtype)
            eps *= 2
    if L is None:
        raise RuntimeError('Could not Cholesky decompose the matrix')
    return L
//  [python]
def robust_cholesky(matrix: torch.Tensor) -> torch.Tensor:
    """
    Implements a Cholesky decomposition.
    If the Cholesky decomposition fails, it is retried with increasing added jitter on the diagonal.
    :param matrix: Symmetric positive semi-definite matrix to factorize.
    :return: Approximate cholesky factor L such that (approximately) LL^T = matrix
    """
    eps = 1e-5 * matrix.trace() / matrix.shape[-1]
    L = None
    for i in range(10):
        try:
            L = torch.linalg.cholesky(matrix)
            break
        except RuntimeError:
            print('Increasing jitter for Cholesky decomposition', flush=True)
            matrix += eps * torch.eye(matrix.shape[-1], device=matrix.device, dtype=matrix.dtype)
            eps *= 2
    if L is None:
        raise RuntimeError('Could not Cholesky decompose the matrix')
    return L
//  [python]
def learn(self, enc_images, enc_segmentations,orig_segmentations=None):
        """
        """
        B, N, C, H, W = enc_segmentations.size()
        B, N, D, H, W = enc_images.size()
        E = 1
        enc_segmentations = enc_segmentations.reshape(B,N,E,C//E,H,W).permute(0,2,1,3,4,5).reshape(B*E,N,C//E,H,W)
        enc_images = enc_images.reshape(B,N,E,D//E,H,W).permute(0,2,1,3,4,5).reshape(B*E,N,D//E,H,W)

        x_s, y_s = self._sample_points_strided_grid(enc_images, enc_segmentations)
        B, S, _ = x_s.size()
        sigma_noise = self.sigma_noise * torch.eye(S, device=x_s.device)[None,:,:]
        K_ss = self.kernel(x_s, x_s) + sigma_noise
        L = torch.cholesky(K_ss)
        alpha = self.tri_solve(L.permute(0,2,1), self.tri_solve(L, y_s), lower=False)
        return L, alpha, x_s
//  [python]
def learn(self, enc_images, enc_segmentations,orig_segmentations=None):
        """
        """
        B, N, C, H, W = enc_segmentations.size()
        B, N, D, H, W = enc_images.size()
        E = 1
        enc_segmentations = enc_segmentations.reshape(B,N,E,C//E,H,W).permute(0,2,1,3,4,5).reshape(B*E,N,C//E,H,W)
        enc_images = enc_images.reshape(B,N,E,D//E,H,W).permute(0,2,1,3,4,5).reshape(B*E,N,D//E,H,W)

        x_s, y_s = self._sample_points_strided_grid(enc_images, enc_segmentations)
        B, S, _ = x_s.size()
        sigma_noise = self.sigma_noise * torch.eye(S, device=x_s.device)[None,:,:]
        K_ss = self.kernel(x_s, x_s) + sigma_noise
        L = torch.linalg.cholesky(K_ss)#torch.cholesky(K_ss)
        alpha = self.tri_solve(L.permute(0,2,1), self.tri_solve(L, y_s), lower=False)
        return L, alpha, x_s
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
    """Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
    Args:
        :attr:`A` (Tensor):
            The tensor to compute the Cholesky decomposition of
        :attr:`upper` (bool, optional):
            See torch.cholesky
        :attr:`out` (Tensor, optional):
            See torch.cholesky
        :attr:`jitter` (float, optional):
            The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
            as 1e-6 (float) or 1e-8 (double)
    """
    try:
        if A.dim() == 2:
            L = torch.cholesky(A, upper=upper, out=out)
            return L
        else:
            L_list = []
            for idx in range(A.shape[0]):
                L = torch.cholesky(A[idx], upper=upper, out=out)
                L_list.append(L)
            return torch.stack(L_list, dim=0)
    except:
        isnan = torch.isnan(A)
        if isnan.any():
            raise NanError(
                f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
            )

        if jitter is None:
            jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
        Aprime = A.clone()
        jitter_prev = 0
        for i in range(8):
            jitter_new = jitter * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                if Aprime.dim() == 2:
                    L = torch.cholesky(Aprime, upper=upper, out=out)
                    warnings.warn(
                        f"A not p.d., added jitter of {jitter_new} to the diagonal",
                        RuntimeWarning,
                    )
                    return L
                else:
                    L_list = []
                    for idx in range(Aprime.shape[0]):
                        L = torch.cholesky(Aprime[idx], upper=upper, out=out)
                        L_list.append(L)
                    warnings.warn(
                        f"A not p.d., added jitter of {jitter_new} to the diagonal",
                        RuntimeWarning,
                    )
                    return torch.stack(L_list, dim=0)
            except:
                continue
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
    """Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
    Args:
        :attr:`A` (Tensor):
            The tensor to compute the Cholesky decomposition of
        :attr:`upper` (bool, optional):
            See torch.cholesky
        :attr:`out` (Tensor, optional):
            See torch.cholesky
        :attr:`jitter` (float, optional):
            The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
            as 1e-6 (float) or 1e-8 (double)
    """
    try:
        if A.dim() == 2:
            L = torch.linalg.cholesky(A, upper=upper, out=out)
            return L
        else:
            L_list = []
            for idx in range(A.shape[0]):
                L = torch.linalg.cholesky(A[idx], upper=upper, out=out)
                L_list.append(L)
            return torch.stack(L_list, dim=0)
    except:
        isnan = torch.isnan(A)
        if isnan.any():
            raise NanError(
                f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
            )

        if jitter is None:
            jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
        Aprime = A.clone()
        jitter_prev = 0
        for i in range(8):
            jitter_new = jitter * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                if Aprime.dim() == 2:
                    L = torch.linalg.cholesky(Aprime, upper=upper, out=out)
                    warnings.warn(
                        f"A not p.d., added jitter of {jitter_new} to the diagonal",
                        RuntimeWarning,
                    )
                    return L
                else:
                    L_list = []
                    for idx in range(Aprime.shape[0]):
                        L = torch.linalg.cholesky(Aprime[idx], upper=upper, out=out)
                        L_list.append(L)
                    warnings.warn(
                        f"A not p.d., added jitter of {jitter_new} to the diagonal",
                        RuntimeWarning,
                    )
                    return torch.stack(L_list, dim=0)
            except:
                continue
//  [python]
def get_reparametrization_parameters(self, sigmas_init, eps = 1e-4):
    """
    Compute the paramters for the reparamterization of the sigmas, such that \sigma = L_lower_init @ L_lower_init.T + torch.exp(log_diag_init) where log_diag_init is a diagonal matrix 
    """

    L_lower_init = torch.cholesky(sigmas_init)
    log_diag_init = torch.log(torch.ones(self.n_topics, self.n_dims)*1e-4)  # initialize diag = (1,...,1)*eps, such that only a small value is added to the diagonal

    return L_lower_init, log_diag_init
//  [python]
def get_reparametrization_parameters(self, sigmas_init, eps = 1e-4):
    """
    Compute the paramters for the reparamterization of the sigmas, such that \sigma = L_lower_init @ L_lower_init.T + torch.exp(log_diag_init) where log_diag_init is a diagonal matrix 
    """

    L_lower_init = torch.linalg.cholesky(sigmas_init)
    log_diag_init = torch.log(torch.ones(self.n_topics, self.n_dims)*1e-4)  # initialize diag = (1,...,1)*eps, such that only a small value is added to the diagonal

    return L_lower_init, log_diag_init