jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x, x_var=0.):
        """
        Forward pass of the CIGP module.

        Args:
            x (torch.Tensor or list): Input tensor or list of tensors.
            x_var (float): Variance of the input tensor.

        Returns:
            tuple: Tuple containing the predicted mean and variance.
        """
        x = self.check_single_tensor(x)

        if self.train_x is None:
            MFGP_LOG.e("gp model model hasn't been trained. predict failed")
            return None
    
        with torch.no_grad():
            Sigma = self.kernel(self.train_x, self.train_x) + JITTER * torch.eye(self.train_x.size(0), device=list(self.parameters())[0].device)
            _noise = self.noise_box.get()
            Sigma = Sigma + _noise.pow(-1) * torch.eye(self.train_x.size(0), device=list(self.parameters())[0].device)

            L = torch.linalg.cholesky(Sigma)
            kx = self.kernel(self.train_x, x)
            LinvKx, _ = torch.triangular_solve(kx, L, upper = False)
            # LinvKx = torch.linalg.solve_triangular(kx, L, upper = False)      # torch recommand but not support nonsquare matrix

            u = kx.t() @ torch.cholesky_solve(self.train_y, L)

            var_diag = self.kernel(x, x).diag().view(-1, 1) - (LinvKx**2).sum(dim = 0).view(-1, 1)
            _noise = self.noise_box.get()
            var_diag = var_diag + _noise.pow(-1)
            var_diag = var_diag.expand_as(u)

            var_diag = var_diag + x_var

        return u, var_diag
//  [python]
def fasterprune(self, gpt, blocksize=128, percdamp=0.01):
        import transformers

        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)  # te default is lower triangle
        H = torch.cholesky_inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H

        M = self.M
        N = self.N

        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]
            mask1 = torch.zeros_like(W1) == 1

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

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

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

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

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

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

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
        if torch.cuda.is_available():
            torch.cuda.synchronize()

        if isinstance(module, transformers.Conv1D):
            W = W.t()
        module.weight.data = W.reshape(module.weight.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

        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)  # te default is lower triangle
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H, upper=True)
        Hinv = H

        M = self.M
        N = self.N

        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]
            mask1 = torch.zeros_like(W1) == 1

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

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

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

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

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

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

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
        if torch.cuda.is_available():
            torch.cuda.synchronize()

        if isinstance(module, transformers.Conv1D):
            W = W.t()
        module.weight.data = W.reshape(module.weight.shape).to(dtype=module.weight.data.dtype)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
//  [python]
def test_gamma_gaussian_tensordot(
    dot_dims, x_batch_shape, x_dim, x_rank, y_batch_shape, y_dim, y_rank
):
    x_rank = min(x_rank, x_dim)
    y_rank = min(y_rank, y_dim)
    x = random_gamma_gaussian(x_batch_shape, x_dim, x_rank)
    y = random_gamma_gaussian(y_batch_shape, y_dim, y_rank)
    na = x_dim - dot_dims
    nb = dot_dims
    nc = y_dim - dot_dims
    try:
        torch.cholesky(x.precision[..., na:, na:] + y.precision[..., :nb, :nb])
    except RuntimeError:
        pytest.skip("Cannot marginalize the common variables of two Gaussians.")

    z = gamma_gaussian_tensordot(x, y, dot_dims)
    assert z.dim() == x_dim + y_dim - 2 * dot_dims

    # We make these precision matrices positive definite to test the math
    x.precision = x.precision + 3 * torch.eye(x.dim())
    y.precision = y.precision + 3 * torch.eye(y.dim())
    z = gamma_gaussian_tensordot(x, y, dot_dims)
    # compare against broadcasting, adding, and marginalizing
    precision = pad(x.precision, (0, nc, 0, nc)) + pad(y.precision, (na, 0, na, 0))
    info_vec = pad(x.info_vec, (0, nc)) + pad(y.info_vec, (na, 0))
    covariance = torch.inverse(precision)
    loc = (
        covariance.matmul(info_vec.unsqueeze(-1)).squeeze(-1)
        if info_vec.size(-1) > 0
        else info_vec
    )
    z_covariance = torch.inverse(z.precision)
    z_loc = z_covariance.matmul(
        z.info_vec.view(z.info_vec.shape + (int(z.dim() > 0),))
    ).sum(-1)
    assert_close(loc[..., :na], z_loc[..., :na])
    assert_close(loc[..., x_dim:], z_loc[..., na:])
    assert_close(covariance[..., :na, :na], z_covariance[..., :na, :na])
    assert_close(covariance[..., :na, x_dim:], z_covariance[..., :na, na:])
    assert_close(covariance[..., x_dim:, :na], z_covariance[..., na:, :na])
    assert_close(covariance[..., x_dim:, x_dim:], z_covariance[..., na:, na:])

    s = torch.randn(z.batch_shape).exp()
    # Assume a = c = 0, integrate out b
    num_samples = 200000
    scale = 10
    # generate samples in [-10, 10]
    value_b = torch.rand((num_samples,) + z.batch_shape + (nb,)) * scale - scale / 2
    value_x = pad(value_b, (na, 0))
    value_y = pad(value_b, (0, nc))
    expect = torch.logsumexp(
        x.log_density(value_x, s) + y.log_density(value_y, s), dim=0
    )
    expect += math.log(scale**nb / num_samples)
    actual = z.log_density(torch.zeros(z.batch_shape + (z.dim(),)), s)
    assert_close(actual.clamp(max=10.0), expect.clamp(max=10.0), atol=0.1, rtol=0.1)
