jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def exacteig(A: LinearOperator, neig: int,
             mode: str, M: Optional[LinearOperator]) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Eigendecomposition using explicit matrix construction.
    No additional option for this method.

    Warnings
    --------
    * As this method construct the linear operators explicitly, it might requires
      a large memory.
    """
    Amatrix = A.fullmatrix()  # (*BA, q, q)
    if M is None:
        # evals, evecs = torch.linalg.eigh(Amatrix, eigenvectors=True)  # (*BA, q), (*BA, q, q)
        evals, evecs = degen_symeig.apply(Amatrix)  # (*BA, q, q)
        return _take_eigpairs(evals, evecs, neig, mode)
    else:
        Mmatrix = M.fullmatrix()  # (*BM, q, q)

        # M decomposition to make A symmetric
        # it is done this way to make it numerically stable in avoiding
        # complex eigenvalues for (near-)degenerate case
        L = torch.cholesky(Mmatrix)  # (*BM, q, q)
        Linv = torch.inverse(L)  # (*BM, q, q)
        LinvT = Linv.transpose(-2, -1).conj()  # (*BM, q, q)
        A2 = torch.matmul(Linv, torch.matmul(Amatrix, LinvT))  # (*BAM, q, q)

        # calculate the eigenvalues and eigenvectors
        # (the eigvecs are normalized in M-space)
        # evals, evecs = torch.linalg.eigh(A2, eigenvectors=True)  # (*BAM, q, q)
        evals, evecs = degen_symeig.apply(A2)  # (*BAM, q, q)
        evals, evecs = _take_eigpairs(evals, evecs, neig, mode)  # (*BAM, neig) and (*BAM, q, neig)
        evecs = torch.matmul(LinvT, evecs)
        return evals, evecs
//  [python]
def exacteig(A: LinearOperator, neig: int,
             mode: str, M: Optional[LinearOperator]) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Eigendecomposition using explicit matrix construction.
    No additional option for this method.

    Warnings
    --------
    * As this method construct the linear operators explicitly, it might requires
      a large memory.
    """
    Amatrix = A.fullmatrix()  # (*BA, q, q)
    if M is None:
        # evals, evecs = torch.linalg.eigh(Amatrix, eigenvectors=True)  # (*BA, q), (*BA, q, q)
        evals, evecs = degen_symeig.apply(Amatrix)  # (*BA, q, q)
        return _take_eigpairs(evals, evecs, neig, mode)
    else:
        Mmatrix = M.fullmatrix()  # (*BM, q, q)

        # M decomposition to make A symmetric
        # it is done this way to make it numerically stable in avoiding
        # complex eigenvalues for (near-)degenerate case
        L = torch.linalg.cholesky(Mmatrix)  # (*BM, q, q)
        Linv = torch.inverse(L)  # (*BM, q, q)
        LinvT = Linv.transpose(-2, -1).conj()  # (*BM, q, q)
        A2 = torch.matmul(Linv, torch.matmul(Amatrix, LinvT))  # (*BAM, q, q)

        # calculate the eigenvalues and eigenvectors
        # (the eigvecs are normalized in M-space)
        # evals, evecs = torch.linalg.eigh(A2, eigenvectors=True)  # (*BAM, q, q)
        evals, evecs = degen_symeig.apply(A2)  # (*BAM, q, q)
        evals, evecs = _take_eigpairs(evals, evecs, neig, mode)  # (*BAM, neig) and (*BAM, q, neig)
        evecs = torch.matmul(LinvT, evecs)
        return evals, evecs
//  [python]
def normal_samples(mean, var, n_samples, generator=None):
    """Produce samples from a batch of Normal distributions either parameterized
    by a diagonal or full covariance given by `var`.

    Parameters
    ----------
    mean : torch.Tensor
        `(batch_size, output_dim)`
    var : torch.Tensor
        (co)variance of the Normal distribution
        `(batch_size, output_dim, output_dim)` or `(batch_size, output_dim)`
    generator : torch.Generator
        random number generator
    """
    assert mean.ndim == 2, 'Invalid input shape of mean, should be 2-dimensional.'
    _, output_dim = mean.shape
    randn_samples = torch.randn((output_dim, n_samples), device=mean.device, generator=generator)
    
    if mean.shape == var.shape:
        # diagonal covariance
        scaled_samples = var.sqrt().unsqueeze(-1) * randn_samples.unsqueeze(0)
        return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
    elif mean.shape == var.shape[:2] and var.shape[-1] == mean.shape[1]:
        # full covariance
        scale = torch.cholesky(var)
        scaled_samples = torch.matmul(scale, randn_samples.unsqueeze(0))  # expand batch dim
        return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
    else:
        raise ValueError('Invalid input shapes.')
