jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def expected_sufficient_statistics(self):
        '''
        stats = (
            S * mu,
            S,
            tr(S * mu * mu^T),
            ln |S|
        )

        E[stats] = (
            v * W * m,
            v * W,
            (D/k) + tr(v * W * m * m^T),
            ( \sum_i psi(.5 * (v + 1 - i)) ) + D * ln 2 + ln |W|
        )
        '''
        mean, scale, = self.params.mean, self.params.scale
        scale_matrix, dof = self.params.scale_matrix, self.params.dof
        dim = mean.shape[-1]
        mean_size = mean.shape if len(mean.shape) > 0 else 1
        if len(mean_size) == 1:
            mean, scale = mean.view(1, -1), scale.view(1, 1)
            scale_matrix, dof = scale_matrix.view(1, dim, dim), dof.view(1, 1)

        idxs = torch.arange(1, dim + 1, dtype=mean.dtype, device=mean.device)
        L = torch.linalg.cholesky(scale_matrix, upper=False)
        logdet = 2 * torch.log(L[:, range(dim), range(dim)]).sum(dim=-1,
                                                                 keepdim=True)
        mean_quad = mean[:, :, None] * mean[:, None, :]
        exp_prec = dof[:, :, None] * scale_matrix
        retval = torch.cat([
            torch.matmul(exp_prec, mean[:, :, None]).reshape(-1, dim),
            exp_prec.reshape(-1, dim ** 2),
            (dim / scale) + _batch_trace(exp_prec, mean_quad, keepdim=True),
            (torch.digamma(.5 * (dof + 1 - idxs)).sum(dim=-1, keepdim=True) \
                + dim * math.log(2) + logdet)
        ], dim=-1)

        if len(mean_size) == 1:
            return retval.view(-1)
        return retval
//  [python]
def precompute_laplacian(self):
		"""Precompute edge weights and Laplacian-Beltrami operator"""

		N = self.num_verts_per_mesh()
		p = self.verts_padded()  # undeformed mesh

		w_padded = self.get_cot_padded()

		L_padded = torch.zeros_like(w_padded)
		L_inv_padded = torch.zeros_like(w_padded)
		for n in range(len(self)):
			w = w_padded[n]
			# diagonal L_ii = sum(w_i). non-diagonal L_ij = -w_ij
			L_padded[n] = torch.diag(torch.sum(w, dim=0)) - w
			L_inv_padded[n] = torch.inverse(torch.cholesky(L_padded[n]))

		self.precomputed_params["L_padded"] = L_padded
		self.precomputed_params["w_padded"] = w_padded
		self.precomputed_params["L_inv_padded"] = L_inv_padded
//  [python]
def precompute_laplacian(self):
		"""Precompute edge weights and Laplacian-Beltrami operator"""

		N = self.num_verts_per_mesh()
		p = self.verts_padded()  # undeformed mesh

		w_padded = self.get_cot_padded()

		L_padded = torch.zeros_like(w_padded)
		L_inv_padded = torch.zeros_like(w_padded)
		for n in range(len(self)):
			w = w_padded[n]
			# diagonal L_ii = sum(w_i). non-diagonal L_ij = -w_ij
			L_padded[n] = torch.diag(torch.sum(w, dim=0)) - w
			L_inv_padded[n] = torch.linalg.cholesky(L_padded[n]).inverse()

		self.precomputed_params["L_padded"] = L_padded
		self.precomputed_params["w_padded"] = w_padded
		self.precomputed_params["L_inv_padded"] = L_inv_padded
//  [python]
def __init__(self, horizon=None, d_action=None, seed=0, mean=None, scale_tril=None,
                 covariance_matrix=None,
                 tensor_args={'device':"cpu", 'dtype':torch.float32}, fixed_samples=False,
                 filter_coeffs=None, **kwargs):
        self.tensor_args = tensor_args
        self.horizon = horizon
        self.d_action = d_action
        self.seed_val = seed
        self.Z = torch.zeros(self.horizon * self.d_action, **tensor_args)
        self.scale_tril = None
        if(scale_tril is None and covariance_matrix is not None):
            self.scale_tril = torch.cholesky(covariance_matrix)
        self.covariance_matrix = covariance_matrix
        self.fixed_samples = fixed_samples
        self.samples = None
        self.sample_shape = 0
        self.filter_coeffs = filter_coeffs
        self.ndims = horizon * d_action
        self.stomp_matrix, self.stomp_scale_tril = get_stomp_cov(self.horizon, self.d_action, tensor_args=self.tensor_args)
//  [python]
def __init__(self, horizon=None, d_action=None, seed=0, mean=None, scale_tril=None,
                 covariance_matrix=None,
                 tensor_args={'device':"cpu", 'dtype':torch.float32}, fixed_samples=False,
                 filter_coeffs=None, **kwargs):
        self.tensor_args = tensor_args
        self.horizon = horizon
        self.d_action = d_action
        self.seed_val = seed
        self.Z = torch.zeros(self.horizon * self.d_action, **tensor_args)
        self.scale_tril = None
        if(scale_tril is None and covariance_matrix is not None):
            self.scale_tril = torch.linalg.cholesky(covariance_matrix)
        self.covariance_matrix = covariance_matrix
        self.fixed_samples = fixed_samples
        self.samples = None
        self.sample_shape = 0
        self.filter_coeffs = filter_coeffs
        self.ndims = horizon * d_action
        self.stomp_matrix, self.stomp_scale_tril = get_stomp_cov(self.horizon, self.d_action, tensor_args=self.tensor_args)
//  [python]
def posterior_params(self, x):
        log_weights = torch.log(self.module.soft_max(self.module.soft_weights))
        T = self.module.covars[None, :, :, :] + x[1][:, None, :, :]
        
        w = log_weights + dist.MultivariateNormal(
            loc=self.module.means, covariance_matrix=T
        ).log_prob(x[0][:, None, :])
        p_weights = w - torch.logsumexp(w, axis=1)[:, None]
        
        L_t = torch.cholesky(T)
        T_inv = torch.cholesky_solve(torch.eye(self.d, device=self.device), L_t)
        
        diff = x[0][:, None, :] - self.module.means
        T_prod = torch.matmul(T_inv, diff[:, :, :, None])
        p_means = self.module.means + torch.matmul(
            self.module.covars,
            T_prod
        ).squeeze()
        
        p_covars = self.module.covars - torch.matmul(
            self.module.covars,
            torch.matmul(T_inv, self.module.covars)
        )
        return p_weights, p_means, p_covars
