def model():
     with pyro.iarange("num_particles", 10, dim=-3):
         with pyro.iarange("components", 2, dim=-1):
             p = pyro.sample("p", dist.Beta(torch.tensor(1.1), torch.tensor(1.1)))
             assert p.shape == torch.Size((10, 1, 2))
         with pyro.iarange("data", data.shape[0], dim=-2):
             pyro.sample("obs", dist.Bernoulli(p), obs=data)

 def guide(subsample):
     loc = pyro.param("loc", lambda: torch.zeros(len(data), requires_grad=True))
     scale = pyro.param("scale", lambda: torch.tensor([1.0], requires_grad=True))
     with pyro.iarange("particles", num_particles):
         with pyro.iarange("data", len(data), subsample_size, subsample) as ind:
             loc_ind = loc[ind].unsqueeze(-1).expand(-1, num_particles)
             pyro.sample("z", Normal(loc_ind, scale))

 def model():
     with pyro.iarange("particles", num_particles):
         pyro.sample("x", dist.Bernoulli(p).expand_by([num_particles]))
         with pyro.iarange("outer", outer_dim):
             pyro.sample("y", dist.Bernoulli(p).expand_by([outer_dim, num_particles]))
             for i in pyro.irange("inner", inner_dim):
                 pyro.sample("z_{}".format(i), dist.Bernoulli(p).expand_by([outer_dim, num_particles]))

 def guide():
     q = pyro.param("q")
     with pyro.iarange("particles", num_particles):
         pyro.sample("y", dist.Bernoulli(q).expand_by([num_particles]),
                     infer={"enumerate": enumerate1})
         with pyro.iarange("iarange", iarange_dim):
             pyro.sample("z", dist.Bernoulli(q).expand_by([iarange_dim, num_particles]),
                         infer={"enumerate": enumerate2})

 def model(loc, cov):
     x = pyro.param("x", torch.randn(2))
     y = pyro.param("y", torch.randn(3, 2))
     z = pyro.param("z", torch.randn(4, 2).abs(), constraint=constraints.greater_than(-1))
     pyro.sample("obs_x", dist.MultivariateNormal(loc, cov), obs=x)
     with pyro.iarange("y_iarange", 3):
         pyro.sample("obs_y", dist.MultivariateNormal(loc, cov), obs=y)
     with pyro.iarange("z_iarange", 4):
         pyro.sample("obs_z", dist.MultivariateNormal(loc, cov), obs=z)

    def model():
        particles_iarange = pyro.iarange("particles", num_particles, dim=-2)
        data_iarange = pyro.iarange("data", len(data), dim=-1)

        pyro.sample("nuisance_a", Normal(0, 1))
        with particles_iarange, data_iarange:
            z = pyro.sample("z", Normal(0, 1))
        pyro.sample("nuisance_b", Normal(2, 3))
        with data_iarange, particles_iarange:
            pyro.sample("x", Normal(z, 1), obs=data)
        pyro.sample("nuisance_c", Normal(4, 5))

def xy_model():
    d = dist.Bernoulli(0.5)
    x_axis = pyro.iarange('x_axis', 2, dim=-1)
    y_axis = pyro.iarange('y_axis', 3, dim=-2)
    pyro.sample('b', d)
    with x_axis:
        pyro.sample('bx', d.expand_by([2]))
    with y_axis:
        pyro.sample('by', d.expand_by([3, 1]))
    with x_axis, y_axis:
        pyro.sample('bxy', d.expand_by([3, 2]))

def iarange_reuse_model_guide(include_obs=True, dim1=3, dim2=2):
    p0 = torch.tensor(math.exp(-0.40 - include_obs * 0.2), requires_grad=True)
    p1 = torch.tensor(math.exp(-0.33 - include_obs * 0.1), requires_grad=True)
    pyro.sample("a1", dist.Bernoulli(p0 * p1))
    my_iarange1 = pyro.iarange("iarange1", dim1)
    my_iarange2 = pyro.iarange("iarange2", dim2)
    with my_iarange1 as ind1:
        with my_iarange2 as ind2:
            pyro.sample("c1", dist.Bernoulli(p1).expand_by([len(ind2), len(ind1)]))
    pyro.sample("b1", dist.Bernoulli(p0 * p1))
    with my_iarange2 as ind2:
        with my_iarange1 as ind1:
            c2 = pyro.sample("c2", dist.Bernoulli(p1).expand_by([len(ind2), len(ind1)]))
            if include_obs:
                pyro.sample("obs", dist.Bernoulli(c2), obs=torch.ones(c2.size()))

