jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x_cat:Tensor, x_cont:Tensor):
        # covariance between combinaisons of samples
        cov_train_test = self.kernel.matrix((self.train_input_cat, self.train_input_cont), (x_cat, x_cont))
        diag_cov_test_test = self.kernel((x_cat, x_cont), (x_cat, x_cont))
        # cholesky decompositions (accelerate solving of linear systems)
        L_train_train, output_weights = self.memoized_cholesky_decomposition()
        (L_test, _) = torch.triangular_solve(cov_train_test, L_train_train, upper=False)
        # predicted mean
        mean = torch.mm(L_test.t(), output_weights) + self.prior(x_cat, x_cont)
        # predicted std
        var_noise = self.std_noise * self.std_noise
        std_scale = self.std_scale.abs()
        covar = (diag_cov_test_test - torch.sum(L_test**2, dim=0)).abs().unsqueeze(dim=1) # abs against negative variance
        stdev = torch.sqrt(covar + var_noise) * std_scale + 1e-10 # epsilon to insure we are strictly above 0
        # puts the std in a member of the mean
        # to make it accesible from the loss function 
        # while being able to use fastai predict pipeline which has strong preconceptions on the output tensor shape
        mean.stdev = stdev
        return mean
//  [python]
def forward(self, x_cat:Tensor, x_cont:Tensor):
        # covariance between combinaisons of samples
        cov_train_test = self.kernel.matrix((self.train_input_cat, self.train_input_cont), (x_cat, x_cont))
        diag_cov_test_test = self.kernel((x_cat, x_cont), (x_cat, x_cont))
        # cholesky decompositions (accelerate solving of linear systems)
        L_train_train, output_weights = self.memoized_cholesky_decomposition()
        (L_test, _) = torch.linalg.solve_triangular(cov_train_test, L_train_train, upper=False)
        # predicted mean
        mean = torch.mm(L_test.t(), output_weights) + self.prior(x_cat, x_cont)
        # predicted std
        var_noise = self.std_noise * self.std_noise
        std_scale = self.std_scale.abs()
        covar = (diag_cov_test_test - torch.sum(L_test**2, dim=0)).abs().unsqueeze(dim=1) # abs against negative variance
        stdev = torch.sqrt(covar + var_noise) * std_scale + 1e-10 # epsilon to insure we are strictly above 0
        # puts the std in a member of the mean
        # to make it accesible from the loss function 
        # while being able to use fastai predict pipeline which has strong preconceptions on the output tensor shape
        mean.stdev = stdev
        return mean
