def _gnn_approx(
        self,
        params: Params,
        train_data_loader: DataLoader,
        mutable: Optional[Mutable] = None,
        calib_params: Optional[CalibParams] = None,
        calib_mutable: Optional[CalibMutable] = None,
        which_params: Optional[Tuple[List[AnyKey, ...]]] = None,
        factorization: str = "diagonal",
        verbose: bool = True,
    ) -> Params:
        """
        Estimate a standard deviation for each parameter using a diagonal Generalized Gauss-Newton Hessian
        approximation.

        Parameters
        ----------
        params : Params
            The random parameters of the probabilistic model.
        train_data_loader: DataLoader
            A training data loader.
        mutable: Optional[Mutable]
            Mutable objects.
        calib_params : Optional[CalibParams]
            The calibration parameters of the probabilistic model.
        calib_mutable : Optional[CalibMutable] = None
            The calibration mutable objects used to evaluate the calibrators.
        which_params : Optional[Tuple[List[AnyKey, ...]]]
            Sequences of keys indicating which parameters to compute the Hessian upon.
        factorization: str = "diagonal"
            Factorization of the GGN approximation. Currently, only "diagonal" is supported.
        verbose: bool
            Whether to log the training progress.

        Returns
        -------
        Params
            An estimate of the likelihood standard deviation for each random parameter.
        """
        rav, unravel = ravel_pytree(
            tuple([nested_get(params, keys) for keys in which_params])
            if which_params
            else params
        )

        def get_params_from_rav(_rav):
            unrav = unravel(_rav)
            return (
                FrozenDict(
                    nested_set(
                        params.unfreeze(),
                        which_params,
                        unrav,
                    )
                )
                if which_params
                else unrav
            )

        def apply_calib_model_manager(_params, _batch_inputs):
            outputs = self.joint.likelihood.model_manager.apply(
                _params, _batch_inputs, mutable=mutable, train=False
            )
            outputs = self.joint.likelihood.output_calib_manager.apply(
                params=(
                    calib_params["output_calibrator"]
                    if calib_params is not None
                    else None
                ),
                mutable=(
                    calib_mutable["output_calibrator"]
                    if calib_mutable is not None
                    else None
                ),
                outputs=outputs,
            )
            return outputs

        apply_fn = lambda _rav, x: apply_calib_model_manager(
            get_params_from_rav(_rav), x
        ).squeeze(0)
        vjp_fn = lambda x: vjp(lambda _rav: apply_fn(_rav, x), rav)[1]

        def eig_hess_fn(vars):
            hess = hessian(
                lambda __o: self.joint.likelihood.prob_output_layer.log_prob(
                    __o, vars[1]
                )
            )(vars[0])
            return jnp.linalg.eigh(hess)

        def compute_hess_batch(_batch_inputs, _batch_targets):
            lam, z = lax.map(
                eig_hess_fn,
                (apply_calib_model_manager(params, _batch_inputs), _batch_targets),
            )
            ztj = lax.map(
                lambda v: lax.map(vjp_fn(v[0]), v[1].T),
                (tree_map(lambda x: x[:, None], _batch_inputs), z),
            )[0]
            if factorization == "diagonal":
                return -jnp.sum(lam[:, :, None] * ztj**2, (0, 1))
            raise ValueError(
                f"`factorization={factorization}` not recognized. Currently, only "
                f"`factorization='diagonal'` is supported."
            )

        n_gpu_devices = len([d for d in devices() if d.platform == "gpu"])
        if n_gpu_devices > 0:
            train_data_loader = DeviceDimensionAugmentedLoader(train_data_loader)
            compute_hess_batch = pmap(compute_hess_batch, axis_name="batch")
        else:
            compute_hess_batch = jit(compute_hess_batch)

        h = 0.0
        for i, (batch_inputs, batch_targets) in enumerate(train_data_loader):
            if verbose:
                logging.info(f"Hessian approximation for batch {i + 1}.")
            h += compute_hess_batch(batch_inputs, batch_targets)

        if n_gpu_devices > 0:
            h = jnp.sum(h, 0)
        return unravel(h)

        Estimate a standard deviation for each parameter using a diagonal Generalized Gauss-Newton Hessian
        approximation.

        Parameters
        ----------
        params : Params
            The random parameters of the probabilistic model.
        train_data_loader: DataLoader
            A training data loader.
        mutable: Optional[Mutable]
            Mutable objects.
        calib_params : Optional[CalibParams]
            The calibration parameters of the probabilistic model.
        calib_mutable : Optional[CalibMutable] = None
            The calibration mutable objects used to evaluate the calibrators.
        which_params : Optional[Tuple[List[AnyKey, ...]]]
            Sequences of keys indicating which parameters to compute the Hessian upon.
        factorization: str = "diagonal"
            Factorization of the GGN approximation. Currently, only "diagonal" is supported.
        verbose: bool
            Whether to log the training progress.

        Returns
        -------
        Params
            An estimate of the likelihood standard deviation for each random parameter.
        
fortuna/prob_model/posterior/laplace/laplace_posterior.py