//  [python]
def posterior_params(self, x):
        log_weights = torch.log(self.module.soft_max(self.module.soft_weights))
        T = self.module.covars[None, :, :, :] + x[1][:, None, :, :]
        
        w = log_weights + dist.MultivariateNormal(
            loc=self.module.means, covariance_matrix=T
        ).log_prob(x[0][:, None, :])
        p_weights = w - torch.logsumexp(w, axis=1)[:, None]
        
        L_t = torch.linalg.cholesky(T)
        T_inv = torch.cholesky_solve(torch.eye(self.d, device=self.device), L_t)
        
        diff = x[0][:, None, :] - self.module.means
        T_prod = torch.matmul(T_inv, diff[:, :, :, None])
        p_means = self.module.means + torch.matmul(
            self.module.covars,
            T_prod
        ).squeeze()
        
        p_covars = self.module.covars - torch.matmul(
            self.module.covars,
            torch.matmul(T_inv, self.module.covars)
        )
        return p_weights, p_means, p_covars
//  [python]
def __init__(self, horizon=0, d_action=0, n_knots=0, degree=3, seed=0, tensor_args={'device':"cpu", 'dtype':torch.float32}, sample_method='halton',
                 covariance_matrix = None, **kwargs):
        self.ndims = n_knots * d_action
        self.n_knots = n_knots
        self.horizon = horizon
        self.d_action = d_action
        self.tensor_args = tensor_args
        self.seed_val = seed
        self.degree = degree
        self.sample_method = sample_method
        self.Z = torch.zeros(self.ndims, **tensor_args)
        if(covariance_matrix is None):
            self.cov_matrix = torch.eye(self.ndims, **tensor_args)
        self.scale_tril = torch.cholesky(self.cov_matrix.to(dtype=torch.float32)).to(**tensor_args)
        self.mvn = MultivariateNormal(loc=self.Z, scale_tril=self.scale_tril, )
//  [python]
def __init__(self, horizon=0, d_action=0, n_knots=0, degree=3, seed=0, tensor_args={'device':"cpu", 'dtype':torch.float32}, sample_method='halton',
                 covariance_matrix = None, **kwargs):
        self.ndims = n_knots * d_action
        self.n_knots = n_knots
        self.horizon = horizon
        self.d_action = d_action
        self.tensor_args = tensor_args
        self.seed_val = seed
        self.degree = degree
        self.sample_method = sample_method
        self.Z = torch.zeros(self.ndims, **tensor_args)
        if(covariance_matrix is None):
            self.cov_matrix = torch.eye(self.ndims, **tensor_args)
        self.scale_tril = torch.linalg.cholesky(self.cov_matrix.to(dtype=torch.float32)).to(**tensor_args)
        self.mvn = MultivariateNormal(loc=self.Z, scale_tril=self.scale_tril, )
//  [python]
def rsample(self, sample_shape=torch.Size([]), lowrank=False):
        # eps = torch.randn(*self.loc.shape, 1)
        # cov_svd = self.covariance_matrix.svd()
        # s_sqrt = torch.diag_embed(cov_svd.S.sqrt())
        # a = torch.bmm(cov_svd.U, torch.bmm(s_sqrt, cov_svd.V))
        # out = self.loc + torch.bmm(a, eps)[..., 0]
        cov_svd = self.cov_matrix.svd()
        s_sqrt = torch.diag_embed(cov_svd.S.sqrt())
        a = torch.bmm(cov_svd.U, torch.bmm(s_sqrt, cov_svd.V))
        pcov_chol = torch.cholesky(self.pcov_matrix, upper=True)

        if sample_shape == torch.Size([]) or sample_shape is None:
            eps = torch.randn(self.batch_shape, self.cov_matrix.shape[1], self.pcov_matrix.shape[1])
            out = self.loc + torch.bmm(a, torch.bmm(eps, pcov_chol))
        else:
            samples = []
            for i in range(sample_shape):
                eps = torch.randn(self.batch_shape, self.cov_matrix.shape[1], self.pcov_matrix.shape[1])
                samples.append(self.loc + torch.bmm(a, torch.bmm(eps, pcov_chol)))
            out = torch.stack(samples, dim=0)
        return out
//  [python]
def rsample(self, sample_shape=torch.Size([]), lowrank=False):
        cov_svd = self.cov_matrix.svd()
        s_sqrt = torch.diag_embed(cov_svd.S.sqrt())
        a = torch.bmm(cov_svd.U, torch.bmm(s_sqrt, cov_svd.V))
        pcov_chol = torch.linalg.cholesky(self.pcov_matrix, upper=True)

        if sample_shape == torch.Size([]) or sample_shape is None:
            eps = torch.randn(self.batch_shape, self.cov_matrix.shape[1], self.pcov_matrix.shape[1])
            out = self.loc + torch.bmm(a, torch.bmm(eps, pcov_chol))
        else:
            samples = []
            for i in range(sample_shape):
                eps = torch.randn(self.batch_shape, self.cov_matrix.shape[1], self.pcov_matrix.shape[1])
                samples.append(self.loc + torch.bmm(a, torch.bmm(eps, pcov_chol)))
            out = torch.stack(samples, dim=0)
        return out
//  [python]
def forward(self, x):
        #x = self.relu(x)
        #x = x @ x.mT + 1e-3
        x = torch.cholesky(x)
        return x
//  [python]
def forward(self, x):
        x = torch.linalg.cholesky(x)
        return x
