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.