def sample(
        self,
        rng: Optional[jax.Array] = None,
        input_shape: Optional[Tuple[int, ...]] = None,
        inputs_loader: Optional[InputsLoader] = None,
        inputs: Optional[Array] = None,
        **kwargs,
    ) -> JointState:
        """
        Sample from the posterior distribution. Either `input_shape` or `_inputs_loader` must be passed.
        Parameters
        ----------
        rng : Optional[jax.Array]
            A random number generator. If not passed, this will be taken from the attributes of this class.
        input_shape: Optional[Tuple[int, ...]]
            Shape of a single input.
        inputs_loader: Optional[InputsLoader]
            Input data loader. If `input_shape` is passed, then `inputs` and `inputs_loader` are ignored.
        inputs: Optional[Array]
            Input variables.
        Returns
        -------
        JointState
            A sample from the posterior distribution.
        """
        if rng is None:
            rng = self.rng.get()
        state = self.state.get()

        _base = self._base
        _architecture = self._architecture
        _unravel = self._unravel

        if state._encoded_which_params is None:
            means = _unravel(
                _architecture.forward(
                    {
                        s: ravel_pytree(
                            {k: v["params"][s] for k, v in state.params.items()}
                        )[0]
                        for s in ["mean", "log_std"]
                    },
                    _base.sample(rng),
                )[0][0]
            )
        else:
            which_params = decode_encoded_tuple_of_lists_of_strings_to_array(
                state._encoded_which_params
            )
            means, log_stds = nested_unpair(
                d=state.params.unfreeze(),
                key_paths=which_params,
                labels=("mean", "log_std"),
            )
            rav_params = {
                k: ravel_pytree([nested_get(d=d, keys=path) for path in which_params])[
                    0
                ]
                for k, d in zip(["mean", "log_std"], [means, log_stds])
            }
            rav_params = _architecture.forward(params=rav_params, u=_base.sample(rng))[
                0
            ][0]

            means = _unravel(rav_params)

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

        Sample from the posterior distribution. Either `input_shape` or `_inputs_loader` must be passed.
        Parameters
        ----------
        rng : Optional[jax.Array]
            A random number generator. If not passed, this will be taken from the attributes of this class.
        input_shape: Optional[Tuple[int, ...]]
            Shape of a single input.
        inputs_loader: Optional[InputsLoader]
            Input data loader. If `input_shape` is passed, then `inputs` and `inputs_loader` are ignored.
        inputs: Optional[Array]
            Input variables.
        Returns
        -------
        JointState
            A sample from the posterior distribution.
        
fortuna/prob_model/posterior/normalizing_flow/advi/advi_posterior.py