def sample(
        self,
        n_target_samples: int,
        params: Params,
        inputs_loader: InputsLoader,
        mutable: Optional[Mutable] = None,
        calib_params: Optional[CalibParams] = None,
        calib_mutable: Optional[CalibMutable] = None,
        return_aux: Optional[List[str]] = None,
        rng: Optional[jax.Array] = None,
        distribute: bool = True,
        **kwargs,
    ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, dict]]:
        """
        Sample target variables from the likelihood function for each input variable.

        Parameters
        ----------
        n_target_samples : int
            The number of samples to draw from the likelihood for each input data point.
        params : Params
            The random parameters of the probabilistic model.
        inputs_loader : InputsLoader
            A loader of input data points.
        mutable : Optional[Mutable]
            The mutable objects used to evaluate the models.
        calib_params : Optional[CalibParams]
            The calibration parameters of the probabilistic model.
        calib_mutable : Optional[CalibMutable]
            The calibration mutable objects used to evaluate the calibrators.
        return_aux : Optional[List[str]]
            The auxiliary objects to return. We support 'outputs'. If this argument is not given, no auxiliary object
            is returned.
        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]]
            The samples of the target variable for each input. If `return_aux` is given, the corresponding auxiliary
            objects are also returned.
        """
        if return_aux is None:
            return_aux = []
        supported_aux = ["outputs"]
        unsupported_aux = [s for s in return_aux if s not in supported_aux]
        if sum(unsupported_aux) > 0:
            raise Exception(
                """The auxiliary objects {} are unknown. Please make sure that all elements of `return_aux`
                            belong to the following list: {}""".format(
                    unsupported_aux, supported_aux
                )
            )

        outputs = self.get_calibrated_outputs(
            params,
            inputs_loader,
            mutable,
            calib_params,
            calib_mutable,
            distribute,
            **kwargs,
        )

        samples = self.prob_output_layer.sample(
            n_target_samples, outputs, rng=rng, **kwargs
        )
        if len(return_aux) > 0:
            return samples, dict(outputs=outputs)
        return samples

        Sample target variables from the likelihood function for each input variable.

        Parameters
        ----------
        n_target_samples : int
            The number of samples to draw from the likelihood for each input data point.
        params : Params
            The random parameters of the probabilistic model.
        inputs_loader : InputsLoader
            A loader of input data points.
        mutable : Optional[Mutable]
            The mutable objects used to evaluate the models.
        calib_params : Optional[CalibParams]
            The calibration parameters of the probabilistic model.
        calib_mutable : Optional[CalibMutable]
            The calibration mutable objects used to evaluate the calibrators.
        return_aux : Optional[List[str]]
            The auxiliary objects to return. We support 'outputs'. If this argument is not given, no auxiliary object
            is returned.
        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]]
            The samples of the target variable for each input. If `return_aux` is given, the corresponding auxiliary
            objects are also returned.
        
fortuna/likelihood/base.py