Random Variables in Aesara

I propose a new approach to represent RandomVariable and PRNG states in Aesara's IR, based on the design of splittable PRNGs. The representation introduces minimal change to the existing RandomVariable interface while being more expressive. It should be easy to transpile to Aesara's current compilation target, and is compatible with the higher-level RandomStream interface.

On counter-based PRNGs:

• Parallel Random Numbers: As easy as 1, 2, 3 (2011)

On splittable PRNG design for functional programs:

• Distributed Random Number Generation (2008)
• Splittable PRNGs Using Cryptographic Hashing (2013)

A good PRNG design satisfies the following conditions:

1. It is expressive: the behavior of the system is predictable by the caller, and allows them to expression any probabilistic program;
2. It makes it possible to build reproducible programs ("seeding");
3. It is embedded in in Aesara's IR;
4. It can be manipulated by Aesara's rewrite system;
5. It can be easily transpiled to current backends;
6. It enables vectorization with generalized universal functions;

This proposal is motivated by two issues that illustrate the shortcomings of the current representation of RandomVariable\s and PRNG states in Aesara:

• In #1036 the use of default_output to hide the PRNG state from the user is causing multiple headaches in the etuplization of RandomVariable\s and unification/reification of expressions with RandomVariable\s. This is the only Op in Aesara that makes use of this property, and to "special-casing" the etuplization logic for RandomVariable\s often appeared as the easiest solution.

• In #66 in AeMCMC, expressing expansions like the following convolution of two normal variables is overly complex:

  from etuples import etuple
  from kanren import var
  
  mu_x, mu_y, sigma2_x, sigma2_y = var(), var(), var(), var()
  
  rng, size, dtype = var(), var(), var()
  X_et = etuple(
      etuplize(at.random.normal),
      rng,
      size,
      dtype,
      etuple(
          etuplize(at.add),
          mu_x,
          mu_y
      ),
      etuple(
          etuplize(at.add),
          sigma2_x,
          sigma2_y,
      )
  )
  
  rng_x, size_x, dtype_x = var(), var(), var()
  rng_y, size_y, dtype_y = var(), var(), var()
  Y_et = etuple(
      etuplize(at.add),
      etuple(
          etuplize(at.random.normal),
          rng_x,
          size_x,
          dtype_x,
          mu_x,
          sigma2_x
      ),
      etuple(
          etuplize(at.random.normal),
          rng_y,
          size_y,
          dtype_y,
          mu_y,
          sigma2_y
      )
  )
  

  It is indeed not clear what the values of rng_x and rng_y should be given the value of rng. A few other application-related shortcomings of the current representation will be given below.

In the following we focus on the symbolic representation of random variables and PRNG states in Aesara's IR. We leave discussions about compilation targets and solution to the previous issues for the end.

If we represent the internal state of the PRNG by the type RandState (short for RandomStateType), the current design of RandomVariable\s can be summarized by the following simplified signature:

RandomVariable :: RandState -> (RandState, TensorVariable)

In other words, RandomVariable\s are responsible for both advancing the state of the PRNG, and producing a random value. This double responsibility is what creates graph dependencies between nodes that have otherwise no data dependency. The following snippet illustrates this:

import aesara
import aesara.tensor as at

rng = at.random.type.RandomStateType()('rng')

rng_x, x_rv = at.random.normal(0, 1, rng=rng, name='x').owner.outputs
rng_y, y_rv = at.random.normal(0, 1, rng=rng_x, name='y').owner.outputs
z_rv = at.random.normal(0, 1, rng=rng_y, name='z')
w_at = x_rv + y_rv + z_rv