//  [python]
def log_det_by_cholesky(matrix):
	"""
	Args:
		matrix: matrix must be a positive define matrix.
				shape [N, C, D, D].
	Ref:
		https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
	"""
	# This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
	# where C is the cholesky decomposition of A.
    
	chol = torch.cholesky(matrix)
	return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def log_det_by_cholesky(matrix):
	"""
	Args:
		matrix: matrix must be a positive define matrix.
				shape [N, C, D, D].
	Ref:
		https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
	"""
	# This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
	# where C is the cholesky decomposition of A.
    
	chol = torch.linalg.cholesky(matrix)
	return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def cross_entropy(p, q):
    d = p.mean.size(-1)
    ret = 0.5 * d * math.log(2 * math.pi)
    ret += q.variance.log().sum(-1)
    q_cov_tril = torch.cholesky(q.covariance_matrix)
    p_cov_tril = torch.cholesky(p.covariance_matrix)
    q_inv_p = torch.triangular_solve(p_cov_tril, q_cov_tril)[0]
    ret += 0.5 * q_inv_p.pow(2).sum(-1).sum(-1)
    Kinv_m = cholesky_solve((p.mean - q.mean).unsqueeze(-1), q_cov_tril)
    ret += 0.5 * torch.sum((p.mean - q.mean) * Kinv_m.squeeze(-1), -1)
    return ret
//  [python]
def cross_entropy(p, q):
    d = p.mean.size(-1)
    ret = 0.5 * d * math.log(2 * math.pi)
    ret += q.variance.log().sum(-1)
    q_cov_tril = torch.linalg.cholesky(q.covariance_matrix)
    p_cov_tril = torch.linalg.cholesky(p.covariance_matrix)
    q_inv_p = torch.triangular_solve(p_cov_tril, q_cov_tril)[0]
    ret += 0.5 * q_inv_p.pow(2).sum(-1).sum(-1)
    Kinv_m = cholesky_solve((p.mean - q.mean).unsqueeze(-1), q_cov_tril)
    ret += 0.5 * torch.sum((p.mean - q.mean) * Kinv_m.squeeze(-1), -1)
    return ret
//  [python]
def laplace_refine(model, X, y, likelihood, prior_prec, n_epochs=1000, lr=1e-3):
    """Laplace in LinNN (GLM) giving posterior mean, and Sigma_chol(d)"""
    J, f = Jacobians(model, X)
    if len(J.shape) == 3:
        J = J.transpose(1, 2)
    theta_star = parameters_to_vector(model.parameters()).detach()
    f_offset = f - (J @ theta_star)
    mu = theta_star.clone().requires_grad_(True)
    opt = Adam([mu], lr=lr)

    losses = list()
    for i in range(n_epochs):
        opt.zero_grad()
        ll_loss = - likelihood.log_likelihood(y, f_offset + J @ mu)
        weight_loss = 0.5 * mu @ mu * prior_prec
        loss = ll_loss + weight_loss
        loss.backward()
        losses.append(loss.detach().cpu().item())
        opt.step()

    Lams = likelihood.Hessian(f_offset + J @ mu.detach())
    if len(J.shape) == 3:
        Hessian = einsum('mkp,mkl,mlq->pq', J, Lams, J, backend='torch')
    else:
        Hessian = (J.T * Lams.reshape(1, -1)) @ J
    Precision = Hessian + prior_prec * torch.eye(len(Hessian), device=Hessian.device)
    Chol = torch.cholesky(Precision)
    Sigma = torch.cholesky_inverse(Chol, upper=False)
    Sigma_chol = torch.cholesky(Sigma, upper=False)
    Sigma_chold = torch.diag(torch.sqrt(1 / torch.diag(Precision)))
    return mu.detach(), Sigma_chol, Sigma_chold, losses
//  [python]
def laplace_refine(model, X, y, likelihood, prior_prec, n_epochs=1000, lr=1e-3):
    """Laplace in LinNN (GLM) giving posterior mean, and Sigma_chol(d)"""
    J, f = Jacobians(model, X)
    if len(J.shape) == 3:
        J = J.transpose(1, 2)
    theta_star = parameters_to_vector(model.parameters()).detach()
    f_offset = f - (J @ theta_star)
    mu = theta_star.clone().requires_grad_(True)
    opt = Adam([mu], lr=lr)

    losses = list()
    for i in range(n_epochs):
        opt.zero_grad()
        ll_loss = - likelihood.log_likelihood(y, f_offset + J @ mu)
        weight_loss = 0.5 * mu @ mu * prior_prec
        loss = ll_loss + weight_loss
        loss.backward()
        losses.append(loss.detach().cpu().item())
        opt.step()

    Lams = likelihood.Hessian(f_offset + J @ mu.detach())
    if len(J.shape) == 3:
        Hessian = einsum('mkp,mkl,mlq->pq', J, Lams, J, backend='torch')
    else:
        Hessian = (J.T * Lams.reshape(1, -1)) @ J
    Precision = Hessian + prior_prec * torch.eye(len(Hessian), device=Hessian.device)
    Chol = torch.linalg.cholesky(Precision)
    Sigma = torch.linalg.solve_triangular(Chol, torch.eye(Chol.shape[0], device=Chol.device), upper=False)
    Sigma_chol = torch.linalg.cholesky(Sigma, upper=False)
    Sigma_chold = torch.diag(torch.sqrt(1 / torch.diag(Precision)))
    return mu.detach(), Sigma_chol, Sigma_chold, losses
//  [python]
def compute_log_likelihood(self, X=None, Y=None):
        """
        This method computes the variational lower bound on the likelihood,
        which is:

            E_{q(F)} [ \log p(Y|F) ] - KL[ q(F) || p(F)]

        with

            q(\\mathbf f) = N(\\mathbf f \\,|\\, \\boldsymbol \\mu, \\boldsymbol \\Sigma)

        """

        assert X is None and Y is None, "{} does not support minibatch mode".format(str(type(self)))
        # Get prior KL.
        KL = kullback_leiblers.gauss_kl_white(self.q_mu.get(), self.q_sqrt.get())

        # Get conditionals
        K = self.kern.K(self.X) + torch.eye(self.num_data, dtype=self.X.dtype, device=self.X.device) * self.jitter_level
        L = torch.cholesky(K, upper=False)

        fmean = torch.matmul(L, self.q_mu) + self.mean_function(self.X)  # NN,ND->ND

        q_sqrt_dnn = kullback_leiblers.batch_tril(self.q_sqrt.get().permute(2, 0, 1))  # D x N x N

        LTA = torch.matmul(L.unsqueeze(0), q_sqrt_dnn)  # D x N x N
        fvar = (LTA**2).sum(2)

        fvar = fvar.t()

        # Get variational expectations.
        var_exp = self.likelihood.variational_expectations(fmean, fvar, self.Y)

        return var_exp.sum() - KL
