def _get_logpdf_fn(
    classwise_mean_features: Array, classwise_cov_features: Array
) -> Callable[[Array], Array]:
    """
    Returns a function to evaluate the log-likelihood of a test sample according to the (fitted) GMM.

    Parameters
    ----------
    classwise_mean_features: Array
        The per-class mean vector of the fitted GMM. The shape of the array is `(num_classes, d)`.
    classwise_cov_features: Array
        The per-class covariance matrix of the fitted GMM. The shape of the array is `(num_classes, d, d)`.

    Returns
    -------
    Callable[[Array], Array]
        A function to evaluate the log-likelihood of a test sample according to the (fitted) GMM.
    """
    for jitter_eps in JITTERS:
        jitter = np.expand_dims(jitter_eps * np.eye(classwise_cov_features.shape[1]), 0)
        gmm_logprob_fn_vmapped = jax.vmap(
            jsp_stats.multivariate_normal.logpdf, in_axes=(None, 0, 0)
        )
        gmm_logprob_fn = lambda x: gmm_logprob_fn_vmapped(
            x, classwise_mean_features, (classwise_cov_features + jitter)
        ).T

        nans = np.isnan(gmm_logprob_fn(classwise_mean_features)).sum()
        if nans > 0:
            logging.info(f"Nans, jittering {jitter_eps}")
            continue
        break

    return gmm_logprob_fn

    Returns a function to evaluate the log-likelihood of a test sample according to the (fitted) GMM.

    Parameters
    ----------
    classwise_mean_features: Array
        The per-class mean vector of the fitted GMM. The shape of the array is `(num_classes, d)`.
    classwise_cov_features: Array
        The per-class covariance matrix of the fitted GMM. The shape of the array is `(num_classes, d, d)`.

    Returns
    -------
    Callable[[Array], Array]
        A function to evaluate the log-likelihood of a test sample according to the (fitted) GMM.
    
fortuna/ood_detection/ddu.py