Add transform handling code

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -255,6 +255,25 @@ 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)
@@ -324,7 +343,6 @@ def recover_marginals(
             other_values = [v for v in m.value_vars if v is not marginalized_value]
 
             # TODO: Handle constants
-            # TODO: Handle transformed variables
 
             joint_logps = vectorize_graph(
                 joint_logp,
@@ -346,7 +364,7 @@ def recover_marginals(
                     on_unused_input="ignore",
                 )
 
-            logvs = [rv_loglike_fn(**vs) for vs in posterior_pts]
+            logvs = [rv_loglike_fn(**self._transform_input(vs)) for vs in posterior_pts]
 
             if include_samples:
                 logps, samples = zip(*logvs)