jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def log_prob(self, z, x, y):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        # inducing points log_prob
        N = y.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        L = torch.linalg.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0]
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.linalg.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        Kdiag = self.get_K_diag(z, x)

        bound = -0.5 * N * np.log(2 * np.pi) - torch.sum(torch.log(torch.diag(LB))) - 0.5 * N * torch.log(
            sigma2) - 0.5 * torch.sum(torch.pow(y, 2)) / sigma2
        bound += 0.5 * torch.sum(torch.pow(c, 2)) - 0.5 * torch.sum(Kdiag) / sigma2 + 0.5 * torch.sum(torch.diag(AAT))
        return bound
//  [python]
def predict(self, z, x, y, z_star, x_star, add_likelihood_variance=False):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        Nstar = z_star.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()

        sigma = torch.sqrt(sigma2)

        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        K_us = self.get_K_without_noise(self.z_u, z_star, self.x_u, x_star)

        L = torch.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        tmp1 = torch.triangular_solve(K_us, L, upper=False)[0]
        tmp2 = torch.triangular_solve(tmp1, LB, upper=False)[0]
        mean = self.intercept + torch.mm(tmp2.t(), c)

        Kdiag = self.get_K_diag(z_star, x_star)
        var = Kdiag + torch.pow(tmp2, 2).sum(dim=0) - torch.pow(tmp1, 2).sum(dim=0)

        if add_likelihood_variance:
            var += self.get_noise_var()

        return mean.reshape(-1), var
//  [python]
def predict(self, z, x, y, z_star, x_star, add_likelihood_variance=False):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        Nstar = z_star.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()

        sigma = torch.sqrt(sigma2)

        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        K_us = self.get_K_without_noise(self.z_u, z_star, self.x_u, x_star)

        L = torch.linalg.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.linalg.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        tmp1 = torch.triangular_solve(K_us, L, upper=False)[0]
        tmp2 = torch.triangular_solve(tmp1, LB, upper=False)[0]
        mean = self.intercept + torch.mm(tmp2.t(), c)

        Kdiag = self.get_K_diag(z_star, x_star)
        var = Kdiag + torch.pow(tmp2, 2).sum(dim=0) - torch.pow(tmp1, 2).sum(dim=0)

        if add_likelihood_variance:
            var += self.get_noise_var()

        return mean.reshape(-1), var
//  [python]
def log_prob(self, z, x, y):
        # inducing points log_prob
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        N = y.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        L = torch.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        Kdiag = self.get_kernel_var().repeat(N)

        bound = -0.5 * N * np.log(2 * np.pi) - torch.sum(torch.log(torch.diag(LB))) - 0.5 * N * torch.log(
            sigma2) - 0.5 * torch.sum(torch.pow(y, 2)) / sigma2
        bound += 0.5 * torch.sum(torch.pow(c, 2)) - 0.5 * torch.sum(Kdiag) / sigma2 + 0.5 * torch.sum(torch.diag(AAT))
        return bound
//  [python]
def log_prob(self, z, x, y):
        # inducing points log_prob
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        N = y.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        L = torch.linalg.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.linalg.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        Kdiag = self.get_kernel_var().repeat(N)

        bound = -0.5 * N * np.log(2 * np.pi) - torch.sum(torch.log(torch.diag(LB))) - 0.5 * N * torch.log(
            sigma2) - 0.5 * torch.sum(torch.pow(y, 2)) / sigma2
        bound += 0.5 * torch.sum(torch.pow(c, 2)) - 0.5 * torch.sum(Kdiag) / sigma2 + 0.5 * torch.sum(torch.diag(AAT))
        return bound
//  [python]
def predict(self, z, x, y, z_star, x_star, add_likelihood_variance=False):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        Nstar = z_star.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)

        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        K_us = self.get_K_without_noise(self.z_u, z_star, self.x_u, x_star)

        L = torch.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        tmp1 = torch.triangular_solve(K_us, L, upper=False)[0]
        tmp2 = torch.triangular_solve(tmp1, LB, upper=False)[0]
        mean = self.intercept + torch.mm(tmp2.t(), c)

        K_ss = self.get_K_without_noise(z_star, z_star, x_star, x_star)
        Kdiag = K_ss.diag()
        var = Kdiag + torch.pow(tmp2, 2).sum(dim=0) - torch.pow(tmp1, 2).sum(dim=0)

        if add_likelihood_variance:
            var += self.get_noise_var()

        return mean.reshape(-1), var
//  [python]
def predict(self, z, x, y, z_star, x_star, add_likelihood_variance=False):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        Nstar = z_star.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)

        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        K_us = self.get_K_without_noise(self.z_u, z_star, self.x_u, x_star)

        L = torch.linalg.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.linalg.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        tmp1 = torch.triangular_solve(K_us, L, upper=False)[0]
        tmp2 = torch.triangular_solve(tmp1, LB, upper=False)[0]
        mean = self.intercept + torch.mm(tmp2.t(), c)

        K_ss = self.get_K_without_noise(z_star, z_star, x_star, x_star)
        Kdiag = K_ss.diag()
        var = Kdiag + torch.pow(tmp2, 2).sum(dim=0) - torch.pow(tmp1, 2).sum(dim=0)

        if add_likelihood_variance:
            var += self.get_noise_var()

        return mean.reshape(-1), var
//  [python]
def log_prob(self, z, x, y):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        # inducing points log_prob

        N = y.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        L = torch.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        K_ff = self.get_K_without_noise(z, z, x, x)
        Kdiag = K_ff.diag()

        bound = -0.5 * N * np.log(2 * np.pi) - torch.sum(torch.log(torch.diag(LB))) - 0.5 * N * torch.log(
            sigma2) - 0.5 * torch.sum(torch.pow(y, 2)) / sigma2
        bound += 0.5 * torch.sum(torch.pow(c, 2)) - 0.5 * torch.sum(Kdiag) / sigma2 + 0.5 * torch.sum(torch.diag(AAT))
        return bound