aesara.dprint(w_at)
# Elemwise{add,no_inplace} [id A]
#  |Elemwise{add,no_inplace} [id B]
#  | |normal_rv{0, (0, 0), floatX, False}.1 [id C] 'x'
#  | | |rng [id D]
#  | | |TensorConstant{[]} [id E]
#  | | |TensorConstant{11} [id F]
#  | | |TensorConstant{0} [id G]
#  | | |TensorConstant{1} [id H]
#  | |normal_rv{0, (0, 0), floatX, False}.1 [id I] 'y'
#  |   |normal_rv{0, (0, 0), floatX, False}.0 [id C]
#  |   |TensorConstant{[]} [id J]
#  |   |TensorConstant{11} [id K]
#  |   |TensorConstant{0} [id L]
#  |   |TensorConstant{1} [id M]
#  |normal_rv{0, (0, 0), floatX, False}.1 [id N] 'z'
#    |normal_rv{0, (0, 0), floatX, False}.0 [id I]
#    |TensorConstant{[]} [id O]
#    |TensorConstant{11} [id P]
#    |TensorConstant{0} [id Q]
#    |TensorConstant{1} [id R]

As we can see in the graph representation, rng_x (id C) is being used as an input to y and rng_y (id I) is being used as an input to z. There is however no data dependency between x, y or z. The intuition that they should not be linked is probably what led to "hiding" these PRNG state outputs so they are not re-used, and the RandomStream interface.

Creating spurious sequential dependencies by threading PRNG states is indeed unsatisfactory from a representation perspective, and unnecessarily complicates the rewrites. It is also problematic for two other reasons:

• Parallelization and Vectorization: Using random variables in user-defined generalized universal functions is going to require a lot of compiler magic to make sure that the random state is updated properly, and the behavior will be completely opaque to the user;
• The fact that callers cannot be intentional about what they do with the random state is limiting. This can be necessary in pratical applications, for instance to implement coupled sampling algorithms in which two algorithms share the same random state.

A natural idea is to simplify the design of RandomVariable\s so that it is only responsible for one thing: generating a random value from a PRNG state. The Op thus creates an Apply node that takes a RandState (using the above notation) as input and outputs a (random) Variable:

RandomVariable :: RandState -> Variable

Providing a RandState to a RandomVariable needs to intentional, and this must be reflected in the user interface. We thus make rng an explicit input of the RandomVariable's __call__ method. This way a user can write:

import aesara.tensor as at

# rng_x, rng_y and rng_z are created before that.
x_rv = at.random.normal(rng_x, 0, 1)
y_rv = at.random.normal(rng_y, 0, 1)
z_rv = at.random.normal(rng_z, 0, 1)

Or, if they want the PRNG state to be shared (silly example, but a legitimate need):

import aesara.tensor as at

# rng_x, rng_y and rng_z are created before that.
x_rv = at.random.normal(rng_x, 0, 1)
y_rv = at.random.normal(rng_x, 0, 1)
z_rv = at.random.normal(rng_x, 0, 1)

This interface presupposes the existence of two operators. First, to build reproducible programs, we need an operator that creates a RandState from a seed, which can be the constructor of RandState itself:

RandState.__init__ :: Seed -> RandState

And then, we need another operator that creates an updated RandomState from a RandomState, so that RandomVariable\s created with these two different states would output different numbers. Let's call it next:

next :: RandomState -> RandomState

We can thus fill in the blanks in the previous code examples:

import aesara.tensor as at

rng = at.random.RandState(0)
rng_y = at.random.next(rng_x)
rng_z = at.random.next(rng_y)

x_rv = at.random.normal(rng_x, 0, 1)
y_rv = at.random.normal(rng_y, 0, 1)
z_rv = at.random.normal(rng_z, 0, 1)

w_at = x_rv + y_rv + z_rv

The code has been specifically formatted to illustrate what we gain from this approach. x_rv, y_rv and z_rv have lost their direct dependency; we could easily execute these three statements in parallel. What we have done implicitly is to create two graphs: the graph between random variables which reflects the dependencies (or lack thereof) on each other's values, and the graph of the updates of the PRNG states. These graphs almost evolve in parallel.

This is similat to what I understand the RandomStream interface does: moving the updates of the PRNG states to the update graphs generated by Aesara's shared variables.

The next operator is however not completely satisfactory. Let us consider a more complex situation, where call is a function that requires a RandomState:

import aesara.tensor as at

rng = at.random.RandState(0)
rng_y = at.random.next(rng)

x_rv = call(rng_x)
y_rv = call(rng_y)
z_at = x_rv + y_rv

