jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_pyro_catch_error(self):
        def potential_fn(z):
            mvn = pyro.distributions.MultivariateNormal(
                loc=torch.zeros(2),
                covariance_matrix=z["K"],
            )
            return mvn.log_prob(torch.zeros(2))

        # Test base case where everything is fine
        z = {"K": torch.eye(2)}
        grads, val = potential_grad(potential_fn, z)
        self.assertAllClose(grads["K"], -0.5 * torch.eye(2))
        norm_mvn = torch.distributions.Normal(0, 1)
        self.assertAllClose(val, 2 * norm_mvn.log_prob(torch.tensor(0.0)))

        # Default behavior should catch the ValueError when trying to instantiate
        # the MVN and return NaN instead
        z = {"K": torch.ones(2, 2)}
        _, val = potential_grad(potential_fn, z)
        self.assertTrue(torch.isnan(val))

        # Default behavior should catch the LinAlgError when peforming a
        # Cholesky decomposition and return NaN instead
        def potential_fn_chol(z):
            return torch.linalg.cholesky(z["K"])

        _, val = potential_grad(potential_fn_chol, z)
        self.assertTrue(torch.isnan(val))

        # Default behavior should not catch other errors
        def potential_fn_rterr_foo(z):
            raise RuntimeError("foo")

        with self.assertRaisesRegex(RuntimeError, "foo"):
            potential_grad(potential_fn_rterr_foo, z)

        # But once we register this specific error then it should
        def catch_runtime_error(e):
            return type(e) is RuntimeError and "foo" in str(e)

        register_exception_handler("foo_runtime", catch_runtime_error)
        _, val = potential_grad(potential_fn_rterr_foo, z)
        self.assertTrue(torch.isnan(val))

        # Unless the error message is different
        def potential_fn_rterr_bar(z):
            raise RuntimeError("bar")

        with self.assertRaisesRegex(RuntimeError, "bar"):
            potential_grad(potential_fn_rterr_bar, z)
//  [python]
def log_det_by_cholesky(matrix):
    """
    Args:
            matrix: matrix must be a positive define matrix.
                            shape [N, C, D, D].
    Ref:
            https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
    """
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.cholesky(matrix)
    return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def log_det_by_cholesky(matrix):
    """
    Args:
            matrix: matrix must be a positive define matrix.
                            shape [N, C, D, D].
    Ref:
            https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
    """
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    # chol = torch.cholesky(matrix)
    chol = torch.linalg.cholesky(matrix, upper=False)
    # return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
    return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False
    ):
        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()
        
        if actorder or self.owq:
            W = W[:, self.ids]
            self.H = self.H[self.ids][:,self.ids]
        
        self.quantizer.find_params(W[:,:self.n_nonout], weight=True)

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

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)
        
        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.inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H
        
        for i1 in range(0, self.n_nonout, blocksize):
            i2 = min(i1 + blocksize, self.n_nonout)
            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]

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

                if groupsize != -1:
                    if (i1 + i) % groupsize == 0:
                        self.quantizer.find_params(W[:, (i1 + i):min((i1 + i + groupsize),(self.columns-self.n_out))], weight=True, num=40)

                q = self.quantizer.quantize(w.unsqueeze(1)).flatten()
                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

            Q[:, i1:i2] = Q1
            Losses[:, i1:i2] = Losses1 / 2

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

        torch.cuda.synchronize()
        print('time %.2f' % (time.time() - tick))
        print('error', torch.sum(Losses).item())
               
        if actorder or self.owq:
            Q[:,self.n_nonout:] = W[:,self.n_nonout:]
            invids = torch.argsort(self.ids)
            Q = Q[:, invids]
        
        if isinstance(self.layer, transformers.Conv1D):
            Q = Q.t()
        
        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False
    ):
        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()
        
        if actorder or self.owq:
            W = W[:, self.ids]
            self.H = self.H[self.ids][:,self.ids]
        
        self.quantizer.find_params(W[:,:self.n_nonout], weight=True)

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

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)
        
        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
        
        for i1 in range(0, self.n_nonout, blocksize):
            i2 = min(i1 + blocksize, self.n_nonout)
            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]

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

                if groupsize != -1:
                    if (i1 + i) % groupsize == 0:
                        self.quantizer.find_params(W[:, (i1 + i):min((i1 + i + groupsize),(self.columns-self.n_out))], weight=True, num=40)

                q = self.quantizer.quantize(w.unsqueeze(1)).flatten()
                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

            Q[:, i1:i2] = Q1
            Losses[:, i1:i2] = Losses1 / 2

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

        torch.cuda.synchronize()
        print('time %.2f' % (time.time() - tick))
        print('error', torch.sum(Losses).item())
               
        if actorder or self.owq:
            Q[:,self.n_nonout:] = W[:,self.n_nonout:]
            invids = torch.argsort(self.ids)
            Q = Q[:, invids]
        
        if isinstance(self.layer, transformers.Conv1D):
            Q = Q.t()
        
        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
//  [python]
def forward(self, inp):
        if self.training:
            with th.no_grad():
                # Ensure input to Cholesky is PD
                cholesky_inp = inp.T @ inp + self.small_number * th.eye(inp.size(1)).type_as(inp)
                ort = th.cholesky(cholesky_inp)
                ort = th.inverse(ort)
            self.update_ort_weights(ort)
        else:
            ort = self.ort_weights

        out = inp @ ort.T
        return out
