jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def low_rank_factor(K, Y, logdet=False, solve=False, jit=1e-6, jitbase=2):
    """
    Inputs: Y, K
    K: a symmetric positive definite (covariance) matrix,
       this will be factored  as K = L @ L.t()
    Y: 1D or 2D
    Returns: Q, logdet, ridge
    ------------------------------------------------------
    The following equality holds:
    Q.t() @ Q = Y.t() @ K.inverse() @ Y
    """
    L, ridge = jitcholesky(K, jit=jit, jitbase=jitbase)
    if len(Y.size()) == 1:
        _1d, _Y = True, Y.view(-1, 1)
    else:
        _1d, _Y = False, Y
    if solve:
        Q, _ = torch.triangular_solve(_Y, L, upper=False)
    else:
        Q = torch.mm(L.inverse(), _Y)
    if logdet:
        ld = 2 * L.diag().log().sum()
    else:
        ld = None
    return Q, ld, ridge
//  [python]
def low_rank_factor(K, Y, logdet=False, solve=False, jit=1e-6, jitbase=2):
    """
    Inputs: Y, K
    K: a symmetric positive definite (covariance) matrix,
       this will be factored  as K = L @ L.t()
    Y: 1D or 2D
    Returns: Q, logdet, ridge
    ------------------------------------------------------
    The following equality holds:
    Q.t() @ Q = Y.t() @ K.inverse() @ Y
    """
    L, ridge = jitcholesky(K, jit=jit, jitbase=jitbase)
    if len(Y.size()) == 1:
        _1d, _Y = True, Y.view(-1, 1)
    else:
        _1d, _Y = False, Y
    if solve:
        Q = torch.linalg.solve_triangular(L, _Y, upper=False)
    else:
        Q = torch.mm(L.inverse(), _Y)
    if logdet:
        ld = 2 * L.diag().log().sum()
    else:
        ld = None
    return Q, ld, ridge
//  [python]
def pca_forward(self, x):
        """normalized tensor"""
        batch_size, feature_dim = x.size()
        f_cov = torch.mm(x.transpose(0, 1), x) / batch_size # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        f_cov_shrink = (1 - self.eps) * f_cov + self.eps * eye
        inv_sqrt = torch.triangular_solve(eye, torch.linalg.cholesky(f_cov_shrink.float()), upper=False)[0]
        inv_sqrt = inv_sqrt.contiguous().view(feature_dim, feature_dim).detach()
        return torch.mm(inv_sqrt, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def pca_forward(self, x):
        """normalized tensor"""
        batch_size, feature_dim = x.size()
        f_cov = torch.mm(x.transpose(0, 1), x) / batch_size # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        f_cov_shrink = (1 - self.eps) * f_cov + self.eps * eye
        inv_sqrt = torch.linalg.solve_triangular(torch.linalg.cholesky(f_cov_shrink.float()), eye, upper=False)
        inv_sqrt = inv_sqrt.contiguous().view(feature_dim, feature_dim).detach()
        return torch.mm(inv_sqrt, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def test_triangular_solve_out_errors_and_warnings(self, device, dtype):
        # dtypes should be safely castable
        a = torch.eye(2, dtype=dtype, device=device)
        b = torch.randn(2, 1, dtype=dtype, device=device)
        out = torch.empty_like(b).to(torch.int)
        clone_a = torch.empty_like(a)
        with self.assertRaisesRegex(RuntimeError, "Expected out tensor to have dtype"):
            torch.triangular_solve(b, a, out=(out, clone_a))

        out = torch.empty_like(b)
        clone_a = clone_a.to(torch.int)
        with self.assertRaisesRegex(RuntimeError, "Expected out tensor to have dtype"):
            torch.triangular_solve(b, a, out=(out, clone_a))

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out = torch.empty(0, dtype=dtype, device=wrong_device)
            clone_a = torch.empty_like(a)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.triangular_solve(b, a, out=(out, clone_a))
            out = torch.empty(0, dtype=dtype, device=device)
            clone_a = torch.empty_like(a).to(wrong_device)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.triangular_solve(b, a, out=(out, clone_a))

        # Trigger the WARN_ONCE deprecation error
        torch.triangular_solve(b, a)

        # if out tensor with wrong shape is passed a warning is given
        with warnings.catch_warnings(record=True) as w:
            out = torch.empty(1, dtype=dtype, device=device)
            clone_a = torch.empty(1, dtype=dtype, device=device)
            # Trigger warning
            torch.triangular_solve(b, a, out=(out, clone_a))
            # Check warning occurs
            self.assertEqual(len(w), 2)
            self.assertTrue("An output with one or more elements was resized" in str(w[0].message))
            self.assertTrue("An output with one or more elements was resized" in str(w[1].message))
//  [python]
def test_triangular_solve_out_errors_and_warnings(self, device, dtype):
        # dtypes should be safely castable
        a = torch.eye(2, dtype=dtype, device=device)
        b = torch.randn(2, 1, dtype=dtype, device=device)
        out = torch.empty_like(b).to(torch.int)
        clone_a = torch.empty_like(a)
        with self.assertRaisesRegex(RuntimeError, "Expected out tensor to have dtype"):
            torch.linalg.solve_triangular(b, a, out=(out, clone_a), upper=False, unitriangular=False)

        out = torch.empty_like(b)
        clone_a = clone_a.to(torch.int)
        with self.assertRaisesRegex(RuntimeError, "Expected out tensor to have dtype"):
            torch.linalg.solve_triangular(b, a, out=(out, clone_a), upper=False, unitriangular=False)

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out = torch.empty(0, dtype=dtype, device=wrong_device)
            clone_a = torch.empty_like(a)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.linalg.solve_triangular(b, a, out=(out, clone_a), upper=False, unitriangular=False)
            out = torch.empty(0, dtype=dtype, device=device)
            clone_a = torch.empty_like(a).to(wrong_device)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.linalg.solve_triangular(b, a, out=(out, clone_a), upper=False, unitriangular=False)

        # Trigger the WARN_ONCE deprecation error
        torch.linalg.solve_triangular(b, a, upper=False, unitriangular=False)

        # if out tensor with wrong shape is passed a warning is given
        with warnings.catch_warnings(record=True) as w:
            out = torch.empty(1, dtype=dtype, device=device)
            clone_a = torch.empty(1, dtype=dtype, device=device)
            # Trigger warning
            torch.linalg.solve_triangular(b, a, out=(out, clone_a), upper=False, unitriangular=False)
            # Check warning occurs
            self.assertEqual(len(w), 2)
            self.assertTrue("An output with one or more elements was resized" in str(w[0].message))
            self.assertTrue("An output with one or more elements was resized" in str(w[1].message))
//  [python]
def gp_cond(u, Kzz, Kzx, Kxx, Lz=None, Lz_Kzx=None):
  '''
  Compute the GP predictive conditional
  p(f|u,x,z)

  μ = Kxz Kzz_inv u
  Σ = Kxx - Kxz Kzz_inv Kzx

  Arguments:
    m: ... x M x 1
    Kzz: ... x M x M
    Kzx: ... x M x N
    Kxx: ... x N x N

  Returns:
    μ: ... x N x 1
    Σ: ... x N x N
  '''
  if Lz is None:
    Lz = cholesky(Kzz)

  Lz_u, _ = torch.triangular_solve(u, Lz, upper=False)

  if Lz_Kzx is None:
    Lz_Kzx, _ = torch.triangular_solve(Kzx, Lz, upper=False)

  μ = torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_u)

  Σ = Kxx - torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_Kzx)

  return μ, Σ
