def update_precision_matrix(
        self, gp_features: Array, gp_logits: Optional[Array], precision_matrix: Array
    ) -> Array:
        """Updates precision matrix given a new batch.

        Parameters
        ----------
          gp_features: Array
            Random features from the new batch, shape (batch_size, hidden_features)
          gp_logits: Optional[Array]
            Predictive logits from the new batch, shape (batch_size, logit_dim).
            Currently only `logit_dim=1` is supported.
          precision_matrix: Array
            The current precision matrix, shape (hidden_features, hidden_features).

        Returns
        -------
        Array
            Updated precision matrix, shape (hidden_features, hidden_features).
        """
        # Computes precision matrix within new batch.
        prob_multiplier = 1.0

        gp_features_adj = jnp.sqrt(prob_multiplier) * gp_features
        batch_prec_mat = jnp.matmul(jnp.transpose(gp_features_adj), gp_features_adj)

        # Updates precision matrix.
        if self.momentum is None:
            # Performs exact update without momentum.
            precision_matrix_updated = precision_matrix + batch_prec_mat
        else:
            batch_size = gp_features.shape[0]
            precision_matrix_updated = (
                self.momentum * precision_matrix
                + (1 - self.momentum) * batch_prec_mat / batch_size
            )
        return precision_matrix_updated
Updates precision matrix given a new batch.

        Parameters
        ----------
          gp_features: Array
            Random features from the new batch, shape (batch_size, hidden_features)
          gp_logits: Optional[Array]
            Predictive logits from the new batch, shape (batch_size, logit_dim).
            Currently only `logit_dim=1` is supported.
          precision_matrix: Array
            The current precision matrix, shape (hidden_features, hidden_features).

        Returns
        -------
        Array
            Updated precision matrix, shape (hidden_features, hidden_features).
        
fortuna/model/utils/random_features.py