//  [python]
def forward(self, inp):
        if self.training:
            with th.no_grad():
                # Ensure input to Cholesky is PD
                cholesky_inp = inp.T @ inp + self.small_number * th.eye(inp.size(1)).type_as(inp)
                ort = th.linalg.cholesky(cholesky_inp)
                ort = th.linalg.inv(ort)
            self.update_ort_weights(ort)
        else:
            ort = self.ort_weights

        out = inp @ ort.T
        return out
//  [python]
def batch_cholesky_inverse(matrix):
    """
    Args: 	matrix, 4-D tensor, [N, C, M, M].
                    matrix must be a symmetric positive define matrix.
    """
    chol_low = torch.cholesky(matrix, upper=False)
    chol_low_inv = batch_low_tri_inv(chol_low)
    return torch.matmul(chol_low_inv.transpose(-2, -1), chol_low_inv)
//  [python]
def batch_cholesky_inverse(matrix):
    """
    Args: 	matrix, 4-D tensor, [N, C, M, M].
                    matrix must be a symmetric positive define matrix.
    """
    # chol_low = torch.cholesky(matrix, upper=False)
    chol_low = torch.linalg.cholesky(matrix, upper=False)
    chol_low_inv = batch_low_tri_inv(chol_low)
    return torch.matmul(chol_low_inv.transpose(-2, -1), chol_low_inv)
//  [python]
def backward(ctx, grad_output):
    batch_size, n_dim = grad_output.shape
    sign_z, = ctx.saved_tensors
    device = grad_output.device
    grad_output = grad_output
    sign_z = sign_z

    S = sign_z != 0
    S[:, -1] = True
    sign_z[:, -1] = 0

    S = S.cuda()
    s_active = torch.zeros(batch_size, n_dim, device=device, dtype=int)
    s_num_active = torch.zeros(batch_size, device=device, dtype=int)
    check_s_active(S, s_active, s_num_active)
    L_S_padded = torch.zeros(batch_size, n_dim, n_dim,
                             device=device, dtype=grad_output.dtype)
    setup_l_s(L_S_padded, s_active, s_num_active)
    LTL_all = L_S_padded.transpose(-1, -2).matmul(L_S_padded)
    pad_ltl_diag(LTL_all, s_num_active)
    grad_u = (grad_output.unsqueeze(-2)
              ).matmul(L_S_padded).squeeze().unsqueeze(-1)

    # Cholesky Solve.
    u_all = torch.cholesky(LTL_all)
    tmp0 = torch.cholesky_solve(grad_u, u_all)
    if ctx.x_need_grad:
      grad_x = L_S_padded.matmul(tmp0).squeeze()
    else:
      grad_x = None
    sign_all = torch.zeros(
        batch_size, n_dim, device=device, dtype=grad_output.dtype)
    setup_sign_all(sign_all, sign_z, s_active, s_num_active)
    grad_lbda = (sign_all.unsqueeze(-2)).matmul(tmp0).squeeze(-1)
    return (grad_x, grad_lbda)
//  [python]
def backward(ctx, grad_output):
    batch_size, n_dim = grad_output.shape
    sign_z, = ctx.saved_tensors
    device = grad_output.device
    grad_output = grad_output
    sign_z = sign_z

    S = sign_z != 0
    S[:, -1] = True
    sign_z[:, -1] = 0

    S = S.cuda()
    s_active = torch.zeros(batch_size, n_dim, device=device, dtype=int)
    s_num_active = torch.zeros(batch_size, device=device, dtype=int)
    check_s_active(S, s_active, s_num_active)
    L_S_padded = torch.zeros(batch_size, n_dim, n_dim,
                             device=device, dtype=grad_output.dtype)
    setup_l_s(L_S_padded, s_active, s_num_active)
    LTL_all = L_S_padded.transpose(-1, -2).matmul(L_S_padded)
    pad_ltl_diag(LTL_all, s_num_active)
    grad_u = (grad_output.unsqueeze(-2)
              ).matmul(L_S_padded).squeeze().unsqueeze(-1)

    # Cholesky Solve.
    u_all = torch.linalg.cholesky(LTL_all)
    tmp0 = torch.cholesky_solve(grad_u, u_all)
    if ctx.x_need_grad:
      grad_x = L_S_padded.matmul(tmp0).squeeze()
    else:
      grad_x = None
    sign_all = torch.zeros(
        batch_size, n_dim, device=device, dtype=grad_output.dtype)
    setup_sign_all(sign_all, sign_z, s_active, s_num_active)
    grad_lbda = (sign_all.unsqueeze(-2)).matmul(tmp0).squeeze(-1)
    return (grad_x, grad_lbda)
//  [python]
def cholesky_factor(self, matrix, jitter_level):
        """
        Compute lower Cholesky factor of given matrix, adding jitter for stability (prevent non-PD error).

        :param matrix: torch.Tensor, matrix subject to Cholesky decomposition
        :param jitter_level: float, jitter added for numerical stability
        :return: torch.Tensor, lower Cholesky factor of input matrix
        """
        jitter = torch.eye(matrix.shape[0], dtype=matrix.dtype, device=matrix.device) * jitter_level
        multiplier = 1.
        while True:
            try:
                L = torch.cholesky(matrix + multiplier * jitter)  # try to compute lower Cholesky factor
                break
            except RuntimeError:
                multiplier *= 2.
                if float(multiplier) == float("inf"):
                    raise RuntimeError("increase to inf jitter")
        return L
