def hand_guide(data):
probs_a = pyro.param("guide_probs_a")
probs_c = pyro.param("guide_probs_c")
a = pyro.sample("a", dist.Categorical(probs_a),
infer={"enumerate": "parallel"})
for i in range(2):
pyro.sample("c_{}".format(i), dist.Categorical(probs_c[a]))
def model(data):
T, N, D = data.shape # time steps, individuals, features
# Gaussian initial distribution.
init_loc = pyro.param("init_loc", torch.zeros(D))
init_scale = pyro.param("init_scale", 1e-2 * torch.eye(D),
constraint=constraints.lower_cholesky)
# Linear dynamics with Gaussian noise.
trans_const = pyro.param("trans_const", torch.zeros(D))
trans_coeff = pyro.param("trans_coeff", torch.eye(D))
noise = pyro.param("noise", 1e-2 * torch.eye(D),
constraint=constraints.lower_cholesky)
obs_plate = pyro.plate("channel", D, dim=-1)
with pyro.plate("data", N, dim=-2):
state = None
for t in range(T):
# Transition.
if t == 0:
loc = init_loc
scale_tril = init_scale
else:
loc = trans_const + funsor.torch.torch_tensordot(trans_coeff, state, 1)
scale_tril = noise
state = pyro.sample("state_{}".format(t),
dist.MultivariateNormal(loc, scale_tril),
infer={"exact": exact})
# Factorial probit likelihood model.
with obs_plate:
pyro.sample("obs_{}".format(t),
dist.Bernoulli(logits=state["channel"]),
obs=data[t])
def auto_guide(data):
probs_a = pyro.param("guide_probs_a")
probs_c = pyro.param("guide_probs_c")
a = pyro.sample("a", dist.Categorical(probs_a),
infer={"enumerate": "parallel"})
with pyro.plate("data", 2, dim=-1):
pyro.sample("c", dist.Categorical(probs_c[a]))
def factorized_guide(reparameterized):
Normal = (dist.Normal if reparameterized else
dist.testing.fakes.NonreparameterizedNormal)
pyro.sample("x0", Normal(pyro.param("q2"), math.sqrt(1.0 / depth)))
for i in range(1, depth):
pyro.sample("x{}".format(i),
Normal(0.0, math.sqrt(float(i + 1) / depth)))
def model_leaf(data, state=0, address=""):
p = pyro.param("p_leaf", torch.ones(10))
pyro.sample(
"leaf_{}".format(address),
dist.Bernoulli(p[state]),
obs=torch.tensor(1.0 if data else 0.0),
)
def model_2(sequences, lengths, args, batch_size=None, include_prior=True):
with ignore_jit_warnings():
num_sequences, max_length, data_dim = map(int, sequences.shape)
assert lengths.shape == (num_sequences, )
assert lengths.max() <= max_length
with handlers.mask(mask=include_prior):
probs_x = pyro.sample(
"probs_x",
dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1))
probs_y = pyro.sample(
"probs_y",
dist.Beta(0.1, 0.9).expand([args.hidden_dim, 2,
data_dim]).to_event(3))
tones_plate = pyro.plate("tones", data_dim, dim=-1)
with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch:
lengths = lengths[batch]
x, y = 0, 0
for t in pyro.markov(range(max_length if args.jit else lengths.max())):
with handlers.mask(mask=(t < lengths).unsqueeze(-1)):
x = pyro.sample("x_{}".format(t),
dist.Categorical(probs_x[x]),
infer={"enumerate": "parallel"})
# Note the broadcasting tricks here: to index probs_y on tensors x and y,
# we also need a final tensor for the tones dimension. This is conveniently
# provided by the plate associated with that dimension.
with tones_plate as tones:
y = pyro.sample("y_{}".format(t),
dist.Bernoulli(probs_y[x, y, tones]),
obs=sequences[batch, t]).long()
def model(reparameterized):
Normal = (dist.Normal if reparameterized else
dist.testing.fakes.NonreparameterizedNormal)
x = pyro.sample("x0", Normal(pyro.param("q2"), math.sqrt(1.0 / depth)))
for i in range(1, depth):
x = pyro.sample("x{}".format(i), Normal(x, math.sqrt(1.0 / depth)))
pyro.sample("y", Normal(x, 1.0), obs=torch.tensor(float(1)))
def model_1(data, history, vectorized):
x_dim = 3
init = pyro.param("init",
lambda: torch.rand(x_dim),
constraint=constraints.simplex)
trans = pyro.param("trans",
lambda: torch.rand((x_dim, x_dim)),
constraint=constraints.simplex)
locs = pyro.param("locs", lambda: torch.rand(x_dim))
x_prev = None
markov_loop = (pyro.vectorized_markov(
name="time", size=len(data), dim=-2, history=history) if vectorized
else pyro.markov(range(len(data)), history=history))
for i in markov_loop:
x_curr = pyro.sample(
"x_{}".format(i),
dist.Categorical(
init if isinstance(i, int) and i < 1 else trans[x_prev]),
)
with pyro.plate("tones", data.shape[-1], dim=-1):
pyro.sample(
"y_{}".format(i),
dist.Normal(Vindex(locs)[..., x_curr], 1.0),
obs=data[i],
)
x_prev = x_curr
def model_2(data, history, vectorized):
x_dim, y_dim = 3, 2
x_init = pyro.param("x_init",
lambda: torch.rand(x_dim),
constraint=constraints.simplex)
x_trans = pyro.param("x_trans",
lambda: torch.rand((x_dim, x_dim)),
constraint=constraints.simplex)
y_init = pyro.param("y_init",
lambda: torch.rand(x_dim, y_dim),
constraint=constraints.simplex)
y_trans = pyro.param(
"y_trans",
lambda: torch.rand((x_dim, y_dim, y_dim)),
constraint=constraints.simplex,
)
x_prev = y_prev = None
markov_loop = (pyro.vectorized_markov(
name="time", size=len(data), dim=-2, history=history) if vectorized
else pyro.markov(range(len(data)), history=history))
for i in markov_loop:
x_curr = pyro.sample(
"x_{}".format(i),
dist.Categorical(
x_init if isinstance(i, int) and i < 1 else x_trans[x_prev]),
)
with pyro.plate("tones", data.shape[-1], dim=-1):
y_curr = pyro.sample(
"y_{}".format(i),
dist.Categorical(y_init[x_curr] if isinstance(i, int) and i < 1
else Vindex(y_trans)[x_curr, y_prev]),
obs=data[i],
)
x_prev, y_prev = x_curr, y_curr
def model():
x_plate = pyro.plate("x_plate",
5,
subsample_size=2 if subsampling else None,
dim=-1)
y_plate = pyro.plate("y_plate",
6,
subsample_size=3 if subsampling else None,
dim=-2)
with pyro.plate("num_particles", 50, dim=-3):
with x_plate:
b = pyro.sample(
"b", dist.Beta(torch.tensor(1.1), torch.tensor(1.1)))
with y_plate:
c = pyro.sample("c", dist.Bernoulli(0.5))
with x_plate, y_plate:
d = pyro.sample("d", dist.Bernoulli(b))
# check shapes
if enumerate_ == "parallel":
assert b.shape == (50, 1, x_plate.subsample_size)
assert c.shape == (2, 1, 1, 1)
assert d.shape == (2, 1, 1, 1, 1)
elif enumerate_ == "sequential":
assert b.shape == (50, 1, x_plate.subsample_size)
assert c.shape in ((), (1, 1, 1)) # both are valid
assert d.shape in ((), (1, 1, 1)) # both are valid
else:
assert b.shape == (50, 1, x_plate.subsample_size)
assert c.shape == (50, y_plate.subsample_size, 1)
assert d.shape == (50, y_plate.subsample_size,
x_plate.subsample_size)
def model_5(data, history, vectorized):
x_dim, y_dim = 3, 2
x_init = pyro.param("x_init",
lambda: torch.rand(x_dim),
constraint=constraints.simplex)
x_init_2 = pyro.param("x_init_2",
lambda: torch.rand(x_dim, x_dim),
constraint=constraints.simplex)
x_trans = pyro.param(
"x_trans",
lambda: torch.rand((x_dim, x_dim, x_dim)),
constraint=constraints.simplex,
)
y_probs = pyro.param("y_probs",
lambda: torch.rand(x_dim, y_dim),
constraint=constraints.simplex)
x_prev = x_prev_2 = None
markov_loop = (pyro.vectorized_markov(
name="time", size=len(data), dim=-2, history=history) if vectorized
else pyro.markov(range(len(data)), history=history))
for i in markov_loop:
if isinstance(i, int) and i == 0:
x_probs = x_init
elif isinstance(i, int) and i == 1:
x_probs = Vindex(x_init_2)[x_prev]
else:
x_probs = Vindex(x_trans)[x_prev_2, x_prev]
x_curr = pyro.sample("x_{}".format(i), dist.Categorical(x_probs))
with pyro.plate("tones", data.shape[-1], dim=-1):
pyro.sample("y_{}".format(i),
dist.Categorical(Vindex(y_probs)[x_curr]),
obs=data[i])
x_prev_2, x_prev = x_prev, x_curr
def model(data=None):
probs_a = torch.tensor([0.45, 0.55])
probs_b = torch.tensor([[0.6, 0.4], [0.4, 0.6]])
probs_c = torch.tensor([[0.75, 0.25], [0.55, 0.45]])
probs_d = torch.tensor([[[0.4, 0.6], [0.3, 0.7]],
[[0.3, 0.7], [0.2, 0.8]]])
b_axis = pyro.plate("b_axis", 2)
c_axis = pyro.plate("c_axis", 2)
a = pyro.sample("a", dist.Categorical(probs_a))
b = [
pyro.sample("b_{}".format(i), dist.Categorical(probs_b[a]))
for i in b_axis
]
c = [
pyro.sample("c_{}".format(j), dist.Categorical(probs_c[a]))
for j in c_axis
]
for i in b_axis:
for j in c_axis:
pyro.sample(
"d_{}_{}".format(i, j),
dist.Categorical(Vindex(probs_d)[b[i], c[j]]),
obs=data[i, j],
)
def model(z=None):
p = pyro.param("p", torch.tensor([0.75, 0.25]))
iz = pyro.sample("z", dist.Categorical(p), obs=z)
z = torch.tensor([0.0, 1.0])[iz]
logger.info("z.shape = {}".format(z.shape))
with pyro.plate("data", 3):
pyro.sample("x", dist.Normal(z, 1.0), obs=data)
def model_0(sequences, lengths, args, batch_size=None, include_prior=True):
assert not torch._C._get_tracing_state()
num_sequences, max_length, data_dim = sequences.shape
with handlers.mask(mask=include_prior):
# Our prior on transition probabilities will be:
# stay in the same state with 90% probability; uniformly jump to another
# state with 10% probability.
probs_x = pyro.sample(
"probs_x",
dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1))
# We put a weak prior on the conditional probability of a tone sounding.
# We know that on average about 4 of 88 tones are active, so we'll set a
# rough weak prior of 10% of the notes being active at any one time.
probs_y = pyro.sample(
"probs_y",
dist.Beta(0.1, 0.9).expand([args.hidden_dim,
data_dim]).to_event(2))
# In this first model we'll sequentially iterate over sequences in a
# minibatch; this will make it easy to reason about tensor shapes.
tones_plate = pyro.plate("tones", data_dim, dim=-1)
for i in pyro.plate("sequences", len(sequences), batch_size):
length = lengths[i]
sequence = sequences[i, :length]
x = 0
for t in pyro.markov(range(length)):
# On the next line, we'll overwrite the value of x with an updated
# value. If we wanted to record all x values, we could instead
# write x[t] = pyro.sample(...x[t-1]...).
x = pyro.sample("x_{}_{}".format(i, t),
dist.Categorical(probs_x[x]),
infer={"enumerate": "parallel"})
with tones_plate:
pyro.sample("y_{}_{}".format(i, t),
dist.Bernoulli(probs_y[x.squeeze(-1)]),
obs=sequence[t])
def model_5(sequences, lengths, args, batch_size=None, include_prior=True):
with ignore_jit_warnings():
num_sequences, max_length, data_dim = map(int, sequences.shape)
assert lengths.shape == (num_sequences, )
assert lengths.max() <= max_length
# Initialize a global module instance if needed.
global tones_generator
if tones_generator is None:
tones_generator = TonesGenerator(args, data_dim)
pyro.module("tones_generator", tones_generator)
with handlers.mask(mask=include_prior):
probs_x = pyro.sample(
"probs_x",
dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1))
with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch:
lengths = lengths[batch]
x = 0
y = torch.zeros(data_dim)
for t in pyro.markov(range(max_length if args.jit else lengths.max())):
with handlers.mask(mask=(t < lengths).unsqueeze(-1)):
x = pyro.sample("x_{}".format(t),
dist.Categorical(probs_x[x]),
infer={"enumerate": "parallel"})
# Note that since each tone depends on all tones at a previous time step
# the tones at different time steps now need to live in separate plates.
with pyro.plate("tones_{}".format(t), data_dim, dim=-1):
y = pyro.sample(
"y_{}".format(t),
dist.Bernoulli(logits=tones_generator(x, y)),
obs=sequences[batch, t])
def model():
locs = pyro.param("locs", torch.randn(3), constraint=constraints.real)
scales = pyro.param("scales",
torch.randn(3).exp(),
constraint=constraints.positive)
p = torch.tensor([0.5, 0.3, 0.2])
x = pyro.sample("x", dist.Categorical(p))
pyro.sample("obs", dist.Normal(locs[x], scales[x]), obs=data)
def model():
with pyro.markov() as m:
with pyro.markov():
with m: # error here
pyro.sample(
"x",
dist.Categorical(torch.ones(4)),
infer={"enumerate": "parallel"},
)
def double_exp_model(data):
k1 = pyro.param("k1", lambda: torch.tensor(0.01), constraint=constraints.positive)
k2 = pyro.param("k2", lambda: torch.tensor(0.05), constraint=constraints.positive)
A = pyro.param("A", lambda: torch.tensor(0.5), constraint=constraints.unit_interval)
k = torch.stack([k1, k2])
with pyro.plate("data", len(data)):
m = pyro.sample("m", dist.Bernoulli(A), infer={"enumerate": "parallel"})
pyro.sample("obs", dist.Exponential(k[m.long()]), obs=data)
def model():
p = pyro.param("p", torch.ones(3, 3))
x = pyro.sample("x", dist.Categorical(p[0]))
y = x
for i in pyro.markov(range(10)):
y = pyro.sample("y_{}".format(i), dist.Categorical(p[y]))
z = y
for j in pyro.markov(range(10)):
z = pyro.sample("z_{}_{}".format(i, j), dist.Categorical(p[z]))
def model():
with pyro.plate("plate", 10, subsample_size=subsample_size, dim=None):
p0 = torch.tensor(0.)
p0 = pyro.subsample(p0, event_dim=0)
assert p0.shape == ()
p = 0.5 * torch.ones(10)
p = pyro.subsample(p, event_dim=0)
assert len(p) == (subsample_size if subsample_size else 10)
pyro.sample("x", dist.Bernoulli(p))
def model():
pyro.sample("w", dist.Bernoulli(0.5), infer={'enumerate': 'parallel'})
with pyro.plate("non_enum", 2):
a = pyro.sample("a", dist.Bernoulli(0.5), infer={'enumerate': None})
p = (1.0 + a.sum(-1)) / (2.0 + a.shape[0]) # introduce dependency of b on a
with pyro.plate("enum_1", 3):
pyro.sample("b", dist.Bernoulli(p), infer={'enumerate': enumerate_})
def model():
p = pyro.param("p", 0.25 * torch.ones(2, 2))
q = pyro.param("q", 0.25 * torch.ones(2))
x_prev = torch.tensor(0)
x_curr = torch.tensor(0)
for t in pyro.markov(range(10), history=history):
probs = p[x_prev, x_curr]
x_prev, x_curr = x_curr, pyro.sample("x_{}".format(t), dist.Bernoulli(probs)).long()
pyro.sample("y_{}".format(t), dist.Bernoulli(q[x_curr]),
obs=torch.tensor(0.))
def model():
x = pyro.sample("x0", dist.Categorical(pyro.param("q0")))
with pyro.plate("local", 3):
for i in range(1, depth):
x = pyro.sample(
"x{}".format(i),
dist.Categorical(pyro.param("q{}".format(i))[..., x, :]))
with pyro.plate("data", 4):
pyro.sample("y",
dist.Bernoulli(pyro.param("qy")[..., x]),
obs=data)
def hmm(data, hidden_dim=10):
transition = 0.3 / hidden_dim + 0.7 * torch.eye(hidden_dim)
means = torch.arange(float(hidden_dim))
states = [0]
for t in pyro.markov(range(len(data))):
states.append(
pyro.sample("states_{}".format(t),
dist.Categorical(transition[states[-1]])))
data[t] = pyro.sample("obs_{}".format(t),
dist.Normal(means[states[-1]], 1.0),
obs=data[t])
return states, data
def model_6(sequences, lengths, args, batch_size=None, include_prior=False):
num_sequences, max_length, data_dim = sequences.shape
assert lengths.shape == (num_sequences, )
assert lengths.max() <= max_length
hidden_dim = args.hidden_dim
if not args.raftery_parameterization:
# Explicitly parameterize the full tensor of transition probabilities, which
# has hidden_dim cubed entries.
probs_x = pyro.param("probs_x",
torch.rand(hidden_dim, hidden_dim, hidden_dim),
constraint=constraints.simplex)
else:
# Use the more parsimonious "Raftery" parameterization of
# the tensor of transition probabilities. See reference:
# Raftery, A. E. A model for high-order markov chains.
# Journal of the Royal Statistical Society. 1985.
probs_x1 = pyro.param("probs_x1",
torch.rand(hidden_dim, hidden_dim),
constraint=constraints.simplex)
probs_x2 = pyro.param("probs_x2",
torch.rand(hidden_dim, hidden_dim),
constraint=constraints.simplex)
mix_lambda = pyro.param("mix_lambda",
torch.tensor(0.5),
constraint=constraints.unit_interval)
# we use broadcasting to combine two tensors of shape (hidden_dim, hidden_dim) and
# (hidden_dim, 1, hidden_dim) to obtain a tensor of shape (hidden_dim, hidden_dim, hidden_dim)
probs_x = mix_lambda * probs_x1 + (1.0 -
mix_lambda) * probs_x2.unsqueeze(-2)
probs_y = pyro.param("probs_y",
torch.rand(hidden_dim, data_dim),
constraint=constraints.unit_interval)
tones_plate = pyro.plate("tones", data_dim, dim=-1)
with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch:
lengths = lengths[batch]
x_curr, x_prev = torch.tensor(0), torch.tensor(0)
# we need to pass the argument `history=2' to `pyro.markov()`
# since our model is now 2-markov
for t in pyro.markov(range(lengths.max()), history=2):
with handlers.mask(mask=(t < lengths).unsqueeze(-1)):
probs_x_t = Vindex(probs_x)[x_prev, x_curr]
x_prev, x_curr = x_curr, pyro.sample(
"x_{}".format(t),
dist.Categorical(probs_x_t),
infer={"enumerate": "parallel"})
with tones_plate:
probs_y_t = probs_y[x_curr.squeeze(-1)]
pyro.sample("y_{}".format(t),
dist.Bernoulli(probs_y_t),
obs=sequences[batch, t])
def model(data, state=0, address=""):
if isinstance(data, bool):
p = pyro.param("p_leaf", torch.ones(10))
pyro.sample("leaf_{}".format(address),
dist.Bernoulli(p[state]),
obs=torch.tensor(1. if data else 0.))
else:
assert isinstance(data, tuple)
p = pyro.param("p_branch", torch.ones(10, 10))
for branch, letter in zip(data, "abcdefg"):
next_state = pyro.sample("branch_{}".format(address + letter),
dist.Categorical(p[state]),
infer={"enumerate": "parallel"})
model(branch, next_state, address + letter)
def model():
p = pyro.param("p", torch.ones(3, 3))
q = pyro.param("q", torch.tensor([0.5, 0.5]))
plate_x = pyro.plate("plate_x",
4,
subsample_size=3 if subsampling else None,
dim=-1)
plate_y = pyro.plate("plate_y",
5,
subsample_size=3 if subsampling else None,
dim=-1)
plate_z = pyro.plate("plate_z",
6,
subsample_size=3 if subsampling else None,
dim=-2)
a = pyro.sample("a", dist.Bernoulli(q[0])).long()
w = 0
for i in pyro.markov(range(4)):
w = pyro.sample("w_{}".format(i), dist.Categorical(p[w]))
with plate_x:
b = pyro.sample("b", dist.Bernoulli(q[a])).long()
x = 0
for i in pyro.markov(range(4)):
x = pyro.sample("x_{}".format(i), dist.Categorical(p[x]))
with plate_y:
c = pyro.sample("c", dist.Bernoulli(q[a])).long()
y = 0
for i in pyro.markov(range(4)):
y = pyro.sample("y_{}".format(i), dist.Categorical(p[y]))
with plate_z:
d = pyro.sample("d", dist.Bernoulli(q[a])).long()
z = 0
for i in pyro.markov(range(4)):
z = pyro.sample("z_{}".format(i), dist.Categorical(p[z]))
with plate_x, plate_z:
# this part is tricky: how do we know to preserve b's dimension?
# also, how do we know how to make b and d have different dimensions?
e = pyro.sample("e",
dist.Bernoulli(q[b if reuse_plate else a])).long()
xz = 0
for i in pyro.markov(range(4)):
xz = pyro.sample("xz_{}".format(i), dist.Categorical(p[xz]))
return a, b, c, d, e
def model():
p = torch.tensor([[0.2, 0.8], [0.1, 0.9]])
xs = [0]
for t in pyro.markov(range(100), history=history):
xs.append(pyro.sample("x_{}".format(t), dist.Categorical(p[xs[-1]])))
assert all(x.dim() <= history + 1 for x in xs[1:])
def model():
p = pyro.param("p", torch.ones(3, 3))
q = pyro.param("q", torch.ones(2))
r = pyro.param("r", torch.ones(3, 2, 4))
x = 0
times = pyro.markov(range(100)) if markov else range(11)
for t in times:
x = pyro.sample("x_{}".format(t), dist.Categorical(p[x]))
y = pyro.sample("y_{}".format(t), dist.Categorical(q))
if use_vindex:
probs = Vindex(r)[x, y]
else:
z_ind = torch.arange(4, dtype=torch.long)
probs = r[x.unsqueeze(-1), y.unsqueeze(-1), z_ind]
pyro.sample("z_{}".format(t), dist.Categorical(probs),
obs=torch.tensor(0.))
def model():
p = torch.tensor([[0.2, 0.8], [0.1, 0.9]])
xs = [0]
for t in pyro.markov(range(10), history=history):
xs.append(pyro.sample("x_{}".format(t), dist.Categorical(p[xs[-1]]),
infer={"enumerate": ("sequential", "parallel")[t % 2]}))
assert all(x.dim() <= history + 1 for x in xs[1:])