def METHOD_NAME(
    self, posterior: AbstractPosterior, dataset: Dataset, key: KeyArray
) -> AbstractPosterior:
    """
    Takes a posterior and corresponding dataset and optimizes the posterior using the
    GPJax `fit` method.
    Args:
        posterior: Posterior being optimized.
        dataset: Dataset used for optimizing posterior.
        key: A JAX PRNG key for generating random numbers.
    Returns:
        Optimized posterior.
    """
    opt_posterior, _ = gpx.fit(
        model=posterior,
        objective=self.optimization_objective,
        train_data=dataset,
        optim=self.optimizer,
        num_iters=self.num_optimization_iters,
        safe=True,
        key=key,
        verbose=False,
    )
    return opt_posterior