Efficient influence function constructed from influence_fn errors if called multiple times

[SamWitty]

chirho.robust.ops.influence_fn at https://github.com/BasisResearch/chirho/blob/staging-robust/chirho/robust/ops.py#L18 return a function for approximating the efficient influence function at a collection of data points. If the returned function is called multiple times it results in the following (somewhat opaque) error:

RuntimeError: !at::functionalization::impl::isFunctionalTensor(base_) INTERNAL ASSERT FAILED at "/Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/FunctionalStorageImpl.cpp":101, please report a bug to PyTorch.

For example, this error can be seen in the current tests by adding multiple calls to the generated function as follows:
https://github.com/BasisResearch/chirho/blob/sw-multiple-eif/tests/robust/test_ops.py#L80

[SamWitty]

@agrawalraj , this is now somewhat blocking #484 , as (our version of) TMLE involves evaluating the influence function on test_points, and then evaluating the influence function on samples drawn from the model and guide. I could come up with a workaround that only evaluated the influence function once by batching these two datasets into a single call, hence why I say "somewhat" blocking.

[eb8680]

I can't reproduce this on my machine with torch==2.1.2. What version of PyTorch were you using?

[SamWitty]

I'm using torch==2.0.1

[SamWitty]

I believe that is consistent with the (implied) requirements of ChiRho via Pyro.

https://github.com/pyro-ppl/pyro/blob/dev/setup.py#L105

[eb8680]

You're right, but I think the easiest short-term fix for this would just be to add a torch>=2.1.0 requirement specifically for chirho.robust.

[SamWitty]

When I upgrade torch to 2.1 I still get the same error. I made this "dummy" PR #494 to test this in a fresh environment with the bumped versions. Hopefully this tells us if/where the problem is.

[SamWitty]

It looks like this error doesn't appear in the CI builds even with torch==2.0.1. I'll see if I can reproduce the error another way...