//  [python]
def cholesky_factor(self, matrix, jitter_level):
        """
        Compute lower Cholesky factor of given matrix, adding jitter for stability (prevent non-PD error).

        :param matrix: torch.Tensor, matrix subject to Cholesky decomposition
        :param jitter_level: float, jitter added for numerical stability
        :return: torch.Tensor, lower Cholesky factor of input matrix
        """
        jitter = torch.eye(matrix.shape[0], dtype=matrix.dtype, device=matrix.device) * jitter_level
        multiplier = 1.
        while True:
            try:
                L = torch.linalg.cholesky(matrix + multiplier * jitter)  # try to compute lower Cholesky factor
                break
            except RuntimeError:
                multiplier *= 2.
                if float(multiplier) == float("inf"):
                    raise RuntimeError("increase to inf jitter")
        return L
//  [python]
def test_ill_conditioned_matrix():
    """Test with an ill-conditioned matrix."""
    # Create an ill-conditioned matrix
    ill_cond_matrix = t.rand((10, 10))
    ill_cond_matrix = ill_cond_matrix @ ill_cond_matrix.T  # Make it symmetric
    ill_cond_matrix[0, 0] = 1e-8  # Introduce ill-conditioning

    # Verify that the matrix really is ill-conditioned
    with pytest.raises(LinAlgError):
        L = t.cholesky(ill_cond_matrix)

    # Test if stable Cholesky decomposition succeeds
    L = stable_cholesky(ill_cond_matrix)
    assert isinstance(L, t.Tensor)
    assert not t.isnan(L).any()
//  [python]
def test_ill_conditioned_matrix():
    """Test with an ill-conditioned matrix."""
    # Create an ill-conditioned matrix
    ill_cond_matrix = t.rand((10, 10))
    ill_cond_matrix = ill_cond_matrix @ ill_cond_matrix.T  # Make it symmetric
    ill_cond_matrix[0, 0] = 1e-8  # Introduce ill-conditioning

    # Verify that the matrix really is ill-conditioned
    with pytest.raises(LinAlgError):
        L = t.linalg.cholesky(ill_cond_matrix)

    # Test if stable Cholesky decomposition succeeds
    L = stable_cholesky(ill_cond_matrix)
    assert isinstance(L, t.Tensor)
    assert not t.isnan(L).any()
//  [python]
def Gaussian_log_likelihood(y, cov, Kinv_method='cholesky3'):
    """
    Compute the log-likelihood of a Gaussian distribution.

    Args:
        y (torch.Tensor): The observed values.
        mean (torch.Tensor): The mean of the Gaussian distribution.
        cov (torch.Tensor): The covariance matrix of the Gaussian distribution.
        Kinv_method (str, optional): The method to compute the inverse of the covariance matrix.
            Defaults to 'cholesky3'.

    Returns:
        torch.Tensor: The log-likelihood of the Gaussian distribution.

    Raises:
        ValueError: If Kinv_method is not 'direct' or 'cholesky'.
    """
    
    # assert if the correct dimension
    assert len(y.shape) == 2 and len(cov.shape) == 2, "y, mean, cov should be 2D tensors"
    
    if Kinv_method == 'cholesky1':
        L = torch.cholesky(cov)
        L_inv = torch.inverse(L)
        K_inv = L_inv.T @ L_inv
        return -0.5 * (y.T @ K_inv @ y + 2 * torch.logdet(cov) + len(y) * np.log(2 * np.pi))
    elif Kinv_method == 'cholesky2':
        L = torch.cholesky(cov)
        gamma = torch.cholesky_solve(y, L)
        return -0.5 * (gamma.T @ gamma + 2 * torch.logdet(cov) + len(y) * np.log(2 * np.pi))
    
    elif Kinv_method == 'cholesky3':
        # fastest implementation so far
        L = torch.cholesky(cov)
        # return -0.5 * (y_use.T @ torch.cholesky_solve(y_use, L) + L.diag().log().sum() + len(x_train) * np.log(2 * np.pi))
        if y.shape[1] > 1:
            Warning('y_use.shape[1] > 1, will treat each column as a sample (for the joint normal distribution) and sum the log-likelihood')
            # 
            # (Alpha ** 2).sum() = (Alpha @ Alpha^T).diag().sum() = \sum_i (Alpha @ Alpha^T)_{ii}
            # 
            y_dim = y.shape[1]
            log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
            gamma = torch.cholesky_solve(y, L, upper = False)
            return - 0.5 * ( (gamma ** 2).sum() + log_det_K * y_dim + len(y) * y_dim * np.log(2 * np.pi) )
        else:
            gamma = torch.cholesky_solve(y, L, upper = False)
            return -0.5 * (gamma.T @ gamma + 2 * L.diag().log().sum() + len(y) * np.log(2 * np.pi))

    elif Kinv_method == 'direct':
        K_inv = torch.inverse(cov)
        return -0.5 * (y.T @ K_inv @ y + 2 * torch.logdet(cov) + len(y) * np.log(2 * np.pi))
    elif Kinv_method == 'torch_distribution_MN1':
        L = torch.cholesky(cov)
        return torch.distributions.MultivariateNormal(y, scale_tril=L).log_prob(y)
    elif Kinv_method == 'torch_distribution_MN2':
        return torch.distributions.MultivariateNormal(y, cov).log_prob(y)
    else:
        raise ValueError('Kinv_method should be either direct or cholesky')
