def credible_set(
        self,
        inputs_loader: InputsLoader,
        n_posterior_samples: int = 30,
        error: float = 0.05,
        rng: Optional[jax.Array] = None,
        distribute: bool = True,
    ) -> jnp.ndarray:
        r"""
        Estimate credible sets for the target variable. This is done by sorting the class probabilities in descending order
        and including classes until the sum > 1-error.

        Parameters
        ----------
        inputs_loader : InputsLoader
            A loader of input data points.
        n_posterior_samples: int
            Number of posterior samples to draw for each input.
        error: float
            The set error. This must be a number between 0 and 1, extremes included. For example,
            `error=0.05` corresponds to a 95% level of credibility.
        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
            A credibility set for each of the inputs.
        """
        p_classes = self.mean(
            inputs_loader=inputs_loader,
            n_posterior_samples=n_posterior_samples,
            rng=rng,
            distribute=distribute,
        )
        n_classes = jnp.shape(p_classes)[1]
        labels = jnp.argsort(p_classes, axis=-1)[:, ::-1]
        p_classes_sorted = jnp.sort(p_classes, axis=-1)[:, ::-1]
        region_classes = jnp.cumsum(p_classes_sorted, axis=1) > (1 - error)

        # Convert CB region into sets
        index_true = jnp.argmax(region_classes, axis=1)  # first index where True
        credible_set = np.zeros(len(index_true), dtype=object)
        for s in np.arange(n_classes):
            idx = np.where(index_true == s)[0]
            credible_set[idx] = labels[idx, : s + 1].tolist()

        return credible_set

        Estimate credible sets for the target variable. This is done by sorting the class probabilities in descending order
        and including classes until the sum > 1-error.

        Parameters
        ----------
        inputs_loader : InputsLoader
            A loader of input data points.
        n_posterior_samples: int
            Number of posterior samples to draw for each input.
        error: float
            The set error. This must be a number between 0 and 1, extremes included. For example,
            `error=0.05` corresponds to a 95% level of credibility.
        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
            A credibility set for each of the inputs.
        
fortuna/prob_model/predictive/classification.py