Metaprob rewrite (another one!)

[alex-lew]

This branch contains a new implementation of Metaprob.

Breakdown of the new code

The Metaprob implementation now consists of the following parts:

Utilities (written in Clojure) metaprob.code-handlers implements helper functions for dealing with Metaprob S-expressions metaprob.trace implements the trace datatype and a number of operations on traces

Core language metaprob.generative-functions implements Metaprob's core abstraction: the generative function. It provides a general constructor make-generative-function for generative functions, and two additional constructs that desugar to make-generative-function calls: the gen macro, and the make-primitive helper.

Libraries metaprob.prelude exposes useful math functions like log and exp, as well as eager, generative-function versions of higher-order functions like map and replicate. metaprob.distributions implements several common distributions as generative functions. metaprob.inference implements the Metaprob standard inference library.

Syntactic transformations metaprob.expander exposes the mp-expand helper function, which authors of Metaprob macros / static analyses can call. It desugars any generative function source code into a small core language consisting only of variables, literals, function calls, quote forms, if expressions, and gen expressions. metaprob.autotrace implements an autotrace macro that relies on mp-expand. (autotrace (gen [] ...)) creates an autotraced generative function.

---

Very quick overview of language changes

Metaprob source files will typically begin with something like

(ns my.project
  (:refer-clojure :exclude [map replicate apply])
  (:require [metaprob.trace :refer :all]
            [metaprob.prelude :refer :all]
            [metaprob.distributions :refer :all]
            [metaprob.generative-functions :refer :all]
            [metaprob.inference :refer :all]))

In their project, Clojure code can be used as usual. To define a generative function, users can:

• Call make-generative-function manually, passing in two arguments:

  • run-in-clojure, a Clojure function that accepts arguments and returns a value
  • make-constrained-generator, a Clojure function that accepts an observation trace and returns a constrained generator (see below for more details).
  • Optionally, the user can pass in a map of arbitrary metadata to attach to the generative function, e.g., {:support [true false]}.

• Call make-primitive, passing in a sampler (fn [& args] ...) and a scorer (fn [val args] ...). The scorer returns the log probability mass or density of val under the distribution on return values induced by (apply sampler args).

• Use the gen macro:

  • The simplest form of the gen macro is (gen [& args] body...). This is a generative function in which no choices are traced.
  • A name can be given to the gen function (e.g., for recursive calls) using either (gen my-name [& args] body...) or (gen {:name my-name} [& args] body...).
  • In general, the user may provide an optional annotations map immediately following the keyword gen and preceding the argument list. These annotations are stored as part of the resulting generative function object's metadata, accessible via Clojure's meta function.
  • Within a gen body, traced random calls can be made in three ways: (at addr other-generative-function arg1 arg2...), (apply-at addr other-generative-function [arg1 arg2...]), and (let-traced [x (some-generative-function arg1 arg2...)] ...).

A generative function can be run like an ordinary Clojure function, but also supports another operation, make-constrained-generator, which accepts an observation trace and returns a constrained generator. A constrained generator is a generative function with the following properties:

• Its inputs are the same inputs as the original generative function from which it was derived.
• It returns a triple [value trace score], such that trace is a complete trace of the original generative function from which it was derived, value is the corresponding return value, and score is a log importance weight.
• As a generative function, it makes a number of traced choices. The requirement is that its trace, when merged with the observation trace, is a supertrace of the trace it returns.

infer-and-score is implemented in terms of make-constrained-generator.

---

The examples in the metaprob.examples namespaces have been updated to use the new syntax, as have all the tests. They pass!

However, we should decide:

• Whether to add new examples / remove some of the old ones.
• Whether to add new tests that give good coverage of Metaprob's language features, and remove some of the old ones.
• How to update the various .md files (documentation) that are out of date.
• Whether to move tutorials and Jupyter support into another repository.

[alex-lew]

It is annoying that source files have to start with

(:refer-clojure :exclude [map replicate apply])

One possible solution: the gen macro itself could insert a (let [map metaprob.eager/map, replicate metaprob.eager/replicate, apply (with-meta apply {:apply? true})] ...) around the body of a function. Outside of generative functions, there's no need for a special special map or replicate, so users could just use the Clojure versions.

@zane @joshuathayer Curious to hear what you all think of this possibility.

[zane]

@alex-lew If the best practice is to replace clojure.core/map with metaprob.eager/map I would prefer that it happen explicitly, in the ns form, rather than implicitly via the gen macro. Otherwise it will be too confusing for people coming to Metaprob from Clojure.

[zane]

@alex-lew If the goal is to provide an eager map could we encourage users to use clojure.core/mapv instead? No :refer needed! One thing to note about clojure.core/mapv is that the output type is always a vector, regardless of what the input type was.

[zane]

@alex-lew Another thing to note is that in ClojureScript, :refer :all is not available. We're going to need to recommend users do one of the following:

1. :refer every "core" Metaprob function individually.
2. Use namespace aliases, for example (:require [metaprob.math :as math]) (math/exp 10 2)

I pretty strongly prefer the second option. It's the Clojure community best-practice [1] [2], and referring each symbol individually would result in ginormous ns boilerplate.

[alex-lew]

@zane thanks for this! some initial responses--

Within a generative function, lazy versions of map, map-indexed, reduce, repeatedly, etc. are dangerous for subtle reasons. I don't think we want Metaprob users to have to reason about when it's ok or not ok to use laziness.

It's interesting that mapv is eager, but I'm not sure that's a solution -- because there's no analogue for map-indexed, reduce, etc.