//  [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, "but got result with dtype Int"):
            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, "but got clone_A with dtype Int"):
            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))

        # 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[-1].message))
            self.assertTrue("An output with one or more elements was resized" in str(w[-2].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, "but got result with dtype Int"):
            torch.linalg.solve_triangular(b, a, out=(out, clone_a))

        out = torch.empty_like(b)
        clone_a = clone_a.to(torch.int)
        with self.assertRaisesRegex(RuntimeError, "but got clone_A with dtype Int"):
            torch.linalg.solve_triangular(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.linalg.solve_triangular(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.linalg.solve_triangular(b, a, out=(out, clone_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.linalg.solve_triangular(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[-1].message))
            self.assertTrue("An output with one or more elements was resized" in str(w[-2].message))
//  [python]
def inverse_mult(self, u, delta, precompute=True):
        """ Computes (I - d A)^-1 u """

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1).unsqueeze(-1)

        if precompute:
            if isinstance(delta, float) and delta in self.backward_cache:
                mat = self.backward_cache[delta]
            else:
                mat = self.precompute_backward(delta) # (n, n) or (..., n, n)

            if len(u.shape) >= len(mat.shape):
                # For memory efficiency, leverage extra batch dimensions
                s = len(u.shape)
                # TODO can make the permutation more efficient by just permuting the last 2 or 3 dim, but need to do more casework
                u = u.permute(list(range(1, s)) + [0])
                x = mat @ u
                x = x.permute([s-1] + list(range(s-1)))
            else:
                x = (mat @ u.unsqueeze(-1))[..., 0]

        else:
            _A = self.I - delta*self.A
            x = torch.triangular_solve(u.unsqueeze(-1), _A, upper=False)[0]
            x = x[..., 0]

        return x
//  [python]
def inverse_mult(self, u, delta, precompute=True):
        """ Computes (I - d A)^-1 u """

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1).unsqueeze(-1)

        if precompute:
            if isinstance(delta, float) and delta in self.backward_cache:
                mat = self.backward_cache[delta]
            else:
                mat = self.precompute_backward(delta) # (n, n) or (..., n, n)

            if len(u.shape) >= len(mat.shape):
                # For memory efficiency, leverage extra batch dimensions
                s = len(u.shape)
                # TODO can make the permutation more efficient by just permuting the last 2 or 3 dim, but need to do more casework
                u = u.permute(list(range(1, s)) + [0])
                x = mat @ u
                x = x.permute([s-1] + list(range(s-1)))
            else:
                x = (mat @ u.unsqueeze(-1))[..., 0]

        else:
            _A = self.I - delta*self.A
            x = torch.linalg.solve_triangular(u.unsqueeze(-1), _A, upper=False)[..., 0]

        return x
//  [python]
def conditional_ND(self, X, full_cov=False, y=None, ind_nan=None):
        """
        q(f|m, S; X, Z)
        @param X:
        @param full_cov:
        @param y:
        @param ind_nan:
        @return:
        """

        _, Lu, Ku_tiled, _ = self.get_prior_matrices()

        Kuf = self.kern(self.inducing_points, X)

        A = torch.triangular_solve(Kuf, Lu, upper=False)[0]


        if not self.white:
            A = torch.triangular_solve(A, Lu.transpose(0,1), upper=True)[0]

        mean = A.transpose(0, 1).matmul(self.q_mu)

        A_tiled = A[None, :, :].repeat([self.num_outputs, 1, 1])
        I = torch.eye(self.num_inducing)[None, :, :].to(self.dtype).to(self.device)

        if self.white:
            SK = -I
        else:
            SK = -Ku_tiled

        if self.q_sqrt_tri is not None:
            qs = torch.zeros((self.num_inducing, self.num_inducing)).to(self.dtype).to(self.device)
            li, lj = torch.tril_indices(self.num_inducing, self.num_inducing, 0)
            qs[li, lj] = self.q_sqrt_tri[0]
            q_sqrt = qs.unsqueeze(0)
            SK +=q_sqrt.matmul(q_sqrt.transpose(1, 2))

        B = SK.matmul(A_tiled)


        if not full_cov:
            # (num_latent, num_X)
            delta_cov = (A_tiled * B).sum(1)
            Kff = self.kern.K_diag(X)


        var = Kff.unsqueeze(0) + delta_cov

        return mean + self.mean_function(X, ind_nan), var.transpose(0,1)
//  [python]
def conditional_ND(self, X, full_cov=False, y=None, ind_nan=None):
        """
        q(f|m, S; X, Z)
        @param X:
        @param full_cov:
        @param y:
        @param ind_nan:
        @return:
        """

        _, Lu, Ku_tiled, _ = self.get_prior_matrices()

        Kuf = self.kern(self.inducing_points, X)

        A = torch.linalg.solve_triangular(Lu, Kuf, upper=False)


        if not self.white:
            A = torch.linalg.solve_triangular(Lu.transpose(0,1), A, upper=True)

        mean = A.transpose(0, 1).matmul(self.q_mu)

        A_tiled = A[None, :, :].repeat([self.num_outputs, 1, 1])
        I = torch.eye(self.num_inducing)[None, :, :].to(self.dtype).to(self.device)

        if self.white:
            SK = -I
        else:
            SK = -Ku_tiled

        if self.q_sqrt_tri is not None:
            qs = torch.zeros((self.num_inducing, self.num_inducing)).to(self.dtype).to(self.device)
            li, lj = torch.tril_indices(self.num_inducing, self.num_inducing, 0)
            qs[li, lj] = self.q_sqrt_tri[0]
            q_sqrt = qs.unsqueeze(0)
            SK +=q_sqrt.matmul(q_sqrt.transpose(1, 2))

        B = SK.matmul(A_tiled)


        if not full_cov:
            # (num_latent, num_X)
            delta_cov = (A_tiled * B).sum(1)
            Kff = self.kern.K_diag(X)


        var = Kff.unsqueeze(0) + delta_cov

        return mean + self.mean_function(X, ind_nan), var.transpose(0,1)
//  [python]
def test_triangular_solve(self):

    def op_fn(b, A, lower, unit_diagonal, transpose_a):
      return A.triangualr_solve(b, True, lower, unit_diagonal, transpose_a)

    def aten_fn(b, A, lower, unit_diagonal, transpose_a):
      return torch.triangular_solve(
          b,
          A,
          upper=not lower,
          unitriangular=unit_diagonal,
          transpose=transpose_a)

    self.runOpBuilderTest(
        'test_triangular_solve',
        [torch.randn(2, 3), torch.randn(2, 2).triu()],
        op_fn,
        aten_fn=aten_fn,
        kwargs={
            'lower': False,
            'unit_diagonal': False,
            'transpose_a': False
        })
//  [python]
def test_triangular_solve(self):

    def op_fn(b, A, lower, unit_diagonal, transpose_a):
      return A.triangular_solve(b, True, lower, unit_diagonal, transpose_a)

    def aten_fn(b, A, lower, unit_diagonal, transpose_a):
      return torch.linalg.solve_triangular(
          b,
          A,
          upper=not lower,
          unitriangular=unit_diagonal,
          left=True,
          transpose=transpose_a)

    self.runOpBuilderTest(
        'test_triangular_solve',
        [torch.randn(2, 3), torch.randn(2, 2).triu()],
        op_fn,
        aten_fn=aten_fn,
        kwargs={
            'lower': False,
            'unit_diagonal': False,
            'transpose_a': False
        })
//  [python]
def predict(self, X_test: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Returns the predictive mean and variance of the objective function at
        the given test points.
        """
        X_test = torch.tensor(X_test, dtype=t_dtype)
        out = self.forward(X_test)

        # Naive implementation:
        #m = torch.matmul(torch.matmul(torch.t(self.e_t[-1]), self.L_t_inv[-1]), torch.t(out))
        m = torch.matmul(
            torch.t(self.e_t[-1]),
            torch.triangular_solve(torch.t(out), self.L_t[-1], upper=False).solution,
        )
        # Remember that we model 1/alpha instead of alpha
        m = (self.beta_t[-1] * self.alpha_t[-1]) * m.reshape((m.shape[1], 1))
        assert (m.shape == (X_test.shape[0], 1))
        if not torch.isfinite(m).all():
            raise ValueError('Infinite predictions %s for input %s' % (m, X_test))
        m = m * self.std_ + self.mean_

        # Naive implementation
        #v = torch.matmul(self.L_inv_t[-1], torch.t(out))
        v = torch.triangular_solve(torch.t(out), self.L_t[-1], upper=False).solution
        # Remember that we model 1/alpha instead of alpha
        v = self.alpha_t[-1] * torch.pow(torch.norm(v, dim=0), 2)
        v = v.reshape((-1, 1))
        assert (v.shape == (X_test.shape[0], 1)), v.shape
        if not torch.isfinite(v).all():
            raise ValueError('Infinite predictions %s for input %s' % (v, X_test))
        v = v * (self.std_ ** 2)

        return m.detach().numpy(), v.detach().numpy()
//  [python]
def predict(self, X_test: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Returns the predictive mean and variance of the objective function at
        the given test points.
        """
        X_test = torch.tensor(X_test, dtype=t_dtype)
        out = self.forward(X_test)

        # Naive implementation:
        #m = torch.matmul(torch.matmul(torch.t(self.e_t[-1]), self.L_t_inv[-1]), torch.t(out))
        m = torch.matmul(
            torch.t(self.e_t[-1]),
            torch.linalg.solve_triangular(self.L_t[-1], torch.t(out), upper=False),
        )
        # Remember that we model 1/alpha instead of alpha
        m = (self.beta_t[-1] * self.alpha_t[-1]) * m.reshape((m.shape[1], 1))
        assert (m.shape == (X_test.shape[0], 1))
        if not torch.isfinite(m).all():
            raise ValueError('Infinite predictions %s for input %s' % (m, X_test))
        m = m * self.std_ + self.mean_

        # Naive implementation
        #v = torch.matmul(self.L_inv_t[-1], torch.t(out))
        v = torch.linalg.solve_triangular(self.L_t[-1], torch.t(out), upper=False)
        # Remember that we model 1/alpha instead of alpha
        v = self.alpha_t[-1] * torch.pow(torch.norm(v, dim=0), 2)
        v = v.reshape((-1, 1))
        assert (v.shape == (X_test.shape[0], 1)), v.shape
        if not torch.isfinite(v).all():
            raise ValueError('Infinite predictions %s for input %s' % (v, X_test))
        v = v * (self.std_ ** 2)

        return m.detach().numpy(), v.detach().numpy()
//  [python]
def __call__(self, x, prior=False):
        if not self.updated_strategy.item() and not prior:
            with torch.no_grad():
                # Get unwhitened p(u)
                prior_function_dist = self(self.inducing_points, prior=True)
                prior_mean = prior_function_dist.loc
                L = self._cholesky_factor(prior_function_dist.lazy_covariance_matrix.add_jitter())

                # Temporarily turn off noise that's added to the mean
                orig_mean_init_std = self._variational_distribution.mean_init_std
                self._variational_distribution.mean_init_std = 0.0

                # Change the variational parameters to be whitened
                variational_dist = self.variational_distribution
                whitened_mean = (
                    torch.triangular_solve((variational_dist.loc - prior_mean).unsqueeze(-1).double(), L, upper=False)[
                        0
                    ]
                    .squeeze(-1)
                    .to(variational_dist.loc.dtype)
                )
                whitened_covar = RootLazyTensor(
                    torch.triangular_solve(
                        variational_dist.lazy_covariance_matrix.root_decomposition().root.evaluate().double(),
                        L,
                        upper=False,
                    )[0].to(variational_dist.loc.dtype)
                )
                whitened_variational_distribution = variational_dist.__class__(whitened_mean, whitened_covar)
                self._variational_distribution.initialize_variational_distribution(whitened_variational_distribution)

                # Reset the random noise parameter of the model
                self._variational_distribution.mean_init_std = orig_mean_init_std

                # Reset the cache
                if hasattr(self, "_memoize_cache"):
                    delattr(self, "_memoize_cache")
                    self._memoize_cache = dict()

                # Mark that we have updated the variational strategy
                self.updated_strategy.fill_(True)

        return super().__call__(x, prior=prior)
//  [python]
def __call__(self, x, prior=False):
        if not self.updated_strategy.item() and not prior:
            with torch.no_grad():
                # Get unwhitened p(u)
                prior_function_dist = self(self.inducing_points, prior=True)
                prior_mean = prior_function_dist.loc
                L = self._cholesky_factor(prior_function_dist.lazy_covariance_matrix.add_jitter())

                # Temporarily turn off noise that's added to the mean
                orig_mean_init_std = self._variational_distribution.mean_init_std
                self._variational_distribution.mean_init_std = 0.0

                # Change the variational parameters to be whitened
                variational_dist = self.variational_distribution
                whitened_mean = (
                    torch.linalg.solve_triangular((variational_dist.loc - prior_mean).unsqueeze(-1).double(), L, upper=False)
                    .squeeze(-1)
                    .to(variational_dist.loc.dtype)
                )
                whitened_covar = RootLazyTensor(
                    torch.linalg.solve_triangular(
                        variational_dist.lazy_covariance_matrix.root_decomposition().root.evaluate().double(),
                        L,
                        upper=False,
                    ).to(variational_dist.loc.dtype)
                )
                whitened_variational_distribution = variational_dist.__class__(whitened_mean, whitened_covar)
                self._variational_distribution.initialize_variational_distribution(whitened_variational_distribution)

                # Reset the random noise parameter of the model
                self._variational_distribution.mean_init_std = orig_mean_init_std

                # Reset the cache
                if hasattr(self, "_memoize_cache"):
                    delattr(self, "_memoize_cache")
                    self._memoize_cache = dict()

                # Mark that we have updated the variational strategy
                self.updated_strategy.fill_(True)

        return super().__call__(x, prior=prior)
//  [python]
def exp_log_full_gaussian(mean1, tril1, mean2, tril2):
        const_term = -0.5 * mean1.shape[0] * np.log(2 * np.pi * np.exp(1))
        log_det_prior = -torch.sum(torch.log(torch.diag(tril2)))
        LpiLq = torch.triangular_solve(
            tril1.double(), tril2.double(), upper=False
        )[0]
        trace_term = -0.5 * torch.sum(LpiLq ** 2)
        mu_diff = mean1 - mean2
        quad_solve = torch.triangular_solve(
            mu_diff.double(), tril2.double(), upper=False
        )[0]
        quad_term = -0.5 * torch.sum(quad_solve ** 2)
        res = const_term + log_det_prior + trace_term + quad_term
        return res
//  [python]
def exp_log_full_gaussian(mean1, tril1, mean2, tril2):
        const_term = -0.5 * mean1.shape[0] * np.log(2 * np.pi * np.exp(1))
        log_det_prior = -torch.sum(torch.log(torch.diag(tril2)))
        LpiLq = torch.linalg.solve_triangular(
            tril1.double(), tril2.double(), upper=False
        )
        trace_term = -0.5 * torch.sum(LpiLq ** 2)
        mu_diff = mean1 - mean2
        quad_solve = torch.linalg.solve_triangular(
            mu_diff.double(), tril2.double(), upper=False
        )
        quad_term = -0.5 * torch.sum(quad_solve ** 2)
        res = const_term + log_det_prior + trace_term + quad_term
        return res
//  [python]
def inverse(self, z):
        L, U = self._create_lower_upper()
        if self.bias is not None: z = z - self.bias
        t, _ = torch.triangular_solve(z.t(), L, upper=False, unitriangular=True)
        t, _ = torch.triangular_solve(t, U, upper=True, unitriangular=False)
        x = t.t()
        return x
//  [python]
def inverse(self, z):
        L, U = self._create_lower_upper()
        if self.bias is not None: z = z - self.bias
        t, _ = torch.linalg.solve_triangular(z.t(), L, upper=False, unitriangular=True)
        t, _ = torch.linalg.solve_triangular(t, U, upper=True, unitriangular=False)
        x = t.t()
        return x
//  [python]
def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None):
        # Compute full prior distribution
        full_inputs = torch.cat([inducing_points, x], dim=-2)
        full_output = self.model.forward(full_inputs)
        full_covar = full_output.lazy_covariance_matrix

        # Covariance terms
        num_induc = inducing_points.size(-2)
        test_mean = full_output.mean[..., num_induc:]
        induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter()
        induc_data_covar = full_covar[..., :num_induc, num_induc:].evaluate()
        data_data_covar = full_covar[..., num_induc:, num_induc:]

        # Compute interpolation terms
        # K_ZZ^{-1/2} K_ZX
        # K_ZZ^{-1/2} \mu_Z
        L = self._cholesky_factor(induc_induc_covar)
        if L.shape != induc_induc_covar.shape:
            # Aggressive caching can cause nasty shape incompatibilies when evaluating with different batch shapes
            del self._memoize_cache["cholesky_factor"]
            L = self._cholesky_factor(induc_induc_covar)
        interp_term = torch.triangular_solve(induc_data_covar.double(), L, upper=False)[0].to(full_inputs.dtype)

        # Compute the mean of q(f)
        # k_XZ K_ZZ^{-1/2} (m - K_ZZ^{-1/2} \mu_Z) + \mu_X
        predictive_mean = (
            torch.matmul(
                interp_term.transpose(-1, -2), (inducing_values - self.prior_distribution.mean).unsqueeze(-1)
            ).squeeze(-1)
            + test_mean
        )

        # Compute the covariance of q(f)
        # K_XX + k_XZ K_ZZ^{-1/2} (S - I) K_ZZ^{-1/2} k_ZX
        middle_term = self.prior_distribution.lazy_covariance_matrix.mul(-1)
        if variational_inducing_covar is not None:
            middle_term = SumLazyTensor(variational_inducing_covar, middle_term)

        if trace_mode.on():
            predictive_covar = (
                data_data_covar.add_jitter(1e-4).evaluate()
                + interp_term.transpose(-1, -2) @ middle_term.evaluate() @ interp_term
            )
        else:
            predictive_covar = SumLazyTensor(
                data_data_covar.add_jitter(1e-4),
                MatmulLazyTensor(interp_term.transpose(-1, -2), middle_term @ interp_term),
            )

        # Return the distribution
        return MultivariateNormal(predictive_mean, predictive_covar)
//  [python]
def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None):
        # Compute full prior distribution
        full_inputs = torch.cat([inducing_points, x], dim=-2)
        full_output = self.model.forward(full_inputs)
        full_covar = full_output.lazy_covariance_matrix

        # Covariance terms
        num_induc = inducing_points.size(-2)
        test_mean = full_output.mean[..., num_induc:]
        induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter()
        induc_data_covar = full_covar[..., :num_induc, num_induc:].evaluate()
        data_data_covar = full_covar[..., num_induc:, num_induc:]

        # Compute interpolation terms
        # K_ZZ^{-1/2} K_ZX
        # K_ZZ^{-1/2} \mu_Z
        L = self._cholesky_factor(induc_induc_covar)
        if L.shape != induc_induc_covar.shape:
            # Aggressive caching can cause nasty shape incompatibilies when evaluating with different batch shapes
            del self._memoize_cache["cholesky_factor"]
            L = self._cholesky_factor(induc_induc_covar)
        interp_term = torch.linalg.solve_triangular(L, induc_data_covar.double(), upper=False).to(full_inputs.dtype)

        # Compute the mean of q(f)
        # k_XZ K_ZZ^{-1/2} (m - K_ZZ^{-1/2} \mu_Z) + \mu_X
        predictive_mean = (
            torch.matmul(
                interp_term.transpose(-1, -2), (inducing_values - self.prior_distribution.mean).unsqueeze(-1)
            ).squeeze(-1)
            + test_mean
        )

        # Compute the covariance of q(f)
        # K_XX + k_XZ K_ZZ^{-1/2} (S - I) K_ZZ^{-1/2} k_ZX
        middle_term = self.prior_distribution.lazy_covariance_matrix.mul(-1)
        if variational_inducing_covar is not None:
            middle_term = SumLazyTensor(variational_inducing_covar, middle_term)

        if trace_mode.on():
            predictive_covar = (
                data_data_covar.add_jitter(1e-4).evaluate()
                + interp_term.transpose(-1, -2) @ middle_term.evaluate() @ interp_term
            )
        else:
            predictive_covar = SumLazyTensor(
                data_data_covar.add_jitter(1e-4),
                MatmulLazyTensor(interp_term.transpose(-1, -2), middle_term @ interp_term),
            )

        # Return the distribution
        return MultivariateNormal(predictive_mean, predictive_covar)
//  [python]
def weight_inverse(self):
        """Cost:
            inverse = O(D^3)
        where:
            D = num of features
        """
        L, U = self._create_lower_upper()
        identity = torch.eye(self.num_features, self.num_features, device=self.unconstrained_upper_diag.device, dtype=self.unconstrained_upper_diag.dtype)
        lower_inverse, _ = torch.triangular_solve(identity, L, upper=False, unitriangular=True)
        weight_inverse, _ = torch.triangular_solve(lower_inverse, U, upper=True, unitriangular=False)
        return weight_inverse
//  [python]
def weight_inverse(self):
        """Cost:
            inverse = O(D^3)
        where:
            D = num of features
        """
        L, U = self._create_lower_upper()
        identity = torch.eye(self.num_features, self.num_features, device=self.unconstrained_upper_diag.device, dtype=self.unconstrained_upper_diag.dtype)
        lower_inverse = torch.linalg.solve_triangular(L, identity, lower=True, unitriangular=True)
        weight_inverse = torch.linalg.solve_triangular(U, lower_inverse, upper=True, unitriangular=False)
        return weight_inverse
//  [python]
def condition(self, x_te, *args, x_tr=None, y_tr=None, **kwargs):
        r""" Calculate the predictive distribution given `x_te`.

        Note
        ----
        Here we allow the speicifaction of sampler but won't actually
        use it here in this version.

        Parameters
        ----------
        x_te : `torch.Tensor`, `shape=(n_te, hidden_dimension)`
            Test input.

        x_tr : `torch.Tensor`, `shape=(n_tr, hidden_dimension)`
             (Default value = None)
             Training input.

        y_tr : `torch.Tensor`, `shape=(n_tr, 1)`
             (Default value = None)
             Test input.

        sampler : `torch.optim.Optimizer` or `pinot.Sampler`
             (Default value = None)
             Sampler.

        Returns
        -------
        distribution : `torch.distributions.Distribution`
            Predictive distribution.

        """

        # get parameters
        (
            k_tr_tr,
            k_te_te,
            k_te_tr,
            k_tr_te,
            l_low,
            alpha,
        ) = self._get_kernel_and_auxiliary_variables(x_tr, y_tr, x_te)

        # compute mean
        # (batch_size_te, 1)
        mean = k_te_tr @ alpha

        # (batch_size_tr, batch_size_te)
        v, _ = torch.triangular_solve(input=k_tr_te, A=l_low, upper=False)

        # (batch_size_te, batch_size_te)
        variance = k_te_te - v.t() @ v

        # ensure symetric
        variance = 0.5 * (variance + variance.t())

        # $ p(y|X) = \int p(y|f)p(f|x) df $
        # variance += torch.exp(self.log_sigma) * torch.eye(
        #         *variance.shape,
        #         device=variance.device)

        # construct noise predictive distribution
        distribution = torch.distributions.multivariate_normal.MultivariateNormal(
            mean.flatten(), variance
        )

        return distribution
//  [python]
def condition(self, x_te, *args, x_tr=None, y_tr=None, **kwargs):
        r""" Calculate the predictive distribution given `x_te`.

        Note
        ----
        Here we allow the speicifaction of sampler but won't actually
        use it here in this version.

        Parameters
        ----------
        x_te : `torch.Tensor`, `shape=(n_te, hidden_dimension)`
            Test input.

        x_tr : `torch.Tensor`, `shape=(n_tr, hidden_dimension)`
             (Default value = None)
             Training input.

        y_tr : `torch.Tensor`, `shape=(n_tr, 1)`
             (Default value = None)
             Test input.

        sampler : `torch.optim.Optimizer` or `pinot.Sampler`
             (Default value = None)
             Sampler.

        Returns
        -------
        distribution : `torch.distributions.Distribution`
            Predictive distribution.

        """

        # get parameters
        (
            k_tr_tr,
            k_te_te,
            k_te_tr,
            k_tr_te,
            l_low,
            alpha,
        ) = self._get_kernel_and_auxiliary_variables(x_tr, y_tr, x_te)

        # compute mean
        # (batch_size_te, 1)
        mean = k_te_tr @ alpha

        # (batch_size_tr, batch_size_te)
        v = torch.linalg.solve_triangular(l_low, k_tr_te, upper=False)

        # (batch_size_te, batch_size_te)
        variance = k_te_te - v.t() @ v

        # ensure symetric
        variance = 0.5 * (variance + variance.t())

        # $ p(y|X) = \int p(y|f)p(f|x) df $
        # variance += torch.exp(self.log_sigma) * torch.eye(
        #         *variance.shape,
        #         device=variance.device)

        # construct noise predictive distribution
        distribution = torch.distributions.multivariate_normal.MultivariateNormal(
            mean.flatten(), variance
        )

        return distribution
//  [python]
def whiten_matrix(self, sigma, eye):  # x [group, dim, dim]
        wm = torch.triangular_solve(
            eye, torch.cholesky(sigma), upper=False
        )[0]
        return wm.permute(0, 2, 1)
//  [python]
def whiten_matrix(self, sigma, eye):  # x [group, dim, dim]
        wm = torch.linalg.solve_triangular(
            eye, torch.cholesky(sigma), upper=False
        )
        return wm.permute(0, 2, 1)
//  [python]
def forward(self, video_id):
        """
        Return whitened embeddings

        :param video_id: id of the embedding [bs,]
        :return: whitened embeddings[bs, embedding_dim]
        """

        x = self.embedding(video_id)
        x_all = self.embedding.weight
        m = x_all.mean(0)
        xn = x - m
        xn_all = x_all - m

        f_cov = torch.mm(xn_all.permute(1, 0), xn_all) / (xn_all.shape[0] - 1)
        eye = torch.eye(self.embedding_dim, device=xn.device)
        f_cov_shrunk = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.triangular_solve(eye, torch.linalg.cholesky(f_cov_shrunk), upper=False)[0]

        return xn @ inv_sqrt.permute(1, 0)
//  [python]
def forward(self, video_id):
        """
        Return whitened embeddings

        :param video_id: id of the embedding [bs,]
        :return: whitened embeddings[bs, embedding_dim]
        """

        x = self.embedding(video_id)
        x_all = self.embedding.weight
        m = x_all.mean(0)
        xn = x - m
        xn_all = x_all - m

        f_cov = torch.mm(xn_all.permute(1, 0), xn_all) / (xn_all.shape[0] - 1)
        eye = torch.eye(self.embedding_dim, device=xn.device)
        f_cov_shrunk = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.linalg.solve_triangular(torch.linalg.cholesky(f_cov_shrunk), eye, upper=False)

        return xn @ inv_sqrt.permute(1, 0)
//  [python]
def inverse_no_cache(self, inputs, full_jacobian=False):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        outputs, _ = torch.triangular_solve(outputs.t(), lower, upper=False, unitriangular=True)
        outputs, _ = torch.triangular_solve(outputs, upper, upper=True, unitriangular=False)
        outputs = outputs.t()

        if full_jacobian:
            jacobian = torch.mm(lower, upper).inverse().unsqueeze(0).expand(outputs.size(0), -1, -1)  # TODO: make this faster
            return outputs, jacobian
        else:
            logabsdet = -self.logabsdet()
            logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])
            return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs, full_jacobian=False):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        outputs, _ = torch.linalg.solve_triangular(outputs.t(), lower, upper=False, unitriangular=True)
        outputs, _ = torch.linalg.solve_triangular(outputs, upper, upper=True, unitriangular=False)
        outputs = outputs.t()

        if full_jacobian:
            jacobian = torch.mm(lower, upper).inverse().unsqueeze(0).expand(outputs.size(0), -1, -1)  # TODO: make this faster
            return outputs, jacobian
        else:
            logabsdet = -self.logabsdet()
            logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])
            return outputs, logabsdet