//  [python]
def compute_log_likelihood(self, X=None, Y=None):
        """
        This method computes the variational lower bound on the likelihood,
        which is:

            E_{q(F)} [ \log p(Y|F) ] - KL[ q(F) || p(F)]

        with

            q(\\mathbf f) = N(\\mathbf f \\,|\\, \\boldsymbol \\mu, \\boldsymbol \\Sigma)

        """

        assert X is None and Y is None, "{} does not support minibatch mode".format(str(type(self)))
        # Get prior KL.
        KL = kullback_leiblers.gauss_kl_white(self.q_mu.get(), self.q_sqrt.get())

        # Get conditionals
        K = self.kern.K(self.X) + torch.eye(self.num_data, dtype=self.X.dtype, device=self.X.device) * self.jitter_level
        L = torch.linalg.cholesky(K, upper=False)

        fmean = torch.matmul(L, self.q_mu) + self.mean_function(self.X)  # NN,ND->ND

        q_sqrt_dnn = kullback_leiblers.batch_tril(self.q_sqrt.get().permute(2, 0, 1))  # D x N x N

        LTA = torch.matmul(L.unsqueeze(0), q_sqrt_dnn)  # D x N x N
        fvar = (LTA**2).sum(2)

        fvar = fvar.t()

        # Get variational expectations.
        var_exp = self.likelihood.variational_expectations(fmean, fvar, self.Y)

        return var_exp.sum() - KL
//  [python]
def forward(self, x):
		self._check_input_dim(x)
		self._check_group_size()

		m = x.mean(0).view(self.num_features, -1).mean(-1).view(1, -1, 1, 1)
		if not self.training and self.track_running_stats: # for inference
			m = self.running_mean
		xn = x - m

		T = xn.permute(1,0,2,3).contiguous().view(self.num_groups, self.group_size,-1)
		f_cov = torch.bmm(T, T.permute(0,2,1)) / T.shape[-1]
		f_cov_shrinked = (1-self.eps) * f_cov + self.eps * torch.eye(self.group_size, out=torch.cuda.FloatTensor() if torch.cuda.is_available() else torch.FloatTensor()).repeat(self.num_groups, 1, 1)

		if not self.training and self.track_running_stats: # for inference
			f_cov_shrinked = (1-self.eps) * self.running_variance + self.eps * torch.eye(self.group_size, out=torch.cuda.FloatTensor() if torch.cuda.is_available() else torch.FloatTensor()).repeat(self.num_groups, 1, 1)

		inv_sqrt = torch.inverse(torch.cholesky(f_cov_shrinked)).contiguous().view(self.num_features, self.group_size, 1, 1)

		decorrelated = conv2d(xn, inv_sqrt, groups = self.num_groups)

		if self.training and self.track_running_stats:
			self.running_mean = torch.add(self.momentum * m.detach(), (1 - self.momentum) * self.running_mean, out=self.running_mean) 
			self.running_variance = torch.add(self.momentum * f_cov.detach(), (1 - self.momentum) * self.running_variance, out=self.running_variance)
			
		return decorrelated
//  [python]
def forward(self, x):
		self._check_input_dim(x)
		self._check_group_size()

		m = x.mean(0).view(self.num_features, -1).mean(-1).view(1, -1, 1, 1)
		if not self.training and self.track_running_stats: # for inference
			m = self.running_mean
		xn = x - m

		T = xn.permute(1,0,2,3).contiguous().view(self.num_groups, self.group_size,-1)
		f_cov = torch.bmm(T, T.permute(0,2,1)) / T.shape[-1]
		f_cov_shrinked = (1-self.eps) * f_cov + self.eps * torch.eye(self.group_size, out=torch.cuda.FloatTensor() if torch.cuda.is_available() else torch.FloatTensor()).repeat(self.num_groups, 1, 1)

		if not self.training and self.track_running_stats: # for inference
			f_cov_shrinked = (1-self.eps) * self.running_variance + self.eps * torch.eye(self.group_size, out=torch.cuda.FloatTensor() if torch.cuda.is_available() else torch.FloatTensor()).repeat(self.num_groups, 1, 1)

		inv_sqrt = torch.inverse(torch.linalg.cholesky(f_cov_shrinked)).contiguous().view(self.num_features, self.group_size, 1, 1)

		decorrelated = conv2d(xn, inv_sqrt, groups = self.num_groups)

		if self.training and self.track_running_stats:
			self.running_mean = torch.add(self.momentum * m.detach(), (1 - self.momentum) * self.running_mean, out=self.running_mean) 
			self.running_variance = torch.add(self.momentum * f_cov.detach(), (1 - self.momentum) * self.running_variance, out=self.running_variance)
			
		return decorrelated
//  [python]
def compute_upper_bound(self):
        num_inducing = self.Z.size(0)
        num_data = self.Y.size(0)

        Kdiag = self.kern.Kdiag(self.X)
        jitter = torch.eye(num_inducing, dtype=self.Z.dtype, device=self.Z.device) * self.jitter_level
        Kuu = self.kern.K(self.Z.get()) + jitter
        Kuf = self.kern.K(self.Z.get(), self.X)

        L = torch.cholesky(Kuu, upper=False)
        LB = torch.cholesky(Kuu + self.likelihood.variance.get() ** -1.0 * torch.matmul(Kuf, Kuf.t()), upper=False)

        LinvKuf, _ = torch.solve(Kuf, L) # could use triangular solve
        # Using the Trace bound, from Titsias' presentation
        c = Kdiag.sum() - (LinvKuf ** 2.0).sum()
        # Kff = self.kern.K(self.X)
        # Qff = tf.matmul(Kuf, LinvKuf, transpose_a=True)

        # Alternative bound on max eigenval:
        # c = tf.reduce_max(tf.reduce_sum(tf.abs(Kff - Qff), 0))
        corrected_noise = self.likelihood.variance.get() + c

        const = -0.5 * num_data * torch.log(2 * float(numpy.pi) * self.likelihood.variance.get())
        logdet = torch.diag(L).log().sum() - torch.diag(LB).log().sum()

        LC = torch.cholesky(Kuu + corrected_noise ** -1.0 * torch.matmul(Kuf, Kuf.t()), upper=False)
        v, _ = torch.solve(corrected_noise ** -1.0 * torch.matmul(Kuf, self.Y), LC)
        quad = -0.5 * corrected_noise ** -1.0 * (self.Y ** 2.0).sum() + 0.5 * (v ** 2.0).sum()

        return const + logdet + quad