//  [python]
def test_gamma_gaussian_tensordot(
    dot_dims, x_batch_shape, x_dim, x_rank, y_batch_shape, y_dim, y_rank
):
    x_rank = min(x_rank, x_dim)
    y_rank = min(y_rank, y_dim)
    x = random_gamma_gaussian(x_batch_shape, x_dim, x_rank)
    y = random_gamma_gaussian(y_batch_shape, y_dim, y_rank)
    na = x_dim - dot_dims
    nb = dot_dims
    nc = y_dim - dot_dims
    try:
        torch.linalg.cholesky(x.precision[..., na:, na:] + y.precision[..., :nb, :nb])
    except RuntimeError:
        pytest.skip("Cannot marginalize the common variables of two Gaussians.")

    z = gamma_gaussian_tensordot(x, y, dot_dims)
    assert z.dim() == x_dim + y_dim - 2 * dot_dims

    # We make these precision matrices positive definite to test the math
    x.precision = x.precision + 3 * torch.eye(x.dim())
    y.precision = y.precision + 3 * torch.eye(y.dim())
    z = gamma_gaussian_tensordot(x, y, dot_dims)
    # compare against broadcasting, adding, and marginalizing
    precision = pad(x.precision, (0, nc, 0, nc)) + pad(y.precision, (na, 0, na, 0))
    info_vec = pad(x.info_vec, (0, nc)) + pad(y.info_vec, (na, 0))
    covariance = torch.inverse(precision)
    loc = (
        covariance.matmul(info_vec.unsqueeze(-1)).squeeze(-1)
        if info_vec.size(-1) > 0
        else info_vec
    )
    z_covariance = torch.inverse(z.precision)
    z_loc = z_covariance.matmul(
        z.info_vec.view(z.info_vec.shape + (int(z.dim() > 0),))
    ).sum(-1)
    assert_close(loc[..., :na], z_loc[..., :na])
    assert_close(loc[..., x_dim:], z_loc[..., na:])
    assert_close(covariance[..., :na, :na], z_covariance[..., :na, :na])
    assert_close(covariance[..., :na, x_dim:], z_covariance[..., :na, na:])
    assert_close(covariance[..., x_dim:, :na], z_covariance[..., na:, :na])
    assert_close(covariance[..., x_dim:, x_dim:], z_covariance[..., na:, na:])

    s = torch.randn(z.batch_shape).exp()
    # Assume a = c = 0, integrate out b
    num_samples = 200000
    scale = 10
    # generate samples in [-10, 10]
    value_b = torch.rand((num_samples,) + z.batch_shape + (nb,)) * scale - scale / 2
    value_x = pad(value_b, (na, 0))
    value_y = pad(value_b, (0, nc))
    expect = torch.logsumexp(
        x.log_density(value_x, s) + y.log_density(value_y, s), dim=0
    )
    expect += math.log(scale**nb / num_samples)
    actual = z.log_density(torch.zeros(z.batch_shape + (z.dim(),)), s)
    assert_close(actual.clamp(max=10.0), expect.clamp(max=10.0), atol=0.1, rtol=0.1)
//  [python]
def sample_xs(self, T, N, eps=0.01):
        """joint
        t: (T)
        xs: (B, N, T, D)
        """
        B, D, device = self.B, self.D, self.device

        t = torch.linspace(eps, 1 - eps, T, device=device)

        ## Compute covariance & its decomposition
        C = self.cov_mtx(t)
        if self.cov_decom == "cholesky":
            A = torch.cholesky(C)
        elif self.cov_decom == "eigh":
            L, Q = torch.linalg.eigh(C)
            A = Q @ torch.diag_embed(L).sqrt()
        assert C.shape == A.shape == (B, T, T)

        ## Compute std_t and mean_t
        noise_t = torch.randn(N, B, T, D, device=device)
        std_t = (A @ noise_t).transpose(0, 1)  # <-- this took most memory
        assert std_t.shape == (B, N, T, D)

        torch.cuda.empty_cache()  # clear out memory

        mean_t = self.mean(t)  # (B, T, D)
        assert mean_t.shape == (B, T, D)

        ## sim-free xs
        xs = mean_t.unsqueeze(1) + std_t
        assert xs.shape == (B, N, T, D)
        return t, xs
//  [python]
def sample_xs(self, T, N, eps=0.01):
        """joint
        t: (T)
        xs: (B, N, T, D)
        """
        B, D, device = self.B, self.D, self.device

        t = torch.linspace(eps, 1 - eps, T, device=device)

        ## Compute covariance & its decomposition
        C = self.cov_mtx(t)
        if self.cov_decom == "cholesky":
            A = torch.linalg.cholesky(C)
        elif self.cov_decom == "eigh":
            L, Q = torch.linalg.eigh(C)
            A = Q @ torch.diag_embed(L).sqrt()
        assert C.shape == A.shape == (B, T, T)

        ## Compute std_t and mean_t
        noise_t = torch.randn(N, B, T, D, device=device)
        std_t = (A @ noise_t).transpose(0, 1)  # <-- this took most memory
        assert std_t.shape == (B, N, T, D)

        torch.cuda.empty_cache()  # clear out memory

        mean_t = self.mean(t)  # (B, T, D)
        assert mean_t.shape == (B, T, D)

        ## sim-free xs
        xs = mean_t.unsqueeze(1) + std_t
        assert xs.shape == (B, N, T, D)
        return t, xs
