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_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_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 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_3(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
hidden_dim = int(args.hidden_dim**0.5) # split between w and x
with handlers.mask(mask=include_prior):
probs_w = pyro.sample(
"probs_w",
dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1).to_event(1))
probs_x = pyro.sample(
"probs_x",
dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1).to_event(1))
probs_y = pyro.sample(
"probs_y",
dist.Beta(0.1, 0.9).expand([hidden_dim, hidden_dim,
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]
w, x = 0, 0
for t in pyro.markov(range(max_length if args.jit else lengths.max())):
with handlers.mask(mask=(t < lengths).unsqueeze(-1)):
w = pyro.sample("w_{}".format(t),
dist.Categorical(probs_w[w]),
infer={"enumerate": "parallel"})
x = pyro.sample("x_{}".format(t),
dist.Categorical(probs_x[x]),
infer={"enumerate": "parallel"})
with tones_plate as tones:
pyro.sample("y_{}".format(t),
dist.Bernoulli(probs_y[w, x, tones]),
obs=sequences[batch, 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 model_8(weeks_data, days_data, history, vectorized):
x_dim, y_dim, w_dim, z_dim = 3, 2, 2, 3
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_probs = pyro.param("y_probs",
lambda: torch.rand(x_dim, y_dim),
constraint=constraints.simplex)
w_init = pyro.param("w_init",
lambda: torch.rand(w_dim),
constraint=constraints.simplex)
w_trans = pyro.param("w_trans",
lambda: torch.rand((w_dim, w_dim)),
constraint=constraints.simplex)
z_probs = pyro.param("z_probs",
lambda: torch.rand(w_dim, z_dim),
constraint=constraints.simplex)
x_prev = None
weeks_loop = (pyro.vectorized_markov(
name="weeks", size=len(weeks_data), dim=-1, history=history)
if vectorized else pyro.markov(range(len(weeks_data)),
history=history))
for i in weeks_loop:
if isinstance(i, int) and i == 0:
x_probs = x_init
else:
x_probs = Vindex(x_trans)[x_prev]
x_curr = pyro.sample("x_{}".format(i), dist.Categorical(x_probs))
pyro.sample(
"y_{}".format(i),
dist.Categorical(Vindex(y_probs)[x_curr]),
obs=weeks_data[i],
)
x_prev = x_curr
w_prev = None
days_loop = (pyro.vectorized_markov(
name="days", size=len(days_data), dim=-1, history=history)
if vectorized else pyro.markov(range(len(days_data)),
history=history))
for j in days_loop:
if isinstance(j, int) and j == 0:
w_probs = w_init
else:
w_probs = Vindex(w_trans)[w_prev]
w_curr = pyro.sample("w_{}".format(j), dist.Categorical(w_probs))
pyro.sample(
"z_{}".format(j),
dist.Categorical(Vindex(z_probs)[w_curr]),
obs=days_data[j],
)
w_prev = w_curr
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(z1=None, z2=None):
p = pyro.param("p", torch.tensor([0.25, 0.75]))
loc = pyro.param("loc", torch.tensor([-1.0, 1.0]))
z1 = pyro.sample("z1", dist.Categorical(p), obs=z1)
with pyro.plate("data[0]", 3):
pyro.sample("x1", dist.Normal(loc[z1], 1.0), obs=data[0])
with pyro.plate("data[1]", 2):
z2 = pyro.sample("z2", dist.Categorical(p), obs=z2)
pyro.sample("x2", dist.Normal(loc[z2], 1.0), obs=data[1])
def model(z1=None, z2=None):
p = pyro.param("p", torch.tensor([[0.25, 0.75], [0.1, 0.9]]))
loc = pyro.param("loc", torch.tensor([-1.0, 1.0]))
z1 = pyro.sample("z1", dist.Categorical(p[0]), obs=z1)
z2 = pyro.sample("z2", dist.Categorical(p[z1]), obs=z2)
logger.info("z1.shape = {}".format(z1.shape))
logger.info("z2.shape = {}".format(z2.shape))
with pyro.plate("data", 3):
pyro.sample("x1", dist.Normal(loc[z1], 1.0), obs=data[0])
pyro.sample("x2", dist.Normal(loc[z2], 1.0), obs=data[1])
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 hand_model():
probs_a = pyro.param("probs_a")
probs_b = pyro.param("probs_b")
probs_c = pyro.param("probs_c")
probs_d = pyro.param("probs_d")
for i in range(2):
a = pyro.sample("a_{}".format(i), dist.Categorical(probs_a),
infer={"enumerate": "parallel"})
pyro.sample("b_{}".format(i), dist.Categorical(probs_b[a]), obs=b_data[i])
for j in range(3):
c = pyro.sample("c_{}".format(j), dist.Categorical(probs_c),
infer={"enumerate": "parallel"})
pyro.sample("d_{}".format(j), dist.Categorical(probs_d[c]), obs=d_data[j])
def auto_model():
probs_a = pyro.param("probs_a")
probs_b = pyro.param("probs_b")
probs_c = pyro.param("probs_c")
probs_d = pyro.param("probs_d")
with pyro.plate("a_axis", 2, dim=-1):
a = pyro.sample("a", dist.Categorical(probs_a),
infer={"enumerate": "parallel"})
pyro.sample("b", dist.Categorical(probs_b[a]), obs=b_data)
with pyro.plate("c_axis", 3, dim=-1):
c = pyro.sample("c", dist.Categorical(probs_c),
infer={"enumerate": "parallel"})
pyro.sample("d", dist.Categorical(probs_d[c]), obs=d_data)
def auto_model(data):
probs_a = pyro.param("model_probs_a")
probs_b = pyro.param("model_probs_b")
probs_c = pyro.param("model_probs_c")
probs_d = pyro.param("model_probs_d")
probs_e = pyro.param("model_probs_e")
a = pyro.sample("a", dist.Categorical(probs_a))
b = pyro.sample("b", dist.Categorical(probs_b[a]),
infer={"enumerate": "parallel"})
with pyro.plate("data", 2, dim=-1):
c = pyro.sample("c", dist.Categorical(probs_c[a]))
d = pyro.sample("d", dist.Categorical(Vindex(probs_d)[b, c]),
infer={"enumerate": "parallel"})
pyro.sample("obs", dist.Categorical(probs_e[d]), obs=data)
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 hand_model(data):
probs_a = pyro.param("model_probs_a")
probs_b = pyro.param("model_probs_b")
probs_c = pyro.param("model_probs_c")
probs_d = pyro.param("model_probs_d")
probs_e = pyro.param("model_probs_e")
a = pyro.sample("a", dist.Categorical(probs_a))
b = pyro.sample("b", dist.Categorical(probs_b[a]),
infer={"enumerate": "parallel"})
for i in range(2):
c = pyro.sample("c_{}".format(i), dist.Categorical(probs_c[a]))
d = pyro.sample("d_{}".format(i),
dist.Categorical(Vindex(probs_d)[b, c]),
infer={"enumerate": "parallel"})
pyro.sample("obs_{}".format(i), dist.Categorical(probs_e[d]), obs=data[i])
def model2():
data = [torch.tensor([-1.0, -1.0, 0.0]), torch.tensor([-1.0, 1.0])]
p = pyro.param("p", torch.tensor([0.25, 0.75]))
loc = pyro.sample("loc", dist.Normal(0, 1).expand([2]).to_event(1))
# FIXME results in infinite loop in transformeddist_to_funsor.
# scale = pyro.sample("scale", dist.LogNormal(0, 1))
z1 = pyro.sample("z1", dist.Categorical(p))
scale = pyro.sample("scale", dist.Normal(torch.tensor([0.0, 1.0])[z1],
1)).exp()
with pyro.plate("data[0]", 3):
pyro.sample("x1", dist.Normal(loc[z1], scale), obs=data[0])
with pyro.plate("data[1]", 2):
z2 = pyro.sample("z2", dist.Categorical(p))
pyro.sample("x2", dist.Normal(loc[z2], scale), obs=data[1])
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 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_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_0(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))
with pyro.plate("sequences", data.shape[0], dim=-3) as sequences:
sequences = sequences[:, None]
x_prev = None
markov_loop = \
pyro.vectorized_markov(name="time", size=data.shape[1], dim=-2, history=history) if vectorized \
else pyro.markov(range(data.shape[1]), 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[2], dim=-1):
pyro.sample("y_{}".format(i),
dist.Normal(Vindex(locs)[..., x_curr], 1.),
obs=Vindex(data)[sequences, i])
x_prev = x_curr
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():
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_6(data, history, vectorized):
x_dim = 3
x_init = pyro.param("x_init",
lambda: torch.rand(x_dim),
constraint=constraints.simplex)
x_trans = pyro.param("x_trans",
lambda: torch.rand((len(data) - 1, 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:
if isinstance(i, int) and i < 1:
x_probs = x_init
elif isinstance(i, int):
x_probs = x_trans[i - 1, x_prev]
else:
x_probs = Vindex(x_trans)[(i - 1)[:, None], 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.Normal(Vindex(locs)[..., x_curr], 1.),
obs=data[i])
x_prev = x_curr
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():
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():
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:])
def model_4(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
hidden_dim = int(args.hidden_dim**0.5) # split between w and x
with handlers.mask(mask=include_prior):
probs_w = pyro.sample(
"probs_w", dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1).to_event(1)
)
probs_x = pyro.sample(
"probs_x",
dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1)
.expand_by([hidden_dim])
.to_event(2),
)
probs_y = pyro.sample(
"probs_y",
dist.Beta(0.1, 0.9).expand([hidden_dim, hidden_dim, 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]
# Note the broadcasting tricks here: we declare a hidden torch.arange and
# ensure that w and x are always tensors so we can unsqueeze them below,
# thus ensuring that the x sample sites have correct distribution shape.
w = x = torch.tensor(0, dtype=torch.long)
for t in pyro.markov(range(max_length if args.jit else lengths.max())):
with handlers.mask(mask=(t < lengths).unsqueeze(-1)):
w = pyro.sample(
"w_{}".format(t),
dist.Categorical(probs_w[w]),
infer={"enumerate": "parallel"},
)
x = pyro.sample(
"x_{}".format(t),
dist.Categorical(Vindex(probs_x)[w, x]),
infer={"enumerate": "parallel"},
)
with tones_plate as tones:
pyro.sample(
"y_{}".format(t),
dist.Bernoulli(probs_y[w, x, tones]),
obs=sequences[batch, t],
)
