Support influence functions for n-ary functionals

[eb8680]

The current implementation of chirho.robust.ops.influence_fn is defined only for functionals of one distribution, but there are many functionals of interest in machine learning, e.g. the KL divergence or mutual information, that are defined over 2 or more distributions and are known to benefit from multivariate versions of one-step correction: https://proceedings.neurips.cc/paper_files/paper/2015/hash/06138bc5af6023646ede0e1f7c1eac75-Abstract.html

Extending our implementation of influence_fn to these cases is conceptually straightforward - we just compute a tuple containing the influence of each individual distribution with the others held fixed.

Here is code illustrating the binary case in terms of the existing unary influence_fn, with guides suppressed as in #477:

def influence2(
    functional: Callable[[Callable[P, Any], Callable[P, Any]], Callable[P, S]],
    model1: Callable[P, Any],
    model2: Callable[P, Any],
    **linearize_kwargs
) -> Callable[P, Tuple[S, S]]:
    def _inf(point1: Point[A], point2: Point[B], *args, **kwargs) -> Tuple[S, S]:
        return (
            influence_fn(model1, lambda m: functional(m, model2))(point1, *args, **kwargs),
            influence_fn(model2, lambda m: functional(model1, m))(point2, *args, **kwargs),
        )
    return _inf

Tasks:

• Change influence_fn and Functional type signatures to support n-ary functionals Make influence_fn a higher-order Functional #492
• Implement general case of new n-ary influence_fn interface