//  [python]
def compute_loss(self, x, y, x_var=0., y_var=0., update_data=False):
        """
        Compute the loss function for the CIGP module.

        Args:
            x (torch.Tensor or list): Input tensor or list of tensors.
            y (torch.Tensor or list): Output tensor or list of tensors.
            x_var (float): Variance of the input tensor.
            y_var (float): Variance of the output tensor.

        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

        Sigma = self.kernel(x, x) + JITTER * torch.eye(x.size(0), device=list(self.parameters())[0].device)
        _noise = self.noise_box.get()
        Sigma = Sigma + _noise.pow(-1) * torch.eye(x.size(0), device=list(self.parameters())[0].device)
        Sigma = Sigma + y_var

        L = torch.cholesky(Sigma)

        gamma = torch.triangular_solve(y, L, upper=False)[0]       # we can use this as an alternative because L is a lower triangular matrix.

        y_num, y_dimension = y.shape
        nll =  0.5 * (gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI, device=list(self.parameters())[0].device)) * y_dimension
        return nll
//  [python]
def compute_loss(self, x, y, x_var=0., y_var=0., update_data=False):
        """
        Compute the loss function for the CIGP module.

        Args:
            x (torch.Tensor or list): Input tensor or list of tensors.
            y (torch.Tensor or list): Output tensor or list of tensors.
            x_var (float): Variance of the input tensor.
            y_var (float): Variance of the output tensor.

        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

        Sigma = self.kernel(x, x) + JITTER * torch.eye(x.size(0), device=list(self.parameters())[0].device)
        _noise = self.noise_box.get()
        Sigma = Sigma + _noise.pow(-1) * torch.eye(x.size(0), device=list(self.parameters())[0].device)
        Sigma = Sigma + y_var

        L = torch.linalg.cholesky(Sigma)

        gamma = L.inverse() @ y       # we can use this as an alternative because L is a lower triangular matrix.

        y_num, y_dimension = y.shape
        nll =  0.5 * (gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI, device=list(self.parameters())[0].device)) * y_dimension
        return nll
//  [python]
def get_kernel_variance(
        self,
        particles: Tensor,
        weights: Tensor,
        samples_per_dim: int = 100,
        kernel_variance_scale: float = 1.0,
    ) -> Tensor:
        """Return kernel variance for a given population of particles and weights."""
        if self.kernel == "gaussian":
            # For variant C, Beaumont et al. 2009, the kernel variance comes from the
            # previous population.
            if self.algorithm_variant == "C":
                # Calculate weighted covariance of particles.
                population_cov = torch.tensor(
                    np.atleast_2d(np.cov(particles, rowvar=False, aweights=weights)),
                    dtype=torch.float32,
                )
                # Make sure variance is nonsingular.
                try:
                    torch.cholesky(kernel_variance_scale * population_cov)
                except RuntimeError:
                    self.logger.warning(
                        """"Singular particle covariance, using unit covariance."""
                    )
                    population_cov = torch.eye(particles.shape[1])
                return kernel_variance_scale * population_cov
            # While for Toni et al. and Sisson et al. it comes from the parameter
            # ranges.
            elif self.algorithm_variant in ("A", "B"):
                particle_ranges = self.get_particle_ranges(
                    particles, weights, samples_per_dim=samples_per_dim
                )
                return kernel_variance_scale * torch.diag(particle_ranges)
            else:
                raise ValueError(f"Variant, '{self.algorithm_variant}' not supported.")
        elif self.kernel == "uniform":
            # Variance spans the range of parameters for every dimension.
            return kernel_variance_scale * self.get_particle_ranges(
                particles, weights, samples_per_dim=samples_per_dim
            )
        else:
            raise ValueError(f"Kernel, '{self.kernel}' not supported.")
//  [python]
def get_kernel_variance(
        self,
        particles: Tensor,
        weights: Tensor,
        samples_per_dim: int = 100,
        kernel_variance_scale: float = 1.0,
    ) -> Tensor:
        """Return kernel variance for a given population of particles and weights."""
        if self.kernel == "gaussian":
            # For variant C, Beaumont et al. 2009, the kernel variance comes from the
            # previous population.
            if self.algorithm_variant == "C":
                # Calculate weighted covariance of particles.
                population_cov = torch.tensor(
                    np.atleast_2d(np.cov(particles, rowvar=False, aweights=weights)),
                    dtype=torch.float32,
                )
                # Make sure variance is nonsingular.
                try:
                    torch.linalg.cholesky(kernel_variance_scale * population_cov)
                except RuntimeError:
                    self.logger.warning(
                        """"Singular particle covariance, using unit covariance."""
                    )
                    population_cov = torch.eye(particles.shape[1])
                return kernel_variance_scale * population_cov
            # While for Toni et al. and Sisson et al. it comes from the parameter
            # ranges.
            elif self.algorithm_variant in ("A", "B"):
                particle_ranges = self.get_particle_ranges(
                    particles, weights, samples_per_dim=samples_per_dim
                )
                return kernel_variance_scale * torch.diag(particle_ranges)
            else:
                raise ValueError(f"Variant, '{self.algorithm_variant}' not supported.")
        elif self.kernel == "uniform":
            # Variance spans the range of parameters for every dimension.
            return kernel_variance_scale * self.get_particle_ranges(
                particles, weights, samples_per_dim=samples_per_dim
            )
        else:
            raise ValueError(f"Kernel, '{self.kernel}' not supported.")