//  [python]
def Gaussian_log_likelihood(y, cov, Kinv_method='cholesky3'):
    """
    Compute the log-likelihood of a Gaussian distribution.

    Args:
        y (torch.Tensor): The observed values.
        mean (torch.Tensor): The mean of the Gaussian distribution.
        cov (torch.Tensor): The covariance matrix of the Gaussian distribution.
        Kinv_method (str, optional): The method to compute the inverse of the covariance matrix.
            Defaults to 'cholesky3'.

    Returns:
        torch.Tensor: The log-likelihood of the Gaussian distribution.

    Raises:
        ValueError: If Kinv_method is not 'direct' or 'cholesky'.
    """
    
    # assert if the correct dimension
    assert len(y.shape) == 2 and len(cov.shape) == 2, "y, mean, cov should be 2D tensors"
    
    if Kinv_method == 'cholesky1':
        L = torch.linalg.cholesky(cov)
        L_inv = torch.inverse(L)
        K_inv = L_inv.T @ L_inv
        return -0.5 * (y.T @ K_inv @ y + 2 * torch.logdet(cov) + len(y) * np.log(2 * np.pi))
    elif Kinv_method == 'cholesky2':
        L = torch.linalg.cholesky(cov)
        gamma = torch.cholesky_solve(y, L)
        return -0.5 * (gamma.T @ gamma + 2 * torch.logdet(cov) + len(y) * np.log(2 * np.pi))
    
    elif Kinv_method == 'cholesky3':
        # fastest implementation so far
        L = torch.linalg.cholesky(cov)
        # return -0.5 * (y_use.T @ torch.cholesky_solve(y_use, L) + L.diag().log().sum() + len(x_train) * np.log(2 * np.pi))
        if y.shape[1] > 1:
            Warning('y_use.shape[1] > 1, will treat each column as a sample (for the joint normal distribution) and sum the log-likelihood')
            # 
            # (Alpha ** 2).sum() = (Alpha @ Alpha^T).diag().sum() = \sum_i (Alpha @ Alpha^T)_{ii}
            # 
            y_dim = y.shape[1]
            log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
            gamma = torch.cholesky_solve(y, L, upper = False)
            return - 0.5 * ( (gamma ** 2).sum() + log_det_K * y_dim + len(y) * y_dim * np.log(2 * np.pi) )
        else:
            gamma = torch.cholesky_solve(y, L, upper = False)
            return -0.5 * (gamma.T @ gamma + 2 * L.diag().log().sum() + len(y) * np.log(2 * np.pi))

    elif Kinv_method == 'direct':
        K_inv = torch.inverse(cov)
        return -0.5 * (y.T @ K_inv @ y + 2 * torch.logdet(cov) + len(y) * np.log(2 * np.pi))
    elif Kinv_method == 'torch_distribution_MN1':
        L = torch.linalg.cholesky(cov)
        return torch.distributions.MultivariateNormal(y, scale_tril=L).log_prob(y)
    elif Kinv_method == 'torch_distribution_MN2':
        return torch.distributions.MultivariateNormal(y, cov).log_prob(y)
    else:
        raise ValueError('Kinv_method should be either direct or cholesky')
//  [python]
def test_solve_least_squares():
    n = 100
    m = 50
    d = 4
    A = torch.randn(n, m)
    B = torch.randn(n, d)
    C = torch.cholesky(A.t() @ A)
    naive_W = torch.cholesky_solve(A.t() @ B, C)
    stable_W = solve_least_squares(A, B)
    with torch.no_grad():
        npt.assert_array_almost_equal(naive_W, stable_W)
//  [python]
def test_solve_least_squares():
    n = 100
    m = 50
    d = 4
    A = torch.randn(n, m)
    B = torch.randn(n, d)
    C = torch.linalg.cholesky(A.t() @ A)
    naive_W = torch.cholesky_solve(A.t() @ B, C)
    stable_W = solve_least_squares(A, B)
    with torch.no_grad():
        npt.assert_array_almost_equal(naive_W, stable_W)
//  [python]
def _solve_linear_system(self, A, B):
        """Solves linear system AX = B.
        If B is a tuple (B1, B2, ...), returns tuple (X1, X2, ...).
        Otherwise returns X.
        """
        B_sizes = None
        # If B is a tuple, concatenate into single tensor:
        if isinstance(B, (tuple, list)):
            B_sizes = list(map(lambda x: x.size(-1), B))
            B = torch.cat(B, dim=-1)
        # Ensure B is 2D (bxmxn):
        if len(B.size()) == 2:
            B = B.unsqueeze(-1)
        try: # Batchwise Cholesky solve
            A_decomp = torch.cholesky(A, upper=False)
            X = torch.cholesky_solve(B, A_decomp, upper=False) # bxmxn
        except RuntimeError: # Revert to loop if batchwise solve fails
            X = torch.zeros_like(B)
            for i in range(A.size(0)):
                try: # Cholesky solve
                    A_decomp = torch.cholesky(A[i, ...], upper=False)
                    X[i, ...] = torch.cholesky_solve(B[i, ...], A_decomp, upper=False) # mxn
                except RuntimeError: # Revert to LU solve
                    X[i, ...] = torch.linalg.solve(A[i, ...], B[i, ...])[0] # mxn
        if B_sizes is not None:
            X = X.split(B_sizes, dim=-1)
        return X