//  [python]
def log_prob(self, z, x, y):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        # inducing points log_prob

        N = y.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        L = torch.linalg.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.linalg.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        K_ff = self.get_K_without_noise(z, z, x, x)
        Kdiag = K_ff.diag()

        bound = -0.5 * N * np.log(2 * np.pi) - torch.sum(torch.log(torch.diag(LB))) - 0.5 * N * torch.log(
            sigma2) - 0.5 * torch.sum(torch.pow(y, 2)) / sigma2
        bound += 0.5 * torch.sum(torch.pow(c, 2)) - 0.5 * torch.sum(Kdiag) / sigma2 + 0.5 * torch.sum(torch.diag(AAT))
        return bound
//  [python]
def log_prob(self, z, y):
        # inducing points log_prob
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
        N = y.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z)
        L = torch.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        K_ff = self.get_K_without_noise(z, z)
        Kdiag = K_ff.diag()

        bound = -0.5 * N * np.log(2 * np.pi) - torch.sum(torch.log(torch.diag(LB))) - 0.5 * N * torch.log(
            sigma2) - 0.5 * torch.sum(torch.pow(y, 2)) / sigma2
        bound += 0.5 * torch.sum(torch.pow(c, 2)) - 0.5 * torch.sum(Kdiag) / sigma2 + 0.5 * torch.sum(torch.diag(AAT))
        return bound
//  [python]
def log_prob(self, z, y):
        # inducing points log_prob
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
        N = y.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z)
        L = torch.linalg.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.linalg.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        K_ff = self.get_K_without_noise(z, z)
        Kdiag = K_ff.diag()

        bound = -0.5 * N * np.log(2 * np.pi) - torch.sum(torch.log(torch.diag(LB))) - 0.5 * N * torch.log(
            sigma2) - 0.5 * torch.sum(torch.pow(y, 2)) / sigma2
        bound += 0.5 * torch.sum(torch.pow(c, 2)) - 0.5 * torch.sum(Kdiag) / sigma2 + 0.5 * torch.sum(torch.diag(AAT))
        return bound
//  [python]
def forward(ctx, A, b):
        u = torch.cholesky(torch.matmul(A.transpose(-1, -2), A), upper=True)
        ret = torch.cholesky_solve(torch.matmul(A.transpose(-1, -2), b), u, upper=True)
        ctx.save_for_backward(u, ret, A, b)
        return ret
//  [python]
def forward(ctx, A, b):
        u = torch.linalg.cholesky(torch.matmul(A.transpose(-1, -2), A))
        ret = torch.linalg.cholesky_solve(torch.matmul(A.transpose(-1, -2), b), u)
        ctx.save_for_backward(u, ret, A, b)
        return ret
//  [python]
def predict(self, z, y, z_star, which_kernels, add_likelihood_variance=False):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
        Nstar = z_star.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, jitter=self.jitter, which_kernels=which_kernels)
        K_uf = self.get_K_without_noise(self.z_u, z, which_kernels=which_kernels)
        K_us = self.get_K_without_noise(self.z_u, z_star, which_kernels=which_kernels)

        L = torch.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        tmp1 = torch.triangular_solve(K_us, L, upper=False)[0]
        tmp2 = torch.triangular_solve(tmp1, LB, upper=False)[0]
        mean = self.intercept + torch.mm(tmp2.t(), c)

        K_ss = self.get_K_without_noise(z_star, z_star)
        Kdiag = K_ss.diag()
        var = Kdiag + torch.pow(tmp2, 2).sum(dim=0) - torch.pow(tmp1, 2).sum(dim=0)

        if add_likelihood_variance:
            var += self.get_noise_var()

        return mean.reshape(-1), var
//  [python]
def predict(self, z, y, z_star, which_kernels, add_likelihood_variance=False):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
        Nstar = z_star.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, jitter=self.jitter, which_kernels=which_kernels)
        K_uf = self.get_K_without_noise(self.z_u, z, which_kernels=which_kernels)
        K_us = self.get_K_without_noise(self.z_u, z_star, which_kernels=which_kernels)

        L = torch.linalg.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0] / sigma
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.linalg.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        tmp1 = torch.triangular_solve(K_us, L, upper=False)[0]
        tmp2 = torch.triangular_solve(tmp1, LB, upper=False)[0]
        mean = self.intercept + torch.mm(tmp2.t(), c)

        K_ss = self.get_K_without_noise(z_star, z_star)
        Kdiag = K_ss.diag()
        var = Kdiag + torch.pow(tmp2, 2).sum(dim=0) - torch.pow(tmp1, 2).sum(dim=0)

        if add_likelihood_variance:
            var += self.get_noise_var()

        return mean.reshape(-1), var
//  [python]
def grow_cov(self, beta, lam):
        Paa = self.Paa
        Sigma = lam * torch.cholesky_inverse(Paa)
        Sigma += beta * torch.eye(Sigma.shape[0])
        Pnew = (1./lam) * torch.cholesky_inverse(Sigma)
        Lmat = torch.cholesky(Pnew)
        tril_indices = torch.tril_indices(row=self.d_act, col=self.d_act, offset=0)
        self.L.data = Lmat[tril_indices[0], tril_indices[1]]
//  [python]
def grow_cov(self, beta, lam):
        Paa = self.Paa
        Sigma = lam * torch.linalg.cholesky_inverse(Paa)
        Sigma += beta * torch.eye(Sigma.shape[0])
        Pnew = (1./lam) * torch.linalg.cholesky_inverse(Sigma)
        Lmat = torch.linalg.cholesky(Pnew)
        tril_indices = torch.tril_indices(row=self.d_act, col=self.d_act, offset=0)
        self.L.data = Lmat[tril_indices[0], tril_indices[1]]
