def epistemic_entropy(
        self,
        inputs_loader: InputsLoader,
        n_posterior_samples: int = 30,
        n_target_samples: int = 30,
        rng: Optional[jax.Array] = None,
        distribute: bool = True,
    ) -> jnp.ndarray:
        r"""
        Estimate the predictive epistemic entropy, that is

        .. math::
            -\mathbb{E}_{Y|x, \mathcal{D}}[\log p(Y|x, \mathcal{D})] +
            \mathbb{E}_{W|\mathcal{D}}[\mathbb{E}_{Y|W, x}[\log p(Y|W, x)]],

        where:
         - :math:`x` is an observed input variable;
         - :math:`Y` is a random target variable;
         - :math:`\mathcal{D}` is the observed training data set;
         - :math:`W` denotes the random model parameters.

         Note that the epistemic entropy above is defined as the difference between the predictive entropy and the
         aleatoric predictive entropy.

        Parameters
        ----------
        inputs_loader : InputsLoader
            A loader of input data points.
        n_posterior_samples : int
            Number of samples to draw from the posterior distribution for each input.
        n_target_samples: int
            Number of target samples to draw for each input.
        rng : Optional[jax.Array]
            A random number generator. If not passed, this will be taken from the attributes of this class.
        distribute: bool
            Whether to distribute computation over multiple devices, if available.

        Returns
        -------
        jnp.ndarray
            An estimate of the predictive epistemic entropy for each input.
        """
        if rng is None:
            rng = self.rng.get()
        key1, *keys = random.split(rng, 1 + n_posterior_samples)

        ensemble_outputs = self.sample_calibrated_outputs(
            inputs_loader=inputs_loader,
            n_output_samples=n_posterior_samples,
            rng=key1,
            distribute=distribute,
        )

        ensemble_target_samples = lax.map(
            lambda variables: self.likelihood.prob_output_layer.sample(
                n_target_samples, variables[0], rng=variables[1]
            ),
            (ensemble_outputs, jnp.array(keys)),
        )

        def fun(i, _curr_sum):
            @vmap
            def _log_pred_fun(target_sample: jnp.ndarray):
                logps = self.likelihood.prob_output_layer.log_prob(
                    ensemble_outputs, target_sample
                )
                return jsp.special.logsumexp(logps, 0) - jnp.log(n_posterior_samples)

            log_preds = _log_pred_fun(ensemble_target_samples[i])
            log_liks = self.likelihood.prob_output_layer.log_prob(
                ensemble_outputs[i], ensemble_target_samples[i]
            )
            _curr_sum -= jnp.mean(log_preds - log_liks, 0)
            return _curr_sum

        curr_sum = fun(0, 0.0)
        curr_sum = lax.fori_loop(1, n_posterior_samples, fun, curr_sum)
        return curr_sum / n_posterior_samples

        Estimate the predictive epistemic entropy, that is

        .. math::
            -\mathbb{E}_{Y|x, \mathcal{D}}[\log p(Y|x, \mathcal{D})] +
            \mathbb{E}_{W|\mathcal{D}}[\mathbb{E}_{Y|W, x}[\log p(Y|W, x)]],

        where:
         - :math:`x` is an observed input variable;
         - :math:`Y` is a random target variable;
         - :math:`\mathcal{D}` is the observed training data set;
         - :math:`W` denotes the random model parameters.

         Note that the epistemic entropy above is defined as the difference between the predictive entropy and the
         aleatoric predictive entropy.

        Parameters
        ----------
        inputs_loader : InputsLoader
            A loader of input data points.
        n_posterior_samples : int
            Number of samples to draw from the posterior distribution for each input.
        n_target_samples: int
            Number of target samples to draw for each input.
        rng : Optional[jax.Array]
            A random number generator. If not passed, this will be taken from the attributes of this class.
        distribute: bool
            Whether to distribute computation over multiple devices, if available.

        Returns
        -------
        jnp.ndarray
            An estimate of the predictive epistemic entropy for each input.
        
fortuna/prob_model/predictive/regression.py