//  [python]
def tri_solve(self, L, b, lower=True):
        return torch.triangular_solve(b, L, upper=not lower)[0]
//  [python]
def tri_solve(self, L, b, lower=True):
        return torch.linalg.solve_triangular(b, L, upper=not lower)
//  [python]
def forward(self, input, weight, logdet=None, reverse=False):
        """
        log-det = log|abs(|W|)| * pixels
        """
        input_dim = input.shape[1]
        assert (input_dim ** 2 - input_dim) / 2 == weight.shape[1]
        assert input.shape[2:] == weight.shape[2:] or input.shape[2:] == (1, 1)

        weight = self.get_weight(weight, input_dim)
        if not reverse:
            z = torch.triangular_solve(input.permute(0, 2, 3, 1).unsqueeze(-1).double(), weight.double(),
                                       unitriangular=True)[0].permute(
                0, 3, 1, 2, 4).squeeze(-1).float().contiguous()
            # z = torch.matmul(weight.double().inverse(), input.permute(0, 2, 3, 1).unsqueeze(-1).double()).permute(0, 3, 1, 2, 4).squeeze(
            #     -1).float().contiguous()
        else:
            z = torch.matmul(weight.double(), input.permute(0, 2, 3, 1).unsqueeze(-1).double()).permute(0, 3, 1, 2,
                                                                                                        4).squeeze(
                -1).float().contiguous()
        return z, logdet
