Add goal for loc-scale transformation

aemcmc/transforms.py

@@ -0,0 +1,46 @@
+import aesara.tensor as at
+from etuples import etuple, etuplize
+from kanren import eq, lall
+from kanren.facts import Relation, fact
+from unification import var
+
+location_scale_family = Relation("loc-scale")
+fact(location_scale_family, at.random.cauchy)
+fact(location_scale_family, at.random.gumbel)
+fact(location_scale_family, at.random.laplace)
+fact(location_scale_family, at.random.logistic)
+fact(location_scale_family, at.random.normal)
+
+
+def location_scale_transform(in_expr, out_expr):
+    """Create a relation to lift and sink scale and location parameters of distributions."""
+
+    # Centered representation
+    rng_lv, size_lv, type_idx_lv = var(), var(), var()
+    mu_lv, sd_lv = var(), var()
+    distribution_lv = var()
+    centered_et = etuple(distribution_lv, rng_lv, size_lv, type_idx_lv, mu_lv, sd_lv)
+
+    # Non-centered representation
+    noncentered_et = etuple(
+        etuplize(at.add),
+        mu_lv,
+        etuple(
+            etuplize(at.mul),
+            sd_lv,
+            etuple(
+                distribution_lv,
+                0.0,
+                1.0,
+                rng=rng_lv,
+                size=size_lv,
+                dtype=type_idx_lv,
+            ),
+        ),
+    )
+
+    return lall(
+        eq(in_expr, centered_et),
+        eq(out_expr, noncentered_et),
+        location_scale_family(distribution_lv),
+    )

tests/test_transforms.py

@@ -0,0 +1,38 @@
+from functools import partial
+
+import aesara.tensor as at
+from aesara.graph.unify import eval_if_etuple
+from kanren import run
+from kanren.graph import reduceo, walko
+from unification import var
+
+from aemcmc.transforms import location_scale_transform
+
+
+def test_normal_scale_loc_transform():
+    """"""
+
+    srng = at.random.RandomStream(0)
+    mu_a_rv = srng.normal(0, 1)
+    sigma_a_rv = srng.halfcauchy(1)
+    a_rv = srng.normal(mu_a_rv, sigma_a_rv, size=(10,))
+
+    mu_b_rv = srng.normal(0, 1)
+    sigma_b_rv = srng.halfcauchy(1)
+    b_rv = srng.normal(mu_b_rv, sigma_b_rv, size=(10))
+
+    mu = a_rv + b_rv
+    sigma_rv = srng.halfcauchy(5.0)
+    Y_rv = srng.cauchy(mu, sigma_rv)
+
+    q_lv = var()
+    (expr_graph,) = run(
+        1, q_lv, walko(partial(reduceo, location_scale_transform), Y_rv, q_lv)
+    )
+    print(expr_graph)
+    Y_nc_rv = eval_if_etuple(expr_graph)
+
+    # Make sure that Y_rv gets replaced with an addition
+    assert Y_nc_rv.owner.op == at.add
+    rhs = Y_nc_rv.owner.inputs[1].owner.inputs[0]
+    assert rhs.owner.op == at.mul