Thanks @fritzo! Sorry for the lack of clarity, I meant a single measurement of the pool values for each item, and then the distribution in each pool depends on the categorical for the item. Your other comments were helpful, and got a bit farther but then the dimensionality of obs still didn’t match:

```
simple_data = TT([[1., 3.], [10., 3.], [1., 8.], [9., 4.], [7., 1.]])
n_pools = simple_data.shape[1]
n_item = simple_data.shape[0]
@config_enumerate
def test_model(data):
pool_size = pyro.sample('pool_size', dist.LogNormal(3., 1.))
print('pool_size', pool_size)
with pyro.plate('item', n_item) as items:
print('items', items)
pool = pyro.sample('pool', dist.Categorical(torch.ones(n_pools)))
print('pool', pool, pool.shape)
rates = (torch.eye(n_pools)[pool] * 0.5 + 0.1) * pool_size
print('rates', rates, rates.shape)
if len(pool.shape) > 1: # hack to skip obs in guide
print('data[pool]', data[pool], data[pool].shape)
obs = pyro.sample('obs', dist.Poisson(rates), obs=data[pool])
print('obs', obs, obs.shape)
test_guide = AutoDiagonalNormal(poutine.block(test_model, expose=['pool_size']))
pyro.clear_param_store()
optim = pyro.optim.Adam({'lr': 0.1, 'betas': [0.8, 0.99]})
elbo = TraceEnum_ELBO(max_plate_nesting=1)
svi = SVI(test_model, test_guide, optim, loss=elbo)
svi.step(simple_data)
# output
('pool_size', tensor(58.5516))
('items', tensor([0, 1, 2, 3, 4]))
('pool', tensor([0, 1, 0, 1, 1]), torch.Size([5]))
('rates', tensor([[35.1310, 5.8552],
[ 5.8552, 35.1310],
[35.1310, 5.8552],
[ 5.8552, 35.1310],
[ 5.8552, 35.1310]]), torch.Size([5, 2]))
('pool_size', tensor(0.6469, grad_fn=<ExpandBackward>))
('items', tensor([0, 1, 2, 3, 4]))
('pool', tensor([[0],
[1]]), torch.Size([2, 1]))
('rates', tensor([[[0.3881, 0.0647]],
[[0.0647, 0.3881]]], grad_fn=<MulBackward0>), torch.Size([2, 1, 2]))
# >> this is not right >>
('data[pool]', tensor([[[ 1., 3.]],
[[10., 3.]]]), torch.Size([2, 1, 2]))
```

So I tried just enumerating the items sequentially and this worked

```
@config_enumerate
def test_model(data):
pool_size = pyro.sample('pool_size', dist.LogNormal(3., 1.))
print('pool_size', pool_size)
for i in pyro.plate('item', n_item):
pool = pyro.sample('pool_%i' % i, dist.Categorical(torch.ones(n_pools)))
print('pool', pool, pool.shape)
rates = (torch.eye(n_pools)[pool] * 0.5 + 0.1) * pool_size
print('rates', rates, rates.shape)
if len(pool.shape) > 1: # hack to skip obs in guide
print('data[pool]', data[pool], data[pool].shape)
obs = pyro.sample('obs_%d' % i, dist.Poisson(rates).to_event(1), obs=data[i])
print('obs', obs, obs.shape)
test_guide = AutoDiagonalNormal(poutine.block(test_model, expose=['pool_size'], hide=['obs']))
pyro.clear_param_store()
optim = pyro.optim.Adam({'lr': 0.1, 'betas': [0.8, 0.99]})
elbo = TraceEnum_ELBO(max_plate_nesting=1)
svi = SVI(test_model, test_guide, optim, loss=elbo)
svi.step(simple_data)
```

This didn’t give any errors, but there were some strange things. The pool shape kept adding dimensions, though IIUC iterating the plate should not add dimensions since they’re conditionally independent?

```
('pool', tensor([[0],
[1]]), torch.Size([2, 1]))
# keeps getting deeper with each iteration
('pool', tensor([[[[[[0]]]]],
[[[[[1]]]]]]), torch.Size([2, 1, 1, 1, 1, 1]))
```

I tried running a hundred steps to optimize loss, then using the code from the Gaussian example using infer_discrete to pull out the pool assignments, but they were random AFAICT - what drives the choice of optimal assignments there?

Thanks for your help, though - this is progress!

Ravi