//  [python]
def gp_cond(u, Kzz, Kzx, Kxx, Lz=None, Lz_Kzx=None):
  '''
  Compute the GP predictive conditional
  p(f|u,x,z)

  μ = Kxz Kzz_inv u
  Σ = Kxx - Kxz Kzz_inv Kzx

  Arguments:
    m: ... x M x 1
    Kzz: ... x M x M
    Kzx: ... x M x N
    Kxx: ... x N x N

  Returns:
    μ: ... x N x 1
    Σ: ... x N x N
  '''
  if Lz is None:
    Lz = cholesky(Kzz)

  Lz_u = torch.linalg.solve_triangular(Lz, u, upper=False)

  if Lz_Kzx is None:
    Lz_Kzx = torch.linalg.solve_triangular(Lz, Kzx, upper=False)

  μ = torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_u)

  Σ = Kxx - torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_Kzx)

  return μ, Σ
//  [python]
def cholseky_inverse(tril: torch.Tensor) -> torch.Tensor:
    r"""Computes batched :math:`(LL^\top)^{-1}`, assuming :math:`L` is a lower-triangular Cholesky
    factor.
    """
    b = eye_like(tril)
    x = torch.triangular_solve(b, tril, upper=False)[0]
    return torch.triangular_solve(x, tril, upper=False, transpose=True)[0]
//  [python]
def cholseky_inverse(tril: torch.Tensor) -> torch.Tensor:
    r"""Computes batched :math:`(LL^\top)^{-1}`, assuming :math:`L` is a lower-triangular Cholesky
    factor.
    """
    b = eye_like(tril)
    x = torch.linalg.solve_triangular(b, tril, upper=False)
    return torch.linalg.solve_triangular(x, tril, upper=False, transpose=True)
//  [python]
def forward(self, inputs, targets=None, train_x=None):
        if train_x is None:
            train_x = inputs

        # Draw samples from the first function
        hidden_covar = rbf_kernel(train_x, train_x, self.ls1, self.os1).add(self.eye.mul(1e-4))
        hidden_chol_covar = psd_safe_cholesky(hidden_covar)
        with pyro.plate("hidden_fns_plate", self.width, dim=-1):
            hidden_dist = pyro.distributions.Normal(self.mean, self.ones).to_event(1)
            hidden_fns = pyro.sample("hidden_fns", hidden_dist).transpose(-1, -2)
            unwhitened_hidden_fns = hidden_chol_covar @ hidden_fns

        final_covar = add_2l_kernel(unwhitened_hidden_fns, unwhitened_hidden_fns, self.ls2, self.os2, self.ls3, self.os3).add(self.eye.mul(1e-4))
        final_chol_covar = psd_safe_cholesky(final_covar)
        with pyro.plate("output_fns_plate", 1, dim=-1):
            final_dist = pyro.distributions.Normal(self.mean, self.ones).to_event(1)
            final_fns = pyro.sample("final_fns", final_dist).transpose(-1, -2)
            unwhitened_final_fns = final_chol_covar @ final_fns

        if torch.equal(train_x, inputs):
            latent_vals = unwhitened_final_fns.squeeze(-1)
        else:
            hidden_cross_covar = rbf_kernel(train_x, inputs, self.ls1, self.os1)
            hidden_interp_term = torch.triangular_solve(hidden_cross_covar, hidden_chol_covar, upper=False)[0]
            pred_hidden_fns = hidden_interp_term.transpose(-1, -2) @ hidden_fns

            final_cross_covar = add_2l_kernel(unwhitened_hidden_fns, pred_hidden_fns, self.ls2, self.os2, self.ls3, self.os3)
            final_interp_term = torch.triangular_solve(final_cross_covar, final_chol_covar, upper=False)[0]
            latent_vals = (final_interp_term.transpose(-1, -2) @ final_fns).squeeze(-1)

        # Observations
        with pyro.plate("output_plate", inputs.size(-2), dim=-1):
            output_dist = pyro.distributions.Normal(latent_vals.squeeze(-1), self.noise.sqrt())
            outputs = pyro.sample("outputs", output_dist, obs=targets)

        return output_dist
//  [python]
def forward(self, inputs, targets=None, train_x=None):
        if train_x is None:
            train_x = inputs

        # Draw samples from the first function
        hidden_covar = rbf_kernel(train_x, train_x, self.ls1, self.os1).add(self.eye.mul(1e-4))
        hidden_chol_covar = psd_safe_cholesky(hidden_covar)
        with pyro.plate("hidden_fns_plate", self.width, dim=-1):
            hidden_dist = pyro.distributions.Normal(self.mean, self.ones).to_event(1)
            hidden_fns = pyro.sample("hidden_fns", hidden_dist).transpose(-1, -2)
            unwhitened_hidden_fns = hidden_chol_covar @ hidden_fns

        final_covar = add_2l_kernel(unwhitened_hidden_fns, unwhitened_hidden_fns, self.ls2, self.os2, self.ls3, self.os3).add(self.eye.mul(1e-4))
        final_chol_covar = psd_safe_cholesky(final_covar)
        with pyro.plate("output_fns_plate", 1, dim=-1):
            final_dist = pyro.distributions.Normal(self.mean, self.ones).to_event(1)
            final_fns = pyro.sample("final_fns", final_dist).transpose(-1, -2)
            unwhitened_final_fns = final_chol_covar @ final_fns

        if torch.equal(train_x, inputs):
            latent_vals = unwhitened_final_fns.squeeze(-1)
        else:
            hidden_cross_covar = rbf_kernel(train_x, inputs, self.ls1, self.os1)
            hidden_interp_term = torch.linalg.solve_triangular(hidden_cross_covar, hidden_chol_covar, upper=False)
            pred_hidden_fns = hidden_interp_term.transpose(-1, -2) @ hidden_fns

            final_cross_covar = add_2l_kernel(unwhitened_hidden_fns, pred_hidden_fns, self.ls2, self.os2, self.ls3, self.os3)
            final_interp_term = torch.linalg.solve_triangular(final_cross_covar, final_chol_covar, upper=False)
            latent_vals = (final_interp_term.transpose(-1, -2) @ final_fns).squeeze(-1)

        # Observations
        with pyro.plate("output_plate", inputs.size(-2), dim=-1):
            output_dist = pyro.distributions.Normal(latent_vals.squeeze(-1), self.noise.sqrt())
            outputs = pyro.sample("outputs", output_dist, obs=targets)

        return output_dist