We can easily find an implementation of call that makes the previous code generate a random state collision:

def call(rng_a):
    a_rv = at.random.normal(rng_a, 0, 1)
    rng_b = at.random.next(rng_a)
    b_rv = at.random.normal(rng_b, 0, 1)
    return a_rv * b_rv

To avoid this kind of issues, we must thus require user-defined functions to return the last PRNG state along the result:

import aesara.tensor as at

def call(rng_a):
    a_rv = at.random.normal(rng_a, 0, 1)
    rng_b = at.random.next(rng_a)
    b_rv = at.random.normal(rng_b, 0, 1)
    return (a_rv * b_rv), at.random.next(rng_b)


rng = at.random.RandState(0)
x_rv, rng_x = call(rng)
y_rv, rng_y = call(rng_x)
z_at = x_rv + y_rv

Threading PRNG state is still necessary to guarantee correctness and the two call functions cannot be called in parallel. The issue arises because, even though we have separated PRNG state update and random value generation, our symbolic structure is still sequential: each RandState has one and only one ancestor. We can of course circumvent this issue knowing how many times next is called within the function, by "jumping" the same number of times to obtain rng_y, but this can quickly become complex (what if call is imported from somewhere else?).

It would make things easier if a RandState\s could have several children, and if each of these child led to separate streams of random number. Let us define the following split operator:

split :: RandState -> (RandState, RandState)

We require that we can never get the same RandState by calling split any number of times on either the left or right returned state. In other words, split should implicitly defines a binary tree in which all the nodes are unique. This can be easily represented by letting RandState holding a number in binary format. The leftmost child state is obtained by appending 0 to the parent's state and the rightmost child state by appending 1:

from typing import NamedTuple


class RandState(NamedTuple):
    key: int
    node_id: int = 0b1


def split(rng):
    left = RandState(rng.key, rng.node_id << 2)
    right = RandState(rng.key, (rng.node_id << 2) + 1)
    return left, right


rng = RandState(0)
l, r = split(rng)
ll, lr = split(l)

print(rng, l, lr)
# RandState(key=0, node_id=1) RandState(key=0, node_id=4) RandState(key=0, node_id=17)

If the generator called by RandomVariable can be made a deterministic function of this binary value, the computations are fully reproducible. We added a key attribute that can be specified by the user at initialization to seed the PRNG state. The tree structure is of course explicit in our graph representation, since l and r depend on rng via the split operator. Nevertheless, we can increment this internal state when building the graph in a way that allows us to compile without traversing the graph.

The next operator we previously defined becomes redundant within this representation. Since its interaction with the split operator would require careful thought we leave it aside in the following. Using the new operator our toy example becomes:

import aesara.tensor as at

rng = at.random.RandState(0)
rng_x, rng_y = at.random.split(rng)

x_rv = at.random.normal(rng_x, 0, 1)
y_rv = at.random.normal(rng_y, 0, 1)
z_at = x_rv + y_rv

Note that the "main" sub-graph that contains random variables, and the PRNG sub-graph are still minimally connected.

Finally, it is also natural to implement the splitn operator represented by:

splitn :: RandState -> Int -> (RandState, ..., RandState)

So we can write the following code:

at.random.split = at.random.Split()

rng = at.random.default_rng()
rng_v, rng_w, rng_x, rng_y = at.random.splitn(rng, 4)

v_rv = at.random.normal(rng_y, 0, 1)
w_rv = at.random.normal(rng_x, 0, 1)
x_rv = at.random.normal(rng_x, 0, 1)
y_rv = at.random.normal(rng_y, 0, 1)
z_at = v_rv + w_rv + x_rv + z_rv

The problems linked to the existence of the default_output attribute disappear since RandomVariable\s do not return PRNG states anymore. The one-to-many difficulty we are facing with the relations between etuplized graphs also disappears with a split operator. Using the example from the beginning we can for instance write:

from etuples import etuple
from kanren import var

mu_x, mu_y, sigma2_x, sigma2_y = var(), var(), var(), var()

