BasisResearch / chirho

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

`ExtractSupports` is not fine grained enough when used with `pyro.deterministic` #559

Open PoorvaGarg opened 3 months ago

PoorvaGarg commented 3 months ago

Consider the following code snippet:

import pyro
import pyro.distributions as dist
from chirho.explainable.handlers import ExtractSupports

def model():
    X = pyro.sample("X", dist.Bernoulli(0.5))
    Y = pyro.deterministic("Y", X)
    return {"X": X, "Y": Y}

with ExtractSupports() as supports:
    model()

print(supports.supports)

Output: {'X': Boolean(), 'Y': IndependentConstraint(Real(), 0)}

Note that the supports extracted for Y says that it has a distribution over real numbers because pyro.deterministic uses a Delta distribution. But as it is obvious from the model, Y has a distribution over Booleans. When ExtractSupports is used downstream for other handlers (such as SearchForExplanation), ExtractSupports currently does not give information that is specific enough.

Either there should be a more involved static analysis to infer the supports of the variables or it should throw a warning that prompts the user to explicitly provide the support.