jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fasterprune(self, gpt, blocksize=128, percdamp=0.01):
        import transformers

        sparsity = self.target_sparsity_ratio
        W = gpt.module.weight.data.clone()
        dev = gpt.dev
        rows = gpt.rows
        columns = gpt.columns
        H = gpt.H
        module = gpt.module
        if isinstance(module, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(module, transformers.Conv1D):
            W = W.t()
        W = W.float()
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros(rows, device=dev)
        damp = percdamp * torch.mean(torch.diag(H))  # λI
        diag = torch.arange(columns, device=dev)
        H[diag, diag] += damp  # H = (X*X.t() + λI)
        H = torch.cholesky(H)  # the default is lower triangle
        H = torch.cholesky_inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H

        for i1 in range(0, columns, blocksize):
            i2 = min(i1 + blocksize, 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]

            tmp = W1**2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
            thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
            mask1 = tmp <= thresh

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]
                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:])
        if torch.cuda.is_available():
            torch.cuda.synchronize()

        if isinstance(module, transformers.Conv1D):
            W = W.t()
        param = module.weight
        param_shape = safe_get_shape(param)
        module.weight.data = W.reshape(param_shape).to(dtype=module.weight.data.dtype)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
//  [python]
def fasterprune(self, gpt, blocksize=128, percdamp=0.01):
        import transformers

        sparsity = self.target_sparsity_ratio
        W = gpt.module.weight.data.clone()
        dev = gpt.dev
        rows = gpt.rows
        columns = gpt.columns
        H = gpt.H
        module = gpt.module
        if isinstance(module, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(module, transformers.Conv1D):
            W = W.t()
        W = W.float()
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros(rows, device=dev)
        damp = percdamp * torch.mean(torch.diag(H))  # λI
        diag = torch.arange(columns, device=dev)
        H[diag, diag] += damp  # H = (X*X.t() + λI)
        H = torch.linalg.cholesky(H)  # the default is lower triangle
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H, upper=True)
        Hinv = H

        for i1 in range(0, columns, blocksize):
            i2 = min(i1 + blocksize, 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]

            tmp = W1**2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
            thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
            mask1 = tmp <= thresh

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]
                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:])
        if torch.cuda.is_available():
            torch.cuda.synchronize()

        if isinstance(module, transformers.Conv1D):
            W = W.t()
        param = module.weight
        param_shape = safe_get_shape(param)
        module.weight.data = W.reshape(param_shape).to(dtype=module.weight.data.dtype)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
//  [python]
def fit(self, X):
        """
        X: training data of shape [n, feature_size]
        """
        n, self.feature_size = X.shape
        X = torch.Tensor(X)
        # the parameter of each gaussian distribution
        self.prior_logit = torch.zeros(self.k)
        self.prior_logit.requires_grad_()
        if self.independent_variance:
            self.log_std = torch.log(X.std(dim=0)).repeat(self.k, 1)
            self.log_std.requires_grad_()
            self.mean = torch.zeros(self.k, self.feature_size)
            self.mean.normal_()
            self.mean.requires_grad_()
        else:
            self.cholesky_inverse_cov = torch.cholesky(torch.cov(X.T)).repeat(self.k, 1, 1)
            self.cholesky_inverse_cov.requires_grad_()
            self.mean = torch.zeros(self.k, self.feature_size)
            self.mean.normal_()
            self.mean.requires_grad_()
        self.optimizer = torch.optim.Adam([self.log_std, self.mean, self.prior_logit], lr=self.learning_rate)

        previous_log_likelihood = -math.inf
        for step in range(self.max_step):
            ##########################################
            # Calculate Likelihood
            ##########################################
            # posterior probability of each sample in each Gaussian model
            # it is exactly the likelihood of parameters including mean, std and prior
            log_likelihood = self.log_likelihood(X, input_tensor=True, return_tensor=True)
            neg_log_likelihood = -log_likelihood.mean()

            if self.verbose:
                if step % 1000 == 0:
                    print('Step', step, ', likelihood is', math.exp(-neg_log_likelihood))

            ##########################################
            # Gradient Descent Step
            ##########################################
            self.optimizer.zero_grad()
            neg_log_likelihood.backward()
            self.optimizer.step()

            # early stopping
            log_likelihood = self.log_likelihood(X, input_tensor=True) 
            if self.verbose:
                print('After step', step, ', likelihood of model parameters is', np.exp(log_likelihood))
            if log_likelihood - previous_log_likelihood < self.epsilon:
                break
            previous_log_likelihood = log_likelihood