//  [python]
def normal_samples(mean, var, n_samples, generator=None):
    """Produce samples from a batch of Normal distributions either parameterized
    by a diagonal or full covariance given by `var`.

    Parameters
    ----------
    mean : torch.Tensor
        `(batch_size, output_dim)`
    var : torch.Tensor
        (co)variance of the Normal distribution
        `(batch_size, output_dim, output_dim)` or `(batch_size, output_dim)`
    generator : torch.Generator
        random number generator
    """
    assert mean.ndim == 2, 'Invalid input shape of mean, should be 2-dimensional.'
    _, output_dim = mean.shape
    randn_samples = torch.randn((output_dim, n_samples), device=mean.device, generator=generator)
    
    if mean.shape == var.shape:
        # diagonal covariance
        scaled_samples = var.sqrt().unsqueeze(-1) * randn_samples.unsqueeze(0)
        return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
    elif mean.shape == var.shape[:2] and var.shape[-1] == mean.shape[1]:
        # full covariance
        scale = torch.linalg.cholesky(var)
        scaled_samples = torch.matmul(scale, randn_samples.unsqueeze(0))  # expand batch dim
        return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
    else:
        raise ValueError('Invalid input shapes.')
//  [python]
def sample_inputs_linalg_cholesky_inverse(op_info, device, dtype, requires_grad=False, **kwargs):
    # Generate Cholesky factors of positive-definite (non-singular) Hermitian (symmetric) matrices
    from torch.testing._internal.common_utils import random_hermitian_pd_matrix
    inputs = (
        torch.zeros(0, 0, dtype=dtype, device=device),  # 0x0 matrix
        torch.zeros(0, 2, 2, dtype=dtype, device=device),  # zero batch of matrices
        random_hermitian_pd_matrix(S, dtype=dtype, device=device),  # single matrix
        random_hermitian_pd_matrix(S, 2, dtype=dtype, device=device),  # batch of matrices
    )
    test_cases = (torch.cholesky(a) for a in inputs)
    out = []
    for a in test_cases:
        a.requires_grad = requires_grad
        out.append(SampleInput(a))
        out.append(SampleInput(a.clone().requires_grad_(requires_grad), kwargs=dict(upper=True)))
    return out
//  [python]
def sample_inputs_linalg_cholesky_inverse(op_info, device, dtype, requires_grad=False, **kwargs):
    # Generate Cholesky factors of positive-definite (non-singular) Hermitian (symmetric) matrices
    from torch.testing._internal.common_utils import random_hermitian_pd_matrix
    inputs = (
        torch.zeros(0, 0, dtype=dtype, device=device),  # 0x0 matrix
        torch.zeros(0, 2, 2, dtype=dtype, device=device),  # zero batch of matrices
        random_hermitian_pd_matrix(S, dtype=dtype, device=device),  # single matrix
        random_hermitian_pd_matrix(S, 2, dtype=dtype, device=device),  # batch of matrices
    )
    test_cases = (torch.linalg.cholesky(a) for a in inputs)
    out = []
    for a in test_cases:
        a.requires_grad = requires_grad
        out.append(SampleInput(a))
        out.append(SampleInput(a.clone().requires_grad_(requires_grad), kwargs=dict(upper=True)))
    return out
//  [python]
def sample_inputs_linalg_cholesky_inverse(op_info, device, dtype, requires_grad=False):
    # Generate Cholesky factors of positive-definite (non-singular) Hermitian (symmetric) matrices
    from torch.testing._internal.common_utils import random_hermitian_pd_matrix
    inputs = (
        torch.zeros(0, 0, dtype=dtype, device=device),  # 0x0 matrix
        torch.zeros(0, 2, 2, dtype=dtype, device=device),  # zero batch of matrices
        random_hermitian_pd_matrix(S, dtype=dtype, device=device),  # single matrix
        random_hermitian_pd_matrix(S, 2, dtype=dtype, device=device),  # batch of matrices
    )
    test_cases = (torch.cholesky(a) for a in inputs)
    out = []
    for a in test_cases:
        a.requires_grad = requires_grad
        out.append(SampleInput(a))
        out.append(SampleInput(a, kwargs=dict(upper=True)))
    return out
//  [python]
def sample_inputs_linalg_cholesky_inverse(op_info, device, dtype, requires_grad=False):
    # Generate Cholesky factors of positive-definite (non-singular) Hermitian (symmetric) matrices
    from torch.testing._internal.common_utils import random_hermitian_pd_matrix
    inputs = (
        torch.zeros(0, 0, dtype=dtype, device=device),  # 0x0 matrix
        torch.zeros(0, 2, 2, dtype=dtype, device=device),  # zero batch of matrices
        random_hermitian_pd_matrix(S, dtype=dtype, device=device),  # single matrix
        random_hermitian_pd_matrix(S, 2, dtype=dtype, device=device),  # batch of matrices
    )
    test_cases = (torch.linalg.cholesky(a) for a in inputs)
    out = []
    for a in test_cases:
        a.requires_grad = requires_grad
        out.append(SampleInput(a))
        out.append(SampleInput(a, kwargs=dict(upper=True)))
    return out