Can you say more about what might be confusing to Clojure programmers about replacing clojure.core/map with metaprob.eager/map inside of gen bodies? I am thinking that metaprob.eager/map would be implemented something like this:

(defn map [& args]
  (doall (apply clojure.core/map args)))

so the only way for it to bite you would be if you were relying on laziness inside your generative function (which we want to avoid).

I do agree that hackily replacing one function with another in a macro is not a particularly clean implementation strategy, though, and would love to find a satisfying alternative.

Re: refer, I could be convinced, but am not yet. Looking at core.logic and anglican, two other Clojure libraries that are trying to be "languages," I see that both encourage users to :use them, exposing all symbols with no qualifiers. In core.logic, e.g., we wouldn't write

(logic/run [q]
  (logic/fresh [x]
    (logic/== x q)))

but rather

(run [q]
  (fresh [x]
     (== x q)))

I can imagine creating a namespace that just exposed names from all the other namespaces, and importing it as something short like m or mp in Clojurescript code -- then you'd use, e.g., mp/trace-value and mp/infer-and-score and mp/importance-resampling. But I don't like the idea of trace/trace-value, gen-func/infer-and-score, and inference/importance-resampling: remembering what function comes from where is a big burden to place on users.

[zane]

I can imagine creating a namespace that just exposed names from all the other namespaces, and importing it as something short like m or mp in Clojurescript code -- remembering what function comes from where is a big burden to place on users.

I was heading there, and agree!

[alex-lew]

Great :) Do you know if there's a way to have gen (the macro) exposed in the same namespace?

So people could do

(ns my.project 
  (:require [metaprob.core :as m :refer [gen]]))

Not sure how to re-expose a macro that's implemented somewhere else.

[zane]

That syntax should work as-is, if I remember, and is how I would do it! 🙂 It would be a little different in ClojureScript (you'd probably use :refer-macros).

[alex-lew]

Ah, yes, that syntax should work, but the part I was wondering about is how metaprob.core (our giant namespace) would expose a gen macro that was actually defined in metaprob.generative-functions :-) (since it couldn't just do (def gen metaprob.generative-functions/gen).)

But I just found https://github.com/ztellman/potemkin, which has an import-vars feature that seems to exist for just this purpose. (I guess there is the question of whether potemkin will work in ClojureScript.)

[zane]

Oh, sorry. I misunderstood!

Yes, I had imagined that we would use Potemkin. But now that I'm looking at it I'm not sure it supports ClojureScript. Will look into it further.

[zane]

You're right that it's not uncommon to refer a few core functions macros from libraries that implement embedded language (like core.logic). gen, letgen, and perhaps infer seem like a prime candidates for this. But in my experience referring to functions via aliases is pretty strongly preferred among Clojure programmers. I would draw the line before referring something like replicate, personally.

[alex-lew]

I think there are two questions:

• where to draw that line in general (i.e., do we include distributions like flip and gaussian and uniform? trace operations like trace-value?)
• what to do about map, reduce, replicate, etc.

In the second, the problem is that people will likely, out of habit, call the Clojure versions of these functions inside a gen form, which will cause hard-to-debug failures -- running infer-and-score might return an empty trace, for instance. How can we effectively prevent this?

[zane]

So, within a generative function, lazy versions of map, map-indexed, reduce, repeatedly, etc. are dangerous for subtle reasons.

I understand, but I don't think rebinding map within gen bodies is the solution. Swapping out the function they intended to use will confuse users who have prior experience with Clojure. They might be surprised to find a different function than the one they intended to use in stack traces, for example.

If we want to take a hard line there it would be better to force the programmer to fix the issue rather than try to fix the problem for them reactively behind the scenes. We could have gen error out if it finds symbols that refer to known lazy functions in its body for example.

[alex-lew]

Hmm, I see. What about inserting a doall in front of all calls to lazy functions? At least within the generated inference code (i.e., this wouldn't take effect when running the function as an ordinary Clojure function).

I think there's maybe a semantic issue here: if you run

(take 1 (take 1000 (repeatedly (fn [] (t (gensym) flip [0.5])))))

how many flips should be in the trace?

[zane]

That's certainly an option! I'm wary of that route because I'm worried that it will cause a divergence between what experienced Clojure programmers think is happening ("This expression will return a lazy sequence!") and what is actually happening ("This expression will return a fully realized list.").

We also can't anticipate the full set of lazy functions that people might use. What if they get used to the behavior of implicitly doall-wrapped functions from clojure.core but then later get bitten by some lazy function that they've written themselves (or that they got from a library)? Better for them to get an error (or warning?) from their first naive use of clojure.core/map and thereafter know to be wary of laziness.

[zane]

@alex-lew and I discussed this a bit more on Slack, and at this stage my preferred way to protect users from the unintended consequences of laziness+tracing are, in order:

1. Providing eager versions of clojure.core functions in a special namespace and encouraging users to use them via an alias. (e.g. (:require [metaprob.tracesafe :as tracesafe] tracesafe/map)
2. Providing functions in an eager namespace per the above, but recommending that users :refer-clojure :only […] and :refer […].
3. Providing functions in an eager namespace per the above with some prefix in the symbol name itself (e.g. map-random).
4. gen Implicitly wraps lazy functions in clojure.core with doall.
5. gen implicitly rebinds lazy functions from clojure.core with eager versions.

[alex-lew]

I think I'd be happy merging this, if @zane and @joshuathayer think we can go ahead. Things to resolve in the future:

• What is the best name for make-constrained-generator?
• How should we enforce eager evaluation of code that makes random choices inside of a gen? (what Zane and I were discussing above).