//  [python]
def compute_upper_bound(self):
        num_inducing = self.Z.size(0)
        num_data = self.Y.size(0)

        Kdiag = self.kern.Kdiag(self.X)
        jitter = torch.eye(num_inducing, dtype=self.Z.dtype, device=self.Z.device) * self.jitter_level
        Kuu = self.kern.K(self.Z.get()) + jitter
        Kuf = self.kern.K(self.Z.get(), self.X)

        L = torch.linalg.cholesky(Kuu, upper=False)
        LB = torch.linalg.cholesky(Kuu + self.likelihood.variance.get() ** -1.0 * torch.matmul(Kuf, Kuf.t()), upper=False)

        LinvKuf, _ = torch.solve(Kuf, L) # could use triangular solve
        # Using the Trace bound, from Titsias' presentation
        c = Kdiag.sum() - (LinvKuf ** 2.0).sum()
        # Kff = self.kern.K(self.X)
        # Qff = tf.matmul(Kuf, LinvKuf, transpose_a=True)

        # Alternative bound on max eigenval:
        # c = tf.reduce_max(tf.reduce_sum(tf.abs(Kff - Qff), 0))
        corrected_noise = self.likelihood.variance.get() + c

        const = -0.5 * num_data * torch.log(2 * float(numpy.pi) * self.likelihood.variance.get())
        logdet = torch.diag(L).log().sum() - torch.diag(LB).log().sum()

        LC = torch.linalg.cholesky(Kuu + corrected_noise ** -1.0 * torch.matmul(Kuf, Kuf.t()), upper=False)
        v, _ = torch.solve(corrected_noise ** -1.0 * torch.matmul(Kuf, self.Y), LC)
        quad = -0.5 * corrected_noise ** -1.0 * (self.Y ** 2.0).sum() + 0.5 * (v ** 2.0).sum()

        return const + logdet + quad
//  [python]
def _common_terms(self):
        num_inducing = self.Z.size(0)
        err = self.Y - self.mean_function(self.X)  # size N x R
        Kdiag = self.kern.Kdiag(self.X)
        Kuf = self.kern.K(self.Z.get(), self.X)
        jitter = torch.eye(num_inducing, dtype=self.Z.dtype, device=self.Z.device) * self.jitter_level
        Kuu = self.kern.K(self.Z.get()) + jitter

        Luu = torch.cholesky(Kuu, upper=False)  # => Luu Luu^T = Kuu
        V, _ = torch.solve(Kuf, Luu)  # => V^T V = Qff = Kuf^T Kuu^-1 Kuf

        diagQff = (V**2).sum(0)
        nu = Kdiag - diagQff + self.likelihood.variance.get()

        B = torch.eye(num_inducing, out=V.data.new()) + torch.matmul(V / nu, V.t())
        L = torch.cholesky(B, upper=False)
        beta = err / nu.unsqueeze(1)  # size N x R
        alpha = torch.matmul(V, beta)  # size N x R

        gamma, _ = torch.solve(alpha, L)  # size N x R

        return err, nu, Luu, L, alpha, beta, gamma
//  [python]
def _common_terms(self):
        num_inducing = self.Z.size(0)
        err = self.Y - self.mean_function(self.X)  # size N x R
        Kdiag = self.kern.Kdiag(self.X)
        Kuf = self.kern.K(self.Z.get(), self.X)
        jitter = torch.eye(num_inducing, dtype=self.Z.dtype, device=self.Z.device) * self.jitter_level
        Kuu = self.kern.K(self.Z.get()) + jitter

        Luu = torch.linalg.cholesky(Kuu, upper=False)  # => Luu Luu^T = Kuu
        V, _ = torch.solve(Kuf, Luu)  # => V^T V = Qff = Kuf^T Kuu^-1 Kuf

        diagQff = (V**2).sum(0)
        nu = Kdiag - diagQff + self.likelihood.variance.get()

        B = torch.eye(num_inducing, out=V.data.new()) + torch.matmul(V / nu, V.t())
        L = torch.linalg.cholesky(B, upper=False)
        beta = err / nu.unsqueeze(1)  # size N x R
        alpha = torch.matmul(V, beta)  # size N x R

        gamma, _ = torch.solve(alpha, L)  # size N x R

        return err, nu, Luu, L, alpha, beta, gamma
//  [python]
def predict_f(self, Xnew, full_cov=False):
        """
        Compute the mean and variance of the latent function at some new points
        Xnew. For a derivation of the terms in here, see the associated SGPR
        notebook.
        """
        num_inducing = self.Z.size(0)
        err = self.Y - self.mean_function(self.X)
        Kuf = self.kern.K(self.Z.get(), self.X)
        jitter = torch.eye(num_inducing, dtype=self.Z.dtype, device=self.Z.device) * self.jitter_level
        Kuu = self.kern.K(self.Z.get()) + jitter
        Kus = self.kern.K(self.Z.get(), Xnew)
        sigma = self.likelihood.variance.get()**0.5
        L = torch.cholesky(Kuu, upper=False)
        A = torch.solve(Kuf, L)[0] / sigma # could use triangular solve here and below
        B = torch.matmul(A,A.t()) + torch.eye(num_inducing, dtype=A.dtype, device=A.device)
        LB = torch.cholesky(B, upper=False)
        Aerr = torch.matmul(A, err)
        c = torch.solve(Aerr, LB)[0] / sigma
        tmp1,_ = torch.solve(Kus, L)
        tmp2,_ = torch.solve(tmp1,LB)
        mean = torch.matmul(tmp2.t(), c)
        if full_cov:
            var = self.kern.K(Xnew) + torch.matmul(tmp2.t(), tmp2) - torch.matmul(tmp1.t(), tmp1)
            var = var.unsqueeze(2).expand(-1, -1, self.Y.size(1))
        else:
            var = self.kern.Kdiag(Xnew) + (tmp2**2).sum(0) - (tmp1**2).sum(0)
            var = var.unsqueeze(1).expand(-1, self.Y.size(1))
        return mean + self.mean_function(Xnew), var