//  [python]
def _reset_cache(self):
		"""Reset the internally stored statistics.

		This method is meant to only be called internally. It resets the
		stored statistics used to update the model parameters as well as
		recalculates the cached values meant to speed up log probability
		calculations.
		"""

		if self._initialized == False:
			return

		self.register_buffer("_w_sum", torch.zeros(self.d, dtype=self.dtype, 
			device=self.device))
		self.register_buffer("_xw_sum", torch.zeros(self.d, dtype=self.dtype,
			device=self.device))

		if self.covariance_type == 'full':
			self.register_buffer("_xxw_sum", torch.zeros(self.d, self.d, 
				dtype=self.dtype, device=self.device))

			if self.covs.sum() > 0.0:
				chol = torch.cholesky(self.covs)
				_inv_cov = torch.triangular_solve(torch.eye(
					len(self.covs), dtype=self.dtype, device=self.device), 
					chol, upper=False)[0].T
				_inv_cov_dot_mu = torch.matmul(self.means, _inv_cov)
				_log_det = -0.5 * torch.slogdet(self.covs)[1]
				_theta = _log_det - 0.5 * (self.d * LOG_2_PI)

				self.register_buffer("_inv_cov", _inv_cov)
				self.register_buffer("_inv_cov_dot_mu", _inv_cov_dot_mu)
				self.register_buffer("_log_det", _log_det)
				self.register_buffer("_theta", _theta)

		elif self.covariance_type in ('diag', 'sphere'):
			self.register_buffer("_xxw_sum", torch.zeros(self.d, 
				dtype=self.dtype, device=self.device))

			if self.covs.sum() > 0.0:
				_log_sigma_sqrt_2pi = -torch.log(torch.sqrt(self.covs) * 
					SQRT_2_PI)
				_inv_two_sigma = 1. / (2 * self.covs)

				self.register_buffer("_log_sigma_sqrt_2pi", _log_sigma_sqrt_2pi)
				self.register_buffer("_inv_two_sigma", _inv_two_sigma)
			
			if any(self.covs < 0):
				raise ValueError("Variances must be positive.")
//  [python]
def _reset_cache(self):
		"""Reset the internally stored statistics.

		This method is meant to only be called internally. It resets the
		stored statistics used to update the model parameters as well as
		recalculates the cached values meant to speed up log probability
		calculations.
		"""

		if self._initialized == False:
			return

		self.register_buffer("_w_sum", torch.zeros(self.d, dtype=self.dtype, 
			device=self.device))
		self.register_buffer("_xw_sum", torch.zeros(self.d, dtype=self.dtype,
			device=self.device))

		if self.covariance_type == 'full':
			self.register_buffer("_xxw_sum", torch.zeros(self.d, self.d, 
				dtype=self.dtype, device=self.device))

			if self.covs.sum() > 0.0:
				chol = torch.linalg.cholesky(self.covs)
				_inv_cov = torch.linalg.solve_triangular(chol, torch.eye(
					len(self.covs), dtype=self.dtype, device=self.device), 
					upper=False).T
				_inv_cov_dot_mu = torch.matmul(self.means, _inv_cov)
				_log_det = -0.5 * torch.linalg.slogdet(self.covs)[1]
				_theta = _log_det - 0.5 * (self.d * LOG_2_PI)

				self.register_buffer("_inv_cov", _inv_cov)
				self.register_buffer("_inv_cov_dot_mu", _inv_cov_dot_mu)
				self.register_buffer("_log_det", _log_det)
				self.register_buffer("_theta", _theta)

		elif self.covariance_type in ('diag', 'sphere'):
			self.register_buffer("_xxw_sum", torch.zeros(self.d, 
				dtype=self.dtype, device=self.device))

			if self.covs.sum() > 0.0:
				_log_sigma_sqrt_2pi = -torch.log(torch.sqrt(self.covs) * 
					SQRT_2_PI)
				_inv_two_sigma = 1. / (2 * self.covs)

				self.register_buffer("_log_sigma_sqrt_2pi", _log_sigma_sqrt_2pi)
				self.register_buffer("_inv_two_sigma", _inv_two_sigma)
			
			if any(self.covs < 0):
				raise ValueError("Variances must be positive.")
//  [python]
def fasterquant(
		self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False
	):
		W = self.layer.weight.data.clone()
		if isinstance(self.layer, nn.Conv2d):
			W = W.flatten(1)
		if isinstance(self.layer, transformers.Conv1D):
			W = W.t()
		W = W.float()

		tick = time.time()

		if 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

		if actorder:
			perm = torch.argsort(torch.diag(H), descending=True)
			W = W[:, perm]
			H = H[perm][:, 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
		
		scale = []
		zero = []
		now_idx = 1

		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)
					
					if ((i1 + i) // groupsize) - now_idx == -1:
						scale.append(self.quantizer.scale)
						zero.append(self.quantizer.zero)
						now_idx += 1

				q = 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))

		torch.cuda.synchronize()
		print('time %.2f' % (time.time() - tick))
		print('error', torch.sum(Losses).item())
		
		if actorder:
			invperm = torch.argsort(perm)
			Q = Q[:, invperm]

		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))
			
		if scale == []:
			scale.append(self.quantizer.scale)
			zero.append(self.quantizer.zero)
		scale = torch.cat(scale,dim=1)
		zero = torch.cat(zero,dim=1)
		return scale,zero
//  [python]
def fasterquant(
		self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False
	):
		W = self.layer.weight.data.clone()
		if isinstance(self.layer, nn.Conv2d):
			W = W.flatten(1)
		if isinstance(self.layer, transformers.Conv1D):
			W = W.t()
		W = W.float()

		tick = time.time()

		if 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

		if actorder:
			perm = torch.argsort(torch.diag(H), descending=True)
			W = W[:, perm]
			H = H[perm][:, 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
		
		scale = []
		zero = []
		now_idx = 1

		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)
					
					if ((i1 + i) // groupsize) - now_idx == -1:
						scale.append(self.quantizer.scale)
						zero.append(self.quantizer.zero)
						now_idx += 1

				q = 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))

		torch.cuda.synchronize()
		print('time %.2f' % (time.time() - tick))
		print('error', torch.sum(Losses).item())
		
		if actorder:
			invperm = torch.argsort(perm)
			Q = Q[:, invperm]

		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))
			
		if scale == []:
			scale.append(self.quantizer.scale)
			zero.append(self.quantizer.zero)
		scale = torch.cat(scale,dim=1)
		zero = torch.cat(zero,dim=1)
		return scale,zero
