BasisResearch / chirho

An experimental language for causal reasoning
https://basisresearch.github.io/chirho/getting_started.html
Apache License 2.0
164 stars 10 forks source link

Separate influence function approximation from `influence_fn` syntax using handlers #532

Closed SamWitty closed 4 months ago

SamWitty commented 4 months ago

This refactoring PR separates out the influence_fn operation, which can be used in a way that is approximation-agnostic, from the effect handlers that implement their approximation.

In particular, this extracts what was previously included in the influence_fn operation and ports that functionality into a new MonteCarloInfluenceEstimator effect handler.

In addition, this PR modifies the tests and some documentation to reflect these changes.

Subsumes #478.

SamWitty commented 4 months ago

@eb8680 , this is ready for review except for linting errors that I'm having trouble addressing on my own. Could you take a look at the lint CI and let me know if you can think of a quick fix?

SamWitty commented 4 months ago

@eb8680 , I took your suggestion to have influence_fn return an effectful operation, rather than have influence_fn itself be effectful. This PR is now ready for review.

SamWitty commented 4 months ago

@eb8680 , I made a pass at the changes you requested. I think it's at least close now. Some things I'd like your feedback on:

  1. Would it better to define an effectful operation influence outside of the definition of influence_fn and influence_functional and then just call it?
  2. Now the linearize computation in the MonteCarloInfluenceEstimator is computed lazily during execution of influence, whereas before it was pre-computed in the influence_functional layer. Does this matter? I'm not sure how else to apply the handler only at the influence call otherwise.