rng, size, dtype = var(), var(), var()
X_et = etuple(
    etuplize(at.random.normal),
    rng,
    size,
    dtype,
    etuple(
        etuplize(at.add),
        mu_x,
        mu_y
    ),
    etuple(
        etuplize(at.add),
        sigma2_x,
        sigma2_y,
    )
)

rng_x, size_x, dtype_x = var(), var(), var()
rng_y, size_y, dtype_y = var(), var(), var()
Y_et = etuple(
    etuplize(at.add),
    etuple(
        etuplize(at.random.normal),
        etuple(
            nth,
            0,
            etuple(
                split,
                rng,
            )
        ),
        size_x,
        dtype_x,
        mu_x,
        sigma2_x
    ),
    etuple(
        etuplize(at.random.normal),
        etuple(
            nth,
            1,
            etuple(
                split,
                rng,
            )
        ),
        size_y,
        dtype_y,
        mu_y,
        sigma2_y
    )
)

Which is guaranteed to be collision-free by construction of the split operator, as long as the PRNG state used by the original normal distribution isn't passed to a split operator somewhere else in the original graph (todo: specify API requirements to guarantee uniqueness of the random numbers).

We would like to be able to control randomness from outside of Aesara. From JAX it is rather clear how that would work:

import aesara

rng = at.random.RNGState()
result = at.random.normal(rng, 0, 1)

# JAX compilation
import jax

fn = aesara.function([rng], result, mode="JAX")
rng_key = jax.random.PRNGKey(0)
fn(rng_key)

But for NumPy/Numba it is not as clear what we should be using.

# Numba compilation
import numpy as np

fn = aesara.function([rng], result, mode="JAX")
rng_key = np.random.default_rng(0)
fn(rng_key)

With the abstractions that are currently available in Aesara/AePPL one wants to define the truncation operator as:

import aesara.tensor as at
import aesara.tensor.random as ar
from aeppl import truncate

x_rv = truncate(ar.normal(0, 1), lower=0)

This reads as defining a new random variable \(X'\) such \(X' = \operatorname{truncate} \circ X\) where \(X \sim \operatorname{Normal}(0, 1)\). As such it is indistinguishable from the existing operator:

y_rv = ar.normal(0, 1)
x_rv = at.clip(y_rv)

Which also reads as defining a new random variable \(X' = \operatorname{clip} \circ X\) with \(X \sim \operatorname{Normal}(0, 1)\). However, we want the truncate operator here to mean something else, namely an operator that takes a probability measure \(\operatorname{Normal}(0, 1)\) and returns a new measure \(\hat{\mathbbm{1}}_{[0, \infty[} \operatorname{Normal}(0, 1)\).

This issue arises because measures are only implicitly defined in Aesara, where we manipulate RandomVariable objects instead. If we want to avoid the kind of ambiguity that we pointed at earlier, we need to treat them as first class citizens. We can thus define:

import aesara.tensor.random as ar

x_m = ar.normal(0, 1)
assert isinstance(x_m, ar.Measure)

where x_m is now a Measure. To bridge the gap with the existing RandomVariable API in Aesara, we define the aesara.tensor.random.sample operator which takes a measure and return an independent random variable that admits this measure every time it is called. sample(rng, \mathcal{F}) takes a rng key \(k \in \Omega\), a Measure object that represents the measure \(\mathbb{P}\) and returns a random variable \(X: \Omega \to E\) where \(E\) is the event space:

sample :: Key -> Measure -> RandomVariable

So to define a normally-distributed random variable one would write:

import aesara.tensor.random as ar

x = ar.sample(rng, ar.normal(0, 1))
assert isinstance(x, ar.RandomVariable)

This is somewhat similar to what is done in the Anglican probabilistic programming language implemented in Clojure. We retain the existing ambiguity that exists between random variables and their realizations, but that does not matter mathematically speaking.

We can keeps something similar to the RandomStream by defining

import aesara.tensor as at
import aesara.tensor.random as ar

srng = ar.RandomStream(0)
x = srng.sample(ar.normal(0, 1))

We will assume in the following that these primitives are available in Aesara.