//  [python]
def _solve_linear_system(self, A, B):
        """Solves linear system AX = B.
        If B is a tuple (B1, B2, ...), returns tuple (X1, X2, ...).
        Otherwise returns X.
        """
        B_sizes = None
        # If B is a tuple, concatenate into single tensor:
        if isinstance(B, (tuple, list)):
            B_sizes = list(map(lambda x: x.size(-1), B))
            B = torch.cat(B, dim=-1)
        # Ensure B is 2D (bxmxn):
        if len(B.size()) == 2:
            B = B.unsqueeze(-1)
        try: # Batchwise Cholesky solve
            A_decomp = torch.linalg.cholesky(A, upper=False)
            X = torch.cholesky_solve(B, A_decomp, upper=False) # bxmxn
        except RuntimeError: # Revert to loop if batchwise solve fails
            X = torch.zeros_like(B)
            for i in range(A.size(0)):
                try: # Cholesky solve
                    A_decomp = torch.linalg.cholesky(A[i, ...], upper=False)
                    X[i, ...] = torch.cholesky_solve(B[i, ...], A_decomp, upper=False) # mxn
                except RuntimeError: # Revert to LU solve
                    X[i, ...] = torch.linalg.solve(A[i, ...], B[i, ...])[0] # mxn
        if B_sizes is not None:
            X = X.split(B_sizes, dim=-1)
        return X
//  [python]
def compute_loss(self, x, y, x_var=0., y_var=0., update_data=False):
        """
        Compute the loss of the module.

        Args:
            x (torch.Tensor or list): Input tensor or list of tensors.
            y (torch.Tensor): Target tensor.
            x_var (float, optional): Variance of the input tensor. Defaults to 0.
            y_var (float, optional): Variance of the target tensor. Defaults to 0.
            update_data (bool, optional): Whether to update the training data. Defaults to False.

        Returns:
            torch.Tensor: Loss value.
        """
        x = self.check_single_tensor(x)
        y = self.check_single_tensor(y)
        assert y.ndim == 2, "y should be 2d tensor"

        # TODO checking if inputs/outputs was changed
        if self.train_x is None:
            self.train_x = x
            self.train_y = y
        elif update_data:
            self.train_x = x
            self.train_y = y

        _noise = self.noise_box.get()

        l1, h1, l2, h2 = self.l1, self.h1, self.l2, self.h2

        m_device = list(self.parameters())[0].device
        sigma = self.kernel(x, x, l1, h1, l2, h2) + _noise.pow(-1) * torch.eye(x.size(0), device=m_device)
        sigma = sigma + JITTER * torch.eye(x.size(0), device=m_device)

        L = torch.cholesky(sigma)
        y_num, y_dimension = y.shape
        Gamma,_ = torch.triangular_solve(y, L, upper = False)
        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 compute_loss(self, x, y, x_var=0., y_var=0., update_data=False):
        """
        Compute the loss of the module.

        Args:
            x (torch.Tensor or list): Input tensor or list of tensors.
            y (torch.Tensor): Target tensor.
            x_var (float, optional): Variance of the input tensor. Defaults to 0.
            y_var (float, optional): Variance of the target tensor. Defaults to 0.
            update_data (bool, optional): Whether to update the training data. Defaults to False.

        Returns:
            torch.Tensor: Loss value.
        """
        x = self.check_single_tensor(x)
        y = self.check_single_tensor(y)
        assert y.ndim == 2, "y should be 2d tensor"

        # TODO checking if inputs/outputs was changed
        if self.train_x is None:
            self.train_x = x
            self.train_y = y
        elif update_data:
            self.train_x = x
            self.train_y = y

        _noise = self.noise_box.get()

        l1, h1, l2, h2 = self.l1, self.h1, self.l2, self.h2

        m_device = list(self.parameters())[0].device
        sigma = self.kernel(x, x, l1, h1, l2, h2) + _noise.pow(-1) * torch.eye(x.size(0), device=m_device)
        sigma = sigma + JITTER * torch.eye(x.size(0), device=m_device)

        L = torch.linalg.cholesky(sigma)
        y_num, y_dimension = y.shape
        Gamma,_ = torch.triangular_solve(y, L, upper = False)
        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) -> torch.Tensor:
        """Return a value proportional to the log likelihood."""
        # model-implied covariance/mean
        sigma, mu = self.implied_sigma_mu()

        if self.fiml:
            # FIML -2 * logL (without constants)
            assert mu is not None
            loss = torch.zeros(1, dtype=mu.dtype, device=mu.device)
            mean_diffs: torch.Tensor = self.data - mu.t()
            for pairs in self.missing_patterns.values():
                # calculate sigma^-1 and logdet(sigma) in batches
                sigmas = torch.stack(
                    [sigma.index_select(0, x[1]).index_select(1, x[1]) for x in pairs],
                )
                L_sigmas = torch.cholesky(sigmas)
                sigmas_logdet = L_sigmas.diagonal(dim1=1, dim2=2).log().sum(dim=1) * 2

                for i, (observations, available) in enumerate(pairs):
                    mean_diff = mean_diffs.index_select(0, observations).index_select(
                        1, available,
                    )
                    loss_current = sigmas_logdet[i] * len(observations) + torch.trace(
                        mean_diff.mm(torch.cholesky_solve(mean_diff.t(), L_sigmas[i])),
                    )
                    loss = loss + loss_current.clamp(min=0.0)

        else:
            # maximum likelihood
            L_sigma = torch.cholesky(sigma)
            loss = L_sigma.diagonal().log().sum() * 2 + torch.trace(
                torch.cholesky_solve(self.sample_covariance, L_sigma),
            )

        return loss
