Marginal probability of single assignment to multiple variables

Misc.
1

Let’s say I have a simple model.

def model():
    z = sample("z", Normal(0, 1))
    return sample("a", Normal(z, 1)), sample("b", Normal(z, 1))

My question is: what’s the simplest way to evaluate the marginal joint probability of a single assignment to a and b? e.g.

P(a = 1, b = 1) = \int_z P(a = 1, b = 1, z = z)dz

For example, I can use trace to find the probability of a joint assignment:

trace(condition(model, {"a": 1, "b": 1})).get_trace().log_prob_sum()

But that doesn’t marginalize over z. I can use EmpiricalMarginal to compute a marginal distribution:

EmpiricalMarginal(Importance(model).run(), sites=["a", "b"])

But this won’t assign probability mass to any individual sample I can bring, since any specific assignment like a: 1, b: 1 won’t likely be sampled.

Also, as an aside—is replay the most appropriate way to evaluate a joint probability using trace, as opposed to condition? If so, are there any facilities for manually constructing Trace objects?

1 Like