//  [python]
def C(self):
        if self.learn_inducing_locations:
            C_out = torch.cholesky(
                self.K(self.tau, self.tau) + self.eps * torch.eye(self.n_tau)
            )
        else:
            C_out = self.C_buffer
        return C_out
//  [python]
def C(self):
        if self.learn_inducing_locations:
            C_out = torch.linalg.cholesky(
                self.K(self.tau, self.tau) + self.eps * torch.eye(self.n_tau)
            )
        else:
            C_out = self.C_buffer
        return C_out
//  [python]
def __init__(self, sample_config: SampleConfig):
        super().__init__(**vars(sample_config))
        self.sample_shape = 0
        self.ndims = self.n_knots * self.d_action
        self.Z = torch.zeros(
            self.ndims, device=self.tensor_args.device, dtype=self.tensor_args.dtype
        )
        if self.covariance_matrix is None:
            self.cov_matrix = torch.eye(
                self.ndims, device=self.tensor_args.device, dtype=self.tensor_args.dtype
            )
        self.scale_tril = torch.cholesky(
            self.cov_matrix.to(dtype=torch.float32, device="cpu")
        ).to(device=self.tensor_args.device, dtype=self.tensor_args.dtype)
        if self.sample_method == "random":
            self.mvn = MultivariateNormal(
                loc=self.Z,
                scale_tril=self.scale_tril,
            )
        if self.sample_method == "halton":
            self.halton_generator = HaltonGenerator(
                self.ndims, seed=self.seed, tensor_args=self.tensor_args
            )
//  [python]
def __init__(self, sample_config: SampleConfig):
        super().__init__(**vars(sample_config))
        self.sample_shape = 0
        self.ndims = self.n_knots * self.d_action
        self.Z = torch.zeros(
            self.ndims, device=self.tensor_args.device, dtype=self.tensor_args.dtype
        )
        if self.covariance_matrix is None:
            self.cov_matrix = torch.eye(
                self.ndims, device=self.tensor_args.device, dtype=self.tensor_args.dtype
            )
        self.scale_tril = torch.linalg.cholesky(
            self.cov_matrix.to(dtype=torch.float32, device="cpu")
        ).to(device=self.tensor_args.device, dtype=self.tensor_args.dtype)
        if self.sample_method == "random":
            self.mvn = MultivariateNormal(
                loc=self.Z,
                scale_tril=self.scale_tril,
            )
        if self.sample_method == "halton":
            self.halton_generator = HaltonGenerator(
                self.ndims, seed=self.seed, tensor_args=self.tensor_args
            )