//  [python]
def negative_log_likelihood(self, hyper=None):
		if hyper is not None:
			param_original = self.model.param_to_vec()
			self.cholesky_update(hyper)
		# cholesky is lower triangular matrix
		mean_vec_sol = torch.triangular_solve(self.mean_vec, self.cholesky, upper=False)[0]
		nll = 0.5 * torch.sum(mean_vec_sol ** 2) + torch.sum(torch.log(torch.diag(self.cholesky))) + 0.5 * self.train_y.size(0) * np.log(2 * np.pi)
		if hyper is not None:
			self.cholesky_update(param_original)
		return nll
//  [python]
def negative_log_likelihood(self, hyper=None):
		if hyper is not None:
			param_original = self.model.param_to_vec()
			self.cholesky_update(hyper)
		# cholesky is lower triangular matrix
		mean_vec_sol = torch.linalg.solve_triangular(self.cholesky, self.mean_vec, upper=False)
		nll = 0.5 * torch.sum(mean_vec_sol ** 2) + torch.sum(torch.log(torch.diag(self.cholesky))) + 0.5 * self.train_y.size(0) * np.log(2 * np.pi)
		if hyper is not None:
			self.cholesky_update(param_original)
		return nll
//  [python]
def linear_marginal_diag(m, S, Kzz, Kzx, Kxx_diag, cache=None):
  '''
  Compute the diagonal of the marginal of
  product of Gaussian densities of the form
  p(z,x) = N(z; m, S)N(y; Az, V) where A = Kxz Kzz_inv.

  This function combines `gp_cond`, `linear_joint`
  and diagonalization into a single efficient function.

  μ = Am + b
  Σ = V + ASA^T

  Arguments:
    m: ... x M x 1
    S: ... x M x M
    Kzx: ... x M x N
    Kzz: ... x M x M
    Kxx_diag: ... x N x N

  Returns:
    μ: ... x N
    Σ: ... x N
  '''

  Lz = cholesky(Kzz)
  Lz_m, _ = torch.triangular_solve(m, Lz, upper=False)
  Lz_Kzx, _ = torch.triangular_solve(Kzx, Lz, upper=False)

  μ = torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_m).squeeze(-1)

  diag1 = Lz_Kzx.pow(2).sum(dim=-2)

  Lz_LS, _ = torch.triangular_solve(cholesky(S), Lz, upper=False)

  diag2 = torch.einsum('...ij,...ik->...jk', Lz_LS, Lz_Kzx).pow(2).sum(dim=-2)

  Σ = Kxx_diag - diag1 + diag2

  if isinstance(cache, dict):
    cache.update(dict(Lz=Lz, Lz_Kzx=Lz_Kzx))

  return μ, Σ
//  [python]
def linear_marginal_diag(m, S, Kzz, Kzx, Kxx_diag, cache=None):
  '''
  Compute the diagonal of the marginal of
  product of Gaussian densities of the form
  p(z,x) = N(z; m, S)N(y; Az, V) where A = Kxz Kzz_inv.

  This function combines `gp_cond`, `linear_joint`
  and diagonalization into a single efficient function.

  μ = Am + b
  Σ = V + ASA^T

  Arguments:
    m: ... x M x 1
    S: ... x M x M
    Kzx: ... x M x N
    Kzz: ... x M x M
    Kxx_diag: ... x N x N

  Returns:
    μ: ... x N
    Σ: ... x N
  '''

  Lz = cholesky(Kzz)
  Lz_m = torch.linalg.solve_triangular(Lz, m, upper=False)
  Lz_Kzx = torch.linalg.solve_triangular(Lz, Kzx, upper=False)

  μ = torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_m).squeeze(-1)

  diag1 = Lz_Kzx.pow(2).sum(dim=-2)

  Lz_LS = torch.linalg.solve_triangular(Lz, cholesky(S), upper=False)

  diag2 = torch.einsum('...ij,...ik->...jk', Lz_LS, Lz_Kzx).pow(2).sum(dim=-2)

  Σ = Kxx_diag - diag1 + diag2

  if isinstance(cache, dict):
    cache.update(dict(Lz=Lz, Lz_Kzx=Lz_Kzx))

  return μ, Σ
//  [python]
def linear_joint(m, S, Kzx, Kzz, V, b, cache=None):
  '''
  Compute product of Gaussian densities of the form
  p(z,x) = N(z; m, S)N(x; Az + b, V) where A = Kxz Kzz_inv.

  Results in N([z,x]; μ, Σ) where
  μ = [m, Am + b]
  Σ = [S, SA^T; AS, V + ASA^T]

  Arguments:
    m: ... x M x 1
    S: ... x M x M
    Kzx: ... x M x N
    Kzz: ... x M x M
    b: ... x N x 1
    V: ... x N x N

  Returns:
    μ: ... x (M + N) x 1
    Σ: ... x (M + N) x (M + N)
  '''

  Lz = cholesky(Kzz)
  Lz_m, _ = torch.triangular_solve(m, Lz, upper=False)
  Lz_Kzx, _ = torch.triangular_solve(Kzx, Lz, upper=False)

  Am = torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_m)

  Lz_S, _ = torch.triangular_solve(S, Lz, upper=False)

  AS = torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_S)
  SAt = torch.einsum('...ij->...ji', AS)

  Lz_SAt, _ = torch.triangular_solve(SAt, Lz, upper=False)

  ASAt = torch.einsum('...ij,...ik->...jk', Lz_SAt, Lz_Kzx)

  μ = torch.cat([m, Am + b], dim=-2)
  Σ = torch.cat([
    torch.cat([S, SAt], dim=-1),
    torch.cat([AS, V + ASAt], dim=-1)
  ], dim=-2)

  if isinstance(cache, dict):
    cache.update(dict(Lz_Kzx=Lz_Kzx, Lz=Lz))

  return μ, Σ
