def init(
        cls,
        params: Params,
        mutable: Optional[Mutable] = None,
        optimizer: Optional[OptaxOptimizer] = None,
        calib_params: Optional[CalibParams] = None,
        calib_mutable: Optional[CalibMutable] = None,
        grad_accumulated: Optional[jnp.ndarray] = None,
        dynamic_scale: Optional[dynamic_scale.DynamicScale] = None,
        **kwargs,
    ) -> Any:
        """
        Initialize a posterior distribution state.

        Parameters
        ----------
        params : Params
            The parameters characterizing an approximation of the posterior distribution.
        optimizer : Optional[OptaxOptimizer]
            An Optax optimizer associated with the posterior state.
        mutable : Optional[Mutable]
            The mutable objects characterizing an approximation of the posterior distribution.
        calib_params : Optional[CalibParams]
            The parameters objects characterizing an approximation of the posterior distribution.
        calib_mutable : Optional[CalibMutable]
            The calibration mutable objects characterizing an approximation of the posterior distribution.
        grad_accumulated : Optional[jnp.ndarray]
            The gradients accumulated in consecutive training steps (used only when `gradient_accumulation_steps > 1`).
        dynamic_scale: Optional[dynamic_scale.DynamicScale]
            Dynamic loss scaling for mixed precision gradients.
        Returns
        -------
        Any
            A posterior distribution state.
        """
        return cls(
            apply_fn=None,
            params=params,
            opt_state=(
                kwargs["opt_state"]
                if optimizer is None and "opt_state" in kwargs
                else optimizer.init(params)
            ),
            mutable=mutable,
            step=kwargs.get("step", 0),
            tx=optimizer,
            calib_params=calib_params,
            calib_mutable=calib_mutable,
            dynamic_scale=dynamic_scale,
            grad_accumulated=grad_accumulated,
            **{
                k: v
                for k, v in kwargs.items()
                if k not in ["opt_state", "apply_fn", "tx", "step"]
            },
        )

        Initialize a posterior distribution state.

        Parameters
        ----------
        params : Params
            The parameters characterizing an approximation of the posterior distribution.
        optimizer : Optional[OptaxOptimizer]
            An Optax optimizer associated with the posterior state.
        mutable : Optional[Mutable]
            The mutable objects characterizing an approximation of the posterior distribution.
        calib_params : Optional[CalibParams]
            The parameters objects characterizing an approximation of the posterior distribution.
        calib_mutable : Optional[CalibMutable]
            The calibration mutable objects characterizing an approximation of the posterior distribution.
        grad_accumulated : Optional[jnp.ndarray]
            The gradients accumulated in consecutive training steps (used only when `gradient_accumulation_steps > 1`).
        dynamic_scale: Optional[dynamic_scale.DynamicScale]
            Dynamic loss scaling for mixed precision gradients.
        Returns
        -------
        Any
            A posterior distribution state.
        
fortuna/prob_model/posterior/state.py