Types for sample spaces and reference measures for random choices

[marcoct]

The semantics of a Gen.Distribution include the sample space and the choice of a reference measure on that space. This choice is currently implicit. It should be made explicit and programmatically retrievable from the Distribution type value. That will allow analyses of generative functions to be written in a way that makes them extensible to handle random choices with arbitrary reference measures, and will clarify the semantics of generative functions (see https://github.com/probcomp/Gen.jl/issues/206 for a related issue).

Initially, sample spaces and reference measures do not need any behaviors (i.e. no functions). They are just types, with their mathematical semantics documented in docstrings. The minimal functionality needed to provide some value (beyond self-documentation and conceptual clarification for users) is to allow equality testing for compatible reference measures (this is required, for example, to divide the densities of random choices in a model and a proposal -- currently this check is based on the underlying Julia type used to store the value, which is indirect and conflates the Julia data type used to store a random choies with its abstract semantics within a Gen generative function).

We should provide an initial type system of sample space and reference measure types for common cases, but it should also be possible for users to write new distributions that use new reference measures that do not have any relationship to the built-in types. Because sample spaces and reference measures do not need any functions, this new system will not introduce any new roadblocks for a user who is constructing a distribution on an exotic sample space. Specifically, they should not be forced to express their exotic type within an elaborate type system for sample spaces and reference measures in order to use their distribution.

Eventually, the types of sample spaces and reference measures for random choice distributions may be merged with one for the return values of generative functions, and for the space of traces of generative functions. But it seems like this can be deferred.

[treigerm]

I just saw this issue and thought that maybe @cscherrer might be interested in reading this since he's working on MeasureTheory.jl. (Sorry if you have already been in touch.)

[cscherrer]

Thanks @treigerm ! I did have a very brief correspondence with @alex-lew about similarities in our goals. I would welcome broader community involvement in MeasureTheory.jl, especially for PPL applications. I'm not developing it to "own" the work in any sense, rather there was just a gap in existing capability someone needed to fill. I'd love for it to become more of a collective effort.

[marcoct]

Thanks @treigerm and @cscherrer. I think what I'm proposing in this issue is less ambitious than MeasureTheory.jl, in that this issue focuses just on labeling Gen distributions with some token for their reference measure (e.g. Lebesgue). The primary purpose is for documentation, and with equality testing of these tokens, some very limited analyses (checking that reference measures are the same) might also be possible.

I think a Julia library that provides a symbolic domain for more sophisticated reasoning about probabilistic programs in measure-theoretic terms would be useful, and such a library could overlap with MeasureTheory.jl significantly. Alex implemented a type system for reasoning about absolute continuity of composite measures for our POPL paper in Haskell. Doing that in Julia, and making it inter-operate with Gen, is a big project that we haven't committed to at this point, but would be interesting.

With this issue, I'm basically looking for a minimal and very incremental way of making the semantics of Gen distributions more explicit, to better document what is already implemented in Gen, and in part to provision for possible future work on more complete systems for reasoning about measure-theoretic semantics of Gen programs.

[marcoct]

But it might be useful to speculate as to how Gen might interface with a more complete system for symbolic (and measure-theoretic) reasoning about probabilistic programs. I think that one design pattern would be for Gen generative functions (and distributions, which are semantically special cases) to optionally implement methods that return the symbolic representation. So you can take a Gen model, convert it into the symbolic domain, and then do analyses in the symbolic domain. Generally, I think that probabilistic programming systems will benefit from having multiple intermediate representations designed for different tasks, and transformations between them. (A key ongoing next step for Gen is making Gen's static IR and combinators more explicitly into a "generative IR" that is the basis for various transformations and analyses). In a conversation with @alex-lew, he also pointed out that it might be possible to transform back from the symbolic domain into the type of executable / generative representation that we currently use in Gen.

I think that recognizing the need for multiple different representations, and transformations between them, might be helpful for understanding how different packages in the Julia probabilistic programming ecosystem can co-evolve. For example, one could make the argument that since every Gen generative function does represent a measure (as well as some other information), Gen generative functions should be an subtype of a MeasureTheory.jl type that expresses the pure mathematical notion of a measure. But I think that instead, allowing for an optional transformation into a MeasureTheory.jl type (e.g. perhaps in a third package GenToMeasureTheory.jl or MeasureTheoryToGen.jl that depends on both) is a more scalable approach that decouples the type systems and allows for much more flexibility. IIRC, GenSoss.jl is similar in its organizing principle.

Curious to hear your thoughts @cscherrer and @alex-lew .

[cscherrer]

I completely agree. I recently added symlogdensity to Soss.jl. Then I can expand this and do some compile-time constant folding. Here's a little writeup: https://informativeprior.com/blog/2021/01-25-symbolic-simplification/

With these optimizations, it can make more sense to integrate Gen's programming model a little more tightly into Soss. We should be able to get to where it's easy to use Gen-style inference, Turing back-ends, or custom generated code.

I think there's also a lot of potential in the recent work on e-graphs, for example as implemented in https://github.com/0x0f0f0f/Metatheory.jl

I think the SymbolicUtils stuff is moving in this direction as well, but I haven't seen any specifics.

[treigerm]

Thanks @cscherrer and @marcoct for elaborating on this more, this is very helpful in understanding the different use case that exist!