Given that all parameters are trained jointly at inference time and a single sample of z is supposed to encode ALL inputs and outputs for a given puzzle (I think), I don't quite understand the role of the latent z here. Feels like μ and Σ could be absorbed into θ (and an additional variance parameter).
Although if they partition z such that each section corresponds to one input and run f_θ through those sections iteratively, then I guess it makes sense.