def compute_predictive_covariance(
        self, gp_features: Array, precision_matrix: nn.Variable, diagonal_only: bool
    ) -> Array:
        """
        Computes the predictive covariance.

        Approximates the Gaussian process posterior using random features.
        Given training random feature :math:`\mathbf{\Phi_{tr}}` (num_train, num_hidden) and testing
        random feature :math:`\mathbf{\Phi_{ts}}` (batch_size, num_hidden). The predictive covariance
        matrix is computed as (assuming Gaussian likelihood):
        :math:`s * \mathbf{\Phi_{ts}}(\mathbf{I}*s + \mathbf{\Phi_{tr}}^{T}*\mathbf{\Phi_{tr}})^{-1}\mathbf{\Phi_{tr}}^{^T}`

        where :math:`s` is the ridge factor to be used for stablizing the inverse, and \mathbf{I} is
        the identity matrix with shape (num_hidden, num_hidden). The above
        description is formal only: the actual implementation uses a Cholesky
        factorization of the covariance matrix.

        Parameters
        ----------
        gp_features: Array
            The random feature of testing data to be used for computing the covariance matrix.
            Shape (batch_size, gp_hidden_size).
        precision_matrix: nn.Variable
            The model's precision matrix.
        diagonal_only: bool
            Whether to return only the diagonal elements of the predictive covariance matrix
            (i.e., the predictive variances).

        Returns
        -------
        Array
            The predictive variances of shape (batch_size, ) if `diagonal_only=True`,
            otherwise the predictive covariance matrix of shape (batch_size, batch_size).
        """
        chol = linalg.cholesky(precision_matrix.value)
        chol_t_cov_feature_product = linalg.triangular_solve(
            chol, gp_features.T, left_side=True, lower=True
        )

        if diagonal_only:
            # Compute diagonal element only, shape (batch_size, ).
            gp_covar = jnp.square(chol_t_cov_feature_product).sum(0)
        else:
            # Compute full covariance matrix, shape (batch_size, batch_size).
            gp_covar = chol_t_cov_feature_product.T @ chol_t_cov_feature_product

        return self.ridge_penalty * gp_covar

        Computes the predictive covariance.

        Approximates the Gaussian process posterior using random features.
        Given training random feature :math:`\mathbf{\Phi_{tr}}` (num_train, num_hidden) and testing
        random feature :math:`\mathbf{\Phi_{ts}}` (batch_size, num_hidden). The predictive covariance
        matrix is computed as (assuming Gaussian likelihood):
        :math:`s * \mathbf{\Phi_{ts}}(\mathbf{I}*s + \mathbf{\Phi_{tr}}^{T}*\mathbf{\Phi_{tr}})^{-1}\mathbf{\Phi_{tr}}^{^T}`

        where :math:`s` is the ridge factor to be used for stablizing the inverse, and \mathbf{I} is
        the identity matrix with shape (num_hidden, num_hidden). The above
        description is formal only: the actual implementation uses a Cholesky
        factorization of the covariance matrix.

        Parameters
        ----------
        gp_features: Array
            The random feature of testing data to be used for computing the covariance matrix.
            Shape (batch_size, gp_hidden_size).
        precision_matrix: nn.Variable
            The model's precision matrix.
        diagonal_only: bool
            Whether to return only the diagonal elements of the predictive covariance matrix
            (i.e., the predictive variances).

        Returns
        -------
        Array
            The predictive variances of shape (batch_size, ) if `diagonal_only=True`,
            otherwise the predictive covariance matrix of shape (batch_size, batch_size).
        
fortuna/model/utils/random_features.py