//  [python]
def fit(self, X):
        """
        X: training data of shape [n, feature_size]
        """
        n, self.feature_size = X.shape
        X = torch.Tensor(X)
        # the parameter of each gaussian distribution
        self.prior_logit = torch.zeros(self.k)
        self.prior_logit.requires_grad_()
        if self.independent_variance:
            self.log_std = torch.log(X.std(dim=0)).repeat(self.k, 1)
            self.log_std.requires_grad_()
            self.mean = torch.zeros(self.k, self.feature_size)
            self.mean.normal_()
            self.mean.requires_grad_()
        else:
            self.cholesky_inverse_cov = torch.linalg.cholesky(torch.cov(X.T)).repeat(self.k, 1, 1)
            self.cholesky_inverse_cov.requires_grad_()
            self.mean = torch.zeros(self.k, self.feature_size)
            self.mean.normal_()
            self.mean.requires_grad_()
        self.optimizer = torch.optim.Adam([self.log_std, self.mean, self.prior_logit], lr=self.learning_rate)

        previous_log_likelihood = -math.inf
        for step in range(self.max_step):
            ##########################################
            # Calculate Likelihood
            ##########################################
            # posterior probability of each sample in each Gaussian model
            # it is exactly the likelihood of parameters including mean, std and prior
            log_likelihood = self.log_likelihood(X, input_tensor=True, return_tensor=True)
            neg_log_likelihood = -log_likelihood.mean()

            if self.verbose:
                if step % 1000 == 0:
                    print('Step', step, ', likelihood is', math.exp(-neg_log_likelihood))

            ##########################################
            # Gradient Descent Step
            ##########################################
            self.optimizer.zero_grad()
            neg_log_likelihood.backward()
            self.optimizer.step()

            # early stopping
            log_likelihood = self.log_likelihood(X, input_tensor=True) 
            if self.verbose:
                print('After step', step, ', likelihood of model parameters is', np.exp(log_likelihood))
            if log_likelihood - previous_log_likelihood < self.epsilon:
                break
            previous_log_likelihood = log_likelihood
//  [python]
def __cholesky(self):
        return torch.cholesky(self.__covariance())
//  [python]
def __cholesky(self):
        return torch.linalg.cholesky(self.__covariance())
//  [python]
def multivariate_stats(x):
    """
    Given a time series x, estimate the mean (mu) and the square root of
    the covariance (sigma) for that time series.
    Inputs:
      x: tensor of shape (tensor(mb_size, channels)) containing the sequence values
    Outputs:
      mu: tensor of shape: (channels,) containing the mean estimates
      sigma: tensor of shape: ((channels, channels) containing an estimate of
      the lower Cholesky factor of the covariance matrix.

      TODO: To improve numerical stability, Use an SVD to compute the Cholesky
      factor rather than the naive formula.
    """
    # Create tensor version of x in case it isn't already
    if not isinstance(x, torch.Tensor):
        x = torch.tensor(x)

    mb_size, channels = x.shape
    mu = torch.mean(x, dim=0)
    error = x - mu.unsqueeze(0).expand((mb_size, channels))
    # error is mb_size x channels
    error1 = error.unsqueeze(2)
    # error1 represents e (mb_size, channels, 1)
    error2 = error.unsqueeze(1)
    # error2 represents e^T (mb_size, 1, channels)
    cov = torch.mean(torch.matmul(error1, error2), dim=0)
    # cov is (channels, channels)

    # Return cholesky factor
    sigma = torch.cholesky(cov)
    return mu, sigma