//  [python]
def linear_joint(m, S, Kzx, Kzz, V, b, cache=None):
  '''
  Compute product of Gaussian densities of the form
  p(z,x) = N(z; m, S)N(x; Az + b, V) where A = Kxz Kzz_inv.

  Results in N([z,x]; μ, Σ) where
  μ = [m, Am + b]
  Σ = [S, SA^T; AS, V + ASA^T]

  Arguments:
    m: ... x M x 1
    S: ... x M x M
    Kzx: ... x M x N
    Kzz: ... x M x M
    b: ... x N x 1
    V: ... x N x N

  Returns:
    μ: ... x (M + N) x 1
    Σ: ... x (M + N) x (M + N)
  '''

  Lz = cholesky(Kzz)
  Lz_m = torch.linalg.solve_triangular(Lz, m, upper=False)
  Lz_Kzx = torch.linalg.solve_triangular(Lz, Kzx, upper=False)

  Am = torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_m)

  Lz_S = torch.linalg.solve_triangular(Lz, S, upper=False)

  AS = torch.einsum('...ij,...ik->...jk', Lz_Kzx, Lz_S)
  SAt = torch.einsum('...ij->...ji', AS)

  Lz_SAt = torch.linalg.solve_triangular(Lz, SAt, upper=False)

  ASAt = torch.einsum('...ij,...ik->...jk', Lz_SAt, Lz_Kzx)

  μ = torch.cat([m, Am + b], dim=-2)
  Σ = torch.cat([
    torch.cat([S, SAt], dim=-1),
    torch.cat([AS, V + ASAt], dim=-1)
  ], dim=-2)

  if isinstance(cache, dict):
    cache.update(dict(Lz_Kzx=Lz_Kzx, Lz=Lz))

  return μ, Σ
//  [python]
def predict(self, pred_x, hyper=None, verbose=False):
		if hyper is not None:
			param_original = self.model.param_to_vec()
			self.cholesky_update(hyper)

		k_pred_train = self.model.kernel(pred_x, self.train_x)
		k_pred = self.model.kernel(pred_x, diagonal=True)
		#print(k_pred_train)
		#print(self.mean_vec)
		# cholesky is lower triangular matrix
		chol_solver = torch.triangular_solve(torch.cat([k_pred_train.t(), self.mean_vec], 1), self.cholesky, upper=False)[0]
		chol_solve_k = chol_solver[:, :-1]
		chol_solve_y = chol_solver[:, -1:]

		pred_mean = torch.mm(chol_solve_k.t(), chol_solve_y) + self.model.mean(pred_x)
		pred_quad = (chol_solve_k ** 2).sum(0).view(-1, 1)
		pred_var = k_pred - pred_quad

		if verbose:
			numerically_stable = (pred_var >= 0).all()
			zero_pred_var = (pred_var <= 0).all()

		if hyper is not None:
			self.cholesky_update(param_original)

		if verbose:
			return pred_mean, pred_var.clamp(min=1e-8), numerically_stable, zero_pred_var
		else:
			return pred_mean, pred_var.clamp(min=1e-8)
//  [python]
def predict(self, pred_x, hyper=None, verbose=False):
		if hyper is not None:
			param_original = self.model.param_to_vec()
			self.cholesky_update(hyper)

		k_pred_train = self.model.kernel(pred_x, self.train_x)
		k_pred = self.model.kernel(pred_x, diagonal=True)
		#print(k_pred_train)
		#print(self.mean_vec)
		# cholesky is lower triangular matrix
		chol_solver = torch.linalg.solve_triangular(torch.cat([k_pred_train.t(), self.mean_vec], 1), self.cholesky, upper=False)
		chol_solve_k = chol_solver[:, :-1]
		chol_solve_y = chol_solver[:, -1:]

		pred_mean = torch.mm(chol_solve_k.t(), chol_solve_y) + self.model.mean(pred_x)
		pred_quad = (chol_solve_k ** 2).sum(0).view(-1, 1)
		pred_var = k_pred - pred_quad

		if verbose:
			numerically_stable = (pred_var >= 0).all()
			zero_pred_var = (pred_var <= 0).all()

		if hyper is not None:
			self.cholesky_update(param_original)

		if verbose:
			return pred_mean, pred_var.clamp(min=1e-8), numerically_stable, zero_pred_var
		else:
			return pred_mean, pred_var.clamp(min=1e-8)
//  [python]
def forward(self, input):
        "Compute covariance"
        if self.demean:
            if self.mu is None:
                temp = input - torch.mean(input, 0)
            else:
                temp = input - self.mu
        else:
            temp = input

        if self.R is None:
            cov = temp.transpose(1, 0) @ temp / temp.shape[0]  # compute covariance matrix
            cov = (cov + cov.transpose(1, 0)) / 2 + 1e-5 * torch.eye(cov.shape[0], device=self.device)
            R = torch.cholesky(cov)  # returns the lower cholesky matrix
        else:
            R = self.R
        Y, _ = torch.triangular_solve(temp.transpose(1, 0), R, upper=False)
        return Y.transpose(1, 0)
//  [python]
def forward(self, input):
        "Compute covariance"
        if self.demean:
            if self.mu is None:
                temp = input - torch.mean(input, 0)
            else:
                temp = input - self.mu
        else:
            temp = input

        if self.R is None:
            cov = temp.transpose(1, 0) @ temp / temp.shape[0]  # compute covariance matrix
            cov = (cov + cov.transpose(1, 0)) / 2 + 1e-5 * torch.eye(cov.shape[0], device=self.device)
            R = torch.cholesky(cov)  # returns the lower cholesky matrix
        else:
            R = self.R
        Y = torch.linalg.solve_triangular(R, temp.transpose(1, 0), upper=False)
        return Y.transpose(1, 0)
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N + KDN)
            logabsdet = O(D)
        where:
            K = num of householder transforms
            D = num of features
            N = num of inputs
        """
        upper = self._create_upper()
        outputs = inputs - self.bias
        outputs, _ = self.orthogonal.inverse(
            outputs
        )  # Ignore logabsdet since we know it's zero.
        outputs, _ = torch.triangular_solve(outputs.t(), upper, upper=True)
        outputs = outputs.t()
        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * torch.ones(outputs.shape[0])
        return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N + KDN)
            logabsdet = O(D)
        where:
            K = num of householder transforms
            D = num of features
            N = num of inputs
        """
        upper = self._create_upper()
        outputs = inputs - self.bias
        outputs, _ = self.orthogonal.inverse(
            outputs
        )  # Ignore logabsdet since we know it's zero.
        outputs, _ = torch.linalg.solve_triangular(outputs.t(), upper, upper=True)
        outputs = outputs.t()
        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * torch.ones(outputs.shape[0])
        return outputs, logabsdet
