This indeed helps. The following code works as expected:
import torch
import pyro.infer
import pyro.distributions as dist
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
def geometric(p, t=None):
if t is None:
t = 0
x = pyro.sample("x_{}".format(t), pyro.distributions.Bernoulli(p))
return torch.tensor(0) if x else 1 + geometric(p, t + 1)
# make the marginal distributions on return values explicit:
g = pyro.infer.Importance(geometric, num_samples=1000)
marginal = pyro.infer.EmpiricalMarginal(g.run(0.6)) # geometric(0.6)
samples = marginal.sample([1000])
labels, counts = np.unique(samples, return_counts=True, axis=0)
probs = counts/np.sum(counts)
plt.bar(range(len(labels)), probs, align='center', tick_label=[str(w) for w in labels])
My original version of geometric was based on the version in part 1 of the tutorial. It did not return a tensor, so I didn’t realize that that was a problem.
I read over tutorial part 1 and 2, but still found them somewhat confusing, perhaps because I had previously worked through PMC. For example, part 1 defines the geometric distribution as an example of the flexibility of Pyro, then later says “Pyro stochastic functions are universal, i.e., they can be used to represent any computable probability distribution.” I interpreted this to mean that I could use geometric in the same way as a built-in distribution, such as by writing:
A = pyro.sample("A", geometric(0.6))
But, of course, that doesn’t work. This raises the question of how to condition on the results of a recursive stochastic function, such as geometric, or how to use abstraction, such as by defining a flip() operator to match that in Church/WebPPL.
I’m not familiar with PyTorch, so I’ll definitely check out the tutorials before going much further with Pyro.