//  [python]
def _sample_approx_posterior_mog(
        self, num_samples, x: Tensor, batch_size: int
    ) -> Tensor:
        r"""Sample from the approximate posterior.

        Args:
            num_samples: Desired number of samples.
            x: Conditioning context for posterior $p(\theta|x)$.
            batch_size: Batch size for sampling.

        Returns:
            Samples from the approximate mixture of Gaussians posterior.
        """

        # Compute the mixture components of the posterior.
        logits_p, m_p, prec_p = self._posthoc_correction(x)

        # Compute the precision factors which represent the upper triangular matrix
        # of the cholesky decomposition of the prec_p.
        prec_factors_p = torch.cholesky(prec_p, upper=True)

        assert logits_p.ndim == 2
        assert m_p.ndim == 3
        assert prec_p.ndim == 4
        assert prec_factors_p.ndim == 4

        # Replicate to use batched sampling from pyknos.
        if batch_size is not None and batch_size > 1:
            logits_p = logits_p.repeat(batch_size, 1)
            m_p = m_p.repeat(batch_size, 1, 1)
            prec_factors_p = prec_factors_p.repeat(batch_size, 1, 1, 1)

        # Get (optionally z-scored) MoG samples.
        theta = MultivariateGaussianMDN.sample_mog(
            num_samples, logits_p, m_p, prec_factors_p
        )

        embedded_context = self._neural_net._embedding_net(x)
        if embedded_context is not None:
            # Merge the context dimension with sample dimension in order to
            # apply the transform.
            theta = torchutils.merge_leading_dims(theta, num_dims=2)
            embedded_context = torchutils.repeat_rows(
                embedded_context, num_reps=num_samples
            )

        theta, _ = self._neural_net._transform.inverse(theta, context=embedded_context)

        if embedded_context is not None:
            # Split the context dimension from sample dimension.
            theta = torchutils.split_leading_dim(theta, shape=[-1, num_samples])

        return theta
//  [python]
def _sample_approx_posterior_mog(
        self, num_samples, x: Tensor, batch_size: int
    ) -> Tensor:
        r"""Sample from the approximate posterior.

        Args:
            num_samples: Desired number of samples.
            x: Conditioning context for posterior $p(\theta|x)$.
            batch_size: Batch size for sampling.

        Returns:
            Samples from the approximate mixture of Gaussians posterior.
        """

        # Compute the mixture components of the posterior.
        logits_p, m_p, prec_p = self._posthoc_correction(x)

        # Compute the precision factors which represent the upper triangular matrix
        # of the cholesky decomposition of the prec_p.
        prec_factors_p = torch.linalg.cholesky(prec_p, upper=True)

        assert logits_p.ndim == 2
        assert m_p.ndim == 3
        assert prec_p.ndim == 4
        assert prec_factors_p.ndim == 4

        # Replicate to use batched sampling from pyknos.
        if batch_size is not None and batch_size > 1:
            logits_p = logits_p.repeat(batch_size, 1)
            m_p = m_p.repeat(batch_size, 1, 1)
            prec_factors_p = prec_factors_p.repeat(batch_size, 1, 1, 1)

        # Get (optionally z-scored) MoG samples.
        theta = MultivariateGaussianMDN.sample_mog(
            num_samples, logits_p, m_p, prec_factors_p
        )

        embedded_context = self._neural_net._embedding_net(x)
        if embedded_context is not None:
            # Merge the context dimension with sample dimension in order to
            # apply the transform.
            theta = torchutils.merge_leading_dims(theta, num_dims=2)
            embedded_context = torchutils.repeat_rows(
                embedded_context, num_reps=num_samples
            )

        theta, _ = self._neural_net._transform.inverse(theta, context=embedded_context)

        if embedded_context is not None:
            # Split the context dimension from sample dimension.
            theta = torchutils.split_leading_dim(theta, shape=[-1, num_samples])

        return theta