//  [python]
def tril_inverse(tril_matrix: torch.Tensor) -> torch.Tensor:
    """Invert a lower-triangular matrix.

    Args:
        tril_matrix (torch.Tensor): Lower-triangular matrix to invert. Shape should be
            `(*, matrix_dim, matrix_dim)`.

    Returns:
        torch.Tensor: Inverted matrix. Shape should be `(*, matrix_dim, matrix_dim)`.
    """
    assert tril_matrix.shape[-1] == tril_matrix.shape[-2], "Input must be square!"
    matrix_dim = tril_matrix.shape[-1]
    batch_dims = tril_matrix.shape[:-2]

    # Get an identity matrix
    identity = torch.eye(tril_matrix.shape[-1], device=tril_matrix.device)

    # View identity w/ batch dimensions
    identity = identity.reshape(
        (1,) * len(batch_dims) + (matrix_dim, matrix_dim)
    ).expand(tril_matrix.shape)

    # Invert with triangular solve
    inverse = torch.triangular_solve(identity, tril_matrix, upper=False).solution
    assert inverse.shape == tril_matrix.shape
    return inverse
//  [python]
def tril_inverse(tril_matrix: torch.Tensor) -> torch.Tensor:
    """Invert a lower-triangular matrix.

    Args:
        tril_matrix (torch.Tensor): Lower-triangular matrix to invert. Shape should be
            `(*, matrix_dim, matrix_dim)`.

    Returns:
        torch.Tensor: Inverted matrix. Shape should be `(*, matrix_dim, matrix_dim)`.
    """
    assert tril_matrix.shape[-1] == tril_matrix.shape[-2], "Input must be square!"
    matrix_dim = tril_matrix.shape[-1]
    batch_dims = tril_matrix.shape[:-2]

    # Get an identity matrix
    identity = torch.eye(tril_matrix.shape[-1], device=tril_matrix.device)

    # View identity w/ batch dimensions
    identity = identity.reshape(
        (1,) * len(batch_dims) + (matrix_dim, matrix_dim)
    ).expand(tril_matrix.shape)

    # Invert with triangular solve
    inverse = torch.linalg.solve_triangular(identity, tril_matrix, upper=False)
    assert inverse.shape == tril_matrix.shape
    return inverse
//  [python]
def logp(self, vals: edic, conds: edic=edic()) -> tc.Tensor: # log_probs
        # [shape_bat, shape_var], [shape_bat, shape_cond] -> [shape_bat]
        meanval, std_trilval = self._meanfn(conds), self._std_trilfn(conds)
        centered_vals_ = (vals[self.name] - meanval).unsqueeze(-1)
        normalized_vals = tc.triangular_solve(centered_vals_, std_trilval, upper=False)[0].squeeze(-1)
        quads = reduce_last(tc.sum, normalized_vals ** 2, len(self.shape))
        half_log_det = reduce_last(tc.sum, std_trilval.diagonal(dim1=-2, dim2=-1).log(), len(self.shape))
        return -.5 * quads - half_log_det - self._log_const
//  [python]
def logp(self, vals: edic, conds: edic=edic()) -> tc.Tensor: # log_probs
        # [shape_bat, shape_var], [shape_bat, shape_cond] -> [shape_bat]
        meanval, std_trilval = self._meanfn(conds), self._std_trilfn(conds)
        centered_vals_ = (vals[self.name] - meanval).unsqueeze(-1)
        normalized_vals = tc.linalg.solve_triangular(centered_vals_, std_trilval, upper=False).squeeze(-1)
        quads = reduce_last(tc.sum, normalized_vals ** 2, len(self.shape))
        half_log_det = reduce_last(tc.sum, std_trilval.diagonal(dim1=-2, dim2=-1).log(), len(self.shape))
        return -.5 * quads - half_log_det - self._log_const
//  [python]
def inverse(self, y, **kwargs):
        L, U = self.get_LU()
        y = torch.triangular_solve(self.P.T @ (y - self.bias).unsqueeze(-1), L, upper=False)[0]
        x = torch.triangular_solve(y, U, upper=True)[0].squeeze(-1)
        ljd = -self.log_diag.expand_as(x)
        return x, ljd
//  [python]
def inverse(self, y, **kwargs):
        L, U = self.get_LU()
        y = torch.linalg.solve_triangular(self.P.T @ (y - self.bias).unsqueeze(-1), L, upper=False)
        x = torch.linalg.solve_triangular(y, U, upper=True).squeeze(-1)
        ljd = -self.log_diag.expand_as(x)
        return x, ljd
//  [python]
def test_triangular_solve(self, device, dtype):
        for (k, n), (upper, unitriangular, transpose) in product(zip([2, 3, 5], [3, 5, 7]),
                                                                 product([True, False], repeat=3)):
            b, A = self.triangular_solve_test_helper((n, n), (n, k), upper,
                                                     unitriangular, device, dtype)
            x = torch.triangular_solve(b, A, upper=upper, unitriangular=unitriangular, transpose=transpose)[0]
            if transpose:
                self.assertLessEqual(b.dist(A.t().mm(x)), 4e-12)
            else:
                self.assertLessEqual(b.dist(A.mm(x)), 4e-12)
//  [python]
def test_triangular_solve(self, device, dtype):
        for (k, n), (upper, unitriangular, transpose) in product(zip([2, 3, 5], [3, 5, 7]),
                                                                 product([True, False], repeat=3)):
            b, A = self.triangular_solve_test_helper((n, n), (n, k), upper,
                                                     unitriangular, device, dtype)
            x = torch.linalg.solve_triangular(b, A, upper=upper, unitriangular=unitriangular, transpose=transpose)
            if transpose:
                self.assertLessEqual(b.dist(A.t().mm(x)), 4e-12)
            else:
                self.assertLessEqual(b.dist(A.mm(x)), 4e-12)
//  [python]
def forward(self, x: torch.Tensor) ->torch.Tensor:
        """Performs whitening using the Cholesky decomposition.

        Args:
            x (torch.Tensor): a batch or slice of projected features.

        Returns:
            torch.Tensor: a batch or slice of whitened features.
        """
        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.output_dim, -1).mean(-1).view(1, -1, 1, 1)
        xn = x - m
        T = xn.permute(1, 0, 2, 3).contiguous().view(self.output_dim, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)
        eye = torch.eye(self.output_dim).type(f_cov.type())
        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye
        inv_sqrt = torch.triangular_solve(eye, torch.linalg.cholesky(f_cov_shrinked), upper=False)[0]
        inv_sqrt = inv_sqrt.contiguous().view(self.output_dim, self.output_dim, 1, 1)
        decorrelated = conv2d(xn, inv_sqrt)
        return decorrelated.squeeze(2).squeeze(2)