//  [python]
def forward(self) -> torch.Tensor:
        """Return a value proportional to the log likelihood."""
        # model-implied covariance/mean
        sigma, mu = self.implied_sigma_mu()

        if self.fiml:
            # FIML -2 * logL (without constants)
            assert mu is not None
            loss = torch.zeros(1, dtype=mu.dtype, device=mu.device)
            mean_diffs: torch.Tensor = self.data - mu.t()
            for pairs in self.missing_patterns.values():
                # calculate sigma^-1 and logdet(sigma) in batches
                sigmas = torch.stack(
                    [sigma.index_select(0, x[1]).index_select(1, x[1]) for x in pairs],
                )
                L_sigmas = torch.linalg.cholesky(sigmas)
                sigmas_logdet = L_sigmas.diagonal(dim1=1, dim2=2).log().sum(dim=1) * 2

                for i, (observations, available) in enumerate(pairs):
                    mean_diff = mean_diffs.index_select(0, observations).index_select(
                        1, available,
                    )
                    loss_current = sigmas_logdet[i] * len(observations) + torch.trace(
                        mean_diff.mm(torch.cholesky_solve(mean_diff.t(), L_sigmas[i])),
                    )
                    loss = loss + loss_current.clamp(min=0.0)

        else:
            # maximum likelihood
            L_sigma = torch.linalg.cholesky(sigma)
            loss = L_sigma.diagonal().log().sum() * 2 + torch.trace(
                torch.cholesky_solve(self.sample_covariance, L_sigma),
            )

        return loss