//  [python]
def forward(self, x_train, y_train, x_test):
        
        if isinstance(y_train, list):
            y_train_var = y_train[1]
            y_train = y_train[0]
        else:
            y_train_var = None
        Sigma = self.kernel(x_train, x_train) + self.log_beta.exp().pow(-1) * torch.eye(x_train.size(0)) \
            + JITTER * torch.eye(x_train.size(0))
        
        kx = self.kernel(x_train, x_test)
        L = torch.cholesky(Sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        # option 1
        mean = kx.t() @ torch.cholesky_solve(y_train, L)  # torch.linalg.cholesky()
        
        var = self.kernel(x_test, x_test) - LinvKx.t() @ LinvKx

        # add the noise uncertainty
        var = var + self.log_beta.exp().pow(-1)
        # if y_train_var is not None:
        #     var = var + y_train_var.diag()* torch.eye(x_test.size(0))

        return mean, var
//  [python]
def forward(self, x_train, y_train, x_test):
        
        if isinstance(y_train, list):
            y_train_var = y_train[1]
            y_train = y_train[0]
        else:
            y_train_var = None
        Sigma = self.kernel(x_train, x_train) + self.log_beta.exp().pow(-1) * torch.eye(x_train.size(0)) \
            + JITTER * torch.eye(x_train.size(0))
        
        kx = self.kernel(x_train, x_test)
        L = torch.linalg.cholesky(Sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        # option 1
        mean = kx.t() @ torch.cholesky_solve(y_train, L)  # torch.linalg.cholesky()
        
        var = self.kernel(x_test, x_test) - LinvKx.t() @ LinvKx

        # add the noise uncertainty
        var = var + self.log_beta.exp().pow(-1)
        # if y_train_var is not None:
        #     var = var + y_train_var.diag()* torch.eye(x_test.size(0))

        return mean, var
//  [python]
def variational_loss(self, x, y, noise, mu_hat, A_hat):
        """
        Computes L_H for the data in the current batch.
        :param x: auxiliary data for current batch (batch, 1 + 1 + M)
        :param y: mean vector for current latent channel, output of the encoder network (batch, 1)
        :param noise: variance vector for current latent channel, output of the encoder network (batch, 1)
        :param mu_hat:
        :param A_hat:
        :return: sum_term, KL_term (variational loss = sum_term + KL_term)  (1,)
        """
        b = x.shape[0]
        m = self.inducing_index_points.shape[0]

        K_mm = self.kernel_matrix(self.inducing_index_points, self.inducing_index_points) # (m,m)
        K_mm_inv = torch.linalg.inv(_add_diagonal_jitter(K_mm, self.jitter)) # (m,m)

        K_nn = self.kernel_matrix(x, x, x_inducing=False, y_inducing=False, diag_only=True) # (b)

        K_nm = self.kernel_matrix(x, self.inducing_index_points, x_inducing=False)  # (b, m)
        K_mn = torch.transpose(K_nm, 0, 1)

#        S = A_hat

        # KL term
        mean_vector = torch.matmul(K_nm, torch.matmul(K_mm_inv, mu_hat))

        K_mm_chol = torch.cholesky(_add_diagonal_jitter(K_mm, self.jitter))
        S_chol = torch.cholesky(_add_diagonal_jitter(A_hat, self.jitter))
        K_mm_log_det = 2 * torch.sum(torch.log(torch.diagonal(K_mm_chol)))
        S_log_det = 2 * torch.sum(torch.log(torch.diagonal(S_chol)))

        KL_term = 0.5 * (K_mm_log_det - S_log_det - m +
                             torch.trace(torch.matmul(K_mm_inv, A_hat)) +
                             torch.sum(mu_hat * torch.matmul(K_mm_inv, mu_hat)))

        # diag(K_tilde), (b, )
        precision = 1 / noise

        K_tilde_terms = precision * (K_nn - torch.diagonal(torch.matmul(K_nm, torch.matmul(K_mm_inv, K_mn))))

        # k_i \cdot k_i^T, (b, m, m)
        lambda_mat = torch.matmul(K_nm.unsqueeze(2), torch.transpose(K_nm.unsqueeze(2), 1, 2))

        # K_mm_inv \cdot k_i \cdot k_i^T \cdot K_mm_inv, (b, m, m)
        lambda_mat = torch.matmul(K_mm_inv, torch.matmul(lambda_mat, K_mm_inv))

        # Trace terms, (b,)
        trace_terms = precision * torch.einsum('bii->b', torch.matmul(A_hat, lambda_mat))

        # L_3 sum part, (1,)
        L_3_sum_term = -0.5 * (torch.sum(K_tilde_terms) + torch.sum(trace_terms) +
                                torch.sum(torch.log(noise)) + b * np.log(2 * np.pi) +
                                torch.sum(precision * (y - mean_vector) ** 2))

        return L_3_sum_term, KL_term
//  [python]
def variational_loss(self, x, y, noise, mu_hat, A_hat):
        """
        Computes L_H for the data in the current batch.
        :param x: auxiliary data for current batch (batch, 1 + 1 + M)
        :param y: mean vector for current latent channel, output of the encoder network (batch, 1)
        :param noise: variance vector for current latent channel, output of the encoder network (batch, 1)
        :param mu_hat:
        :param A_hat:
        :return: sum_term, KL_term (variational loss = sum_term + KL_term)  (1,)
        """
        b = x.shape[0]
        m = self.inducing_index_points.shape[0]

        K_mm = self.kernel_matrix(self.inducing_index_points, self.inducing_index_points) # (m,m)
        K_mm_inv = torch.linalg.inv(_add_diagonal_jitter(K_mm, self.jitter)) # (m,m)

        K_nn = self.kernel_matrix(x, x, x_inducing=False, y_inducing=False, diag_only=True) # (b)

        K_nm = self.kernel_matrix(x, self.inducing_index_points, x_inducing=False)  # (b, m)
        K_mn = torch.transpose(K_nm, 0, 1)

#        S = A_hat

        # KL term
        mean_vector = torch.matmul(K_nm, torch.matmul(K_mm_inv, mu_hat))

        K_mm_chol = torch.linalg.cholesky(_add_diagonal_jitter(K_mm, self.jitter))
        S_chol = torch.linalg.cholesky(_add_diagonal_jitter(A_hat, self.jitter))
        K_mm_log_det = 2 * torch.sum(torch.log(torch.diagonal(K_mm_chol)))
        S_log_det = 2 * torch.sum(torch.log(torch.diagonal(S_chol)))

        KL_term = 0.5 * (K_mm_log_det - S_log_det - m +
                             torch.trace(torch.matmul(K_mm_inv, A_hat)) +
                             torch.sum(mu_hat * torch.matmul(K_mm_inv, mu_hat)))

        # diag(K_tilde), (b, )
        precision = 1 / noise

        K_tilde_terms = precision * (K_nn - torch.diagonal(torch.matmul(K_nm, torch.matmul(K_mm_inv, K_mn))))

        # k_i \cdot k_i^T, (b, m, m)
        lambda_mat = torch.matmul(K_nm.unsqueeze(2), torch.transpose(K_nm.unsqueeze(2), 1, 2))

        # K_mm_inv \cdot k_i \cdot k_i^T \cdot K_mm_inv, (b, m, m)
        lambda_mat = torch.matmul(K_mm_inv, torch.matmul(lambda_mat, K_mm_inv))

        # Trace terms, (b,)
        trace_terms = precision * torch.einsum('bii->b', torch.matmul(A_hat, lambda_mat))

        # L_3 sum part, (1,)
        L_3_sum_term = -0.5 * (torch.sum(K_tilde_terms) + torch.sum(trace_terms) +
                                torch.sum(torch.log(noise)) + b * np.log(2 * np.pi) +
                                torch.sum(precision * (y - mean_vector) ** 2))

        return L_3_sum_term, KL_term
//  [python]
def negative_log_likelihood(self, x_train, y_train):
        if isinstance(y_train, list):
            y_train_var = y_train[1]
            y_train = y_train[0]
        else:
            y_train_var = None
        y_num, y_dimension = y_train.shape
        Sigma = self.kernel(x_train, x_train) + self.log_beta.exp().pow(-1) * torch.eye(
            x_train.size(0)) + JITTER * torch.eye(x_train.size(0))
        if y_train_var is not None:
            Sigma = Sigma + y_train_var.diag()* torch.eye(x_train.size(0))
        L = torch.cholesky(Sigma)
        #option 1 (use this if torch supports)
        Gamma,_ = torch.triangular_solve(y_train, L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return -nll
//  [python]
def negative_log_likelihood(self, x_train, y_train):
        if isinstance(y_train, list):
            y_train_var = y_train[1]
            y_train = y_train[0]
        else:
            y_train_var = None
        y_num, y_dimension = y_train.shape
        Sigma = self.kernel(x_train, x_train) + self.log_beta.exp().pow(-1) * torch.eye(
            x_train.size(0)) + JITTER * torch.eye(x_train.size(0))
        if y_train_var is not None:
            Sigma = Sigma + y_train_var.diag()* torch.eye(x_train.size(0))
        L = torch.linalg.cholesky(Sigma)
        #option 1 (use this if torch supports)
        Gamma,_ = torch.triangular_solve(y_train, L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return -nll
//  [python]
def test_multivariate_normal_properties(self):
        loc = torch.randn(5)
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
        m = MultivariateNormal(loc=loc, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0]))
        self.assertEqual(m.scale_tril, torch.cholesky(m.covariance_matrix))
//  [python]
def test_multivariate_normal_properties(self):
        loc = torch.randn(5)
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
        m = MultivariateNormal(loc=loc, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0]))
        self.assertEqual(m.scale_tril, torch.linalg.cholesky(m.covariance_matrix))
