def FIM(logits_fn, theta):
    """ Fisher information matrix for a probability model log p(y|theta)
        where p is a classifier probability distribution. Averages over batch dimensions.

        Args:
            logit_fn function that maps parameters to logits of shape (*, n_classes)
            theta (array_like): parameter vector to eval Fisher at

        Returns:
            Hessian(KL(p(y|theta')||p(y|theta))) (w.r.t. theta)
    """
    xnp = get_library_fns(theta.dtype)
    probs = xnp.softmax(logits_fn(theta), axis=-1)

    def entropy(theta):
        log_probs = xnp.log_softmax(logits_fn(theta), axis=-1)
        return -xnp.sum(probs * log_probs, axis=-1).mean()
 Fisher information matrix for a probability model log p(y|theta)
        where p is a classifier probability distribution. Averages over batch dimensions.

        Args:
            logit_fn function that maps parameters to logits of shape (*, n_classes)
            theta (array_like): parameter vector to eval Fisher at

        Returns:
            Hessian(KL(p(y|theta')||p(y|theta))) (w.r.t. theta)
    
cola/ops/operators.py