def sample(
        self,
        inputs_loader: InputsLoader,
        n_target_samples: int = 1,
        return_aux: Optional[List[str]] = None,
        rng: Optional[jax.Array] = None,
        distribute: bool = True,
        **kwargs,
    ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Dict[str, jnp.ndarray]]]:
        r"""
        Sample from an approximation of the predictive distribution for each input data point, that is

        .. math::
            y^{(i)}\sim p(\cdot|x, \mathcal{D}),

        where:
         - :math:`x` is an observed input variable;
         - :math:`\mathcal{D}` is the observed training data set;
         - :math:`y^{(i)}` is a sample of the target variable for the input :math:`x`.

        Parameters
        ----------
        inputs_loader : InputsLoader
            A loader of input data points.
        n_target_samples : int
            Number of target samples to sample for each input data point.
        return_aux : Optional[List[str]]
            Return auxiliary objects. We currently support 'outputs'.
        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
        -------
        Union[jnp.ndarray, Tuple[jnp.ndarray, Dict[str, jnp.ndarray]]]
            Samples for each input data point. Optionally, an auxiliary object is returned.
        """
        if not rng:
            rng = self.rng.get()

        def fun(_inputs):
            return self._batched_sample(
                _inputs, n_target_samples, return_aux, rng, **kwargs
            )

        if distribute:
            inputs_loader = DeviceDimensionAugmentedLoader(inputs_loader)
            fun = pmap(fun)
            if return_aux is None or len(return_aux) == 0:
                return jnp.concatenate(
                    [
                        self.likelihood._unshard_array(fun(inputs))
                        for inputs in inputs_loader
                    ],
                    1,
                )
            else:
                samples, aux_outputs = [], []
                for inputs in inputs_loader:
                    _samples, _aux = fun(inputs)
                    samples.append(self.likelihood._unshard_array(_samples))
                    if "outputs" in _aux:
                        aux_outputs.append(
                            self.likelihood._unshard_array(_aux["outputs"])
                        )
                samples = jnp.concatenate(samples, axis=0)
                aux = dict()
                if "outputs" in aux:
                    aux["outputs"] = jnp.concatenate(aux_outputs, axis=0)
                return samples, aux
        else:
            fun = jit(fun)
            if return_aux is None or len(return_aux) == 0:
                return jnp.concatenate([fun(inputs) for inputs in inputs_loader], 1)
            else:
                samples, aux_outputs = [], []
                for inputs in inputs_loader:
                    _samples, _aux = fun(inputs)
                    samples.append(_samples)
                    if "outputs" in _aux:
                        aux_outputs.append(_aux["outputs"])
                samples = jnp.concatenate(samples, axis=0)
                aux = dict()
                if "outputs" in aux:
                    aux["outputs"] = jnp.concatenate(aux_outputs, axis=0)
                return samples, aux

        Sample from an approximation of the predictive distribution for each input data point, that is

        .. math::
            y^{(i)}\sim p(\cdot|x, \mathcal{D}),

        where:
         - :math:`x` is an observed input variable;
         - :math:`\mathcal{D}` is the observed training data set;
         - :math:`y^{(i)}` is a sample of the target variable for the input :math:`x`.

        Parameters
        ----------
        inputs_loader : InputsLoader
            A loader of input data points.
        n_target_samples : int
            Number of target samples to sample for each input data point.
        return_aux : Optional[List[str]]
            Return auxiliary objects. We currently support 'outputs'.
        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
        -------
        Union[jnp.ndarray, Tuple[jnp.ndarray, Dict[str, jnp.ndarray]]]
            Samples for each input data point. Optionally, an auxiliary object is returned.
        
fortuna/prob_model/predictive/base.py