//  [python]
def multivariate_stats(x):
    """
    Given a time series x, estimate the mean (mu) and the square root of
    the covariance (sigma) for that time series.
    Inputs:
      x: tensor of shape (tensor(mb_size, channels)) containing the sequence values
    Outputs:
      mu: tensor of shape: (channels,) containing the mean estimates
      sigma: tensor of shape: ((channels, channels) containing an estimate of
      the lower Cholesky factor of the covariance matrix.

      TODO: To improve numerical stability, Use an SVD to compute the Cholesky
      factor rather than the naive formula.
    """
    # Create tensor version of x in case it isn't already
    if not isinstance(x, torch.Tensor):
        x = torch.tensor(x)

    mb_size, channels = x.shape
    mu = torch.mean(x, dim=0)
    error = x - mu.unsqueeze(0).expand((mb_size, channels))
    # error is mb_size x channels
    error1 = error.unsqueeze(2)
    # error1 represents e (mb_size, channels, 1)
    error2 = error.unsqueeze(1)
    # error2 represents e^T (mb_size, 1, channels)
    cov = torch.mean(torch.matmul(error1, error2), dim=0)
    # cov is (channels, channels)

    # Return cholesky factor
    sigma = torch.linalg.cholesky(cov)
    return mu, sigma
//  [python]
def log_likelihood(self, x_train, y_train, Kinv_method='cholesky3'):
        """
        Computes the log-likelihood of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

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

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        if isinstance(y_train, list):
            y_train_var = y_train[1]
            y_train = y_train[0]
        else:
            y_train_var = None

        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        if y_train_var is not None:
            K = K + y_train_var 
        
        if Kinv_method == 'cholesky1':
            L = torch.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            return -0.5 * (y_train.T @ K_inv @ y_train + 2 * torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky2':
            L = torch.cholesky(K)
            gamma = torch.cholesky_solve(y_train, L)
            return -0.5 * ((gamma.T @ gamma).sum() + 2 * torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky3':
            L = torch.cholesky(K)
            if y_train.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_train.shape[1]
                log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
                gamma = torch.cholesky_solve(y_train, L, upper = False)
                return - 0.5 * ( (gamma ** 2).sum() + log_det_K * y_dim + len(y_train) * y_dim * np.log(2 * np.pi) )
            else:
                gamma = torch.cholesky_solve(y_train, L, upper = False)
                return -0.5 * (gamma.T @ gamma + 2 * L.diag().log().sum() + len(y_train) * np.log(2 * np.pi))
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            return -0.5 * (y_train.T @ K_inv @ y_train + 2 * torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'torch_distribution_MN1':
            L = torch.cholesky(K)
            return torch.distributions.MultivariateNormal(y_train, scale_tril=L).log_prob(y_train)
        elif Kinv_method == 'torch_distribution_MN2':
            return torch.distributions.MultivariateNormal(y_train, K).log_prob(y_train)
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
//  [python]
def log_likelihood(self, x_train, y_train, Kinv_method='cholesky3'):
        """
        Computes the log-likelihood of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

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

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        if isinstance(y_train, list):
            y_train_var = y_train[1]
            y_train = y_train[0]
        else:
            y_train_var = None

        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        if y_train_var is not None:
            K = K + y_train_var 
        
        if Kinv_method == 'cholesky1':
            L = torch.linalg.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            return -0.5 * (y_train.T @ K_inv @ y_train + 2 * torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky2':
            L = torch.linalg.cholesky(K)
            gamma = torch.cholesky_solve(y_train, L)
            return -0.5 * ((gamma.T @ gamma).sum() + 2 * torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky3':
            L = torch.linalg.cholesky(K)
            if y_train.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_train.shape[1]
                log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
                gamma = torch.cholesky_solve(y_train, L, upper = False)
                return - 0.5 * ( (gamma ** 2).sum() + log_det_K * y_dim + len(y_train) * y_dim * np.log(2 * np.pi) )
            else:
                gamma = torch.cholesky_solve(y_train, L, upper = False)
                return -0.5 * (gamma.T @ gamma + 2 * L.diag().log().sum() + len(y_train) * np.log(2 * np.pi))
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            return -0.5 * (y_train.T @ K_inv @ y_train + 2 * torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'torch_distribution_MN1':
            L = torch.linalg.cholesky(K)
            return torch.distributions.MultivariateNormal(y_train, scale_tril=L).log_prob(y_train)
        elif Kinv_method == 'torch_distribution_MN2':
            return torch.distributions.MultivariateNormal(y_train, K).log_prob(y_train)
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
//  [python]
def fasterquant(self, blocksize=256, percdamp=.01, groupsize=-1):
        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() # Shape: Out X In

        tick = time.time()

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

        M = torch.zeros_like(W) + 1
        E = (torch.zeros_like(W))[:, :blocksize]

        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

        print("Starting pruning: ", end="")

        for i in range(0, self.columns, blocksize):
            i2 = min(i + blocksize, self.columns)
            count = i2 - i
            assert count == blocksize, (
                count, blocksize, i, i2, self.layer.weight.data.shape)

            for ji in range(count):
                j = ji + i

                if ji == 0:
                    # Determine the weights for pruning mask selection for the next column
                    copy_linear = torch.nn.Linear(blocksize, W.shape[0]).to(W.device)
                    copy_linear.weight.data = W[:, j:j+blocksize] ** 2
                    copy_linear.weight.data /= torch.diag(Hinv).unsqueeze(0)[:, j:j+blocksize]

                    prune.l1_unstructured(copy_linear, name='weight', amount=self.amount_prune)
                    print(f"{j}:{j+blocksize}", end = ", ")
                    M[:, j:j+blocksize] = copy_linear.weight_mask
                E[:, j-i] = (1 - M[:, j]) * (W[:, j] / Hinv[j, j])
                W[:, j:i+blocksize] -= E[:, j-i].unsqueeze(1) * Hinv[j, j:i+blocksize].unsqueeze(0)

            W[:, i+blocksize:] -= E.matmul(Hinv[i:i+blocksize, i+blocksize:]) # Keep as is

        if DEVICE.type != "cpu":
            torch.cuda.synchronize()
        print('\nPrune time %.2f' % (time.time() - tick))
        self.new_weight_with_mask = (W, M)
        self.layer.weight.data = (self.layer.weight.data * M).to(self.layer.weight.data.dtype)
//  [python]
def fasterquant(self, blocksize=256, percdamp=.01, groupsize=-1):
        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() # Shape: Out X In

        tick = time.time()

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

        M = torch.zeros_like(W) + 1
        E = (torch.zeros_like(W))[:, :blocksize]

        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

        print("Starting pruning: ", end="")

        for i in range(0, self.columns, blocksize):
            i2 = min(i + blocksize, self.columns)
            count = i2 - i
            assert count == blocksize, (
                count, blocksize, i, i2, self.layer.weight.data.shape)

            for ji in range(count):
                j = ji + i

                if ji == 0:
                    # Determine the weights for pruning mask selection for the next column
                    copy_linear = torch.nn.Linear(blocksize, W.shape[0]).to(W.device)
                    copy_linear.weight.data = W[:, j:j+blocksize] ** 2
                    copy_linear.weight.data /= torch.diag(Hinv).unsqueeze(0)[:, j:j+blocksize]

                    prune.l1_unstructured(copy_linear, name='weight', amount=self.amount_prune)
                    print(f"{j}:{j+blocksize}", end = ", ")
                    M[:, j:j+blocksize] = copy_linear.weight_mask
                E[:, j-i] = (1 - M[:, j]) * (W[:, j] / Hinv[j, j])
                W[:, j:i+blocksize] -= E[:, j-i].unsqueeze(1) * Hinv[j, j:i+blocksize].unsqueeze(0)

            W[:, i+blocksize:] -= E.matmul(Hinv[i:i+blocksize, i+blocksize:]) # Keep as is

        if DEVICE.type != "cpu":
            torch.cuda.synchronize()
        print('\nPrune time %.2f' % (time.time() - tick))
        self.new_weight_with_mask = (W, M)
        self.layer.weight.data = (self.layer.weight.data * M).to(self.layer.weight.data.dtype)
//  [python]
def fasterquant(self, blocksize=128, percdamp=.01, groupsize=-1):
        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 not self.quantizer.ready():
            self.quantizer.find_params(W, 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.cholesky_inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H

        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]

            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):(i1 + i + groupsize)], weight=True)

                q = quant.quantize(
                    w.unsqueeze(1), self.quantizer.scale, self.quantizer.zero, self.quantizer.maxq
                ).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:])

            if DEBUG:
                self.layer.weight.data[:, :i2] = Q[:, :i2]
                self.layer.weight.data[:, i2:] = W[:, i2:]
                print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
                print(torch.sum(Losses))

        if DEVICE.type != 'cpu':
            torch.cuda.synchronize()
        print('time %.2f' % (time.time() - tick))
        print('error', torch.sum(Losses).item())

        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)
        if DEBUG:
            print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