//  [python]
def forward(self, input, weight, logdet=None, reverse=False):
        """
        log-det = log|abs(|W|)| * pixels
        """
        input_dim = input.shape[1]
        assert (input_dim ** 2 - input_dim) / 2 == weight.shape[1]
        assert input.shape[2:] == weight.shape[2:] or input.shape[2:] == (1, 1)

        weight = self.get_weight(weight, input_dim)
        if not reverse:
            z = torch.linalg.solve_triangular(weight.double(), input.permute(0, 2, 3, 1).unsqueeze(-1).double(), upper=True, unitriangular=True)[0].permute(
                0, 3, 1, 2, 4).squeeze(-1).float().contiguous()
            # z = torch.matmul(weight.double().inverse(), input.permute(0, 2, 3, 1).unsqueeze(-1).double()).permute(0, 3, 1, 2, 4).squeeze(
            #     -1).float().contiguous()
        else:
            z = torch.matmul(weight.double(), input.permute(0, 2, 3, 1).unsqueeze(-1).double()).permute(0, 3, 1, 2,
                                                                                                        4).squeeze(
                -1).float().contiguous()
        return z, logdet
//  [python]
def test_triangular_solve(self, device, dtype):
        ks = [0, 1, 3]
        ns = [0, 5]
        for k, n, (upper, unitriangular, transpose) in itertools.product(ks, ns,
                                                                         itertools.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.assertEqual(b, np.matmul(A.t().cpu(), x.cpu()))
            else:
                self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
//  [python]
def test_triangular_solve(self, device, dtype):
        ks = [0, 1, 3]
        ns = [0, 5]
        for k, n, (upper, unitriangular, transpose) in itertools.product(ks, ns,
                                                                         itertools.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.assertEqual(b, np.matmul(A.t().cpu(), x.cpu()))
            else:
                self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
//  [python]
def _inducing_inv_root(self):
        if not self.training and hasattr(self, "_cached_kernel_inv_root"):
            return self._cached_kernel_inv_root
        else:
            chol = psd_safe_cholesky(self._inducing_mat, upper=True, jitter=settings.cholesky_jitter.value())
            eye = torch.eye(chol.size(-1), device=chol.device, dtype=chol.dtype)
            inv_root = torch.triangular_solve(eye, chol)[0]

            res = inv_root
            if not self.training:
                self._cached_kernel_inv_root = res
            return res
//  [python]
def _inducing_inv_root(self):
        if not self.training and hasattr(self, "_cached_kernel_inv_root"):
            return self._cached_kernel_inv_root
        else:
            chol = psd_safe_cholesky(self._inducing_mat, upper=True, jitter=settings.cholesky_jitter.value())
            eye = torch.eye(chol.size(-1), device=chol.device, dtype=chol.dtype)
            inv_root = torch.linalg.solve_triangular(chol, eye, upper=True)

            res = inv_root
            if not self.training:
                self._cached_kernel_inv_root = res
            return res
//  [python]
def forward(self, x):
        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.num_features, -1).mean(-1).view(1, -1, 1, 1)
        if not self.training and self.track_running_stats:  # for inference
            m = self.running_mean
        xn = x - m

        T = xn.permute(1, 0, 2, 3).contiguous().view(self.num_features, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)

        eye = torch.eye(self.num_features).type(f_cov.type())

        if not self.training and self.track_running_stats:  # for inference
            f_cov = self.running_variance

        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.triangular_solve(
            eye, torch.cholesky(f_cov_shrinked), upper=False
        )[0]
        inv_sqrt = inv_sqrt.contiguous().view(
            self.num_features, self.num_features, 1, 1
        )

        decorrelated = conv2d(xn, inv_sqrt)

        if self.training and self.track_running_stats:
            self.running_mean = torch.add(
                self.momentum * m.detach(),
                (1 - self.momentum) * self.running_mean,
                out=self.running_mean,
            )
            self.running_variance = torch.add(
                self.momentum * f_cov.detach(),
                (1 - self.momentum) * self.running_variance,
                out=self.running_variance,
            )

        return decorrelated.squeeze(2).squeeze(2)
//  [python]
def forward(self, x):
        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.num_features, -1).mean(-1).view(1, -1, 1, 1)
        if not self.training and self.track_running_stats:  # for inference
            m = self.running_mean
        xn = x - m

        T = xn.permute(1, 0, 2, 3).contiguous().view(self.num_features, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)

        eye = torch.eye(self.num_features).type(f_cov.type())

        if not self.training and self.track_running_stats:  # for inference
            f_cov = self.running_variance

        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.linalg.solve_triangular(
            eye, torch.cholesky(f_cov_shrinked), upper=False
        )
        inv_sqrt = inv_sqrt.contiguous().view(
            self.num_features, self.num_features, 1, 1
        )

        decorrelated = conv2d(xn, inv_sqrt)

        if self.training and self.track_running_stats:
            self.running_mean = torch.add(
                self.momentum * m.detach(),
                (1 - self.momentum) * self.running_mean,
                out=self.running_mean,
            )
            self.running_variance = torch.add(
                self.momentum * f_cov.detach(),
                (1 - self.momentum) * self.running_variance,
                out=self.running_variance,
            )

        return decorrelated.squeeze(2).squeeze(2)