//  [python]
def predict_f(self, Xnew, full_cov=False):
        """
        Compute the mean and variance of the latent function at some new points
        Xnew. For a derivation of the terms in here, see the associated SGPR
        notebook.
        """
        num_inducing = self.Z.size(0)
        err = self.Y - self.mean_function(self.X)
        Kuf = self.kern.K(self.Z.get(), self.X)
        jitter = torch.eye(num_inducing, dtype=self.Z.dtype, device=self.Z.device) * self.jitter_level
        Kuu = self.kern.K(self.Z.get()) + jitter
        Kus = self.kern.K(self.Z.get(), Xnew)
        sigma = self.likelihood.variance.get()**0.5
        L = torch.linalg.cholesky(Kuu, upper=False)
        A = torch.solve(Kuf, L)[0] / sigma # could use triangular solve here and below
        B = torch.matmul(A,A.t()) + torch.eye(num_inducing, dtype=A.dtype, device=A.device)
        LB = torch.linalg.cholesky(B, upper=False)
        Aerr = torch.matmul(A, err)
        c = torch.solve(Aerr, LB)[0] / sigma
        tmp1,_ = torch.solve(Kus, L)
        tmp2,_ = torch.solve(tmp1,LB)
        mean = torch.matmul(tmp2.t(), c)
        if full_cov:
            var = self.kern.K(Xnew) + torch.matmul(tmp2.t(), tmp2) - torch.matmul(tmp1.t(), tmp1)
            var = var.unsqueeze(2).expand(-1, -1, self.Y.size(1))
        else:
            var = self.kern.Kdiag(Xnew) + (tmp2**2).sum(0) - (tmp1**2).sum(0)
            var = var.unsqueeze(1).expand(-1, self.Y.size(1))
        return mean + self.mean_function(Xnew), var
//  [python]
def compute_log_likelihood(self, X=None, Y=None):
        """
        For a derivation of the terms in here, see the associated
        SGPR notebook.
        """

        assert X is None and Y is None, "{} does not support minibatch mode".format(str(type(self)))

        num_inducing = self.Z.size(0)
        num_data = self.Y.size(0)
        output_dim = self.Y.size(1)

        err = self.Y - self.mean_function(self.X)
        Kdiag = self.kern.Kdiag(self.X)
        Kuf = self.kern.K(self.Z.get(), self.X)
        jitter = torch.eye(num_inducing, dtype=self.Z.dtype, device=self.Z.device) * self.jitter_level
        Kuu = self.kern.K(self.Z.get()) + jitter
        L = torch.cholesky(Kuu, upper=False)
        sigma = self.likelihood.variance.get()**0.5

        # Compute intermediate matrices
        A = torch.solve(Kuf, L)[0] / sigma  # could use triangular solve
        AAT = torch.matmul(A, A.t())
        B = AAT + torch.eye(num_inducing, dtype=AAT.dtype, device=AAT.device)
        LB = torch.cholesky(B, upper=False)
        Aerr = torch.matmul(A, err)
        c = torch.solve(Aerr, LB)[0] / sigma # could use triangular solve

        # compute log marginal bound
        bound = -0.5 * num_data * output_dim * float(numpy.log(2 * numpy.pi))
        bound += -output_dim * torch.sum(torch.log(torch.diag(LB)))
        bound = bound - 0.5 * num_data * output_dim * torch.log(self.likelihood.variance.get())
        bound += -0.5 * torch.sum(err**2) / self.likelihood.variance.get()
        bound += 0.5 * torch.sum(c**2)
        bound += -0.5 * output_dim * torch.sum(Kdiag) / self.likelihood.variance.get()
        bound += 0.5 * output_dim * torch.sum(torch.diag(AAT))

        return bound
//  [python]
def compute_log_likelihood(self, X=None, Y=None):
        """
        For a derivation of the terms in here, see the associated
        SGPR notebook.
        """

        assert X is None and Y is None, "{} does not support minibatch mode".format(str(type(self)))

        num_inducing = self.Z.size(0)
        num_data = self.Y.size(0)
        output_dim = self.Y.size(1)

        err = self.Y - self.mean_function(self.X)
        Kdiag = self.kern.Kdiag(self.X)
        Kuf = self.kern.K(self.Z.get(), self.X)
        jitter = torch.eye(num_inducing, dtype=self.Z.dtype, device=self.Z.device) * self.jitter_level
        Kuu = self.kern.K(self.Z.get()) + jitter
        L = torch.linalg.cholesky(Kuu)
        sigma = self.likelihood.variance.get()**0.5

        # Compute intermediate matrices
        A = torch.solve(Kuf, L)[0] / sigma  # could use triangular solve
        AAT = torch.matmul(A, A.t())
        B = AAT + torch.eye(num_inducing, dtype=AAT.dtype, device=AAT.device)
        LB = torch.linalg.cholesky(B)
        Aerr = torch.matmul(A, err)
        c = torch.solve(Aerr, LB)[0] / sigma # could use triangular solve

        # compute log marginal bound
        bound = -0.5 * num_data * output_dim * float(numpy.log(2 * numpy.pi))
        bound += -output_dim * torch.sum(torch.log(torch.diag(LB)))
        bound = bound - 0.5 * num_data * output_dim * torch.log(self.likelihood.variance.get())
        bound += -0.5 * torch.sum(err**2) / self.likelihood.variance.get()
        bound += 0.5 * torch.sum(c**2)
        bound += -0.5 * output_dim * torch.sum(Kdiag) / self.likelihood.variance.get()
        bound += 0.5 * output_dim * torch.sum(torch.diag(AAT))

        return bound
