Find horsehoe priors in an Aesara graph

Remember that the horseshoe prior is mathematically defined as:

We would like to be able to tell whether the model that the graph implements contains a horseshoe prior, and if so return the variables that represent \(\lambda\) and \(\tau\). We start by define a horseshoe pattern against which we are going to try to match the model, using unification's logic variables and etuples' emulation of S-expressions:

from unification import var
from etuples import etuple, etuplize

horseshoe_1_lv, horseshoe_2_lv = var('horseshoe_1'), var('horsehoe_2')
zero_lv = var('zero')
horseshoe_pattern = etuple(
    etuplize(at.random.normal),
    var(),
    var(),
    var(),
    zero_lv,
    etuple(
        etuplize(at.mul),
        horseshoe_1_lv,
        horseshoe_2_lv)
)

We can then unify this pattern with the subgraph beta_rv using pythological's unification package:

from unification import unify
from IPython.lib.pretty import pprint

s = unify(horseshoe_pattern, etuplize(beta_rv))
pprint(s)

{~_32: RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FCCBFFD8740>),
 ~_33: TensorConstant{[]},
 ~_34: TensorConstant{11},
 ~zero: TensorConstant{0},
 ~horseshoe_1: e(
   e(
     aesara.tensor.random.basic.HalfCauchyRV,
     'halfcauchy',
     0,
     (0, 0),
     'floatX',
     False),
   RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FCCBFFD8F20>),
   TensorConstant{(1,) of 1},
   TensorConstant{11},
   TensorConstant{1},
   TensorConstant{1.0}),
 ~horsehoe_2: e(
   e(
     aesara.tensor.random.basic.HalfCauchyRV,
     'halfcauchy',
     0,
     (0, 0),
     'floatX',
     False),
   RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FCCBFFCB820>),
   TensorConstant{(1,) of 10},
   TensorConstant{11},
   TensorConstant{1},
   TensorConstant{1.0})}

The pattern we used in the previous section only checks that the graphical structure matches the horseshoe prior, but not that the correct distributions and correct parameters are used. However, we can only identify an aesara graph with the mathematical model if structure and distributions match (or if the pattern represents the normal form of the model under a set of rewrite rules, but we'll keep things simple for now).

Let us thus improve the design by defining a unify_horseshoe function that not only checks the basic graphical structure but also the shapes and distributions. This function behaves similarly to unify; it returns the random variables (lambda_rv, tau_rv) when the subgraph is identified, and False if unification is impossible.

from aesara.graph.unify import eval_if_etuple
from etuples import etuple, etuplize
from unification import unify, var

def unify_horseshoe(graph):
    horseshoe_1_lv, horseshoe_2_lv = var('horseshoe_1'), var('horsehoe_2')
    zero_lv = var('zero')
    horseshoe_pattern = etuple(
        etuplize(at.random.normal),
        var(),
        var(),
        var(),
        zero_lv,
        etuple(
            etuplize(at.mul),
            horseshoe_1_lv,
            horseshoe_2_lv)
    )

    s = unify(graph, horseshoe_pattern)
    if s is False:
        return False

    # Check that horseshoe_1 was unified with a half-cauchy distributed RV
    halfcauchy_1 = eval_if_etuple(s[horseshoe_1_lv])
    if halfcauchy_1.owner is None or not isinstance(
        halfcauchy_1.owner.op, type(at.random.halfcauchy)
    ):
         return False

    # Check that horseshoe_2 was unified with a half-cauchy distributed RV
    halfcauchy_2 = eval_if_etuple(s[horseshoe_2_lv])
    if halfcauchy_2.owner is None or not isinstance(
        halfcauchy_2.owner.op, type(at.random.halfcauchy)
    ):
        return False
    # Check that at least one of the RVs is a scalar
    if halfcauchy_1.type.shape == (1,):
        lmbda_rv = halfcauchy_2
        tau_rv = halfcauchy_1
    elif halfcauchy_2.type.shape == (1,):
        lmbda_rv = halfcauchy_1
        tau_rv = halfcauchy_2
    else:
        return false

    return (lmbda_rv, tau_rv)

Again we check that we can unify the subgraph beta_rv:

print(unify_horseshoe(beta_rv))

(halfcauchy_rv{0, (0, 0), floatX, False}.out, halfcauchy_rv{0, (0, 0), floatX, False}.out)

unify_horseshoe will only work if the pattern matches the whole graph. Indeed if we try to unify the horsehoe pattern with Y_rv we get:

print(unify_horseshoe(Y_rv))

False

To identify subgraphs we thus need to walk through the graph (here breadth-first) and attempt unification at each step:

from aesara.graph.basic import walk
from aesara.tensor.random.op import RandomVariable

def expand(var):
    if var.owner:
        return var.owner.inputs
    else:
        return

for node in walk([Y_rv], expand, bfs=True):
    try:
        if isinstance(node.owner.op, RandomVariable):
            s = unify_horseshoe(node)
            if s:
                break
    except AttributeError:
        continue

We can check that \(\tau\) and \(\lambda\) have been correctly identified:

pprint(s)

(halfcauchy_rv{0, (0, 0), floatX, False}.out,
 halfcauchy_rv{0, (0, 0), floatX, False}.out)

Being able to unify a pattern with a subgraph is a (small) first step towards being able to assign sampling steps to random variables in an arbitrary graph. In the following post we will show how we can automatically build a Gibbs sampler for (sub)graphs that represent a horseshoe prior.