//  [python]
def forward(self, x: torch.Tensor) ->torch.Tensor:
        """Performs whitening using the Cholesky decomposition.

        Args:
            x (torch.Tensor): a batch or slice of projected features.

        Returns:
            torch.Tensor: a batch or slice of whitened features.
        """
        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.output_dim, -1).mean(-1).view(1, -1, 1, 1)
        xn = x - m
        T = xn.permute(1, 0, 2, 3).contiguous().view(self.output_dim, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)
        eye = torch.eye(self.output_dim).type(f_cov.type())
        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye
        inv_sqrt = torch.linalg.solve_triangular(torch.linalg.cholesky(f_cov_shrinked), eye, upper=False)
        inv_sqrt = inv_sqrt.contiguous().view(self.output_dim, self.output_dim, 1, 1)
        decorrelated = conv2d(xn, inv_sqrt)
        return decorrelated.squeeze(2).squeeze(2)
//  [python]
def test_triangular_solve_batched_many_batches(self, device, dtype):
        for upper, transpose, unitriangular in product([True, False], repeat=3):
            b, A = self.triangular_solve_test_helper((256, 256, 5, 5), (5, 1),
                                                     upper, unitriangular, device, dtype)
            x, _ = torch.triangular_solve(b, A,
                                          upper=upper, transpose=transpose, unitriangular=unitriangular)
            if transpose:
                A = A.transpose(-2, -1)
            self.assertEqual(torch.matmul(A, x), b.expand(A.shape[:-2] + (5, 1)))

            b, A = self.triangular_solve_test_helper((3, 3), (512, 512, 3, 1),
                                                     upper, unitriangular, device, dtype)
            x, _ = torch.triangular_solve(b, A, upper=upper, transpose=transpose,
                                          unitriangular=unitriangular)
            if transpose:
                A = A.transpose(-2, -1)
            self.assertEqual(torch.matmul(A, x), b)
//  [python]
def test_triangular_solve_batched_many_batches(self, device, dtype):
        for upper, transpose, unitriangular in product([True, False], repeat=3):
            b, A = self.triangular_solve_test_helper((256, 256, 5, 5), (5, 1),
                                                     upper, unitriangular, device, dtype)
            x = torch.linalg.solve_triangular(b, A, upper=upper, transpose=transpose, unitriangular=unitriangular)
            if transpose:
                A = A.transpose(-2, -1)
            self.assertEqual(torch.matmul(A, x), b.expand(A.shape[:-2] + (5, 1)))

            b, A = self.triangular_solve_test_helper((3, 3), (512, 512, 3, 1),
                                                     upper, unitriangular, device, dtype)
            x = torch.linalg.solve_triangular(b, A, upper=upper, transpose=transpose,
                                          unitriangular=unitriangular)
            if transpose:
                A = A.transpose(-2, -1)
            self.assertEqual(torch.matmul(A, x), b)
//  [python]
def predict(self, test_mean, train_test_cov, test_test_cov):
        if self.USE_LU:
            predictive_mean  = test_mean + ( train_test_cov.T @ self.alpha ).view(*test_mean.shape)
            predictive_covar = test_test_cov - train_test_cov.T @ torch.lu_solve( train_test_cov, *self.LU )
        else:
            v = torch.triangular_solve(train_test_cov, self.L, upper=False).solution
            predictive_mean  = test_mean + ( train_test_cov.T @ self.alpha ).view(*test_mean.shape)
            predictive_covar = test_test_cov - v.T @ v
        return MultitaskMultivariateNormal(predictive_mean, predictive_covar)
//  [python]
def predict(self, test_mean, train_test_cov, test_test_cov):
        if self.USE_LU:
            predictive_mean  = test_mean + ( train_test_cov.T @ self.alpha ).view(*test_mean.shape)
            predictive_covar = test_test_cov - train_test_cov.T @ torch.lu_solve( train_test_cov, *self.LU )
        else:
            v = torch.linalg.solve_triangular(train_test_cov, self.L, upper=False)
            predictive_mean  = test_mean + ( train_test_cov.T @ self.alpha ).view(*test_mean.shape)
            predictive_covar = test_test_cov - v.T @ v
        return MultitaskMultivariateNormal(predictive_mean, predictive_covar)
//  [python]
def forward(self, x: Tensor) -> Tensor:
        return torch.triangular_solve(
            (x - self.mu).t(), self.L, upper=False
        ).solution.t()
//  [python]
def forward(self, x: Tensor) -> Tensor:
        return torch.linalg.solve_triangular(
            (x - self.mu).t(), self.L, upper=False
        ).t()
//  [python]
def _update_precond_dense_scale(Ql, qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update (dense, scaling) preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql), where
    dX and dG have shape (M, N)
    Ql has shape (M, M)
    qr has shape (1, N)
    qr is the diagonal part of Qr
    dX is perturbation of (matrix) parameter
    dG is perturbation of (matrix) gradient
    step: update step size
    _tiny: an offset to avoid division by zero
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(qr)

    rho = torch.sqrt(max_l / max_r)
    Ql /= rho
    qr *= rho

    A = Ql.mm(dG * qr)
    Bt = torch.triangular_solve(dX / qr, Ql, upper=True, transpose=True)[0]

    grad1 = torch.triu(A.mm(A.t()) - Bt.mm(Bt.t()))
    grad2 = torch.sum(A * A, dim=0, keepdim=True) - torch.sum(Bt * Bt, dim=0, keepdim=True)

    step1 = step / (torch.max(torch.abs(grad1)) + _tiny)
    step2 = step / (torch.max(torch.abs(grad2)) + _tiny)

    return Ql - step1 * grad1.mm(Ql), qr - step2 * grad2 * qr
//  [python]
def _update_precond_dense_scale(Ql, qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update (dense, scaling) preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql), where
    dX and dG have shape (M, N)
    Ql has shape (M, M)
    qr has shape (1, N)
    qr is the diagonal part of Qr
    dX is perturbation of (matrix) parameter
    dG is perturbation of (matrix) gradient
    step: update step size
    _tiny: an offset to avoid division by zero
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(qr)

    rho = torch.sqrt(max_l / max_r)
    Ql /= rho
    qr *= rho

    A = Ql.mm(dG * qr)
    Bt = torch.linalg.solve_triangular(Ql, dX / qr, upper=True, left=False)

    grad1 = torch.triu(A.mm(A.t()) - Bt.mm(Bt.t()))
    grad2 = torch.sum(A * A, dim=0, keepdim=True) - torch.sum(Bt * Bt, dim=0, keepdim=True)

    step1 = step / (torch.max(torch.abs(grad1)) + _tiny)
    step2 = step / (torch.max(torch.abs(grad2)) + _tiny)

    return Ql - step1 * grad1.mm(Ql), qr - step2 * grad2 * qr