//  [python]
def fasterquant(self, blocksize=128, percdamp=.01, groupsize=-1):
        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 not self.quantizer.ready():
            self.quantizer.find_params(W, 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.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]

            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):(i1 + i + groupsize)], weight=True)

                q = quant.quantize(
                    w.unsqueeze(1), self.quantizer.scale, self.quantizer.zero, self.quantizer.maxq
                ).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:])

            if DEBUG:
                self.layer.weight.data[:, :i2] = Q[:, :i2]
                self.layer.weight.data[:, i2:] = W[:, i2:]
                print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
                print(torch.sum(Losses))

        if DEVICE.type != 'cpu':
            torch.cuda.synchronize()
        print('time %.2f' % (time.time() - tick))
        print('error', torch.sum(Losses).item())

        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)
        if DEBUG:
            print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
//  [python]
def pre_train(self):
        opt = torch.optim.Adam(self.k.parameters(),lr=1e-1)
        for i in tqdm.tqdm(range(self.its)):
            opt.zero_grad()
            L=torch.cholesky(self.k(self.Z).evaluate()+self.eye*1e-2)
            loss= -self.objective(L)
            print(loss)
            loss.backward()
            opt.step()
        return self.k
//  [python]
def pre_train(self):
        opt = torch.optim.Adam(self.k.parameters(),lr=1e-1)
        for i in tqdm.tqdm(range(self.its)):
            opt.zero_grad()
            L=torch.linalg.cholesky(self.k(self.Z).evaluate()+self.eye*1e-2)
            loss= -self.objective(L)
            print(loss)
            loss.backward()
            opt.step()
        return self.k