//  [python]
def unscented(
    kappa: float = 1.0,
    unscented_mean: bool = True,
) -> Projection:
    """Configures unscented projection for further use

    Args:
        kappa (float, optional): kappa parameter for the unscented projection. Defaults to 1.0.
        unscented_mean (bool, optional): use the unscented mean as mean. Defaults to True.

    Returns:
        Projection: Callable that does the unscented transformation
    """

    def _unscented(
        points: torch.Tensor,
        K_inv: torch.Tensor,
        covariances: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Perfoms the unscented transformation of a covariance from the image plane of a pinhole camera onto the unit sphere in 3D. Accounts for the rank deficiency of the covariance (max rank = 2). Defaults to linear projection if not covariances are provided.

        Args:
            point (torch.Tensor): image point in homogenous coordinates (...xNx3)
            K_inv (torch.Tensor): inverse camera matrix (...x3x3)
            covariance (Optional[torch.Tensor]): covariances corresponding to the points (...xNx3x3), default = None

        Returns:
            Tuple[torch.Tensor, Optional[torch.Tensor]]: projected points (...xNx3) and covariance matrices (...xNx3x3) (if given)
        """
        if covariances is None:
            return linear(points, K_inv, covariances)

        n = 2

        chol_covs = torch.cholesky(covariances[..., :2, :2])

        weights = torch.zeros((2 * n + 1), dtype=torch.float64)
        sample_points = []
        weights[0] = kappa / (n + kappa)
        sample_points.append(points)
        for i in range(n):
            sample_points.append(points + torch.nn.functional.pad(chol_covs[..., :, i], (0, 1), value=0.0))
            weights[i + 1] = 0.5 / (n + kappa)
            sample_points.append(points - torch.nn.functional.pad(chol_covs[..., :, i], (0, 1), value=0.0))
            weights[i + 1 + n] = 0.5 / (n + kappa)
        sample_points = torch.stack(sample_points, dim=-1)

        transformed_points = torch.einsum("ij,...jk->...ik", K_inv, sample_points)
        transformed_points = transformed_points / torch.norm(transformed_points, dim=-2)[..., None, :]

        means = torch.sum(weights * transformed_points, dim=-1)

        difference = transformed_points - means[..., :, None]

        covs = torch.einsum("...ij,j,...kj->...ik", difference, weights, difference)

        if not unscented_mean:
            # use the projected point as the mean.
            means = transformed_points[:, 0]

        return means, covs

    return _unscented
//  [python]
def unscented(
    kappa: float = 1.0,
    unscented_mean: bool = True,
) -> Projection:
    """Configures unscented projection for further use

    Args:
        kappa (float, optional): kappa parameter for the unscented projection. Defaults to 1.0.
        unscented_mean (bool, optional): use the unscented mean as mean. Defaults to True.

    Returns:
        Projection: Callable that does the unscented transformation
    """

    def _unscented(
        points: torch.Tensor,
        K_inv: torch.Tensor,
        covariances: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Perfoms the unscented transformation of a covariance from the image plane of a pinhole camera onto the unit sphere in 3D. Accounts for the rank deficiency of the covariance (max rank = 2). Defaults to linear projection if not covariances are provided.

        Args:
            point (torch.Tensor): image point in homogenous coordinates (...xNx3)
            K_inv (torch.Tensor): inverse camera matrix (...x3x3)
            covariance (Optional[torch.Tensor]): covariances corresponding to the points (...xNx3x3), default = None

        Returns:
            Tuple[torch.Tensor, Optional[torch.Tensor]]: projected points (...xNx3) and covariance matrices (...xNx3x3) (if given)
        """
        if covariances is None:
            return linear(points, K_inv, covariances)

        n = 2

        chol_covs = torch.linalg.cholesky(covariances[..., :2, :2])

        weights = torch.zeros((2 * n + 1), dtype=torch.float64)
        sample_points = []
        weights[0] = kappa / (n + kappa)
        sample_points.append(points)
        for i in range(n):
            sample_points.append(points + torch.nn.functional.pad(chol_covs[..., :, i], (0, 1), value=0.0))
            weights[i + 1] = 0.5 / (n + kappa)
            sample_points.append(points - torch.nn.functional.pad(chol_covs[..., :, i], (0, 1), value=0.0))
            weights[i + 1 + n] = 0.5 / (n + kappa)
        sample_points = torch.stack(sample_points, dim=-1)

        transformed_points = torch.einsum("ij,...jk->...ik", K_inv, sample_points)
        transformed_points = transformed_points / torch.norm(transformed_points, dim=-2)[..., None, :]

        means = torch.sum(weights * transformed_points, dim=-1)

        difference = transformed_points - means[..., :, None]

        covs = torch.einsum("...ij,j,...kj->...ik", difference, weights, difference)

        if not unscented_mean:
            # use the projected point as the mean.
            means = transformed_points[:, 0]

        return means, covs

    return _unscented
//  [python]
def __init__(self, mu, sigma, low, high, policy_state_shape=None):
        """
        Constructor.

        Args:
            mu (Regressor): the regressor representing the mean w.r.t. the state;
            sigma (torch.tensor): a square positive definite matrix representing the covariance matrix. The size of this
                matrix must be n x n, where n is the action dimensionality;
            low (torch.tensor): a vector containing the minimum action for each component;
            high (torch.tensor): a vector containing the maximum action for each component.

        """
        super().__init__(policy_state_shape)

        self._approximator = mu
        self._predict_params = dict()
        self._chol_sigma = torch.cholesky(sigma)
        self._low = torch.as_tensor(low)
        self._high = torch.as_tensor(high)

        self._add_save_attr(
            _approximator='mushroom',
            _predict_params='pickle',
            _chol_sigma='torch',
            _low='torch',
            _high='torch'
        )
//  [python]
def __init__(self, mu, sigma, low, high, policy_state_shape=None):
        """
        Constructor.

        Args:
            mu (Regressor): the regressor representing the mean w.r.t. the state;
            sigma (torch.tensor): a square positive definite matrix representing the covariance matrix. The size of this
                matrix must be n x n, where n is the action dimensionality;
            low (torch.tensor): a vector containing the minimum action for each component;
            high (torch.tensor): a vector containing the maximum action for each component.

        """
        super().__init__(policy_state_shape)

        self._approximator = mu
        self._predict_params = dict()
        self._chol_sigma = torch.linalg.cholesky(sigma)
        self._low = torch.as_tensor(low)
        self._high = torch.as_tensor(high)

        self._add_save_attr(
            _approximator='mushroom',
            _predict_params='pickle',
            _chol_sigma='torch',
            _low='torch',
            _high='torch'
        )
//  [python]
def fit(self, features, logits, targets):
        self.global_mean = features.mean(0, keepdim=True)
        center_f = features - self.global_mean
        n, dim = features.size()

        if self.covariance_type == "diagonal":
            self.W = 1. / center_f.std(0, keepdim=True)
            transf_f = center_f * self.W
        else:
            reg = torch.eye(dim) * 1e-6
            self.W = (center_f.t() @ center_f).div(n - 1).add(reg).inverse()
            self.W = torch.cholesky(self.W)
            transf_f = center_f @ self.W

        self.mu = torch.stack([
            transf_f[targets == t].mean(0) for t in targets.unique()
        ])
//  [python]
def fit(self, features, logits, targets):
        self.global_mean = features.mean(0, keepdim=True)
        center_f = features - self.global_mean
        n, dim = features.size()

        if self.covariance_type == "diagonal":
            self.W = 1. / center_f.std(0, keepdim=True)
            transf_f = center_f * self.W
        else:
            reg = torch.eye(dim) * 1e-6
            self.W = (center_f.t() @ center_f).div(n - 1).add(reg).inverse()
            self.W = torch.linalg.cholesky(self.W)
            transf_f = center_f @ self.W

        self.mu = torch.stack([
            transf_f[targets == t].mean(0) for t in targets.unique()
        ])
//  [python]
def fit(self, features, logits, targets):
        n, dim = features.size()
        reg = torch.eye(dim) * 1e-6

        # parameters for MD_0
        self.m0 = features.mean(0, keepdim=True)
        center0 = features - self.m0
        self.W0 = (center0.t() @ center0).div(n).add(reg).inverse()
        self.W0 = torch.cholesky(self.W0)
        self.m0 @= self.W0

        # parameters for MD_k
        self.mk = torch.stack([
            features[targets == t].mean(0) for t in targets.unique()
        ])
        centerk = features - self.mk[targets]
        self.Wk = (centerk.t() @ centerk).div(n).add(reg).inverse()
        self.Wk = torch.cholesky(self.Wk)
        self.mk @= self.Wk
//  [python]
def fit(self, features, logits, targets):
        n, dim = features.size()
        reg = torch.eye(dim) * 1e-6

        # parameters for MD_0
        self.m0 = features.mean(0, keepdim=True)
        center0 = features - self.m0
        self.W0 = (center0.t() @ center0).div(n).add(reg).inverse()
        self.W0 = torch.linalg.cholesky(self.W0)
        self.m0 @= self.W0

        # parameters for MD_k
        self.mk = torch.stack([
            features[targets == t].mean(0) for t in targets.unique()
        ])
        centerk = features - self.mk[targets]
        self.Wk = (centerk.t() @ centerk).div(n).add(reg).inverse()
        self.Wk = torch.linalg.cholesky(self.Wk)
        self.mk @= self.Wk
//  [python]
def Eigenvalues(self, kpoints, time_symm=True, unit="Hartree",if_eigvec=False):
        """ using the tight-binding H and S matrix calculate eigenvalues at kpoints.
        
        Args:
            kpoints: the k-kpoints used to calculate the eigenvalues.
        Note: must have the BondHBlock and BondSBlock 
        """
        hkmat = self.hs_block_R2k(kpoints=kpoints, HorS='H', time_symm=time_symm)
        if not self.use_orthogonal_basis:
            skmat =  self.hs_block_R2k(kpoints=kpoints, HorS='S', time_symm=time_symm)
        else:
            skmat = torch.eye(hkmat.shape[1], dtype=self.cdtype).unsqueeze(0).repeat(hkmat.shape[0], 1, 1)

        if self.use_orthogonal_basis:
            Heff = hkmat
        else:
            chklowt = th.cholesky(skmat)
            chklowtinv = th.inverse(chklowt)
            Heff = (chklowtinv @ hkmat @ th.transpose(chklowtinv,dim0=1,dim1=2).conj())
        # the factor 13.605662285137 * 2 from Hartree to eV.
        # eigks = th.linalg.eigvalsh(Heff) * 13.605662285137 * 2
        if if_eigvec:
            eigks, eigvec = th.linalg.eigh(Heff)
        else:
            eigks = th.linalg.eigvalsh(Heff)
        
        if unit == "Hartree":
            factor = 13.605662285137 * 2
        elif unit == "eV":
            factor = 1.0
        elif unit == "Ry":
            factor = 13.605662285137
        else:
            log.error("The unit name is not correct !")
            raise ValueError
        eigks = eigks * factor
        # Qres = Q.detach()
        # else:
        #     chklowt = np.linalg.cholesky(skmat)
        #     chklowtinv = np.linalg.inv(chklowt)
        #     Heff = (chklowtinv @ hkmat @ np.transpose(chklowtinv,(0,2,1)).conj())
        #     eigks = np.linalg.eigvalsh(Heff) * 13.605662285137 * 2
        #     Qres = 0
        
        if if_eigvec:
            return eigks, eigvec
        else:
            return eigks, None
//  [python]
def Eigenvalues(self, kpoints, time_symm=True, unit="Hartree",if_eigvec=False):
        """ using the tight-binding H and S matrix calculate eigenvalues at kpoints.
        
        Args:
            kpoints: the k-kpoints used to calculate the eigenvalues.
        Note: must have the BondHBlock and BondSBlock 
        """
        hkmat = self.hs_block_R2k(kpoints=kpoints, HorS='H', time_symm=time_symm)
        if not self.use_orthogonal_basis:
            skmat =  self.hs_block_R2k(kpoints=kpoints, HorS='S', time_symm=time_symm)
        else:
            skmat = torch.eye(hkmat.shape[1], dtype=self.cdtype).unsqueeze(0).repeat(hkmat.shape[0], 1, 1)

        if self.use_orthogonal_basis:
            Heff = hkmat
        else:
            chklowt = th.linalg.cholesky(skmat)
            chklowtinv = th.linalg.inv(chklowt)
            Heff = (chklowtinv @ hkmat @ th.transpose(chklowtinv,dim0=1,dim1=2).conj())
        # the factor 13.605662285137 * 2 from Hartree to eV.
        # eigks = th.linalg.eigvalsh(Heff) * 13.605662285137 * 2
        if if_eigvec:
            eigks, eigvec = th.linalg.eigh(Heff)
        else:
            eigks = th.linalg.eigvalsh(Heff)
        
        if unit == "Hartree":
            factor = 13.605662285137 * 2
        elif unit == "eV":
            factor = 1.0
        elif unit == "Ry":
            factor = 13.605662285137
        else:
            log.error("The unit name is not correct !")
            raise ValueError
        eigks = eigks * factor
        # Qres = Q.detach()
        # else:
        #     chklowt = np.linalg.cholesky(skmat)
        #     chklowtinv = np.linalg.inv(chklowt)
        #     Heff = (chklowtinv @ hkmat @ np.transpose(chklowtinv,(0,2,1)).conj())
        #     eigks = np.linalg.eigvalsh(Heff) * 13.605662285137 * 2
        #     Qres = 0
        
        if if_eigvec:
            return eigks, eigvec
        else:
            return eigks, None
//  [python]
def forward(ctx, H, b):
        return_chol = False
        # don't crash training if cholesky decomp fails
        U = None
        U_inv = None
        H_inv = None
        try:
            U = torch.cholesky(H)
            if not return_chol:
                xs = torch.cholesky_solve(b, U)
                ctx.save_for_backward(U, xs)
                ctx.failed = False
            else:
                identity = torch.eye(U.shape[1], device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1)
                U_inv = torch.triangular_solve(identity, U, upper=False)[0]
                U_conj_inv = torch.triangular_solve(identity, torch.transpose(U, -1, -2), upper=True)[0]
                H_inv = U_conj_inv @ U_inv
                xs = H_inv @ b

            ctx.save_for_backward(U, xs)
            ctx.failed = False
        except Exception as e:
            print(e)
            ctx.failed = True
            xs = torch.zeros_like(b)

        return xs, U, U_inv, H_inv
//  [python]
def forward(ctx, H, b):
        return_chol = False
        # don't crash training if cholesky decomp fails
        U = None
        U_inv = None
        H_inv = None
        try:
            U = torch.linalg.cholesky(H)
            if not return_chol:
                xs = torch.cholesky_solve(b, U)
                ctx.save_for_backward(U, xs)
                ctx.failed = False
            else:
                identity = torch.eye(U.shape[1], device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1)
                U_inv = torch.triangular_solve(identity, U, upper=False)[0]
                U_conj_inv = torch.triangular_solve(identity, torch.transpose(U, -1, -2), upper=True)[0]
                H_inv = U_conj_inv @ U_inv
                xs = H_inv @ b

            ctx.save_for_backward(U, xs)
            ctx.failed = False
        except Exception as e:
            print(e)
            ctx.failed = True
            xs = torch.zeros_like(b)

        return xs, U, U_inv, H_inv
//  [python]
def layer_transform(self, layer, name):
        self.qparams = {}
        self.columns = self.layers_cache[name]["columns"]
        W = layer.weight.data.clone()
        if isinstance(layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(layer, transformers.Conv1D):
            W = W.t()

        W = W.float()

        tick = time.time()

        self.groups = [None] * (self.columns // self.wquantizer.group_size)

        H = self.layers_cache[name]["H"]
        del self.layers_cache[name]["H"]

        if self.actorder:
            self.perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, self.perm]
            H = H[self.perm][:, self.perm]
            self.invperm = torch.argsort(self.perm)
            layer.register_buffer("buf_perm", self.perm)
            layer.register_buffer("buf_invperm", self.invperm)

        dead = torch.diag(H) == 0
        if self.percdamp > 0:
            damp = self.percdamp * abs(torch.diag(H)).mean()
            diag = torch.arange(self.columns, device=self.dev)
            H[diag, diag] += damp
            del diag
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros_like(W)
        tmp = torch.zeros_like(W)

        H = torch.cholesky(H)
        H = torch.inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H
        mask = torch.zeros_like(W, dtype=torch.bool)
        self.weight_transform(W, Hinv, Losses, tmp, mask)

        torch.cuda.synchronize()
        logger.info(f"time {time.time() - tick}")
        logger.info(f"error {torch.sum(Losses).item()}")

        if self.actorder:
            tmp = tmp[:, self.invperm]
            mask = mask[:, self.invperm]

        if isinstance(layer, transformers.Conv1D):
            tmp = tmp.t()
            mask = mask.t()

        assert layer.weight.shape == tmp.shape
        layer.weight.data = tmp

        logger.info(f"tmp {tmp}")
        logger.info(f"outliers {torch.sum(mask)} / {mask.numel()}")

        if self.wquantizer.granularity == "per_group":
            self.set_model_qparams(layer)
            layer.register_buffer("buf_mask", mask.float().to_sparse())
//  [python]
def layer_transform(self, layer, name):
        self.qparams = {}
        self.columns = self.layers_cache[name]["columns"]
        W = layer.weight.data.clone()
        if isinstance(layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(layer, transformers.Conv1D):
            W = W.t()

        W = W.float()

        tick = time.time()

        self.groups = [None] * (self.columns // self.wquantizer.group_size)

        H = self.layers_cache[name]["H"]
        del self.layers_cache[name]["H"]

        if self.actorder:
            self.perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, self.perm]
            H = H[self.perm][:, self.perm]
            self.invperm = torch.argsort(self.perm)
            layer.register_buffer("buf_perm", self.perm)
            layer.register_buffer("buf_invperm", self.invperm)

        dead = torch.diag(H) == 0
        if self.percdamp > 0:
            damp = self.percdamp * abs(torch.diag(H)).mean()
            diag = torch.arange(self.columns, device=self.dev)
            H[diag, diag] += damp
            del diag
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros_like(W)
        tmp = torch.zeros_like(W)

        H = torch.linalg.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H, upper=True)
        Hinv = H
        mask = torch.zeros_like(W, dtype=torch.bool)
        self.weight_transform(W, Hinv, Losses, tmp, mask)

        torch.cuda.synchronize()
        logger.info(f"time {time.time() - tick}")
        logger.info(f"error {torch.sum(Losses).item()}")

        if self.actorder:
            tmp = tmp[:, self.invperm]
            mask = mask[:, self.invperm]

        if isinstance(layer, transformers.Conv1D):
            tmp = tmp.t()
            mask = mask.t()

        assert layer.weight.shape == tmp.shape
        layer.weight.data = tmp

        logger.info(f"tmp {tmp}")
        logger.info(f"outliers {torch.sum(mask)} / {mask.numel()}")

        if self.wquantizer.granularity == "per_group":
            self.set_model_qparams(layer)
            layer.register_buffer("buf_mask", mask.float().to_sparse())
//  [python]
def faster_quant(self, H, W):
        percdamp = self.percdamp
        blocksize = self.blocksize
        groupsize = self.groupsize
        orig_dtype = W.dtype
        W = W.detach().float()
        rows, columns = W.shape[0], W.shape[1]
        device = W.device

        if groupsize == -1:
            cur_qparams = self.get_qparams_func(W)
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros_like(W)
        DQ = torch.zeros_like(W)

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

        all_qparams = []
        for i1 in range(0, columns, blocksize):
            i2 = min(i1 + blocksize, columns)
            count = i2 - i1
            W1 = W[:, i1:i2].clone()
            DQ1 = 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 and (i1 + i) % groupsize == 0:  # start of new group
                    cur_qparams = self.get_qparams_func(
                        W[:, (i1 + i) : (i1 + i + groupsize)]
                    )
                    all_qparams.append(cur_qparams)

                q = self.quantize_func(w.unsqueeze(1), cur_qparams).flatten()
                dq = self.dequantize_func(q.unsqueeze(1), cur_qparams).flatten()

                DQ1[:, i] = dq
                Losses1[:, i] = (w - dq) ** 2 / d**2

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

            DQ[:, i1:i2] = DQ1
            Losses[:, i1:i2] = Losses1 / 2

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

        torch.cuda.synchronize()

        if all_qparams == []:
            all_qparams.append(cur_qparams)

        # convert a list of qparams objects into a single one. enerally by
        # concatenating a bunch of n,1 scale/zeros tensors into a n,num_groups tensor
        all_qparams = self.combine_qparams_list_func(all_qparams)
        Q = self.quantize_func(DQ, all_qparams)
        return Q, DQ.to(orig_dtype), all_qparams
//  [python]
def faster_quant(self, H, W):
        percdamp = self.percdamp
        blocksize = self.blocksize
        groupsize = self.groupsize
        orig_dtype = W.dtype
        W = W.detach().float()
        rows, columns = W.shape[0], W.shape[1]
        device = W.device

        if groupsize == -1:
            cur_qparams = self.get_qparams_func(W)
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros_like(W)
        DQ = torch.zeros_like(W)

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

        all_qparams = []
        for i1 in range(0, columns, blocksize):
            i2 = min(i1 + blocksize, columns)
            count = i2 - i1
            W1 = W[:, i1:i2].clone()
            DQ1 = 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 and (i1 + i) % groupsize == 0:  # start of new group
                    cur_qparams = self.get_qparams_func(
                        W[:, (i1 + i) : (i1 + i + groupsize)]
                    )
                    all_qparams.append(cur_qparams)

                q = self.quantize_func(w.unsqueeze(1), cur_qparams).flatten()
                dq = self.dequantize_func(q.unsqueeze(1), cur_qparams).flatten()

                DQ1[:, i] = dq
                Losses1[:, i] = (w - dq) ** 2 / d**2

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

            DQ[:, i1:i2] = DQ1
            Losses[:, i1:i2] = Losses1 / 2

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

        torch.cuda.synchronize()

        if all_qparams == []:
            all_qparams.append(cur_qparams)

        # convert a list of qparams objects into a single one. enerally by
        # concatenating a bunch of n,1 scale/zeros tensors into a n,num_groups tensor
        all_qparams = self.combine_qparams_list_func(all_qparams)
        Q = self.quantize_func(DQ, all_qparams)
        return Q, DQ.to(orig_dtype), all_qparams
//  [python]
def set_cov(self, cov=None):
        self.chol_cov = torch.cholesky(cov or self.cov0)
//  [python]
def set_cov(self, cov=None):
        self.chol_cov = torch.linalg.cholesky(cov or self.cov0)