//  [python]
def eKxz_parallel(self, Z, Xmean, Xcov):
        # TODO: add test
        """Parallel implementation (needs more space, but less time)
        Refer to GPflow implementation

        Args:
            Args:
            Z (Variable): m x q inducing input
            Xmean (Variable): n x q mean of input X
            Xcov (Varible): posterior covariance of X
                two sizes are accepted:
                    n x q x q: each q(x_i) has full covariance
                    n x q: each q(x_i) has diagonal covariance (uncorrelated),
                        stored in each row
        Returns:
            (Variable): n x m
        """

        # Revisit later, check for backward support for n-D tensor
        n = Xmean.size(0)
        q = Xmean.size(1)
        m = Z.size(0)
        if Xcov.dim() == 2:
            # from flattered diagonal to full matrix
            cov = Variable(th.Tensor(n, q, q).type(float_type))
            for i in range(Xmean.size(0)):
                cov[i] = Xcov[i].diag()
            Xcov = cov
            del cov
        length_scales = self.length_scales.transform()
        Lambda = length_scales.pow(2).diag().unsqueeze(0).expand_as(Xcov)
        L = cholesky(Lambda + Xcov)
        xz = Xmean.unsqueeze(2).expand(n, q, m) - Z.unsqueeze(0).expand(n, q, m)
        Lxz = th.triangular_solve(xz, L, upper=False)[0]
        half_log_dets = L.diag().log().sum(1) - length_scales.log().sum().expand(n)

        return self.variance.transform().expand(n, m) * th.exp(
            -0.5 * Lxz.pow(2).sum(1) - half_log_dets.expand(n, m)
        )
//  [python]
def eKxz_parallel(self, Z, Xmean, Xcov):
        # TODO: add test
        """Parallel implementation (needs more space, but less time)
        Refer to GPflow implementation

        Args:
            Args:
            Z (Variable): m x q inducing input
            Xmean (Variable): n x q mean of input X
            Xcov (Varible): posterior covariance of X
                two sizes are accepted:
                    n x q x q: each q(x_i) has full covariance
                    n x q: each q(x_i) has diagonal covariance (uncorrelated),
                        stored in each row
        Returns:
            (Variable): n x m
        """

        # Revisit later, check for backward support for n-D tensor
        n = Xmean.size(0)
        q = Xmean.size(1)
        m = Z.size(0)
        if Xcov.dim() == 2:
            # from flattered diagonal to full matrix
            cov = Variable(th.Tensor(n, q, q).type(float_type))
            for i in range(Xmean.size(0)):
                cov[i] = Xcov[i].diag()
            Xcov = cov
            del cov
        length_scales = self.length_scales.transform()
        Lambda = length_scales.pow(2).diag().unsqueeze(0).expand_as(Xcov)
        L = cholesky(Lambda + Xcov)
        xz = Xmean.unsqueeze(2).expand(n, q, m) - Z.unsqueeze(0).expand(n, q, m)
        Lxz = th.linalg.solve_triangular(xz, L, upper=False)
        half_log_dets = L.diag().log().sum(1) - length_scales.log().sum().expand(n)

        return self.variance.transform().expand(n, m) * th.exp(
            -0.5 * Lxz.pow(2).sum(1) - half_log_dets.expand(n, m)
        )
//  [python]
def log_prob(self, x):
        D = self.loc.shape[-1]
        y = torch.triangular_solve((x - self.loc).unsqueeze(-1), self.scale_tril, upper=False)[0].squeeze(-1)
        maha = (y * y).sum(-1)
        return -.5 * (self.df + D) * torch.log1p(maha / self.df) - self._log_normalizer()
//  [python]
def log_prob(self, x):
        D = self.loc.shape[-1]
        y = torch.linalg.solve_triangular((x - self.loc).unsqueeze(-1), self.scale_tril, upper=False).squeeze(-1)
        maha = (y * y).sum(-1)
        return -.5 * (self.df + D) * torch.log1p(maha / self.df) - self._log_normalizer()
//  [python]
def _update_precond_norm_dense(ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update (normalization, dense) Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql), where
    dX and dG have shape (M, N)
    ql has shape (2, M)
    Qr has shape (N, N)
    ql[0] is the diagonal part of Ql
    ql[1,0:-1] is the last column of Ql, excluding the last entry
    dX is perturbation of (matrix) parameter
    dG is perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    # make sure that Ql and Qr have similar dynamic range
    max_l = torch.max(ql[0])
    max_r = torch.max(torch.diag(Qr))
    rho = torch.sqrt(max_l / max_r)
    ql /= rho
    Qr *= rho

    # refer to https://arxiv.org/abs/1512.04202 for details
    A = ql[0:1].t() * dG + ql[1:].t().mm(dG[-1:])  # Ql*dG
    A = A.mm(Qr.t())

    Bt = dX / ql[0:1].t()
    Bt[-1:] -= (ql[1:] / (ql[0:1] * ql[0, -1])).mm(dX)
    Bt = torch.triangular_solve(Bt.t(), Qr, upper=True, transpose=True)[0].t()

    grad1_diag = torch.sum(A * A, dim=1) - torch.sum(Bt * Bt, dim=1)
    grad1_bias = A[:-1].mm(A[-1:].t()) - Bt[:-1].mm(Bt[-1:].t())
    grad1_bias = torch.cat([torch.squeeze(grad1_bias), grad1_bias.new_zeros(1)])

    step1 = step / (torch.max(torch.max(torch.abs(grad1_diag)),
                              torch.max(torch.abs(grad1_bias))) + _tiny)
    new_ql0 = ql[0] - step1 * grad1_diag * ql[0]
    new_ql1 = ql[1] - step1 * (grad1_diag * ql[1] + ql[0, -1] * grad1_bias)

    grad2 = torch.triu(A.t().mm(A) - Bt.t().mm(Bt))
    step2 = step / (torch.max(torch.abs(grad2)) + _tiny)

    return torch.stack((new_ql0, new_ql1)), Qr - step2 * grad2.mm(Qr)