//  [python]
def test_multivariate_standardize_tfm():
    d = 5
    tril_ind = torch.tril_indices(d, d)
    Q = torch.zeros(d, d)
    Q[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
    Sigma = Q @ Q.t() + torch.eye(d) * 1e-3
    L = torch.cholesky(Sigma)
    x = torch.randn(1000, d) @ L.t() + torch.arange(d)
    tfm = MultivariateStandardizeTransform(x)
    y = tfm(x)
    with torch.no_grad():
        npt.assert_array_almost_equal(y.mean(0), torch.zeros(d))
        npt.assert_array_almost_equal(sample_covariance(y), torch.eye(d))
        npt.assert_array_almost_equal(tfm.inverse(y), x)
//  [python]
def test_multivariate_standardize_tfm():
    d = 5
    tril_ind = torch.tril_indices(d, d)
    Q = torch.zeros(d, d)
    Q[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
    Sigma = Q @ Q.t() + torch.eye(d) * 1e-3
    L = torch.linalg.cholesky(Sigma)
    x = torch.randn(1000, d) @ L.t() + torch.arange(d)
    tfm = MultivariateStandardizeTransform(x)
    y = tfm(x)
    with torch.no_grad():
        npt.assert_array_almost_equal(y.mean(0), torch.zeros(d))
        npt.assert_array_almost_equal(sample_covariance(y), torch.eye(d))
        npt.assert_array_almost_equal(tfm.inverse(y), x)
//  [python]
def solve_trust_region_2d(B, g, Delta):
    """Solve a general trust-region problem in 2 dimensions.
    The problem is reformulated as a 4th order algebraic equation,
    the solution of which is found by numpy.roots.
    """
    try:
        L = torch.cholesky(B)
        p = - torch.cholesky_solve(g.unsqueeze(1), L).squeeze(1)
        if p.dot(p) <= Delta**2:
            return p, True
    except RuntimeError as exc:
        if not 'cholesky' in exc.args[0]:
            raise

    # move things to numpy
    device = B.device
    dtype = B.dtype
    B = B.data.cpu().numpy()
    g = g.data.cpu().numpy()
    Delta = float(Delta)

    a = B[0, 0] * Delta**2
    b = B[0, 1] * Delta**2
    c = B[1, 1] * Delta**2
    d = g[0] * Delta
    f = g[1] * Delta

    coeffs = np.array([-b + d, 2 * (a - c + f), 6 * b, 2 * (-a + c + f), -b - d])
    t = np.roots(coeffs)  # Can handle leading zeros.
    t = np.real(t[np.isreal(t)])

    p = Delta * np.vstack((2 * t / (1 + t**2), (1 - t**2) / (1 + t**2)))
    value = 0.5 * np.sum(p * B.dot(p), axis=0) + np.dot(g, p)
    p = p[:, np.argmin(value)]

    # convert back to torch
    p = torch.tensor(p, device=device, dtype=dtype)

    return p, False
//  [python]
def solve_trust_region_2d(B, g, Delta):
    """Solve a general trust-region problem in 2 dimensions.
    The problem is reformulated as a 4th order algebraic equation,
    the solution of which is found by numpy.roots.
    """
    try:
        L = torch.linalg.cholesky(B)
        p = - torch.cholesky_solve(g.unsqueeze(1), L).squeeze(1)
        if p.dot(p) <= Delta**2:
            return p, True
    except RuntimeError as exc:
        if not 'cholesky' in exc.args[0]:
            raise

    # move things to numpy
    device = B.device
    dtype = B.dtype
    B = B.data.cpu().numpy()
    g = g.data.cpu().numpy()
    Delta = float(Delta)

    a = B[0, 0] * Delta**2
    b = B[0, 1] * Delta**2
    c = B[1, 1] * Delta**2
    d = g[0] * Delta
    f = g[1] * Delta

    coeffs = np.array([-b + d, 2 * (a - c + f), 6 * b, 2 * (-a + c + f), -b - d])
    t = np.roots(coeffs)  # Can handle leading zeros.
    t = np.real(t[np.isreal(t)])

    p = Delta * np.vstack((2 * t / (1 + t**2), (1 - t**2) / (1 + t**2)))
    value = 0.5 * np.sum(p * B.dot(p), axis=0) + np.dot(g, p)
    p = p[:, np.argmin(value)]

    # convert back to torch
    p = torch.tensor(p, device=device, dtype=dtype)

    return p, False
//  [python]
def cholesky(mat):
    return torch.cholesky(mat)
//  [python]
def cholesky(mat):
    return torch.linalg.cholesky(mat)
//  [python]
def sample_parameters(self) -> torch.Tensor:
        L = torch.cholesky(self._U)
        return self._R + L.matmul(torch.randn_like(self._R))
//  [python]
def sample_parameters(self) -> torch.Tensor:
        L = torch.linalg.cholesky(self._U)
        return self._R + L.matmul(torch.randn_like(self._R))
//  [python]
def model(self):
        self.set_mode("model")

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

        zero_loc = self.X.new_zeros(self.X.size(0))
        f_loc = zero_loc + self.mean_function(self.X)
        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 model(self):
        self.set_mode("model")

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

        zero_loc = self.X.new_zeros(self.X.size(0))
        f_loc = zero_loc + self.mean_function(self.X)
        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 full_mixture_posterior_predictive(model, test_x, trace_hyper):
     
      ''' Returns the mixture posterior predictive multivariate normal '''
     
      # Make predictions by feeding model through likelihood
     
      list_of_y_pred_dists = []
      
      for i in range(len(trace_hyper)):
          
         hyper_sample = trace_hyper[i]
         
         ## Training mode for overwriting hypers 
         model.train()
         model.likelihood.train()
         
         if hyper_sample['sig_n']**2 < 1e-4:
             hyper_sample['sig_n'] = 0.01
             
         model.likelihood.noise_covar.noise = hyper_sample['sig_n']**2
         model.base_covar_module.outputscale = hyper_sample['sig_f']**2
         model.base_covar_module.base_kernel.lengthscale = hyper_sample['ls']
         model.covar_module.inducing_points = torch.nn.Parameter(torch.Tensor(hyper_sample['Z']).double())
         
         with torch.no_grad():
             
              ## Testing mode for computing the posterior predictive 
              model.eval()
              model.likelihood.eval()
              pred = model.likelihood(model(test_x))
              
              try:
                    chol = torch.cholesky(pred.covariance_matrix + torch.eye(len(test_x))*1e-5)
                    list_of_y_pred_dists.append(pred)
              except RuntimeError:
                   print('Not psd for sample ' + str(i))
     
      return list_of_y_pred_dists
//  [python]
def full_mixture_posterior_predictive(model, test_x, trace_hyper):
     
      ''' Returns the mixture posterior predictive multivariate normal '''
     
      # Make predictions by feeding model through likelihood
     
      list_of_y_pred_dists = []
      
      for i in range(len(trace_hyper)):
          
         hyper_sample = trace_hyper[i]
         
         ## Training mode for overwriting hypers 
         model.train()
         model.likelihood.train()
         
         if hyper_sample['sig_n']**2 < 1e-4:
             hyper_sample['sig_n'] = 0.01
             
         model.likelihood.noise_covar.noise = hyper_sample['sig_n']**2
         model.base_covar_module.outputscale = hyper_sample['sig_f']**2
         model.base_covar_module.base_kernel.lengthscale = hyper_sample['ls']
         model.covar_module.inducing_points = torch.nn.Parameter(torch.Tensor(hyper_sample['Z']).double())
         
         with torch.no_grad():
             
              ## Testing mode for computing the posterior predictive 
              model.eval()
              model.likelihood.eval()
              pred = model.likelihood(model(test_x))
              
              try:
                    chol = torch.linalg.cholesky(pred.covariance_matrix + torch.eye(len(test_x))*1e-5)
                    list_of_y_pred_dists.append(pred)
              except RuntimeError:
                   print('Not psd for sample ' + str(i))
     
      return list_of_y_pred_dists
//  [python]
def random_mvn(batch_shape, dim, requires_grad=False):
    rank = dim + dim
    loc = torch.randn(batch_shape + (dim,), requires_grad=requires_grad)
    cov = torch.randn(batch_shape + (dim, rank))
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.cholesky(cov)
    scale_tril.requires_grad_(requires_grad)
    return dist.MultivariateNormal(loc, scale_tril=scale_tril)
//  [python]
def random_mvn(batch_shape, dim, requires_grad=False):
    rank = dim + dim
    loc = torch.randn(batch_shape + (dim,), requires_grad=requires_grad)
    cov = torch.randn(batch_shape + (dim, rank))
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.linalg.cholesky(cov)
    scale_tril.requires_grad_(requires_grad)
    return dist.MultivariateNormal(loc, scale_tril=scale_tril)
//  [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:
        L = torch.cholesky(A, upper=upper, out=out)
        return L
    except RuntimeError as e:
        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(10):
            jitter_new = jitter * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                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
            except RuntimeError:
                continue
        # return torch.randn_like(A).tril()
        raise e
//  [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:
        L = torch.linalg.cholesky(A, upper=upper, out=out)
        return L
    except RuntimeError as e:
        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(10):
            jitter_new = jitter * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                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
            except RuntimeError:
                continue
        # return torch.randn_like(A).tril()
        raise e
//  [python]
def test_covariance_matrix(self):
        for d in self.distributions:
            cov = d.covariance_matrix
            self.assertTrue(isinstance(cov, Tensor))

            # testing for symmetry
            self.assertAllClose(cov, cov.mT)

            # testing for positive-definiteness
            ispd = False
            try:
                torch.cholesky(cov)
                ispd = True
            except RuntimeError:
                pass
            self.assertTrue(ispd)

            # checking that linear operator to tensor conversion
            # leads to same covariance matrix
            xcov_linop = DenseLinearOperator(d.cross_covariance_matrix)
            usn_linop = UnifiedSkewNormal(
                trunc=d.trunc, gauss=d.gauss, cross_covariance_matrix=xcov_linop
            )
            cov_linop = usn_linop.covariance_matrix
            self.assertTrue(isinstance(cov_linop, Tensor))
            self.assertAllClose(cov, cov_linop)
//  [python]
def test_covariance_matrix(self):
        for d in self.distributions:
            cov = d.covariance_matrix
            self.assertTrue(isinstance(cov, Tensor))

            # testing for symmetry
            self.assertAllClose(cov, cov.mT)

            # testing for positive-definiteness
            ispd = False
            try:
                torch.linalg.cholesky(cov)
                ispd = True
            except RuntimeError:
                pass
            self.assertTrue(ispd)

            # checking that linear operator to tensor conversion
            # leads to same covariance matrix
            xcov_linop = DenseLinearOperator(d.cross_covariance_matrix)
            usn_linop = UnifiedSkewNormal(
                trunc=d.trunc, gauss=d.gauss, cross_covariance_matrix=xcov_linop
            )
            cov_linop = usn_linop.covariance_matrix
            self.assertTrue(isinstance(cov_linop, Tensor))
            self.assertAllClose(cov, cov_linop)
//  [python]
def sample(self, X_in: torch.Tensor, **kwargs) -> torch.Tensor:
        """Sample current layer neuron values given upper layer neuron values. When
        self._sample_strat is LayerSampleStrat.MAP and LayerSampleStrat.P_PRED, sample from
        p(X_out| argmax_W p(W), self._Sigma, ...) and E[p(X_out| W, self._Sigma, ...)].
        NOTE: Treats elements in batch independently.

        Args:
            X_in (torch.Tensor): Upper layer neuron values of shape <d_batch x d_in>.

        Returns:
            torch.Tensor: Sampled lower layer neuron values of shape <d_batch x d_out>.
        """
        if self._update_strat == LayerUpdateStrat.KERNEL:
            # HACK: Create a new class for kernel stuff later
            mus, Sigmas = kernel_posterior_params(X_train=self._X_train, Y_train=self._Y_train,
                                                  X_test=X_in, kernel=self._kernel_type,
                                                  kernel_params=self._kernel_params,
                                                  Sigma_trtr_inv=self._X_train_kernel)
            L = torch.cholesky(Sigmas)
            return mus + L.matmul(torch.randn_like(mus))

        marginal_mean = self.predict(X_in=X_in)
        if self._sample_strat == LayerSampleStrat.MAP:
            marginal_Sigma = self._Sigma
        elif self._sample_strat == LayerSampleStrat.P_PRED:
            Z_in = self._f(X_in)
            marginal_Sigma = self._Sigma + Z_in.matmul(self._U).matmul(Z_in.T).diag()
            marginal_Sigma = marginal_Sigma.unsqueeze(-1)
        else:
            raise NotImplementedError()
        dist = pdists.Normal(torch.zeros_like(marginal_mean), marginal_Sigma ** 0.5)
        obs = None if kwargs.get('X_obs') is None else (kwargs['X_obs']-marginal_mean)
        sample = marginal_mean + pyro.sample(self._name, dist, obs=obs)
        return sample