Symbolic Conditioning of Arrays in Probabilistic Programs

Probabilistic programming systems make machine learning methods modular by automating inference. Recent work by Shan and Ramsey (2017) makes inference itself modular by automating the common component of conditioning. Their work introduces a symbolic program transformation that treats conditioning generally via the measure-theoretic notion of disintegration. This technique, however, is limited to conditioning a single scalar variable. As a step towards tackling realistic machine learning applications, we have extended the disintegration algorithm to symbolically condition arrays in probabilistic programs. The extended algorithm implements lifted disintegration, where repetition is treated symbolically and without unrolling loops. The technique uses a language of index variables for tracking expressions at various array levels. We find that the method works well for arbitrarily-sized arrays of independent random choices, with the conditioning step taking time linear in the number of indices needed to select an element.