You should not be using a scan to compute pi, as noted above. Broadcast tensor operations should always be preferred to scans for optimization reasons. You can also dispense with the transformation of μ and σ into μlog and σlog, and instead let the log normal distribution handle the transformation of the parameters. I also used init = 'jitter+adapt_diag_grad when I ran pm.sample, although this might not matter after you make the other changes.
1 Like