Implement method for fast evals of (un)transformed RV shapes

Author: michaelosthege

pymc/model.py

@@ -62,7 +62,14 @@
 from pymc.distributions.transforms import Transform
 from pymc.exceptions import ImputationWarning, SamplingError, ShapeError
 from pymc.math import flatten_list
-from pymc.util import UNSET, WithMemoization, get_var_name, treedict, treelist
+from pymc.util import (
+    UNSET,
+    WithMemoization,
+    get_transformed_name,
+    get_var_name,
+    treedict,
+    treelist,
+)
 from pymc.vartypes import continuous_types, discrete_types, typefilter
 
 __all__ = [
@@ -1603,6 +1610,33 @@ def update_start_vals(self, a: Dict[str, np.ndarray], b: Dict[str, np.ndarray]):
 
         a.update({k: v for k, v in b.items() if k not in a})
 
+    def eval_rv_shapes(self) -> Dict[str, Tuple[int, ...]]:
+        """Evaluates shapes of untransformed AND transformed free variables.
+
+        Returns
+        -------
+        shapes : dict
+            Maps untransformed and transformed variable names to shape tuples.
+        """
+        names = []
+        outputs = []
+        for rv in self.free_RVs:
+            rv_var = self.rvs_to_values[rv]
+            transform = getattr(rv_var.tag, "transform", None)
+            if transform is not None:
+                names.append(get_transformed_name(rv.name, transform))
+                outputs.append(transform.forward(rv, rv).shape)
+            names.append(rv.name)
+            outputs.append(rv.shape)
+        f = aesara.function(
+            inputs=[],
+            outputs=outputs,
+            givens=[(obs, obs.tag.observations) for obs in self.observed_RVs],
+            mode=aesara.compile.mode.FAST_COMPILE,
+            on_unused_input="ignore",
+        )
+        return {name: tuple(shape) for name, shape in zip(names, f())}
+
     def check_start_vals(self, start):
         r"""Check that the starting values for MCMC do not cause the relevant log probability
         to evaluate to something invalid (e.g. Inf or NaN)

pymc/tests/test_model.py

@@ -594,6 +594,28 @@ def test_soft_update_parent(self):
         assert_almost_equal(start["interv_interval__"], test_point["interv_interval__"])
 
 
+class TestShapeEvaluation:
+    def test_eval_rv_shapes(self):
+        with pm.Model(
+            coords={
+                "city": ["Sydney", "Las Vegas", "Düsseldorf"],
+            }
+        ) as pmodel:
+            pm.Data("budget", [1, 2, 3, 4], dims="year")
+            pm.Normal("untransformed", size=(1, 2))
+            pm.Uniform("transformed", size=(7,))
+            obs = pm.Uniform("observed", size=(3,), observed=[0.1, 0.2, 0.3])
+            pm.LogNormal("lognorm", mu=at.log(obs))
+            pm.Normal("from_dims", dims=("city", "year"))
+        shapes = pmodel.eval_rv_shapes()
+        assert shapes["untransformed"] == (1, 2)
+        assert shapes["transformed"] == (7,)
+        assert shapes["transformed_interval__"] == (7,)
+        assert shapes["lognorm"] == (3,)
+        assert shapes["lognorm_log__"] == (3,)
+        assert shapes["from_dims"] == (3, 4)
+
+
 class TestCheckStartVals(SeededTest):
     def setup_method(self):
         super().setup_method()