//  [python]
def build_cache(self):
        """
        Builds a cache of computations that uniquely define a sample from posterior process
        1. Generate and fix parameters of Fourier features
        2. Generate and fix inducing posterior sample
        3. Intermediate computations based on the inducing sample for pathwise update
        """
        # generate parameters required for the Fourier feature maps
        self.rff_weights = sample_normal((self.S, self.D_out))  # (S,D_out)
        self.rff_omega = self.kern.sample_freq(self.S)  # (D_in,S) or (D_in,S,D_out)
        phase_shape = (1, self.S, self.D_out) if self.dimwise else (1, self.S)
        self.rff_phase = sample_uniform(phase_shape) * 2 * np.pi  # (S,D_out)

        # generate sample from the inducing posterior
        inducing_val = self.sample_inducing()  # (M,D)

        # compute th term nu = k(Z,Z)^{-1}(u-f(Z)) in whitened form of inducing variables
        # equation (13) from http://proceedings.mlr.press/v119/wilson20a/wilson20a.pdf
        Ku = self.kern.K(self.inducing_loc())  # (M,M) or (D,M,M)
        Lu = torch.cholesky(Ku + torch.eye(self.M) * jitter)  # (M,M) or (D,M,M)
        u_prior = self.rff_forward(self.inducing_loc())  # (M,D)

        if not self.dimwise:
            nu = torch.triangular_solve(u_prior, Lu, upper=False)[0]  # (M,D)
            nu = torch.triangular_solve((inducing_val - nu),
                                        Lu.T, upper=True)[0]  # (M,D)
        else:
            nu = torch.triangular_solve(u_prior.T.unsqueeze(2), Lu, upper=False)[0]  # (D,M,1)
            nu = torch.triangular_solve((inducing_val.T.unsqueeze(2) - nu),
                                        Lu.permute(0, 2, 1), upper=True)[0]  # (D,M,1)
        self.nu = nu
//  [python]
def build_cache(self):
        """
        Builds a cache of computations that uniquely define a sample from posterior process
        1. Generate and fix parameters of Fourier features
        2. Generate and fix inducing posterior sample
        3. Intermediate computations based on the inducing sample for pathwise update
        """
        # generate parameters required for the Fourier feature maps
        self.rff_weights = sample_normal((self.S, self.D_out))  # (S,D_out)
        self.rff_omega = self.kern.sample_freq(self.S)  # (D_in,S) or (D_in,S,D_out)
        phase_shape = (1, self.S, self.D_out) if self.dimwise else (1, self.S)
        self.rff_phase = sample_uniform(phase_shape) * 2 * np.pi  # (S,D_out)

        # generate sample from the inducing posterior
        inducing_val = self.sample_inducing()  # (M,D)

        # compute th term nu = k(Z,Z)^{-1}(u-f(Z)) in whitened form of inducing variables
        # equation (13) from http://proceedings.mlr.press/v119/wilson20a/wilson20a.pdf
        Ku = self.kern.K(self.inducing_loc())  # (M,M) or (D,M,M)
        Lu = torch.linalg.cholesky(Ku + torch.eye(self.M) * jitter)  # (M,M) or (D,M,M)
        u_prior = self.rff_forward(self.inducing_loc())  # (M,D)

        if not self.dimwise:
            nu = torch.triangular_solve(u_prior, Lu, upper=False)[0]  # (M,D)
            nu = torch.triangular_solve((inducing_val - nu),
                                        Lu.T, upper=True)[0]  # (M,D)
        else:
            nu = torch.triangular_solve(u_prior.T.unsqueeze(2), Lu, upper=False)[0]  # (D,M,1)
            nu = torch.triangular_solve((inducing_val.T.unsqueeze(2) - nu),
                                        Lu.permute(0, 2, 1), upper=True)[0]  # (D,M,1)
        self.nu = nu
//  [python]
def log_det_by_cholesky(matrix):
    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):
    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 forward(self, x):
        x = torch.stack(torch.chunk(x, 4, 1), 1).permute(1, 0, 2, 3, 4)
        axes, d = (1, *range(3, x.dim())), x.shape[0]
        shape = 1, x.shape[2], *([1] * (x.dim() - 3))

        if self.training:
            mean = x.mean(dim=axes)
            if self.running_mean is not None:
                with torch.no_grad():
                    self.running_mean = self.momentum * self.running_mean + \
                                        (1.0 - self.momentum) * mean
        else:
            mean = self.running_mean

        x = x - mean.reshape(d, *shape)

        if self.training:
            perm = x.permute(2, 0, *axes).flatten(2, -1)
            cov = torch.matmul(perm, perm.transpose(-1, -2)) / perm.shape[-1]

            if self.running_cov is not None:
                with torch.no_grad():
                    self.running_cov = self.momentum * self.running_cov + \
                                       (1.0 - self.momentum) * cov

        else:
            cov = self.running_cov

        ell = torch.cholesky(cov + self.eye, upper=True)
        soln = torch.triangular_solve(
            x.unsqueeze(-1).permute(*range(1, x.dim()), 0, -1),
            ell.reshape(*shape, d, d)
        )

        wht = soln.solution.squeeze(-1)
        z = torch.stack(torch.unbind(wht, dim=-1), dim=0)

        if self.affine:
            weight = self.weight.view(4, 4, *shape)
            scaled = torch.stack([
                z[0] * weight[0, 0] + z[1] * weight[0, 1] + z[2] * weight[0, 2] + z[3] * weight[0, 3],
                z[0] * weight[1, 0] + z[1] * weight[1, 1] + z[2] * weight[1, 2] + z[3] * weight[1, 3],
                z[0] * weight[2, 0] + z[1] * weight[2, 1] + z[2] * weight[2, 2] + z[3] * weight[2, 3],
                z[0] * weight[3, 0] + z[1] * weight[3, 1] + z[2] * weight[3, 2] + z[3] * weight[3, 3],
            ], dim=0)
            z = scaled + self.bias.reshape(4, *shape)

        z = torch.cat(torch.chunk(z, 4, 0), 2).squeeze()

        return Q(z)
