What are packed tensors in Pyro?

Can anyone provide some background on what packed tensors are? The Trace object has a method pack_tensors that is called when enumeration is required (in TraceEnum_ELBO). This creates a packed item in the stochastic nodes but there not a lot of information about what they are and why you need them.

Think of it as a squeezed tensor with named dimensions.

x = torch.rand(3, 1, 2)
x_packed = pack(x, dim_to_symbol={0: "a", 1: "b", 2: "c"})
assert x_packed.shape == (3, 2)
assert x_packed._pyro_dims == "ac"

It allows to track named dimensions like batch dim or enumerated dim in the inference algorithm.

Thanks @ordabayev! This helps a lot.

I have noticed in the code that b, d, f… are used for enumerated dimensions and a, c, e… are use for the rest. Is this correct?

And do I understand correctly that packing is used only in TraceEnum_ELBO because there is no need to track enumerated dimensions in Trace_ELBO?

As far as I know there are no specific dims reserved for enumerated dims so they should change from model to model.

That’s right.