//  [python]
def test_triangular_solve(self):
        def _test_with_size(A_dims, B_dims):
            A = torch.rand(*A_dims).requires_grad_()
            b = torch.rand(*B_dims).requires_grad_()

            for upper, transpose, unitriangular in product((True, False), repeat=3):
                def func(A, b):
                    return torch.triangular_solve(b, A, upper, transpose, unitriangular)

                gradcheck(func, [A, b])
                gradgradcheck(func, [A, b])

        _test_with_size((3, 3), (3, 4))
        _test_with_size((3, 3), (3, 2))
        _test_with_size((2, 3, 3), (2, 3, 4))
        _test_with_size((2, 3, 3), (2, 3, 2))
//  [python]
def test_triangular_solve(self):
        def _test_with_size(A_dims, B_dims):
            A = torch.rand(*A_dims).requires_grad_()
            b = torch.rand(*B_dims).requires_grad_()

            for upper, transpose, unitriangular in product((True, False), repeat=3):
                def func(A, b):
                    return torch.linalg.solve_triangular(b, A, upper=upper, transpose=transpose, unitriangular=unitriangular)

                gradcheck(func, [A, b])
                gradgradcheck(func, [A, b])

        _test_with_size((3, 3), (3, 4))
        _test_with_size((3, 3), (3, 2))
        _test_with_size((2, 3, 3), (2, 3, 4))
        _test_with_size((2, 3, 3), (2, 3, 2))
