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_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 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_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(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 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(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(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(): p = torch.tensor(0.5, requires_grad=True) with pyro.plate("plate_outer", 5, dim=plate_dims[0]): pyro.sample("x", dist.Bernoulli(p)) with pyro.plate("plate_inner_1", 6, dim=plate_dims[1]): pyro.sample("y", dist.Bernoulli(p)) with pyro.plate("plate_inner_2", 7, dim=plate_dims[2]): pyro.sample("z", dist.Bernoulli(p)) with pyro.plate("plate_inner_3", 8, dim=plate_dims[3]): pyro.sample("q", 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(): 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 model(): d = dist.Categorical(p) context1 = pyro.plate("outer", outer_dim, dim=-1) context2 = pyro.plate("inner", inner_dim, dim=-2) pyro.sample("w", d) with context1: pyro.sample("x", d) with context2: pyro.sample("y", d) with context1, context2: pyro.sample("z", d)
def guide(): d = dist.Categorical(pyro.param("q")) context1 = pyro.plate("outer", outer_dim, dim=-1) context2 = pyro.plate("inner", inner_dim, dim=-2) pyro.sample("w", d, infer={"enumerate": "parallel"}) with context1: pyro.sample("x", d, infer={"enumerate": "parallel"}) with context2: pyro.sample("y", d, infer={"enumerate": "parallel"}) with context1, context2: pyro.sample("z", d, infer={"enumerate": "parallel"})
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 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 model(): pyro.sample("w", dist.Bernoulli(0.5), infer={'enumerate': 'parallel'}) inner_plate = pyro.plate("plate", 10, subsample_size=4 if subsampling else None) for i in pyro.plate( "iplate", 10, subsample=torch.arange(3) if subsampling else None): pyro.sample("y_{}".format(i), dist.Bernoulli(0.5)) with inner_plate: pyro.sample("x_{}".format(i), dist.Bernoulli(0.5), infer={'enumerate': enumerate_})
def model(): x_plate = pyro.plate("x_plate", 10, dim=-1) y_plate = pyro.plate("y_plate", 11, dim=-2) q = pyro.param("q", torch.tensor([0.5, 0.5])) pyro.sample("a", dist.Bernoulli(0.5)) with x_plate: b = pyro.sample("b", dist.Bernoulli(0.5)).long() with y_plate: # Note that it is difficult to check that c does not depend on b. c = pyro.sample("c", dist.Bernoulli(0.5)).long() with x_plate, y_plate: pyro.sample("d", dist.Bernoulli(Vindex(q)[b] if reuse_plate else 0.5)) assert c.shape != b.shape or enumerate_ == "sequential"
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 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_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_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(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_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 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_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_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], )
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(): 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(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)) with c_axis: c = pyro.sample("c", dist.Categorical(probs_c[a])) for i in b_axis: b_i = pyro.sample("b_{}".format(i), dist.Categorical(probs_b[a])) with c_axis: pyro.sample("d_{}".format(i), dist.Categorical(Vindex(probs_d)[b_i, c]), obs=data[i])