Register the conjugates relation in a database

aesara-devs · Jun 7, 2022 · 87d6f75 · 87d6f75
1 parent 249e65a
commit 87d6f75
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/aemcmc/conjugates.py b/aemcmc/conjugates.py
@@ -1,8 +1,12 @@
 import aesara.tensor as at
+from aesara.graph.kanren import KanrenRelationSub
+from aesara.graph.optdb import OptimizationDatabase
 from etuples import etuple, etuplize
 from kanren import eq, lall
 from unification import var
 
+conjugatesdb = OptimizationDatabase()
+
 
 def beta_binomial_conjugateo(observed_rv_expr, posterior_expr):
     r"""Produce a goal that represents the application of Bayes theorem
@@ -64,3 +68,6 @@ def beta_binomial_conjugateo(observed_rv_expr, posterior_expr):
         eq(observed_rv_expr[1], y_lv),
         eq(posterior_expr, p_posterior_et),
     )
+
+
+conjugatesdb.register("beta_binomial", KanrenRelationSub(beta_binomial_conjugateo))