//  [python]
def inv_matmul(self, right_tensor: Tensor, left_tensor: Optional[Tensor] = None) -> Tensor:
        if isinstance(self._tensor, NonLazyTensor):
            res = torch.triangular_solve(right_tensor, self.evaluate(), upper=self.upper).solution
        elif isinstance(self._tensor, BatchRepeatLazyTensor):
            res = self._tensor.base_lazy_tensor.inv_matmul(right_tensor, left_tensor)
            # TODO: Proper broadcasting
            res = res.expand(self._tensor.batch_repeat + res.shape[-2:])
        else:
            # TODO: Can we be smarter here?
            res = self._tensor.inv_matmul(right_tensor=right_tensor, left_tensor=left_tensor)
        if left_tensor is not None:
            res = left_tensor @ res
        return res
//  [python]
def inv_matmul(self, right_tensor: Tensor, left_tensor: Optional[Tensor] = None) -> Tensor:
        if isinstance(self._tensor, NonLazyTensor):
            res = torch.linalg.solve_triangular(self.evaluate(), right_tensor, upper=self.upper)
        elif isinstance(self._tensor, BatchRepeatLazyTensor):
            res = self._tensor.base_lazy_tensor.inv_matmul(right_tensor, left_tensor)
            # TODO: Proper broadcasting
            res = res.expand(self._tensor.batch_repeat + res.shape[-2:])
        else:
            # TODO: Can we be smarter here?
            res = self._tensor.inv_matmul(right_tensor=right_tensor, left_tensor=left_tensor)
        if left_tensor is not None:
            res = left_tensor @ res
        return res
