.WIP

Author: ricardoV94

pymc_experimental/model/marginal/distributions.py

@@ -39,7 +39,8 @@ def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> tuple[int, ...]:
         return (0, 1)
     elif isinstance(op, Categorical):
         [p_param] = dist_params
-        return tuple(range(pt.get_vector_length(p_param)))
+        [p_param_length] = constant_fold([p_param.shape[-1]])
+        return tuple(range(p_param_length))
     elif isinstance(op, DiscreteUniform):
         lower, upper = constant_fold(dist_params)
         return tuple(np.arange(lower, upper + 1))
@@ -60,7 +61,7 @@ def _add_reduce_batch_dependent_logps(
     for dependent_logp in dependent_logps:
         dbcast = dependent_logp.type.broadcastable
         dim_diff = len(dbcast) - len(mbcast)
-        mbcast_aligned = (True,) * dim_diff + mbcast
+        mbcast_aligned = mbcast + (True,) * dim_diff
         vbcast_axis = [i for i, (m, v) in enumerate(zip(mbcast_aligned, dbcast)) if m and not v]
         reduced_logps.append(dependent_logp.sum(vbcast_axis))
     return pt.add(*reduced_logps)
@@ -79,6 +80,8 @@ def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
     inner_rv_values = dict(zip(inner_rvs, values))
     marginalized_vv = marginalized_rv.clone()
     rv_values = inner_rv_values | {marginalized_rv: marginalized_vv}
+    print("")
+    print("Inner conditional logp call >> ")
     logps_dict = conditional_logp(rv_values=rv_values, **kwargs)
 
     # Reduce logp dimensions corresponding to broadcasted variables

pymc_experimental/model/marginal/marginal_model.py

@@ -13,11 +13,12 @@
 from pymc.distributions.transforms import Chain
 from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Model
-from pymc.pytensorf import compile_pymc, constant_fold
+from pymc.pytensorf import compile_pymc, constant_fold, toposort_replace
 from pymc.util import RandomState, _get_seeds_per_chain, treedict
 from pytensor.graph import FunctionGraph, clone_replace
+from pytensor.graph.basic import truncated_graph_inputs, Constant, ancestors
 from pytensor.graph.replace import vectorize_graph
-from pytensor.tensor import TensorVariable
+from pytensor.tensor import TensorVariable, extract_constant
 from pytensor.tensor.special import log_softmax
 
 __all__ = ["MarginalModel", "marginalize"]
@@ -544,52 +545,45 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     if not dependent_rvs:
         raise ValueError(f"No RVs depend on marginalized RV {rv_to_marginalize}")
 
-    ndim_supp = max({rv.owner.op.ndim_supp for rv in dependent_rvs})
-
     marginalized_rv_input_rvs = find_conditional_input_rvs([rv_to_marginalize], all_rvs)
     other_direct_rv_ancestors = [
         rv
         for rv in find_conditional_input_rvs(dependent_rvs, all_rvs)
         if rv is not rv_to_marginalize
     ]
 
-    # If the marginalized RV has multiple dimensions, check that graph between
-    # marginalized RV and dependent RVs does not mix information from batch dimensions
-    # (otherwise logp would require enuremating over all combinations of batch dimension values)
-    if any(not bcast for bcast in rv_to_marginalize.type.broadcastable):
-        # When there are batch dimensions, we call `batch_dims_subgraph` to make sure these are not mixed
-        dependent_rvs_dims = subgraph_dim_connection(
+    if all (rv_to_marginalize.type.broadcastable):
+        ndim_supp = max([dependent_rv.type.ndim for dependent_rv in dependent_rvs])
+    else:
+        # If the marginalized RV has multiple dimensions, check that graph between
+        # marginalized RV and dependent RVs does not mix information from batch dimensions
+        # (otherwise logp would require enumerating over all combinations of batch dimension values)
+        dependent_rvs_dim_connections = subgraph_dim_connection(
             rv_to_marginalize, other_direct_rv_ancestors, dependent_rvs
         )
+        # dependent_rvs_dim_connections = subgraph_dim_connection(
+        #     rv_to_marginalize, other_inputs, dependent_rvs
+        # )
 
-        # Cr
+        ndim_supp = max((dependent_rv.type.ndim - rv_to_marginalize.type.ndim) for dependent_rv in dependent_rvs)
 
-        if any(len(dim) > 1 for dim in dependent_rvs_dims):
+        if any(len(dim) > 1 for rv_dim_connections in dependent_rvs_dim_connections for dim in rv_dim_connections):
             raise NotImplementedError("Multiple dimensions are mixed")
 
-        # We further check that any extra batch dimensions of dependent RVs beyond those implied by the MarginalizedRV
-        # show up on the left, so that collapsing logic in logp can be more straightforward.
+        # We further check that:
+        # 1) Dimensions of dependent RVs are aligned with those of the marginalized RV
+        # 2) Any extra batch dimensions of dependent RVs beyond those implied by the MarginalizedRV
+        # show up on the right, so that collapsing logic in logp can be more straightforward.
         # This also ensures the MarginalizedRV still behaves as an RV itself
         marginal_batch_ndim = rv_to_marginalize.owner.op.batch_ndim(rv_to_marginalize.owner)
         marginal_batch_dims = tuple((i,) for i in range(marginal_batch_ndim))
-        for dependent_rv, dependent_rv_batch_dims in zip(dependent_rvs, dependent_rvs_dims):
-            extra_batch_ndim = (
-                dependent_rv.type.ndim - marginal_batch_ndim - dependent_rv.owner.op.ndim_supp
-            )
-            valid_dependent_batch_dims = (((),) * extra_batch_ndim) + marginal_batch_dims
+        for dependent_rv, dependent_rv_batch_dims in zip(dependent_rvs, dependent_rvs_dim_connections):
+            extra_batch_ndim = dependent_rv.type.ndim - marginal_batch_ndim
+            valid_dependent_batch_dims = marginal_batch_dims + (((),) * extra_batch_ndim)
             if dependent_rv_batch_dims != valid_dependent_batch_dims:
                 raise NotImplementedError(
-                    "Any extra batch dimensions introduced by dependent RVs must be "
-                    "on the left of dimensions introduced by the marginalized RV"
-                )
-
-        for dependent_rv, dependent_rv_batch_dims in zip(dependent_rvs, dependent_rvs_dims):
-            shared_batch_dims = [
-                batch_dim for batch_dim in dependent_rv_batch_dims if batch_dim is not None
-            ]
-            if shared_batch_dims != sorted(shared_batch_dims):
-                raise NotImplementedError(
-                    "Shared batch dimensions between marginalized RV and dependent RVs must be aligned positionally"
+                    "Any extra dimensions introduced by dependent RVs must appear to the right of dimensions "
+                    "introduced by the marginalized RV."
                 )
 
     input_rvs = [*marginalized_rv_input_rvs, *other_direct_rv_ancestors]
@@ -598,7 +592,22 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     outputs = rvs_to_marginalize
     # We are strict about shared variables in SymbolicRandomVariables
     inputs = input_rvs + collect_shared_vars(rvs_to_marginalize, blockers=input_rvs)
-
+    # inputs = [
+    #     inp
+    #     for rv in rvs_to_marginalize  # should be toposort
+    #     for inp in rv.owner.inputs
+    #     if not(all(isinstance(a, Constant) for a in ancestors([inp], blockers=all_rvs)))
+    # ]
+    # inputs = [
+    #     inp for inp in truncated_graph_inputs(outputs, ancestors_to_include=inputs)
+    #     if not (all(isinstance(a, Constant) for a in ancestors([inp], blockers=all_rvs)))
+    # ]
+    # inputs = truncated_graph_inputs(outputs, ancestors_to_include=[
+    # #     inp
+    # #     for output in outputs
+    # #     for inp in output.owner.inputs
+    # # ])
+    # inputs = [inp for inp in inputs if not isinstance(constant_fold([inp], raise_not_constant=False)[0], Constant | np.ndarray)]
     if isinstance(rv_to_marginalize.owner.op, DiscreteMarkovChain):
         marginalize_constructor = DiscreteMarginalMarkovChainRV
     else:
@@ -611,6 +620,14 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     )
 
     marginalized_rvs = marginalization_op(*inputs)
-    fgraph.replace_all(tuple(zip(rvs_to_marginalize, marginalized_rvs)))
+    print()
+    import pytensor
+    pytensor.dprint(marginalized_rvs, print_type=True)
+    fgraph.replace_all(reversed(tuple(zip(rvs_to_marginalize, marginalized_rvs))))
+    # assert 0
+    # fgraph.dprint()
+    # assert 0
+    # toposort_replace(fgraph, tuple(zip(rvs_to_marginalize, marginalized_rvs)))
+    # assert 0
     return rvs_to_marginalize, marginalized_rvs
 

tests/model/marginal/test_graph_analysis.py

@@ -39,6 +39,9 @@ def test_subtensor(self):
         [dims] = subgraph_dim_connection(inp, [], [valid_out])
         assert dims == ((), (2,))
 
+    def test_advanced_subtensor(self):
+        raise NotImplementedError()
+
     def test_elemwise(self):
         inp = pt.zeros(shape=(5, 5))
 
@@ -93,5 +96,3 @@ def test_symbolic_random_variable(self):
         [dims] = subgraph_dim_connection(inp, [], [out])
         assert dims == ((0, 2), (1, 2))
 
-    def test_advanced_indexing(self):
-        raise NotImplementedError()

tests/model/marginal/test_marginal_model.py

@@ -133,36 +133,52 @@ def test_rv_dependent_multiple_marginalized_rvs():
     np.testing.assert_allclose(np.exp(logp({"z": 2})), 0.1 * 0.3)
 
 
-def test_nested_marginalized_rvs():
+@pytest.mark.parametrize("batched", (False, True))
+def test_nested_marginalized_rvs(batched):
     """Test that marginalization works when there are nested marginalized RVs"""
 
-    with MarginalModel() as m:
-        sigma = pm.HalfNormal("sigma")
+    def build_model(build_batched: bool) -> MarginalModel:
+        idx_shape =  (3,) if build_batched else ()
+        sub_idx_shape = (3, 5) if build_batched else (5,)
 
-        idx = pm.Bernoulli("idx", p=0.75)
-        dep = pm.Normal("dep", mu=pt.switch(pt.eq(idx, 0), -1000.0, 1000.0), sigma=sigma)
+        with MarginalModel() as m:
+            sigma = pm.HalfNormal("sigma")
 
-        sub_idx = pm.Bernoulli("sub_idx", p=pt.switch(pt.eq(idx, 0), 0.15, 0.95), shape=(5,))
-        sub_dep = pm.Normal("sub_dep", mu=dep + sub_idx * 100, sigma=sigma, shape=(5,))
+            idx = pm.Bernoulli("idx", p=0.75, shape=idx_shape)
+            dep = pm.Normal("dep", mu=pt.switch(pt.eq(idx, 0), -1000.0, 1000.0), sigma=sigma)
 
-    ref_logp_fn = m.compile_logp(vars=[idx, dep, sub_idx, sub_dep])
+            sub_idx_p = pt.switch(pt.eq(idx, 0), 0.15, 0.95)
+            if build_batched:
+                sub_idx_p = sub_idx_p[:, None]
+                dep = dep[:, None]
+            sub_idx = pm.Bernoulli("sub_idx", p=sub_idx_p, shape=sub_idx_shape)
+            sub_dep = pm.Normal("sub_dep", mu=dep + sub_idx * 100, sigma=sigma)
 
-    with pytest.warns(UserWarning, match="There are multiple dependent variables"):
-        m.marginalize([idx, sub_idx])
+        return m
 
-    assert set(m.marginalized_rvs) == {idx, sub_idx}
+    m = build_model(build_batched=batched)
+    with pytest.warns(UserWarning, match="There are multiple dependent variables"):
+        m.marginalize(["idx", "sub_idx"])
+    assert sorted(m.name for m in m.marginalized_rvs) == ["idx", "sub_idx"]
+    return
 
     # Test logp
-    test_point = m.initial_point()
-    test_point["dep"] = 1000
-    test_point["sub_dep"] = np.full((5,), 1000 + 100)
+    ref_m = build_model(build_batched=False)
+    ref_logp_fn = ref_m.compile_logp(vars=[ref_m["idx"], ref_m["dep"], ref_m["sub_idx"], ref_m["sub_dep"]])
 
+    test_point = ref_m.initial_point()
+    test_point["dep"] = np.full_like(test_point["dep"], 1000)
+    test_point["sub_dep"] = np.full_like(test_point["sub_dep"], 1000 + 100)
     ref_logp = [
         ref_logp_fn({**test_point, **{"idx": idx, "sub_idx": np.array(sub_idxs)}})
         for idx in (0, 1)
         for sub_idxs in itertools.product((0, 1), repeat=5)
     ]
-    logp = m.compile_logp(vars=[dep, sub_dep])(test_point)
+
+    test_point = m.initial_point()
+    test_point["dep"] = np.full_like(test_point["dep"], 1000)
+    test_point["sub_dep"] = np.full_like(test_point["sub_dep"], 1000 + 100)
+    logp = m.compile_logp(vars=[m["dep"], m["sub_dep"]])(test_point)
 
     np.testing.assert_almost_equal(
         logp,
@@ -615,8 +631,8 @@ def test_change_point_model_sampling(self, disaster_model):
 
     def test_k_censored_clusters_model(self):
         def build_model(batch: bool) -> MarginalModel:
-            data = np.array([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]).T
-            nobs = data.shape[-1]
+            data = np.array([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])
+            nobs = data.shape[0]
             n_clusters = 5
             coords = {
                 "cluster": range(n_clusters),
@@ -641,17 +657,17 @@ def build_model(batch: bool) -> MarginalModel:
                     initval=np.linspace(-1, 1, n_clusters),
                 )
                 mu_y = pm.Normal("mu_y", dims=["cluster"])
-                mu = pm.math.concatenate([mu_x[None], mu_y[None]], axis=0)  # (ndim, cluster)
+                mu = pm.math.stack([mu_x, mu_y], axis=-1)  # (cluster, ndim)
 
                 sigma = pm.HalfNormal("sigma")
 
                 y = pm.Censored(
                     "y",
-                    dist=pm.Normal.dist(mu[:, idx], sigma),
+                    dist=pm.Normal.dist(mu[idx, :], sigma),
                     lower=-3,
                     upper=3,
                     observed=data,
-                    dims=["ndim", "obs"],
+                    dims=["obs", "ndim"],
                 )
 
             return m