//  [python]
def forward(self, x):
        x = torch.stack(torch.chunk(x, 4, 1), 1).permute(1, 0, 2, 3, 4)
        axes, d = (1, *range(3, x.dim())), x.shape[0]
        shape = 1, x.shape[2], *([1] * (x.dim() - 3))

        if self.training:
            mean = x.mean(dim=axes)
            if self.running_mean is not None:
                with torch.no_grad():
                    self.running_mean = self.momentum * self.running_mean + \
                                        (1.0 - self.momentum) * mean
        else:
            mean = self.running_mean

        x = x - mean.reshape(d, *shape)

        if self.training:
            perm = x.permute(2, 0, *axes).flatten(2, -1)
            cov = torch.matmul(perm, perm.transpose(-1, -2)) / perm.shape[-1]

            if self.running_cov is not None:
                with torch.no_grad():
                    self.running_cov = self.momentum * self.running_cov + \
                                       (1.0 - self.momentum) * cov

        else:
            cov = self.running_cov

        ell = torch.linalg.cholesky(cov + self.eye, upper=True)
        soln = torch.triangular_solve(
            x.unsqueeze(-1).permute(*range(1, x.dim()), 0, -1),
            ell.reshape(*shape, d, d)
        )

        wht = soln.solution.squeeze(-1)
        z = torch.stack(torch.unbind(wht, dim=-1), dim=0)

        if self.affine:
            weight = self.weight.view(4, 4, *shape)
            scaled = torch.stack([
                z[0] * weight[0, 0] + z[1] * weight[0, 1] + z[2] * weight[0, 2] + z[3] * weight[0, 3],
                z[0] * weight[1, 0] + z[1] * weight[1, 1] + z[2] * weight[1, 2] + z[3] * weight[1, 3],
                z[0] * weight[2, 0] + z[1] * weight[2, 1] + z[2] * weight[2, 2] + z[3] * weight[2, 3],
                z[0] * weight[3, 0] + z[1] * weight[3, 1] + z[2] * weight[3, 2] + z[3] * weight[3, 3],
            ], dim=0)
            z = scaled + self.bias.reshape(4, *shape)

        z = torch.cat(torch.chunk(z, 4, 0), 2).squeeze()

        return Q(z)
//  [python]
def build_conditional(self, x, full_cov=False):
        """
        Calculates conditional distribution q(f(x)) = N(m(x), Sigma(x))
            where  m(x) = k(x,Z)k(Z,Z)^{-1}u, k(x,x)
                   Sigma(x) = k(x,Z)k(Z,Z)^{-1}(S-K(Z,Z))k(Z,Z)^{-1}k(Z,x))

        @param x: input tensor (N,D)
        @param full_cov: if True, returns full Sigma(x) else returns only diagonal
        @return: m(x), Sigma(x)
        """
        Ku = self.kern.K(self.inducing_loc())  # (M,M) or (D,M,M)
        Lu = torch.cholesky(Ku + torch.eye(self.M) * jitter)  # (M,M) or (D,M,M)
        Kuf = self.kern.K(self.inducing_loc(), x)  # (M,N) or (D,M,N)
        A = torch.triangular_solve(Kuf, Lu, upper=False)[0]  # (M,M)@(M,N) --> (M,N) or (D,M,M)@(D,M,N) --> (D,M,N)

        Us_sqrt = self.Us_sqrt().T[:, :, None] if self.q_diag else self.Us_sqrt()  # (D,M,1) or (D,M,M)
        SK = (Us_sqrt @ Us_sqrt.permute(0, 2, 1)) - torch.eye(Ku.shape[1]).unsqueeze(0)  # (D,M,M)
        B = torch.einsum('dme, den->dmn' if self.dimwise else 'dmi, in->dmn', SK, A)  # (D,M,N)

        if full_cov:
            delta_cov = torch.einsum('dme, dmn->den' if self.dimwise else 'me, dmn->den', A, B)  # (D,M,N)
            Kff = self.kern.K(x) if self.dimwise else self.kern.K(x).unsqueeze(0)  # (1,N,N) or (D,N,N)
        else:
            delta_cov = ((A if self.dimwise else A.unsqueeze(0)) * B).sum(1)  # (D,N)
            if self.dimwise:
                Kff = torch.diagonal(self.kern.K(x), dim1=1, dim2=2)  # (N,)
            else:
                Kff = torch.diagonal(self.kern.K(x), dim1=0, dim2=1)  # (D,N)

        var = Kff + delta_cov
        mean = torch.einsum('dmn, md->nd' if self.dimwise else 'mn, md->nd', A, self.Um())
        return mean, var.T
