def __call__(
        self,
        gp_features: Array,
        gp_logits: Optional[Array] = None,
        diagonal_only: bool = True,
    ) -> Optional[Array]:
        """
        Updates the precision matrix and computes the predictive covariance.

        NOTE:
        The precision matrix will be updated only during training (i.e., when
        `self.collection_name` are in the list of mutable variables). The covariance
        matrix will be computed only during inference to avoid repeated calls to the
        (expensive) `linalg.inv` op.

        Parameters
        ----------
        gp_features: Array
            The nd-array of random fourier features, shape (batch_size, ..., hidden_features).
        gp_logits: Optional[Array]
            The nd-array of predictive logits, shape (batch_size, ..., logit_dim).
            Cannot be None.
        diagonal_only: bool
            Whether to return only the diagonal elements of the predictive covariance matrix (i.e., the predictive variance).

        Returns
        -------
        Optional[Array]
          The predictive variances of shape (batch_size, ) if diagonal_only=True,
          otherwise the predictive covariance matrix of shape
          (batch_size, batch_size).
        """
        gp_features = jnp.asarray(gp_features, self.dtype)

        # Flatten GP features and logits to 2-d, by doing so we treat all the
        # non-final dimensions as the batch dimensions.
        gp_features = jnp.reshape(gp_features, [-1, self.hidden_features])

        if gp_logits is not None:
            gp_logits = jnp.asarray(gp_logits, self.dtype)
            gp_logits = jnp.reshape(gp_logits, [gp_features.shape[0], -1])

        precision_matrix = self.variable(
            self.collection_name,
            "precision_matrix",
            lambda: self.initial_precision_matrix(),
        )  # pylint: disable=unnecessary-lambda

        # Updates the precision matrix during training.
        initializing = self.is_mutable_collection("params")
        training = self.is_mutable_collection(self.collection_name)

        if training and not initializing:
            precision_matrix.value = self.update_precision_matrix(
                gp_features, gp_logits, precision_matrix.value
            )

        # Computes covariance matrix during inference.
        if not training:
            return self.compute_predictive_covariance(
                gp_features, precision_matrix, diagonal_only
            )

        Updates the precision matrix and computes the predictive covariance.

        NOTE:
        The precision matrix will be updated only during training (i.e., when
        `self.collection_name` are in the list of mutable variables). The covariance
        matrix will be computed only during inference to avoid repeated calls to the
        (expensive) `linalg.inv` op.

        Parameters
        ----------
        gp_features: Array
            The nd-array of random fourier features, shape (batch_size, ..., hidden_features).
        gp_logits: Optional[Array]
            The nd-array of predictive logits, shape (batch_size, ..., logit_dim).
            Cannot be None.
        diagonal_only: bool
            Whether to return only the diagonal elements of the predictive covariance matrix (i.e., the predictive variance).

        Returns
        -------
        Optional[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