//  [python]
def _update_precond_norm_dense(ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update (normalization, dense) Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql), where
    dX and dG have shape (M, N)
    ql has shape (2, M)
    Qr has shape (N, N)
    ql[0] is the diagonal part of Ql
    ql[1,0:-1] is the last column of Ql, excluding the last entry
    dX is perturbation of (matrix) parameter
    dG is perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    # make sure that Ql and Qr have similar dynamic range
    max_l = torch.max(ql[0])
    max_r = torch.max(torch.diag(Qr))
    rho = torch.sqrt(max_l / max_r)
    ql /= rho
    Qr *= rho

    # refer to https://arxiv.org/abs/1512.04202 for details
    A = ql[0:1].t() * dG + ql[1:].t().mm(dG[-1:])  # Ql*dG
    A = A.mm(Qr.t())

    Bt = dX / ql[0:1].t()
    Bt[-1:] -= (ql[1:] / (ql[0:1] * ql[0, -1])).mm(dX)
    Bt = torch.linalg.solve_triangular(Qr, Bt.t(), upper=True, transpose=True).t()

    grad1_diag = torch.sum(A * A, dim=1) - torch.sum(Bt * Bt, dim=1)
    grad1_bias = A[:-1].mm(A[-1:].t()) - Bt[:-1].mm(Bt[-1:].t())
    grad1_bias = torch.cat([torch.squeeze(grad1_bias), grad1_bias.new_zeros(1)])

    step1 = step / (torch.max(torch.max(torch.abs(grad1_diag)),
                              torch.max(torch.abs(grad1_bias))) + _tiny)
    new_ql0 = ql[0] - step1 * grad1_diag * ql[0]
    new_ql1 = ql[1] - step1 * (grad1_diag * ql[1] + ql[0, -1] * grad1_bias)

    grad2 = torch.triu(A.t().mm(A) - Bt.t().mm(Bt))
    step2 = step / (torch.max(torch.abs(grad2)) + _tiny)

    return torch.stack((new_ql0, new_ql1)), Qr - step2 * grad2.mm(Qr)
//  [python]
def forward(self, inputs, targets=None, train_x=None):
        if train_x is None:
            train_x = inputs

        # Draw samples from the first function
        hidden_covar = rbf_kernel(train_x, train_x, self.ls1, self.os1).add(self.eye.mul(1e-4))
        hidden_chol_covar = psd_safe_cholesky(hidden_covar)
        with pyro.plate("hidden_fns_plate", self.width, dim=-1):
            hidden_dist = pyro.distributions.Normal(self.mean, self.ones).to_event(1)
            hidden_fns = pyro.sample("hidden_fns", hidden_dist).transpose(-1, -2)
            unwhitened_hidden_fns = hidden_chol_covar @ hidden_fns

        final_covar = add_rbf_kernel(unwhitened_hidden_fns, unwhitened_hidden_fns, self.ls2, self.os2).add(self.eye.mul(1e-4))
        final_chol_covar = psd_safe_cholesky(final_covar)
        with pyro.plate("output_fns_plate", 1, dim=-1):
            final_dist = pyro.distributions.Normal(self.mean, self.ones).to_event(1)
            final_fns = pyro.sample("final_fns", final_dist).transpose(-1, -2)
            unwhitened_final_fns = final_chol_covar @ final_fns

        if torch.equal(train_x, inputs):
            latent_vals = unwhitened_final_fns.squeeze(-1)
        else:
            hidden_cross_covar = rbf_kernel(train_x, inputs, self.ls1, self.os1)
            hidden_interp_term = torch.triangular_solve(hidden_cross_covar, hidden_chol_covar, upper=False)[0]
            pred_hidden_fns = hidden_interp_term.transpose(-1, -2) @ hidden_fns

            final_cross_covar = add_rbf_kernel(unwhitened_hidden_fns, pred_hidden_fns, self.ls2, self.os2)
            final_interp_term = torch.triangular_solve(final_cross_covar, final_chol_covar, upper=False)[0]
            latent_vals = (final_interp_term.transpose(-1, -2) @ final_fns).squeeze(-1)

        # Observations
        with pyro.plate("output_plate", inputs.size(-2), dim=-1):
            output_dist = pyro.distributions.Normal(latent_vals.squeeze(-1), self.noise.sqrt())
            outputs = pyro.sample("outputs", output_dist, obs=targets)

        return output_dist
//  [python]
def forward(self, inputs, targets=None, train_x=None):
        if train_x is None:
            train_x = inputs

        # Draw samples from the first function
        hidden_covar = rbf_kernel(train_x, train_x, self.ls1, self.os1).add(self.eye.mul(1e-4))
        hidden_chol_covar = psd_safe_cholesky(hidden_covar)
        with pyro.plate("hidden_fns_plate", self.width, dim=-1):
            hidden_dist = pyro.distributions.Normal(self.mean, self.ones).to_event(1)
            hidden_fns = pyro.sample("hidden_fns", hidden_dist).transpose(-1, -2)
            unwhitened_hidden_fns = hidden_chol_covar @ hidden_fns

        final_covar = add_rbf_kernel(unwhitened_hidden_fns, unwhitened_hidden_fns, self.ls2, self.os2).add(self.eye.mul(1e-4))
        final_chol_covar = psd_safe_cholesky(final_covar)
        with pyro.plate("output_fns_plate", 1, dim=-1):
            final_dist = pyro.distributions.Normal(self.mean, self.ones).to_event(1)
            final_fns = pyro.sample("final_fns", final_dist).transpose(-1, -2)
            unwhitened_final_fns = final_chol_covar @ final_fns

        if torch.equal(train_x, inputs):
            latent_vals = unwhitened_final_fns.squeeze(-1)
        else:
            hidden_cross_covar = rbf_kernel(train_x, inputs, self.ls1, self.os1)
            hidden_interp_term = torch.linalg.solve_triangular(hidden_cross_covar, hidden_chol_covar, upper=False)
            pred_hidden_fns = hidden_interp_term.transpose(-1, -2) @ hidden_fns

            final_cross_covar = add_rbf_kernel(unwhitened_hidden_fns, pred_hidden_fns, self.ls2, self.os2)
            final_interp_term = torch.linalg.solve_triangular(final_cross_covar, final_chol_covar, upper=False)
            latent_vals = (final_interp_term.transpose(-1, -2) @ final_fns).squeeze(-1)

        # Observations
        with pyro.plate("output_plate", inputs.size(-2), dim=-1):
            output_dist = pyro.distributions.Normal(latent_vals.squeeze(-1), self.noise.sqrt())
            outputs = pyro.sample("outputs", output_dist, obs=targets)

        return output_dist
//  [python]
def get_alpha(self):
        if self.x_train is None or self.y_train is None:
            raise RuntimeError("Call forward at least once before calling predict.")
        with torch.autograd.no_grad():
            L, A, AAT, LB, c = self._calc_intermediate(self.x_train, self.y_train)
            tmp1 = torch.triangular_solve(c, LB, upper=False, transpose=True).solution
            tmp2 = torch.triangular_solve(tmp1, L, upper=False, transpose=True).solution
            return tmp2