def gmm_batch_guide(data):
    with pyro.iarange("data", len(data)) as batch:
        n = len(batch)
        probs = pyro.param("probs", torch.tensor(torch.ones(n, 1) * 0.6, requires_grad=True))
        probs = torch.cat([probs, 1 - probs], dim=1)
        z = pyro.sample("z", dist.OneHotCategorical(probs))
        assert z.shape[-2:] == (n, 2)

def gmm_batch_guide(data):
    with pyro.iarange("data", len(data)) as batch:
        n = len(batch)
        ps = pyro.param("ps", Variable(torch.ones(n, 1) * 0.6, requires_grad=True))
        ps = torch.cat([ps, 1 - ps], dim=1)
        z = pyro.sample("z", dist.Categorical(ps))
        assert z.size() == (n, 2)

    def guide(self, xs, ys=None):
        """
        The guide corresponds to the following:
        q(y|x) = categorical(alpha(x))              # infer digit from an image
        q(z|x,y) = normal(mu(x,y),sigma(x,y))       # infer handwriting style from an image and the digit
        mu, sigma are given by a neural network `encoder_z`
        alpha is given by a neural network `encoder_y`
        :param xs: a batch of scaled vectors of pixels from an image
        :param ys: (optional) a batch of the class labels i.e.
                   the digit corresponding to the image(s)
        :return: None
        """
        # inform Pyro that the variables in the batch of xs, ys are conditionally independent
        with pyro.iarange("independent"):

            # if the class label (the digit) is not supervised, sample
            # (and score) the digit with the variational distribution
            # q(y|x) = categorical(alpha(x))
            if ys is None:
                alpha = self.encoder_y.forward(xs)
                ys = pyro.sample("y", dist.categorical, alpha)

            # sample (and score) the latent handwriting-style with the variational
            # distribution q(z|x,y) = normal(mu(x,y),sigma(x,y))
            mu, sigma = self.encoder_z.forward([xs, ys])
            zs = pyro.sample("z", dist.normal, mu, sigma)   # noqa: F841

 def guide(subsample_size):
     mu = pyro.param("mu", lambda: Variable(torch.zeros(len(data)), requires_grad=True))
     sigma = pyro.param("sigma", lambda: Variable(torch.ones(1), requires_grad=True))
     with pyro.iarange("data", len(data), subsample_size) as ind:
         mu = mu[ind]
         sigma = sigma.expand(subsample_size)
         pyro.sample("z", dist.Normal(mu, sigma, reparameterized=reparameterized))

def iarange_model(subsample_size):
    loc = torch.zeros(20)
    scale = torch.ones(20)
    with pyro.iarange('iarange', 20, subsample_size) as batch:
        pyro.sample("x", dist.Normal(loc[batch], scale[batch]))
        result = list(batch.data)
    return result

 def guide(self, x):
     # register PyTorch module `encoder` with Pyro
     pyro.module("encoder", self.encoder)
     with pyro.iarange("data", x.size(0)):
         # use the encoder to get the parameters used to define q(z|x)
         z_loc, z_scale = self.encoder.forward(x)
         # sample the latent code z
         pyro.sample("latent", dist.Normal(z_loc, z_scale).independent(1))

 def model(self, data):
     decoder = pyro.module('decoder', self.vae_decoder)
     z_mean, z_std = torch.zeros([data.size(0), 20]), torch.ones([data.size(0), 20])
     with pyro.iarange('data', data.size(0)):
         z = pyro.sample('latent', Normal(z_mean, z_std).independent(1))
         img = decoder.forward(z)
         pyro.sample('obs',
                     Bernoulli(img).independent(1),
                     obs=data.reshape(-1, 784))

    def model(data):
        weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(K)))
        locs = pyro.sample('locs', dist.Normal(0, 10).expand_by([K]).independent(1))
        scale = pyro.sample('scale', dist.LogNormal(0, 1))

        with pyro.iarange('data', len(data)):
            weights = weights.expand(torch.Size((len(data),)) + weights.shape)
            assignment = pyro.sample('assignment', dist.Categorical(weights))
            pyro.sample('obs', dist.Normal(locs[assignment], scale), obs=data)