//  [python]
def ortho_norm(weight):
    wtw = torch.mm(weight.t(), weight) + 1e-4 * torch.eye(weight.shape[1]).to(weight.device)
    L = torch.cholesky(wtw)
    weight_ortho = torch.mm(weight, L.inverse().t())
    return weight_ortho
//  [python]
def ortho_norm(weight):
    wtw = torch.mm(weight.t(), weight) + 1e-4 * torch.eye(weight.shape[1]).to(weight.device)
    L = torch.linalg.cholesky(wtw)
    weight_ortho = torch.mm(weight, L.inverse().t())
    return weight_ortho
//  [python]
def log_pf(self) -> Tuple[torch.Tensor, torch.Tensor]:
        # mu_a: (1, num_replicas, 1, num_components, num_variables)
        # mu_b: (1, num_replicas, num_components, 1, num_variables)
        mu = self.mu
        mu_a, mu_b = mu.unsqueeze(dim=-3), mu.unsqueeze(dim=-2)
        # cov_a: (1, num_replicas, 1, num_components, num_variables, num_variables)
        # cov_b: (1, num_replicas, num_components, 1, num_variables, num_variables)
        cov = self.__covariance()
        cov_a, cov_b = cov.unsqueeze(dim=-4), cov.unsqueeze(dim=-3)
        pairwise_cov = cov_a + cov_b
        pairwise_scale_tril = torch.cholesky(pairwise_cov)

        # sq_norm_dist: (1, num_replicas, num_components, num_components)
        sq_norm_dist = batch_mahalanobis(pairwise_scale_tril, mu_a - mu_b)
        # log_det_cov: (1, num_replicas, num_components, num_components)
        log_det_cov = torch.sum(safelog(pairwise_scale_tril.diagonal(dim1=-2, dim2=-1)), dim=-1)

        # log_z: (1, num_replicas, 1, num_components, num_components)
        log_z = -0.5 * (self.mu.shape[-1] * self.log_two_pi + sq_norm_dist) - log_det_cov
        log_z = log_z.unsqueeze(dim=-3)
        return log_z, torch.ones_like(log_z, requires_grad=False)