//  [python]
def test_multivariate_normal_sample(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        mean = torch.randn(3, requires_grad=True)
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.cholesky(cov).requires_grad_()

        self._check_sampler_sampler(MultivariateNormal(mean, cov),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, cov={})'.format(mean, cov),
                                    multivariate=True)
        self._check_sampler_sampler(MultivariateNormal(mean, precision_matrix=prec),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, atol={})'.format(mean, prec),
                                    multivariate=True)
        self._check_sampler_sampler(MultivariateNormal(mean, scale_tril=scale_tril),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, scale_tril={})'.format(mean, scale_tril),
                                    multivariate=True)
//  [python]
def test_multivariate_normal_sample(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        mean = torch.randn(3, requires_grad=True)
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov).requires_grad_()

        self._check_sampler_sampler(MultivariateNormal(mean, cov),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, cov={})'.format(mean, cov),
                                    multivariate=True)
        self._check_sampler_sampler(MultivariateNormal(mean, precision_matrix=prec),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, atol={})'.format(mean, prec),
                                    multivariate=True)
        self._check_sampler_sampler(MultivariateNormal(mean, scale_tril=scale_tril),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, scale_tril={})'.format(mean, scale_tril),
                                    multivariate=True)
//  [python]
def test_multivariate_normal_log_prob(self):
        mean = torch.randn(3, requires_grad=True)
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.cholesky(cov).requires_grad_()

        # check that logprob values match scipy logpdf,
        # and that covariance and scale_tril parameters are equivalent
        dist1 = MultivariateNormal(mean, cov)
        dist2 = MultivariateNormal(mean, precision_matrix=prec)
        dist3 = MultivariateNormal(mean, scale_tril=scale_tril)
        ref_dist = scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy())

        x = dist1.sample((10,))
        expected = ref_dist.logpdf(x.numpy())

        self.assertEqual(0.0, np.mean((dist1.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)
        self.assertEqual(0.0, np.mean((dist2.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)
        self.assertEqual(0.0, np.mean((dist3.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)

        # Double-check that batched versions behave the same as unbatched
        mean = torch.randn(5, 3, requires_grad=True)
        tmp = torch.randn(5, 3, 10)
        cov = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()

        dist_batched = MultivariateNormal(mean, cov)
        dist_unbatched = [MultivariateNormal(mean[i], cov[i]) for i in range(mean.size(0))]

        x = dist_batched.sample((10,))
        batched_prob = dist_batched.log_prob(x)
        unbatched_prob = torch.stack([dist_unbatched[i].log_prob(x[:, i]) for i in range(5)]).t()

        self.assertEqual(batched_prob.shape, unbatched_prob.shape)
        self.assertEqual(0.0, (batched_prob - unbatched_prob).abs().max(), atol=1e-3, rtol=0)
//  [python]
def test_multivariate_normal_log_prob(self):
        mean = torch.randn(3, requires_grad=True)
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov).requires_grad_()

        # check that logprob values match scipy logpdf,
        # and that covariance and scale_tril parameters are equivalent
        dist1 = MultivariateNormal(mean, cov)
        dist2 = MultivariateNormal(mean, precision_matrix=prec)
        dist3 = MultivariateNormal(mean, scale_tril=scale_tril)
        ref_dist = scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy())

        x = dist1.sample((10,))
        expected = ref_dist.logpdf(x.numpy())

        self.assertEqual(0.0, np.mean((dist1.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)
        self.assertEqual(0.0, np.mean((dist2.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)
        self.assertEqual(0.0, np.mean((dist3.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)

        # Double-check that batched versions behave the same as unbatched
        mean = torch.randn(5, 3, requires_grad=True)
        tmp = torch.randn(5, 3, 10)
        cov = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()

        dist_batched = MultivariateNormal(mean, cov)
        dist_unbatched = [MultivariateNormal(mean[i], cov[i]) for i in range(mean.size(0))]

        x = dist_batched.sample((10,))
        batched_prob = dist_batched.log_prob(x)
        unbatched_prob = torch.stack([dist_unbatched[i].log_prob(x[:, i]) for i in range(5)]).t()

        self.assertEqual(batched_prob.shape, unbatched_prob.shape)
        self.assertEqual(0.0, (batched_prob - unbatched_prob).abs().max(), atol=1e-3, rtol=0)
//  [python]
def test_wishart_sample(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        df = (torch.rand([], requires_grad=True) + 1) * 3
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.cholesky(cov).requires_grad_()

        self._check_sampler_sampler(Wishart(df, cov),
                                    scipy.stats.wishart(df.item(), cov.detach().numpy()),
                                    'Wishart(df={}, covariance_matrix={})'.format(df, cov),
                                    multivariate=True)
        self._check_sampler_sampler(Wishart(df, precision_matrix=prec),
                                    scipy.stats.wishart(df.item(), cov.detach().numpy()),
                                    'Wishart(df={}, precision_matrix={})'.format(df, prec),
                                    multivariate=True)
        self._check_sampler_sampler(Wishart(df, scale_tril=scale_tril),
                                    scipy.stats.wishart(df.item(), cov.detach().numpy()),
                                    'Wishart(df={}, scale_tril={})'.format(df, scale_tril),
                                    multivariate=True)
//  [python]
def test_wishart_sample(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        df = (torch.rand([], requires_grad=True) + 1) * 3
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov).requires_grad_()

        self._check_sampler_sampler(Wishart(df, cov),
                                    scipy.stats.wishart(df.item(), cov.detach().numpy()),
                                    'Wishart(df={}, covariance_matrix={})'.format(df, cov),
                                    multivariate=True)
        self._check_sampler_sampler(Wishart(df, precision_matrix=prec),
                                    scipy.stats.wishart(df.item(), cov.detach().numpy()),
                                    'Wishart(df={}, precision_matrix={})'.format(df, prec),
                                    multivariate=True)
        self._check_sampler_sampler(Wishart(df, scale_tril=scale_tril),
                                    scipy.stats.wishart(df.item(), cov.detach().numpy()),
                                    'Wishart(df={}, scale_tril={})'.format(df, scale_tril),
                                    multivariate=True)
//  [python]
def test_wishart_properties(self):
        df = (torch.rand([]) + 1) * 5
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
        m = Wishart(df=df, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0]))
        self.assertEqual(m.scale_tril, torch.cholesky(m.covariance_matrix))
//  [python]
def test_wishart_properties(self):
        df = (torch.rand([]) + 1) * 5
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
        m = Wishart(df=df, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0]))
        self.assertEqual(m.scale_tril, torch.linalg.cholesky(m.covariance_matrix))
//  [python]
def decompose_C(self) -> None:
        """Perform the decomposition C = AA^T using a cholesky decomposition
        Note that traditionally CMA-ES uses the eigendecomposition C = BDDB^-1. In our case,
        we keep track of zs, ys and xs when sampling, so we never need C^-1/2.
        Therefore, a cholesky decomposition is all that is necessary. This generally requires
        O(d^3/3) operations, rather than the more costly O(d^3) operations associated with the eigendecomposition.
        """
        if self.separable:
            self.A = self.C.pow(0.5)
        else:
            self.A = torch.cholesky(self.C)
//  [python]
def decompose_C(self) -> None:
        """Perform the decomposition C = AA^T using a cholesky decomposition
        Note that traditionally CMA-ES uses the eigendecomposition C = BDDB^-1. In our case,
        we keep track of zs, ys and xs when sampling, so we never need C^-1/2.
        Therefore, a cholesky decomposition is all that is necessary. This generally requires
        O(d^3/3) operations, rather than the more costly O(d^3) operations associated with the eigendecomposition.
        """
        if self.separable:
            self.A = self.C.pow(0.5)
        else:
            self.A = torch.linalg.cholesky(self.C)
//  [python]
def test_wishart_log_prob(self):
        df = (torch.rand([], requires_grad=True) + 1) * 10
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.cholesky(cov).requires_grad_()

        # check that logprob values match scipy logpdf,
        # and that covariance and scale_tril parameters are equivalent
        dist1 = Wishart(df, cov)
        dist2 = Wishart(df, precision_matrix=prec)
        dist3 = Wishart(df, scale_tril=scale_tril)
        ref_dist = scipy.stats.wishart(df.item(), cov.detach().numpy())

        x = dist1.sample((10,))
        expected = ref_dist.logpdf(x.transpose(0, 2).numpy())

        self.assertEqual(0.0, np.mean((dist1.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)
        self.assertEqual(0.0, np.mean((dist2.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)
        self.assertEqual(0.0, np.mean((dist3.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)

        # Double-check that batched versions behave the same as unbatched
        df = (torch.rand(5, requires_grad=True) + 1) * 3
        tmp = torch.randn(5, 3, 10)
        cov = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()

        dist_batched = Wishart(df, cov)
        dist_unbatched = [Wishart(df[i], cov[i]) for i in range(df.size(0))]

        x = dist_batched.sample((10,))
        batched_prob = dist_batched.log_prob(x)
        unbatched_prob = torch.stack([dist_unbatched[i].log_prob(x[:, i]) for i in range(5)]).t()

        self.assertEqual(batched_prob.shape, unbatched_prob.shape)
        self.assertEqual(0.0, (batched_prob - unbatched_prob).abs().max(), atol=1e-3, rtol=0)
//  [python]
def test_wishart_log_prob(self):
        df = (torch.rand([], requires_grad=True) + 1) * 10
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov).requires_grad_()

        # check that logprob values match scipy logpdf,
        # and that covariance and scale_tril parameters are equivalent
        dist1 = Wishart(df, cov)
        dist2 = Wishart(df, precision_matrix=prec)
        dist3 = Wishart(df, scale_tril=scale_tril)
        ref_dist = scipy.stats.wishart(df.item(), cov.detach().numpy())

        x = dist1.sample((10,))
        expected = ref_dist.logpdf(x.transpose(0, 2).numpy())

        self.assertEqual(0.0, np.mean((dist1.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)
        self.assertEqual(0.0, np.mean((dist2.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)
        self.assertEqual(0.0, np.mean((dist3.log_prob(x).detach().numpy() - expected)**2), atol=1e-3, rtol=0)

        # Double-check that batched versions behave the same as unbatched
        df = (torch.rand(5, requires_grad=True) + 1) * 3
        tmp = torch.randn(5, 3, 10)
        cov = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()

        dist_batched = Wishart(df, cov)
        dist_unbatched = [Wishart(df[i], cov[i]) for i in range(df.size(0))]

        x = dist_batched.sample((10,))
        batched_prob = dist_batched.log_prob(x)
        unbatched_prob = torch.stack([dist_unbatched[i].log_prob(x[:, i]) for i in range(5)]).t()

        self.assertEqual(batched_prob.shape, unbatched_prob.shape)
        self.assertEqual(0.0, (batched_prob - unbatched_prob).abs().max(), atol=1e-3, rtol=0)
//  [python]
def SigmatoH(Sigmas):
    """
    Transform covariance matrices into logCholesky factors

    Inputs:
        Sigmas : (K,D,D,B) array, B observations of the K cluster covariances

    Outpus:
        Hs : (K,D,D,B) array, B observations of the K cluster logCholesky factors
    """

    idx=np.diag_indices(Sigmas.shape[1])
    Ls=torch.cholesky(torch.moveaxis(Sigmas,3,1))
    Hs=torch.clone(Ls)
    Hs[:,:,idx[0],idx[1]]=torch.log(Ls[:,:,idx[0],idx[1]])
    return torch.moveaxis(Hs,1,3)
//  [python]
def SigmatoH(Sigmas):
    """
    Transform covariance matrices into logCholesky factors

    Inputs:
        Sigmas : (K,D,D,B) array, B observations of the K cluster covariances

    Outpus:
        Hs : (K,D,D,B) array, B observations of the K cluster logCholesky factors
    """

    idx=np.diag_indices(Sigmas.shape[1])
    Ls=torch.linalg.cholesky(torch.moveaxis(Sigmas,3,1))
    Hs=torch.clone(Ls)
    Hs[:,:,idx[0],idx[1]]=torch.log(Ls[:,:,idx[0],idx[1]])
    return torch.moveaxis(Hs,1,3)