Implementing a DSL which desugars to Jaxpr

[femtomc]

I'm toying around with a probabilistic language built on top of JAX, but I'm running into some conceptual roadblocks.

Let's imagine I create a new primitive trace, with usage as follows:

def model(rng_key, v):
    sub_key, v = trace(stats.beta, rng_key, v, 1.0)
    return sub_key, v

here, trace will be given semantics by an external interpreter -- but the fallback will be just to sample from the specified distribution.

Now, for PPL interfaces -- I actually want to be able to construct a Jaxpr and operate on it with an interpreter.

If I try to do this, I'll run into an issue:

jaxpr = jax.make_jaxpr(model)(key, 3.0)

because stats.beta as a module object is not a valid JAX value.

That's okay! So what I want is the ability to construct an interpreter which is similar to make_jaxpr, but instead targets the lower "pure" set of Jaxpr values and implements the semantics I want. For example, let's study simulate:

v, tree = simulate(model)(key, 3.0)

roughly, what simulate should do is the same thing as make_jaxpr, except the primitive trace is treated specially under partial evaluation. Roughly:

sub_key, v = trace(stats.beta, rng_key, v, 1.0)

=>

sub_key, ssub_key = random.split(sub_key)
new = random.beta(ssub_key, v, 1.0)
score = stats.beta.logpdf(new, v, 1.0)

What's the right level to construct this functionality? Is the correct thing to make my own make_jaxpr (?)

[sharadmv]

To answer your question about primitives, if trace is a primitive, it can only take in array valued inputs as main arguments and any metadata can be passed in as params. You could pass in stats.beta as a param like this:

trace_p = jax.core.Primitive("trace")

def trace(dist, key, *args):
  return trace_p.bind(key, *args, dist=dist)

This will bake stats.beta into the jaxpr and a custom interpreter will be able to use that information.

I think Jaxpr interpreters using make_jaxpr are a viable approach for building a PPL on top of JAX. You could look at some existing JAX PPLs as an example. NumPyro implements its own effect-handling system in Python but Oryx uses Jaxpr interpreters.

Specifically, in Oryx you could look at the log_prob interpreter which might be a closer example of what you want to do (i.e. convert a sampling function into a log-density). You could also look at Oryx's guide for probabilistic programming.

[femtomc]

@sharadmv I ended up working this problem out using a slightly different approach (of forwarding bind to the primitive, and using trace as syntactic sugar). In general, I'm parametrizing the interpreter by statically dispatched "primitive handlers" (similar to effects, but zero overhead) -- so I suspect my implementation will be close to Oryx!

I'll take a look -- thanks for the reference.

Separately: for new primitives, is it possible to define multiple return types for def_abstract_eval (?)

For example, I'd like to represent sampling operations as pure functions from (PRNGKey, *args) to (PRNGKey, return) but I'm having a bit of trouble figuring out what types I should return under abstract evaluation.

E.g.

def bernoulli_abstract_eval(key, p):
    ...
    return (key, ShapedArray(...))

[sharadmv]

Separately: for new primitives, is it possible to define multiple return types for def_abstract_eval (?)

You'll need to mark the primitive as having multiple results:

prim = jax.core.Primitive('...')
prim.multiple_results = True