//  [python]
def build_conditional(self, x, full_cov=False):
        """
        Calculates conditional distribution q(f(x)) = N(m(x), Sigma(x))
            where  m(x) = k(x,Z)k(Z,Z)^{-1}u, k(x,x)
                   Sigma(x) = k(x,Z)k(Z,Z)^{-1}(S-K(Z,Z))k(Z,Z)^{-1}k(Z,x))

        @param x: input tensor (N,D)
        @param full_cov: if True, returns full Sigma(x) else returns only diagonal
        @return: m(x), Sigma(x)
        """
        Ku = self.kern.K(self.inducing_loc())  # (M,M) or (D,M,M)
        Lu = torch.linalg.cholesky(Ku + torch.eye(self.M) * jitter)  # (M,M) or (D,M,M)
        Kuf = self.kern.K(self.inducing_loc(), x)  # (M,N) or (D,M,N)
        A = torch.triangular_solve(Kuf, Lu, upper=False)[0]  # (M,M)@(M,N) --> (M,N) or (D,M,M)@(D,M,N) --> (D,M,N)

        Us_sqrt = self.Us_sqrt().T[:, :, None] if self.q_diag else self.Us_sqrt()  # (D,M,1) or (D,M,M)
        SK = (Us_sqrt @ Us_sqrt.permute(0, 2, 1)) - torch.eye(Ku.shape[1]).unsqueeze(0)  # (D,M,M)
        B = torch.einsum('dme, den->dmn' if self.dimwise else 'dmi, in->dmn', SK, A)  # (D,M,N)

        if full_cov:
            delta_cov = torch.einsum('dme, dmn->den' if self.dimwise else 'me, dmn->den', A, B)  # (D,M,N)
            Kff = self.kern.K(x) if self.dimwise else self.kern.K(x).unsqueeze(0)  # (1,N,N) or (D,N,N)
        else:
            delta_cov = ((A if self.dimwise else A.unsqueeze(0)) * B).sum(1)  # (D,N)
            if self.dimwise:
                Kff = torch.diagonal(self.kern.K(x), dim1=1, dim2=2)  # (N,)
            else:
                Kff = torch.diagonal(self.kern.K(x), dim1=0, dim2=1)  # (D,N)

        var = Kff + delta_cov
        mean = torch.einsum('dmn, md->nd' if self.dimwise else 'mn, md->nd', A, self.Um())
        return mean, var.T