//  [python]
def chol(self):
        if not hasattr(self, "_chol"):
            self._chol = LowerTriangularMatrix(t.cholesky(self.A))
        return self._chol
//  [python]
def chol(self):
        if not hasattr(self, "_chol"):
            self._chol = LowerTriangularMatrix(t.linalg.cholesky(self.A))
        return self._chol
//  [python]
def estimate_variance(self, train_loader, var0):
        tau = 1/var0
        criterion = nn.CrossEntropyLoss()
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        # for x, y in tqdm(train_loader):
        for x, y in train_loader:
            x = x.cuda(non_blocking=True)

            output = self(x)
            distribution = torch.distributions.Categorical(logits=output)
            y = distribution.sample()

            loss = criterion(output, y)
            loss.backward()
            opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = len(train_loader.dataset)
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m, device='cuda')
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n, device='cuda')

                # Use cholesky_inverse for numerical stability since U, V SPD
                u = torch.cholesky(U)
                v = torch.cholesky(V)

                U_half = torch.cholesky(torch.cholesky_inverse(u), upper=False)
                V_half = torch.cholesky(torch.cholesky_inverse(v), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['MaskedLinear', 'Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def estimate_variance(self, train_loader, var0):
        tau = 1/var0
        criterion = nn.CrossEntropyLoss()
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        # for x, y in tqdm(train_loader):
        for x, y in train_loader:
            x = x.cuda(non_blocking=True)

            output = self(x)
            distribution = torch.distributions.Categorical(logits=output)
            y = distribution.sample()

            loss = criterion(output, y)
            loss.backward()
            opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = len(train_loader.dataset)
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m, device='cuda')
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n, device='cuda')

                # Use cholesky_inverse for numerical stability since U, V SPD
                u = torch.linalg.cholesky(U)
                v = torch.linalg.cholesky(V)

                U_half = torch.linalg.cholesky(torch.linalg.cholesky_inverse(u), upper=False)
                V_half = torch.linalg.cholesky(torch.linalg.cholesky_inverse(v), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['MaskedLinear', 'Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def get_dv_w_J_e(self, bs_idx, v, Minv, Jac, Jac_v, v_star, mu, cor, DPhi):
        n_cld, d, n_o, n_p, _ = Jac.shape
        n = n_o*n_p
        C = DPhi.shape[-1]
        # calculate A_decom
        J_e_T = DPhi[:,0,:,:,0,:].reshape(n*d, C)
        J_e = J_e_T.t() # 
        Minv_J_e_T = (Minv @ J_e_T.reshape(n, d*C)).reshape(n*d, C)
        J_e_Minv_J_e_T = J_e @ Minv_J_e_T # (C, C)
        M_e = torch.inverse(J_e_Minv_J_e_T) # (C, C) inertia in equality contact space
        factor = torch.eye(n*d).type_as(v) - J_e_T @ M_e @ Minv_J_e_T.t()
        M_hat_inv = (Minv @ factor.reshape(n, d*n*d)).reshape(n*d, n*d)
        A = Jac.reshape(n_cld*d, n*d) @ M_hat_inv @ Jac.reshape(n_cld*d, n*d).t()
        
        if torch.is_tensor(self.reg):
            reg = F.softplus(self.reg)
        else:
            reg = self.reg
        A_decom = torch.cholesky(A+torch.eye(A.shape[0]).type_as(A)*reg, upper=True)
        # make sure v_star is "valid", avoid unbounded cvx problem
        v_star_c = self.get_v_star_c_w_J_e(n_cld, n, d, v_star, Jac.reshape(n_cld*d, n*d), J_e)

        # compression phase impulse
        try:
            impulse = self.solve_compression_impulse(
                A_decom, Jac_v, mu, n_cld, d
            )
        except SolverError as e:
            try:
                impulse = self.solve_compression_impulse_w_penalty(
                    A_decom, Jac_v, mu, n_cld, d
                )
            except SolverError as e:
                # self.save(e, "comp", bs_idx, v, Minv, Jac, Jac_v, v_star, mu, cor, DPhi)
                impulse = torch.zeros(n_cld*d, 1).type_as(v)
        # velocity after compression phase (before retitution phase)
        M_hat_inv_Jac_T = M_hat_inv @ Jac.reshape(n_cld*d, n*d).t() # (n*d, n_cld*d)
        dv_comp = (M_hat_inv_Jac_T @ impulse).reshape(n, d)
        v_prev_r = v[bs_idx] + dv_comp
        # restitution phase impulse
        Jac_v_prev_r = Jac.reshape(n_cld*d, n*d) @ v_prev_r.reshape(n*d, 1) # n_cld*d, 1
        target_impulse = (cor.reshape(-1) * impulse.reshape(n_cld, d)[:, 0]).reshape(n_cld, 1)
        try:
            impulse_star_r = self.solve_restitution_impulse(
                A_decom, Jac_v_prev_r, v_star_c, mu, n_cld, d, target_impulse=target_impulse
            )
        except SolverError as e:
            try:
                impulse_star_r = self.solve_restitution_impulse_w_penalty(
                    A_decom, Jac_v_prev_r, v_star_c, mu, n_cld, d, target_impulse=target_impulse
                )
            except SolverError as e:
                # self.save(e, "rest", bs_idx, v, Minv, Jac, Jac_v, v_star, mu, cor, DPhi)
                impulse_star_r = torch.zeros(n_cld*d, 1).type_as(v)
        # velocity after restitution phase, compensate penetration
        dv_rest = (M_hat_inv_Jac_T @ impulse_star_r).reshape(n, d)
        dv = torch.zeros_like(v)
        dv[bs_idx] = dv_comp + dv_rest
        return dv
//  [python]
def get_dv_w_J_e(self, bs_idx, v, Minv, Jac, Jac_v, v_star, mu, cor, DPhi):
        n_cld, d, n_o, n_p, _ = Jac.shape
        n = n_o*n_p
        C = DPhi.shape[-1]
        # calculate A_decom
        J_e_T = DPhi[:,0,:,:,0,:].reshape(n*d, C)
        J_e = J_e_T.t() # 
        Minv_J_e_T = (Minv @ J_e_T.reshape(n, d*C)).reshape(n*d, C)
        J_e_Minv_J_e_T = J_e @ Minv_J_e_T # (C, C)
        M_e = torch.linalg.inv(J_e_Minv_J_e_T) # (C, C) inertia in equality contact space
        factor = torch.eye(n*d).type_as(v) - J_e_T @ M_e @ Minv_J_e_T.t()
        M_hat_inv = (Minv @ factor.reshape(n, d*n*d)).reshape(n*d, n*d)
        A = Jac.reshape(n_cld*d, n*d) @ M_hat_inv @ Jac.reshape(n_cld*d, n*d).t()
        
        if torch.is_tensor(self.reg):
            reg = F.softplus(self.reg)
        else:
            reg = self.reg
        A_decom = torch.linalg.cholesky(A+torch.eye(A.shape[0]).type_as(A)*reg, upper=True)
        # make sure v_star is "valid", avoid unbounded cvx problem
        v_star_c = self.get_v_star_c_w_J_e(n_cld, n, d, v_star, Jac.reshape(n_cld*d, n*d), J_e)

        # compression phase impulse
        try:
            impulse = self.solve_compression_impulse(
                A_decom, Jac_v, mu, n_cld, d
            )
        except SolverError as e:
            try:
                impulse = self.solve_compression_impulse_w_penalty(
                    A_decom, Jac_v, mu, n_cld, d
                )
            except SolverError as e:
                # self.save(e, "comp", bs_idx, v, Minv, Jac, Jac_v, v_star, mu, cor, DPhi)
                impulse = torch.zeros(n_cld*d, 1).type_as(v)
        # velocity after compression phase (before retitution phase)
        M_hat_inv_Jac_T = M_hat_inv @ Jac.reshape(n_cld*d, n*d).t() # (n*d, n_cld*d)
        dv_comp = (M_hat_inv_Jac_T @ impulse).reshape(n, d)
        v_prev_r = v[bs_idx] + dv_comp
        # restitution phase impulse
        Jac_v_prev_r = Jac.reshape(n_cld*d, n*d) @ v_prev_r.reshape(n*d, 1) # n_cld*d, 1
        target_impulse = (cor.reshape(-1) * impulse.reshape(n_cld, d)[:, 0]).reshape(n_cld, 1)
        try:
            impulse_star_r = self.solve_restitution_impulse(
                A_decom, Jac_v_prev_r, v_star_c, mu, n_cld, d, target_impulse=target_impulse
            )
        except SolverError as e:
            try:
                impulse_star_r = self.solve_restitution_impulse_w_penalty(
                    A_decom, Jac_v_prev_r, v_star_c, mu, n_cld, d, target_impulse=target_impulse
                )
            except SolverError as e:
                # self.save(e, "rest", bs_idx, v, Minv, Jac, Jac_v, v_star, mu, cor, DPhi)
                impulse_star_r = torch.zeros(n_cld*d, 1).type_as(v)
        # velocity after restitution phase, compensate penetration
        dv_rest = (M_hat_inv_Jac_T @ impulse_star_r).reshape(n, d)
        dv = torch.zeros_like(v)
        dv[bs_idx] = dv_comp + dv_rest
        return dv
//  [python]
def prior_kl(self, sample_idxs=None):
        q_mu, q_sqrt, z = self.prms
        assert (q_mu.shape[0] == q_sqrt.shape[0])
        if not self.tied_samples and sample_idxs is not None:
            q_mu = q_mu[sample_idxs]
            q_sqrt = q_sqrt[sample_idxs]
        z = self._expand_z(z)
        e = torch.eye(self.n_inducing).to(q_mu.device)
        if not self.whiten:
            kzz = self.kernel(z, z) + (e * jitter)
            l = torch.cholesky(kzz, upper=False)
        q = MultivariateNormal(q_mu, scale_tril=q_sqrt)
        p_mu = torch.zeros(self.n, self.n_inducing).to(q_mu.device)
        if not self.whiten:
            prior = MultivariateNormal(p_mu, scale_tril=l)
        else:
            prior = MultivariateNormal(p_mu, scale_tril=e)

        return kl_divergence(q, prior)
//  [python]
def prior_kl(self, sample_idxs=None):
        q_mu, q_sqrt, z = self.prms
        assert (q_mu.shape[0] == q_sqrt.shape[0])
        if not self.tied_samples and sample_idxs is not None:
            q_mu = q_mu[sample_idxs]
            q_sqrt = q_sqrt[sample_idxs]
        z = self._expand_z(z)
        e = torch.eye(self.n_inducing).to(q_mu.device)
        if not self.whiten:
            kzz = self.kernel(z, z) + (e * jitter)
            l = torch.linalg.cholesky(kzz, upper=False)
        q = MultivariateNormal(q_mu, scale_tril=q_sqrt)
        p_mu = torch.zeros(self.n, self.n_inducing).to(q_mu.device)
        if not self.whiten:
            prior = MultivariateNormal(p_mu, scale_tril=l)
        else:
            prior = MultivariateNormal(p_mu, scale_tril=e)

        return kl_divergence(q, prior)
//  [python]
def estimate_variance_tensor(self, X, criterion, dist, var0, lik_prec, mb_size=128):
        tau = 1/var0
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}
        
        lik_std = sqrt(1/lik_prec)

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        # for x, y in tqdm(train_loader):
        for i in range(0, X.shape[0], mb_size):
            x = X[i:i+mb_size]

            output = self(x)
            distribution = dist(output, lik_std)
            y = distribution.sample()
            
            loss = criterion(output, y)
            loss.backward()
            opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = X.shape[0]
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m, device=self.device)
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n, device=self.device)

#                 try:
                    # Use cholesky_inverse for numerical stability since U, V SPD
                u = torch.cholesky(U)
                v = torch.cholesky(V)
#                 except Exception as e:
#                     u = torch.cholesky(U + 1*torch.eye(m, device=self.device))
#                     v = torch.cholesky(V)

                U_half = torch.cholesky(torch.cholesky_inverse(u), upper=False)
                V_half = torch.cholesky(torch.cholesky_inverse(v), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['MaskedLinear', 'Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)