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

[PoorvaGarg]

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.