def sample(
        self,
        rng: Optional[jax.Array] = None,
        inputs_loader: Optional[InputsLoader] = None,
        inputs: Optional[Array] = None,
        **kwargs,
    ) -> JointState:
        """
        Sample from the posterior distribution.

        Parameters
        ----------
        rng : Optional[jax.Array]
            A random number generator. If not passed, this will be taken from the attributes of this class.
        inputs_loader: Optional[InputsLoader]
            Input data loader. This or `inputs` is required if the posterior state includes mutable objects.
        inputs: Optional[Array]
            Input variables. This or `inputs_loader` is required if the posterior state includes mutable objects.

        Returns
        -------
        JointState
            A sample from the posterior distribution.
        """
        if rng is None:
            rng = self.rng.get()
        state = self.state.get()
        if state.mutable is not None and inputs_loader is None and inputs is None:
            raise ValueError(
                "The posterior state contains mutable objects. Please pass `inputs_loader` or `inputs`."
            )

        n_params = len(state.mean)
        rank = state.dev.shape[-1]
        which_params = decode_encoded_tuple_of_lists_of_strings_to_array(
            state._encoded_which_params
        )

        unravel = ravel_pytree(
            state.params
            if which_params is None
            else [nested_get(state.params, path) for path in which_params]
        )[1]

        coeff1 = 1 / jnp.sqrt(2)
        coeff2 = coeff1 / jnp.sqrt(rank)

        rng, key1, key2 = random.split(rng, 3)
        z1 = random.normal(key1, shape=(n_params,))
        z2 = random.normal(key2, shape=(rank,))
        if which_params is None:
            state = state.replace(
                params=self._get_sample(
                    mean=state.mean,
                    std=state.std,
                    dev=state.dev,
                    z1=z1,
                    z2=z2,
                    coeff1=coeff1,
                    coeff2=coeff2,
                    unravel=unravel,
                )
            )
        else:
            state = state.replace(
                params=FrozenDict(
                    nested_set(
                        d=state.params.unfreeze(),
                        key_paths=which_params,
                        objs=tuple(
                            self._get_sample(
                                mean=state.mean,
                                std=state.std,
                                dev=state.dev,
                                z1=z1,
                                z2=z2,
                                coeff1=coeff1,
                                coeff2=coeff2,
                                unravel=unravel,
                            )
                        ),
                    )
                )
            )

        if state.mutable:
            if inputs_loader is not None:
                for batch_inputs in inputs_loader:
                    state = state.replace(
                        mutable=self.joint.likelihood.model_manager.apply(
                            state.params,
                            batch_inputs,
                            mutable=state.mutable,
                            train=True,
                            rng=rng,
                        )[1]["mutable"]
                    )
            else:
                state = state.replace(
                    mutable=self.joint.likelihood.model_manager.apply(
                        state.params, inputs, mutable=state.mutable, train=True, rng=rng
                    )[1]["mutable"]
                )

        return JointState(
            params=state.params,
            mutable=state.mutable,
            calib_params=state.calib_params,
            calib_mutable=state.calib_mutable,
        )

        Sample from the posterior distribution.

        Parameters
        ----------
        rng : Optional[jax.Array]
            A random number generator. If not passed, this will be taken from the attributes of this class.
        inputs_loader: Optional[InputsLoader]
            Input data loader. This or `inputs` is required if the posterior state includes mutable objects.
        inputs: Optional[Array]
            Input variables. This or `inputs_loader` is required if the posterior state includes mutable objects.

        Returns
        -------
        JointState
            A sample from the posterior distribution.
        
fortuna/prob_model/posterior/swag/swag_posterior.py