//  [python]
def forward(self, x_te):
        x_te = (x_te - self.x_mean.expand_as(x_te)) / self.x_std.expand_as(x_te)

        sigma = self.kernel_customized(self.x, self.x) + self.log_beta.exp().pow(-1) * torch.eye(self.x.size(0)) \
            + JITTER * torch.eye(self.x.size(0))

        kx = self.kernel_customized(self.x, x_te)
        L = torch.cholesky(sigma)
        l_inv_kx, _ = torch.triangular_solve(kx, L, upper=False)

        # option 1
        mean = kx.t() @ torch.cholesky_solve(self.y, L)  # torch.linalg.cholesky()
        # var_diag = self.log_coe_rbf.exp().expand(n_test, 1) \
        #     - (l_inv_kx**2).sum(dim=0).view(-1, 1)
        var_diag = self.kernel_customized(x_te, x_te).diag().view(-1, 1) - (l_inv_kx**2).sum(dim=0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        mean = mean * self.y_std.expand_as(mean) + self.y_mean.expand_as(mean)
        var_diag = var_diag.expand_as(mean) * self.y_std ** 2
        return mean, var_diag
//  [python]
def forward(self, x_te):
        x_te = (x_te - self.x_mean.expand_as(x_te)) / self.x_std.expand_as(x_te)

        sigma = self.kernel_customized(self.x, self.x) + self.log_beta.exp().pow(-1) * torch.eye(self.x.size(0)) \
            + JITTER * torch.eye(self.x.size(0))

        kx = self.kernel_customized(self.x, x_te)
        L = torch.linalg.cholesky(sigma)
        l_inv_kx, _ = torch.triangular_solve(kx, L, upper=False)

        # option 1
        mean = kx.t() @ torch.linalg.solve_triangular(self.y, L, upper=False)  # torch.linalg.cholesky()
        # var_diag = self.log_coe_rbf.exp().expand(n_test, 1) \
        #     - (l_inv_kx**2).sum(dim=0).view(-1, 1)
        var_diag = self.kernel_customized(x_te, x_te).diag().view(-1, 1) - (l_inv_kx**2).sum(dim=0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        mean = mean * self.y_std.expand_as(mean) + self.y_mean.expand_as(mean)
        var_diag = var_diag.expand_as(mean) * self.y_std ** 2
        return mean, var_diag
//  [python]
def complexity(sigma, c, samples):
    n = sigma.shape[0]
    device = sigma.device
    id = torch.eye(n, device=device)

    assert ( sigma == sigma.t() ).all()
    u = torch.cholesky(sigma)
    inv = torch.cholesky_inverse(u)
    assert (torch.mm(sigma, inv) - id).abs().max() < 1e-03

    nth_root_det = u.diag().pow(2/n).prod()
    inv_trace = inv.diag().sum()
    inv_proj = torch.dot(c, torch.mv(inv, c))

    formula_1 = n/5.0 + nth_root_det * ( (0.5 - 1/math.pi)*inv_trace + inv_proj/math.pi )

    mat = nth_root_det*inv - id

    z = torch.randn((n,samples), device=device).abs()
    cz = torch.mul(c.unsqueeze(1), z)
    contribs = -0.5*(cz * torch.matmul(mat, cz)).sum(dim=0)

    formula_0 = n*math.log(2) + math.log(samples) - torch.logsumexp(contribs, dim=0)

    return formula_0.item(), formula_1.item()
//  [python]
def complexity(sigma, c, samples):
    n = sigma.shape[0]
    device = sigma.device
    id = torch.eye(n, device=device)

    assert ( sigma == sigma.t() ).all()
    u = torch.linalg.cholesky(sigma)
    inv = torch.linalg.cholesky_inverse(u)
    assert (torch.mm(sigma, inv) - id).abs().max() < 1e-03

    nth_root_det = u.diag().pow(2/n).prod()
    inv_trace = inv.diag().sum()
    inv_proj = torch.dot(c, torch.mv(inv, c))

    formula_1 = n/5.0 + nth_root_det * ( (0.5 - 1/math.pi)*inv_trace + inv_proj/math.pi )

    mat = nth_root_det*inv - id

    z = torch.randn((n,samples), device=device).abs()
    cz = torch.mul(c.unsqueeze(1), z)
    contribs = -0.5*(cz * torch.matmul(mat, cz)).sum(dim=0)

    formula_0 = n*math.log(2) + math.log(samples) - torch.logsumexp(contribs, dim=0)

    return formula_0.item(), formula_1.item()
//  [python]
def gp_model(X, y):
    N = X.shape[0]
    
    # Priors for kernel parameters.
    alpha = pyro.sample("alpha", dist.LogNormal(0., 0.1))
    rho = pyro.sample("rho", dist.LogNormal(0., 1.))
    sigma = pyro.sample("sigma", dist.LogNormal(0., 1.))
    
    # Covariance matrix. 
    K = sq_exp_kernel_matrix(X, alpha, rho) + torch.eye(N) * sigma * sigma
    L = torch.cholesky(K)
    
    # Marginal likelihood.
    pyro.sample('obs',
                dist.MultivariateNormal(torch.zeros(N), scale_tril=L),
                obs=y)
//  [python]
def gp_model(X, y):
    N = X.shape[0]
    
    # Priors for kernel parameters.
    alpha = pyro.sample("alpha", dist.LogNormal(0., 0.1))
    rho = pyro.sample("rho", dist.LogNormal(0., 1.))
    sigma = pyro.sample("sigma", dist.LogNormal(0., 1.))
    
    # Covariance matrix. 
    K = sq_exp_kernel_matrix(X, alpha, rho) + torch.eye(N) * sigma * sigma
    L = torch.linalg.cholesky(K)
    
    # Marginal likelihood.
    pyro.sample('obs',
                dist.MultivariateNormal(torch.zeros(N), scale_tril=L),
                obs=y)
//  [python]
def forward(ctx, A, B):
        r"""In the forward pass we solve the matrix least-squares
        problem given the problem data A and B. We cache A, B, theta,
        and u (the cholesky decomposition).

        Args:
            A: p x n Tensor.
            B: p x m Tensor.

        Returns:
            theta^ls: Solution to matrix least-squares problem.
        """
        # compute gram matrix of A, and its QR factorization
        with torch.no_grad():
            u = torch.cholesky(A.t() @ A, upper=True)
            theta = torch.potrs(A.t() @ B, u)

        # cache information needed for backward pass
        ctx.save_for_backward(A, B, theta, u)

        return theta
//  [python]
def forward(ctx, A, B):
        r"""In the forward pass we solve the matrix least-squares
        problem given the problem data A and B. We cache A, B, theta,
        and u (the cholesky decomposition).

        Args:
            A: p x n Tensor.
            B: p x m Tensor.

        Returns:
            theta^ls: Solution to matrix least-squares problem.
        """
        # compute gram matrix of A, and its QR factorization
        with torch.no_grad():
            u = torch.linalg.cholesky(A.t() @ A, upper=True)
            theta = torch.potrs(A.t() @ B, u)

        # cache information needed for backward pass
        ctx.save_for_backward(A, B, theta, u)

        return theta
//  [python]
def solve_linear_system(self, A, b, K, delta = 0.0):
    I = torch.eye(self.N, self.N, device=self.device)
    A_t_K = torch.mm(A.t(), K)
    A_t_A = torch.mm(A_t_K, A)
    LAM = A_t_A + delta*I
    R = torch.mm(A_t_K, b)
    #Solve using cholesky
    l = torch.cholesky(LAM, upper=False)
    z = torch.trtrs(R, l, transpose=False, upper=False)[0]
    dtheta = torch.trtrs(z, l, transpose=True, upper=False)[0]
    return dtheta.view(self.num_traj_states, self.state_dim)
//  [python]
def solve_linear_system(self, A, b, K, delta = 0.0):
    I = torch.eye(self.N, self.N, device=self.device)
    A_t_K = torch.mm(A.t(), K)
    A_t_A = torch.mm(A_t_K, A)
    LAM = A_t_A + delta*I
    R = torch.mm(A_t_K, b)
    #Solve using cholesky
    l = torch.linalg.cholesky(LAM, upper=False)
    z = torch.trtrs(R, l, transpose=False, upper=False)[0]
    dtheta = torch.trtrs(z, l, transpose=True, upper=False)[0]
    return dtheta.view(self.num_traj_states, self.state_dim)
//  [python]
def _stable_cholesky(input, upper=False, out=None, eps=1e-6):
    try:
        return torch.cholesky(input, upper=upper, out=out)
    except RuntimeError:
        input = condition_scm(input, eps)
        return torch.cholesky(input, upper=upper, out=out)
//  [python]
def _stable_cholesky(input, upper=False, out=None, eps=1e-6):
    try:
        return torch.linalg.cholesky(input, upper=upper, out=out)
    except RuntimeError:
        input = condition_scm(input, eps)
        return torch.linalg.cholesky(input, upper=upper, out=out)
//  [python]
def KL_closed(covar_module, train_x, likelihoods, data, mu, log_var):
    """
    Closed form KL divergence.

    :param covar_module: additive kernel (sum of cross-covariances)
    :param train_x: auxiliary covariate information
    :param likelihoods: GPyTorch likelihood model
    :param data: sample measurements
    :param mu: mean of approximating variational distribution
    :param log_var: log variance of approximating variational distribution
    :return: KL divergence between variational distribution and additive GP prior
    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    K1 = covar_module(train_x.to(device), train_x.to(device)).evaluate() \
         + likelihoods.noise * torch.eye(data.shape[0]).to(device)
    v1 = torch.exp(log_var.view(-1)).to(device)
    mu1 = mu.view(-1)
    LK1 = torch.cholesky(K1)
    iK1 = torch.cholesky_solve(torch.eye(data.shape[0],
                                         dtype=torch.double).to(device), LK1).to(device)
    logdet11 = 2 * torch.sum(torch.log(torch.diag(LK1))).to(device)
    qf1 = torch.sum(mu1 * torch.matmul(iK1, mu1)).to(device)
    tr1 = torch.sum(v1 * torch.diag(iK1)).to(device)
    logdet10 = log_var.sum().to(device)
    kld1 = 0.5 * (tr1 + qf1 - data.shape[0] + logdet11 - logdet10)
    return kld1
//  [python]
def KL_closed(covar_module, train_x, likelihoods, data, mu, log_var):
    """
    Closed form KL divergence.

    :param covar_module: additive kernel (sum of cross-covariances)
    :param train_x: auxiliary covariate information
    :param likelihoods: GPyTorch likelihood model
    :param data: sample measurements
    :param mu: mean of approximating variational distribution
    :param log_var: log variance of approximating variational distribution
    :return: KL divergence between variational distribution and additive GP prior
    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    K1 = covar_module(train_x.to(device), train_x.to(device)).evaluate() \
         + likelihoods.noise * torch.eye(data.shape[0]).to(device)
    v1 = torch.exp(log_var.view(-1)).to(device)
    mu1 = mu.view(-1)
    LK1 = torch.linalg.cholesky(K1)
    iK1 = torch.linalg.solve_triangular(LK1, torch.eye(data.shape[0],
                                         dtype=torch.double).to(device), upper=False).to(device)
    logdet11 = 2 * torch.sum(torch.log(torch.diag(LK1))).to(device)
    qf1 = torch.sum(mu1 * torch.matmul(iK1, mu1)).to(device)
    tr1 = torch.sum(v1 * torch.diag(iK1)).to(device)
    logdet10 = log_var.sum().to(device)
    kld1 = 0.5 * (tr1 + qf1 - data.shape[0] + logdet11 - logdet10)
    return kld1
//  [python]
def inverse_by_cholesky(tensor: torch.Tensor, damping: torch.Tensor) -> torch.Tensor:
    damped = tensor + torch.eye(tensor.shape[-1], device=tensor.device, dtype=tensor.dtype) * damping
    cholesky = torch.cholesky(damped)
    return torch.cholesky_inverse(cholesky)
//  [python]
def inverse_by_cholesky(tensor: torch.Tensor, damping: torch.Tensor) -> torch.Tensor:
    damped = tensor + torch.eye(tensor.shape[-1], device=tensor.device, dtype=tensor.dtype) * damping
    cholesky = torch.linalg.cholesky(damped)
    return torch.linalg.cholesky_inverse(cholesky)
//  [python]
def solve_linear_system_batch(self, A, b, K, delta = 0.0):
    I = torch.eye(self.N, self.N, device=self.device).unsqueeze(0).repeat(self.batch_size,1,1)
    # K = K + 1e-8 * torch.eye(K.shape[1], K.shape[2], device=self.device).repeat(self.batch_size,1,1)
    A_t_K = torch.bmm(A.transpose(1,2), K)
    A_t_A = torch.bmm(A_t_K, A)
    LAM = A_t_A + delta*I
    R = torch.bmm(A_t_K, b)
    #Solve using cholesky
    # l = torch.cholesky(LAM, upper=False)
    # z = torch.trtrs(R, l, transpose=False, upper=False)[0]
    # dtheta = torch.trtrs(z, l, transpose=True, upper=False)[0]

    u = torch.cholesky(LAM, upper=True)
    z = torch.bmm(torch.inverse(u.transpose(1,2)), R)
    dtheta = torch.bmm(torch.inverse(u), z) 

    # u = torch.cholesky(LAM, upper=True)
    # laminv = torch.potri(u, upper=True)
    # dtheta = torch.bmm(laminv, R)

    return dtheta.view(self.batch_size, self.num_traj_states, self.state_dim)
//  [python]
def solve_linear_system_batch(self, A, b, K, delta = 0.0):
    I = torch.eye(self.N, self.N, device=self.device).unsqueeze(0).repeat(self.batch_size,1,1)
    # K = K + 1e-8 * torch.eye(K.shape[1], K.shape[2], device=self.device).repeat(self.batch_size,1,1)
    A_t_K = torch.bmm(A.transpose(1,2), K)
    A_t_A = torch.bmm(A_t_K, A)
    LAM = A_t_A + delta*I
    R = torch.bmm(A_t_K, b)
    #Solve using cholesky
    # l = torch.linalg.cholesky(LAM, upper=False)
    # z = torch.linalg.solve_triangular(R, l, upper=False, transpose=False)[0]
    # dtheta = torch.linalg.solve_triangular(z, l, upper=False, transpose=True)[0]

    u = torch.linalg.cholesky(LAM, upper=True)
    z = torch.bmm(torch.linalg.inv(u.transpose(1,2)), R)
    dtheta = torch.bmm(torch.linalg.inv(u), z) 

    # u = torch.linalg.cholesky(LAM, upper=True)
    # laminv = torch.linalg.cholesky_inverse(u, upper=True)
    # dtheta = torch.bmm(laminv, R)

    return dtheta.view(self.batch_size, self.num_traj_states, self.state_dim)
//  [python]
def chol_scale_invert_kron_factor(factor, prior_scale, data_scale, upper=False):
    scaled_factor = data_scale * factor + prior_scale * torch.eye(factor.shape[0]).type(factor.type())
    inv_factor = torch.inverse(scaled_factor)
    chol_inv_factor = torch.cholesky(inv_factor, upper=upper)
    return chol_inv_factor
//  [python]
def chol_scale_invert_kron_factor(factor, prior_scale, data_scale, upper=False):
    scaled_factor = data_scale * factor + prior_scale * torch.eye(factor.shape[0]).type(factor.type())
    inv_factor = torch.inverse(scaled_factor)
    chol_inv_factor = torch.linalg.cholesky(inv_factor, upper=upper)
    return chol_inv_factor
//  [python]
def test_discrete_mvn_log_prob(init_shape, trans_shape, obs_shape, state_dim):
    event_size = 4
    init_logits = torch.randn(init_shape + (state_dim,))
    trans_logits = torch.randn(trans_shape + (state_dim, state_dim))
    loc = torch.randn(obs_shape + (state_dim, event_size))
    cov = torch.randn(obs_shape + (state_dim, event_size, 2 * event_size))
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.cholesky(cov)
    obs_dist = dist.MultivariateNormal(loc, scale_tril=scale_tril)

    actual_dist = DiscreteHMM(init_logits, trans_logits, obs_dist)
    expected_dist = dist.DiscreteHMM(init_logits, trans_logits, obs_dist)
    assert actual_dist.event_shape == expected_dist.event_shape
    assert actual_dist.batch_shape == expected_dist.batch_shape

    batch_shape = broadcast_shape(init_shape + (1,), trans_shape, obs_shape)
    data = obs_dist.expand(batch_shape + (state_dim,)).sample()
    data = data[(slice(None),) * len(batch_shape) + (0,)]
    actual_log_prob = actual_dist.log_prob(data)
    expected_log_prob = expected_dist.log_prob(data)
    assert_close(actual_log_prob, expected_log_prob)
    check_expand(actual_dist, data)
//  [python]
def test_discrete_mvn_log_prob(init_shape, trans_shape, obs_shape, state_dim):
    event_size = 4
    init_logits = torch.randn(init_shape + (state_dim,))
    trans_logits = torch.randn(trans_shape + (state_dim, state_dim))
    loc = torch.randn(obs_shape + (state_dim, event_size))
    cov = torch.randn(obs_shape + (state_dim, event_size, 2 * event_size))
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.linalg.cholesky(cov)
    obs_dist = dist.MultivariateNormal(loc, scale_tril=scale_tril)

    actual_dist = DiscreteHMM(init_logits, trans_logits, obs_dist)
    expected_dist = dist.DiscreteHMM(init_logits, trans_logits, obs_dist)
    assert actual_dist.event_shape == expected_dist.event_shape
    assert actual_dist.batch_shape == expected_dist.batch_shape

    batch_shape = broadcast_shape(init_shape + (1,), trans_shape, obs_shape)
    data = obs_dist.expand(batch_shape + (state_dim,)).sample()
    data = data[(slice(None),) * len(batch_shape) + (0,)]
    actual_log_prob = actual_dist.log_prob(data)
    expected_log_prob = expected_dist.log_prob(data)
    assert_close(actual_log_prob, expected_log_prob)
    check_expand(actual_dist, data)
//  [python]
def optimizer_step(g, H, lambda_=0, mute=False, mask=None, eps=1e-6):
    """One optimization step with Gauss-Newton or Levenberg-Marquardt.

    Args:
        g: batched gradient tensor of size (..., N).
        H: batched hessian tensor of size (..., N, N).
        lambda_: damping factor for LM (use GN if lambda_=0).
        mask: denotes valid elements of the batch (optional).
    """
    if lambda_ is 0:  # noqa
        diag = torch.zeros_like(g)
    else:
        diag = H.diagonal(dim1=-2, dim2=-1) * lambda_
    H = H + diag.clamp(min=eps).diag_embed()

    if mask is not None:
        # make sure that masked elements are not singular
        H = torch.where(mask[..., None, None], H, torch.eye(H.shape[-1]).to(H))
        # set g to 0 to delta is 0 for masked elements
        g = g.masked_fill(~mask[..., None], 0.0)

    H_, g_ = H.cpu(), g.cpu()
    try:
        U = torch.cholesky(H_, upper=True)
    except RuntimeError as e:
        if 'singular U' in str(e):
            if not mute:
                logger.debug('Cholesky decomposition failed, fallback to LU.')
            delta = -torch.solve(g_[..., None], H_)[0][..., 0]
        else:
            raise
    else:
        delta = -torch.cholesky_solve(g_[..., None], U, upper=True)[..., 0]

    return delta.to(H.device)
//  [python]
def optimizer_step(g, H, lambda_=0, mute=False, mask=None, eps=1e-6):
    """One optimization step with Gauss-Newton or Levenberg-Marquardt.

    Args:
        g: batched gradient tensor of size (..., N).
        H: batched hessian tensor of size (..., N, N).
        lambda_: damping factor for LM (use GN if lambda_=0).
        mask: denotes valid elements of the batch (optional).
    """
    if lambda_ is 0:  # noqa
        diag = torch.zeros_like(g)
    else:
        diag = H.diagonal(dim1=-2, dim2=-1) * lambda_
    H = H + diag.clamp(min=eps).diag_embed()

    if mask is not None:
        # make sure that masked elements are not singular
        H = torch.where(mask[..., None, None], H, torch.eye(H.shape[-1]).to(H))
        # set g to 0 to delta is 0 for masked elements
        g = g.masked_fill(~mask[..., None], 0.0)

    H_, g_ = H.cpu(), g.cpu()
    try:
        U = torch.linalg.cholesky(H_)
    except RuntimeError as e:
        if 'singular U' in str(e):
            if not mute:
                logger.debug('Cholesky decomposition failed, fallback to LU.')
            delta = -torch.solve(g_[..., None], H_)[0][..., 0]
        else:
            raise
    else:
        delta = -torch.linalg.solve_triangular(U, g_[..., None], upper=False)[..., 0]

    return delta.to(H.device)