//  [python]
def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        D = self.event_shape[-1]
        logdet = _posdef_logdet(value)
        # TODO: Replace with cholesky_solve once its derivative is implemented
        # prod = torch.cholesky_solve(value, self.scale_tril, upper=False)
        prod_ = torch.triangular_solve(value, self.scale_tril, upper=False)[0]
        prod = torch.triangular_solve(prod_, self.scale_tril, upper=False, transpose=True)[0]
        trace = prod.diagonal(dim1=-2, dim2=-1).sum(-1)
        return .5 * (self.df - D - 1) * logdet - .5 * trace - self.log_normalizer
//  [python]
def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        D = self.event_shape[-1]
        logdet = _posdef_logdet(value)
        # TODO: Replace with cholesky_solve once its derivative is implemented
        # prod = torch.cholesky_solve(value, self.scale_tril, upper=False)
        prod_ = torch.linalg.solve_triangular(value, self.scale_tril, upper=False, left=True)
        prod = torch.linalg.solve_triangular(prod_, self.scale_tril, upper=False, left=False, transpose=True)
        trace = prod.diagonal(dim1=-2, dim2=-1).sum(-1)
        return .5 * (self.df - D - 1) * logdet - .5 * trace - self.log_normalizer
//  [python]
def inverse_using_low_rank_factor(Q, D):
    """
    returns inverse of Q @ Q.T + D.diag()
    algorithm inspired by torch.distributions.LowRankMultivariateNormal
    """
    m = Q.size(-1)
    W = Q.t() / D[None]
    K = torch.matmul(W, Q).contiguous()
    K.view(-1, m * m)[:, :: m + 1] += 1  # add identity matrix to K
    C, _ = jitcholesky(K)  # robust
    A = torch.triangular_solve(W, C, upper=False)[0]
    inv = torch.diag_embed(D.reciprocal()) - torch.matmul(A.transpose(-1, -2), A)
    return inv
//  [python]
def inverse_using_low_rank_factor(Q, D):
    """
    returns inverse of Q @ Q.T + D.diag()
    algorithm inspired by torch.distributions.LowRankMultivariateNormal
    """
    m = Q.size(-1)
    W = Q.t() / D[None]
    K = torch.matmul(W, Q).contiguous()
    K.view(-1, m * m)[:, :: m + 1] += 1  # add identity matrix to K
    C, _ = jitcholesky(K)  # robust
    A = torch.linalg.solve_triangular(C, W, upper=False)
    inv = torch.diag_embed(D.reciprocal()) - torch.matmul(A.transpose(-1, -2), A)
    return inv