def sample(
        self,
        inputs_loader: InputsLoader,
        n_samples: int = 1,
        rng: Optional[jax.Array] = None,
        distribute: bool = True,
    ) -> 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_samples : int
            Number of target samples to sample for each input data point.
        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
            Samples for each input data point.
        """
        state = self.state.get()
        return self.likelihood.sample(
            n_target_samples=n_samples,
            params=state.params,
            inputs_loader=inputs_loader,
            mutable=state.mutable,
            rng=rng,
            distribute=distribute,
        )

        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_samples : int
            Number of target samples to sample for each input data point.
        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
            Samples for each input data point.
        
fortuna/calib_model/predictive/base.py