Don't have to worry about transforms if you never transform

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -255,25 +255,6 @@ def marginalize(self, rvs_to_marginalize: Union[TensorVariable, Sequence[TensorV
         # Raise errors and warnings immediately
         self.clone()._marginalize(user_warnings=True)
 
-    def _transform_input(self, inputs):
-        "Create a function from the untransformed space to the transformed space"
-        transformed_rvs = []
-        transformed_names = []
-
-        for rv in self.free_RVs:
-            transform = self.rvs_to_transforms.get(rv)
-            if transform is None:
-                transformed_rvs.append(rv)
-                transformed_names.append(rv.name)
-            else:
-                transformed_rv = transform.forward(rv, *rv.owner.inputs)
-                transformed_rvs.append(transformed_rv)
-                transformed_names.append(self.rvs_to_values[rv].name)
-
-        fn = self.compile_fn(inputs=self.free_RVs, outs=transformed_rvs)
-        rets = fn(inputs)
-        return dict(zip(transformed_names, rets))
-
     def unmarginalize(self, rvs_to_unmarginalize):
         for rv in rvs_to_unmarginalize:
             self.marginalized_rvs.remove(rv)
@@ -320,8 +301,15 @@ def recover_marginals(
         rv_dict = {}
         rv_dims_dict = {}
 
+        # Disable all transforms
+        model = self.clone()
+        model.rvs_to_transforms = {k: None for k in model.rvs_to_transforms}
+        for rv in model.free_RVs:
+            model.rvs_to_values[rv].name = rv.name
+        var_names = [model.vars_to_clone[rv] for rv in var_names]
+
         for rv in var_names:
-            m = self.clone()
+            m = model.clone()
             rv = m.vars_to_clone[rv]
             m.unmarginalize([rv])
             joint_logp = m.logp()
@@ -343,7 +331,6 @@ def recover_marginals(
             other_values = [v for v in m.value_vars if v is not marginalized_value]
 
             # TODO: Handle constants
-
             joint_logps = vectorize_graph(
                 joint_logp,
                 replace={marginalized_value: rv_domain_tensor},
@@ -364,7 +351,7 @@ def recover_marginals(
                     on_unused_input="ignore",
                 )
 
-            logvs = [rv_loglike_fn(**self._transform_input(vs)) for vs in posterior_pts]
+            logvs = [rv_loglike_fn(**vs) for vs in posterior_pts]
 
             if include_samples:
                 logps, samples = zip(*logvs)

pymc_experimental/tests/model/test_marginal_model.py

@@ -255,7 +255,7 @@ def test_marginalized_change_point_model_sampling(disaster_model):
 
 def test_recover_marginals_basic():
     with MarginalModel() as m:
-        sigma = pm.HalfNormal("sigma", transform=None)
+        sigma = pm.HalfNormal("sigma")
         p = np.array([0.5, 0.2, 0.3])
         k = pm.Categorical("k", p=p)
         mu = np.array([-3.0, 0.0, 3.0])