def _mean_field_logits(
        self, logits: Array, covariance: Optional[Array] = None
    ) -> Array:
        """
        Adjust the model logits s.t. its softmax approximates the posterior mean
        (`Zhiyun L. et al., 2020 <https://arxiv.org/abs/2006.07584>`_).

        Parameters
        ----------
          logits: Array
            A float tensor of shape (batch_size, num_classes).
          covariance: Array
            A float tensor of shape (batch_size, batch_size) or (batch_size,). If None then it
            assumes the covariance is an identity matrix.

        Returns
        ----------
        Array
            The adjusted model logits.
        """
        # Implementation adapted from https://github.com/google/edward2/blob/520e28285e905e0021e49b52b982ee5ea170071c/edward2/tensorflow/layers/utils.py#L379
        if self.mean_field_factor < 0:
            raise ValueError("mean_field_factor cannot be < 0.")

        # Compute standard deviation.
        if covariance is None:
            variances = 1.0
        else:
            variances = jnp.diagonal(covariance) if covariance.ndim == 2 else covariance

        # Compute scaling coefficient for mean-field approximation.
        logits_scale = jnp.sqrt(1.0 + variances * self.mean_field_factor)

        # Cast logits_scale to compatible dimension.
        if len(logits.shape) > 1:
            logits_scale = jnp.expand_dims(logits_scale, axis=-1)

        return logits / logits_scale

        Adjust the model logits s.t. its softmax approximates the posterior mean
        (`Zhiyun L. et al., 2020 <https://arxiv.org/abs/2006.07584>`_).

        Parameters
        ----------
          logits: Array
            A float tensor of shape (batch_size, num_classes).
          covariance: Array
            A float tensor of shape (batch_size, batch_size) or (batch_size,). If None then it
            assumes the covariance is an identity matrix.

        Returns
        ----------
        Array
            The adjusted model logits.
        
fortuna/model/model_manager/classification.py