//  [python]
def log_pf(self) -> Tuple[torch.Tensor, torch.Tensor]:
        # mu_a: (1, num_replicas, 1, num_components, num_variables)
        # mu_b: (1, num_replicas, num_components, 1, num_variables)
        mu = self.mu
        mu_a, mu_b = mu.unsqueeze(dim=-3), mu.unsqueeze(dim=-2)
        # cov_a: (1, num_replicas, 1, num_components, num_variables, num_variables)
        # cov_b: (1, num_replicas, num_components, 1, num_variables, num_variables)
        cov = self.__covariance()
        cov_a, cov_b = cov.unsqueeze(dim=-4), cov.unsqueeze(dim=-3)
        pairwise_cov = cov_a + cov_b
        pairwise_scale_tril = torch.linalg.cholesky(pairwise_cov)

        # sq_norm_dist: (1, num_replicas, num_components, num_components)
        sq_norm_dist = batch_mahalanobis(pairwise_scale_tril, mu_a - mu_b)
        # log_det_cov: (1, num_replicas, num_components, num_components)
        log_det_cov = torch.sum(safelog(pairwise_scale_tril.diagonal(dim1=-2, dim2=-1)), dim=-1)

        # log_z: (1, num_replicas, 1, num_components, num_components)
        log_z = -0.5 * (self.mu.shape[-1] * self.log_two_pi + sq_norm_dist) - log_det_cov
        log_z = log_z.unsqueeze(dim=-3)
        return log_z, torch.ones_like(log_z, requires_grad=False)
//  [python]
def sample_inputs_linalg_cholesky_inverse(
    op_info, device, dtype, requires_grad=False, **kwargs
):
    from torch.testing._internal.common_utils import random_well_conditioned_matrix

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

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

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

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

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

        # generate upper-triangular inputs
        u = l.detach().clone().mT.contiguous().requires_grad_(requires_grad)
        yield SampleInput(u, kwargs=dict(upper=True))
//  [python]
def stable_kalman_gain(H, P, Sigma_pred):
    # Kalman gain, a more stable implementation than naive P @ H^T @ y_sigma^{-1}
    L = torch.cholesky(Sigma_pred)
    K = torch.triangular_solve(H @ P.transpose(1, 2), L, upper=False)[0]
    K = torch.triangular_solve(K, L.transpose(1, 2))[0].transpose(1, 2)
    return K
