Add a .from_pymc adapter.

PiperOrigin-RevId: 597570748

bayeux/_src/bayeux.py

@@ -21,6 +21,7 @@
 from bayeux import vi
 from bayeux._src import shared
 import jax
+import jax.numpy as jnp
 import oryx
 
 _MODULES = (mcmc, optimize, vi)
@@ -151,3 +152,75 @@ def transform_fn(x):
         transform_fn=transform_fn,
         initial_state=initial_state)
 
+  @classmethod
+  def from_pymc(cls, pm_model, initial_state=None):
+    import pymc as pm  # pylint: disable=g-import-not-at-top
+    import pymc.sampling.jax as pm_jax  # pylint: disable=g-import-not-at-top
+
+    uc_model = pm.model.transform.conditioning.remove_value_transforms(pm_model)
+    uc_vars = uc_model.value_vars
+    uc_names = [v.name for v in uc_vars]
+
+    logp = pm_jax.get_jaxified_logp(model=uc_model)
+
+    def identity(x, *_):
+      return x
+
+    def none(*_):
+      return 0.
+
+    # We have different ideas of forward and backward!
+    fwd_transforms = {
+        k.name: identity if v is None else v.backward
+        for k, v in pm_model.rvs_to_transforms.items()}
+    bwd_transforms = {
+        k.name: identity if v is None else v.forward
+        for k, v in pm_model.rvs_to_transforms.items()}
+
+    def forward_transform(pt):
+      return [
+          fwd_transforms[k](v, *([] if v.owner is None else v.owner.inputs))
+          for k, v in zip(uc_names, pt)]
+
+    def backward_transform(pt):
+      return [
+          bwd_transforms[k](v, *([] if v.owner is None else v.owner.inputs))
+          for k, v in zip(uc_names, pt)
+      ]
+
+    ildjs = {
+        k.name: none if v is None else v.log_jac_det
+        for k, v in pm_model.rvs_to_transforms.items()}
+
+    def ildj(pt):
+      return -pm.math.log(
+          pm.math.sum([
+              ildjs[k](v, *([] if v.owner is None else v.owner.inputs))
+              for k, v in zip(uc_names, pt)]))
+
+    fwd = pm_jax.get_jaxified_graph(
+        inputs=uc_vars, outputs=forward_transform(uc_vars))
+    bwd = pm_jax.get_jaxified_graph(
+        inputs=uc_vars, outputs=backward_transform(uc_vars))
+    ildj = pm_jax.get_jaxified_graph(inputs=uc_vars, outputs=[ildj(uc_vars)])
+    def logp_wrap(args):
+      return logp([args[k] for k in uc_names])
+
+    def fwd_wrap(args):
+      ret = fwd(*[args[k] for k in uc_names])
+      return dict(zip(uc_names, ret))
+
+    def bwd_wrap(args):
+      ret = bwd(*[args[k] for k in uc_names])
+      return dict(zip(uc_names, ret))
+
+    def ildj_wrap(args):
+      return ildj(*[args[k] for k in uc_names])[0]
+
+    test_point = {k.name: jnp.ones(k.shape.eval()) for k in uc_model.free_RVs}
+    return cls(log_density=logp_wrap,
+               test_point=test_point,
+               transform_fn=fwd_wrap,
+               inverse_transform_fn=bwd_wrap,
+               inverse_log_det_jacobian=ildj_wrap,
+               initial_state=initial_state)

bayeux/tests/compat_test.py

@@ -65,3 +65,24 @@ def tfp_model():
   bx_model = bx.Model.from_tfp(pinned_model)
   ret = bx_model.optimize.optax_adam(seed=jax.random.key(0))
   assert ret is not None
+
+
+def test_from_pymc():
+  import pymc as pm  # pylint: disable=g-import-not-at-top
+
+  treatment_effects = np.array([28, 8, -3, 7, -1, 1, 18, 12], dtype=np.float32)
+  treatment_stddevs = np.array(
+      [15, 10, 16, 11, 9, 11, 10, 18], dtype=np.float32)
+
+  with pm.Model() as model:
+    avg_effect = pm.Normal('avg_effect', 0., 10.)
+    avg_stddev = pm.HalfNormal('avg_stddev', 10.)
+    school_effects = pm.Normal('school_effects', shape=8)
+    pm.Normal('observed',
+              avg_effect + avg_stddev * school_effects,
+              treatment_stddevs,
+              observed=treatment_effects)
+
+  bx_model = bx.Model.from_pymc(model)
+  ret = bx_model.optimize.optax_adam(seed=jax.random.key(0))
+  assert ret is not None

bayeux/tests/optimize_test.py

@@ -86,6 +86,8 @@ def test_optimizers(method, linear_model):  # pylint: disable=redefined-outer-na
 
   if method not in {
       "optax_adafactor",
+      "optax_adagrad",
+      "optax_sm3",
       "optimistix_chord",
       "optimistix_nelder_mead"}:
     np.testing.assert_allclose(expected, params.w, atol=atol)

pyproject.toml

@@ -42,6 +42,7 @@ dev = [
     "pytest-xdist",
     "pylint>=2.6.0",
     "pyink",
+    "pymc",
 ]