//  [python]
def stable_kalman_gain(H, P, Sigma_pred):
    # Kalman gain, a more stable implementation than naive P @ H^T @ y_sigma^{-1}
    L = torch.linalg.cholesky(Sigma_pred)
    K = torch.triangular_solve(H @ P.transpose(1, 2), L, upper=False)[0]
    K = torch.triangular_solve(K, L.transpose(1, 2))[0].transpose(1, 2)
    return K
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False, static_groups=False
    ):
        W = self.layer.weight.data.clone()
        W = W.float()

        tick = time.time()

        if not self.quantizer.ready():
            self.quantizer.find_params(W)

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

        if static_groups:
            import copy
            groups = []
            for i in range(0, self.columns, groupsize):
                quantizer = copy.deepcopy(self.quantizer)
                quantizer.find_params(W[:, i:(i + groupsize)])
                groups.append(quantizer)

        if actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, perm]
            H = H[perm][:, perm]
            invperm = torch.argsort(perm)

        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.cholesky_inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H

        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]

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

                if groupsize != -1:
                    if not static_groups:
                        if (i1 + i) % groupsize == 0:
                            self.quantizer.find_params(W[:, (i1 + i):(i1 + i + groupsize)])
                    else:
                        idx = i1 + i
                        if actorder:
                            idx = perm[idx]
                        self.quantizer = groups[idx // groupsize]

                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()

        if actorder:
            Q = Q[:, invperm]

        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
        if torch.any(torch.isnan(self.layer.weight.data)):
            logging.warning('NaN in weights')
            import pprint
            pprint.pprint(self.quantizer.bits, self.quantizer.scale, self.quantizer.zero_point)
            raise ValueError('NaN in weights')
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False, static_groups=False
    ):
        W = self.layer.weight.data.clone()
        W = W.float()

        tick = time.time()

        if not self.quantizer.ready():
            self.quantizer.find_params(W)

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

        if static_groups:
            import copy
            groups = []
            for i in range(0, self.columns, groupsize):
                quantizer = copy.deepcopy(self.quantizer)
                quantizer.find_params(W[:, i:(i + groupsize)])
                groups.append(quantizer)

        if actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, perm]
            H = H[perm][:, perm]
            invperm = torch.argsort(perm)

        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.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]

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

                if groupsize != -1:
                    if not static_groups:
                        if (i1 + i) % groupsize == 0:
                            self.quantizer.find_params(W[:, (i1 + i):(i1 + i + groupsize)])
                    else:
                        idx = i1 + i
                        if actorder:
                            idx = perm[idx]
                        self.quantizer = groups[idx // groupsize]

                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()

        if actorder:
            Q = Q[:, invperm]

        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
        if torch.any(torch.isnan(self.layer.weight.data)):
            logging.warning('NaN in weights')
            import pprint
            pprint.pprint(self.quantizer.bits, self.quantizer.scale, self.quantizer.zero_point)
            raise ValueError('NaN in weights')
//  [python]
def init_L(self):
        with torch.no_grad():
            kx = self.k(self.Z, self.X).evaluate()
            self.kzz = self.k(self.Z).evaluate()
        L_set = False
        for reg in [1e-7,1e-6,1e-5,1e-4,1e-3,1e-2]:#,1e-1,1.0,5.0,10.]:
            try:
                with torch.no_grad():
                    chol_L = torch.cholesky(self.kzz + kx @ kx.t() / self.sigma + self.eye*reg)
                    L = torch.cholesky(torch.inverse(chol_L))
                    # L = torch.linalg.cholesky(torch.inverse(self.kzz  + self.eye*self.sigma))
                    self.L = torch.nn.Parameter(L)
                    L_set = True
            except Exception as e:
                gc.collect()
                print('cholesky init error: ',reg)
                print(e)
                    # torch.cuda.empty_cache()
        if not L_set:
            print("Welp this is awkward, it don't want to factorize")

            gc.collect()
            # torch.cuda.empty_cache()
            L = torch.randn_like(self.kzz) * 0.1
            self.L = torch.nn.Parameter(L)
        inv_set=False
        for reg in [1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]:#, 1e-1, 1.0,2.5,5.0,10.]:
            self.reg=reg
            try:
                if not self.parametrize_Z:
                    with torch.no_grad():
                        self.register_buffer('inv_L', torch.cholesky(self.kzz + self.eye * self.reg))
                        self.register_buffer('inv_Z',torch.inverse(self.inv_L))
                    inv_set=True
            except Exception as e:
                gc.collect()
                print(e)
                print('cholesky inv error: ',reg)
                pass
        if not inv_set:
            gc.collect()
            self.parametrize_Z=True
//  [python]
def init_L(self):
        with torch.no_grad():
            kx = self.k(self.Z, self.X).evaluate()
            self.kzz = self.k(self.Z).evaluate()
        L_set = False
        for reg in [1e-7,1e-6,1e-5,1e-4,1e-3,1e-2]:#,1e-1,1.0,5.0,10.]:
            try:
                with torch.no_grad():
                    chol_L = torch.linalg.cholesky(self.kzz + kx @ kx.t() / self.sigma + self.eye*reg)
                    L = torch.linalg.cholesky(torch.cholesky_inverse(chol_L))
                    # L = torch.linalg.cholesky(torch.inverse(self.kzz  + self.eye*self.sigma))
                    self.L = torch.nn.Parameter(L)
                    L_set = True
            except Exception as e:
                gc.collect()
                print('cholesky init error: ',reg)
                print(e)
                    # torch.cuda.empty_cache()
        if not L_set:
            print("Welp this is awkward, it don't want to factorize")

            gc.collect()
            # torch.cuda.empty_cache()
            L = torch.randn_like(self.kzz) * 0.1
            self.L = torch.nn.Parameter(L)
        inv_set=False
        for reg in [1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]:#, 1e-1, 1.0,2.5,5.0,10.]:
            self.reg=reg
            try:
                if not self.parametrize_Z:
                    with torch.no_grad():
                        self.register_buffer('inv_L', torch.linalg.cholesky(self.kzz + self.eye * self.reg))
                        self.register_buffer('inv_Z',torch.cholesky_inverse(self.inv_L))
                    inv_set=True
            except Exception as e:
                gc.collect()
                print(e)
                print('cholesky inv error: ',reg)
                pass
        if not inv_set:
            gc.collect()
            self.parametrize_Z=True
//  [python]
def forward(self,x1,x2=None):
        L = torch.tril(self.L)+self.eye*self.reg
        L = ensure_pos_diag(L)
        Z= self.Z
        if x2 is None:
            kzx = self.k(Z, x1).evaluate()
            t= L.t() @ kzx #L\cdot k(Z,X)

            if self.parametrize_Z:
                kzz = self.k(Z).evaluate()
                chol_z = torch.cholesky(kzz+self.eye*self.reg)
                sol = torch.cholesky_solve(kzx,chol_z)
            else:
                # sol = torch.cholesky_solve(kzx, self.inv_L)
                sol = self.inv_Z @ kzx
            if len(t.shape)==3: #t=[L^T k(Z,X_i),L^T k(Z,X_{i+1}),]
                T_mat = t.permute(0,2,1) @ t #T_mat ok
                inverse_part = kzx.permute(0,2,1)@sol
                out = self.k(x1).evaluate()- inverse_part + T_mat#/self.sigma
                return out
            else:
                T_mat = t.t() @ t
                out =self.k(x1).evaluate()- kzx.t()@sol + T_mat#/self.sigma
                return out
        else:
            kzx_1 = self.k(Z, x1).evaluate()
            kzx_2 = self.k(Z, x2).evaluate()
            t= L.t() @ kzx_2
            t_ = kzx_1.t() @ L
            T_mat = t_ @ t
            if self.parametrize_Z:
                kzz = self.k(Z).evaluate()
                chol_z = torch.cholesky(kzz + self.eye * self.reg)
                sol = torch.cholesky_solve(kzx_2, chol_z)
            else:
                # sol = torch.cholesky_solve(kzx_2, self.inv_L)
                sol = self.inv_Z @ kzx_2
            out=self.k(x1, x2).evaluate() - kzx_1.t() @ sol +T_mat#/self.sigma
            return out
//  [python]
def forward(self,x1,x2=None):
        L = torch.tril(self.L)+self.eye*self.reg
        L = ensure_pos_diag(L)
        Z= self.Z
        if x2 is None:
            kzx = self.k(Z, x1).evaluate()
            t= L.t() @ kzx #L\cdot k(Z,X)

            if self.parametrize_Z:
                kzz = self.k(Z).evaluate()
                chol_z = torch.linalg.cholesky(kzz+self.eye*self.reg)
                sol = torch.cholesky_solve(kzx,chol_z)
            else:
                # sol = torch.cholesky_solve(kzx, self.inv_L)
                sol = self.inv_Z @ kzx
            if len(t.shape)==3: #t=[L^T k(Z,X_i),L^T k(Z,X_{i+1}),]
                T_mat = t.permute(0,2,1) @ t #T_mat ok
                inverse_part = kzx.permute(0,2,1)@sol
                out = self.k(x1).evaluate()- inverse_part + T_mat#/self.sigma
                return out
            else:
                T_mat = t.t() @ t
                out =self.k(x1).evaluate()- kzx.t()@sol + T_mat#/self.sigma
                return out
        else:
            kzx_1 = self.k(Z, x1).evaluate()
            kzx_2 = self.k(Z, x2).evaluate()
            t= L.t() @ kzx_2
            t_ = kzx_1.t() @ L
            T_mat = t_ @ t
            if self.parametrize_Z:
                kzz = self.k(Z).evaluate()
                chol_z = torch.linalg.cholesky(kzz + self.eye * self.reg)
                sol = torch.cholesky_solve(kzx_2, chol_z)
            else:
                # sol = torch.cholesky_solve(kzx_2, self.inv_L)
                sol = self.inv_Z @ kzx_2
            out=self.k(x1, x2).evaluate() - kzx_1.t() @ sol +T_mat#/self.sigma
            return out