.WIP backtrack

Author: ricardoV94

pymc_experimental/model/marginal/graph_analysis.py

@@ -2,6 +2,7 @@
 from itertools import chain, zip_longest
 
 from pymc import SymbolicRandomVariable
+from pymc.distributions.custom import CustomSymbolicDistRV
 from pytensor.compile import SharedVariable
 from pytensor.graph import Constant, Variable, ancestors, graph_inputs
 from pytensor.graph.basic import io_toposort
@@ -112,7 +113,6 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
 
         if not any(inputs_dims):
             # None of the inputs are related to the batch_axes of the marginalized_rv
-            # We could set `()` for everything, but for now that doesn't seem needed
             continue
 
         elif isinstance(node.op, DimShuffle):
@@ -122,7 +122,11 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
             )
             var_dims[node.outputs[0]] = output_dims
 
-        elif isinstance(node.op, SymbolicRandomVariable) and node.op.extended_signature is None:
+        elif (
+            isinstance(node.op, CustomSymbolicDistRV)
+            or isinstance(node.op, SymbolicRandomVariable)
+            and node.op.extended_signature is None
+        ):
             # SymbolicRandomVariables without signature are a wild-card, so we need to introspect the inner graph.
             # MarginalRVs are such a case!
             inner_var_dims = {
@@ -154,23 +158,18 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
                     missing_ndim = op_batch_ndim - (len(param_dims) - param_core_ndim)
                     inputs_dims[param_idx] = ((),) * missing_ndim + param_dims
 
-            # Collapse all core_dims
-            core_dims = tuple(
-                sorted(
-                    chain.from_iterable(
-                        [i for input_dim in inputs_dims for i in input_dim[op_batch_ndim:]]
-                    )
+            if any(any(input_dim[op_batch_ndim:]) for input_dim in inputs_dims):
+                raise ValueError(
+                    f"Use of known dimensions as core dimensions of op {node.op} not supported."
                 )
-            )
+
             batch_dims = _broadcast_dims(
                 tuple(input_dims[:op_batch_ndim] for input_dims in inputs_dims)
             )
-            # Add batch dims to each output_dims
-            batch_dims = tuple(batch_dim + core_dims for batch_dim in batch_dims)
             for out in node.outputs:
                 if isinstance(out.type, TensorType):
                     core_ndim = out.type.ndim - op_batch_ndim
-                    output_dims = batch_dims + (core_dims,) * core_ndim
+                    output_dims = batch_dims + ((),) * core_ndim
                     var_dims[out] = output_dims
 
         elif isinstance(node.op, CAReduce):
@@ -182,8 +181,12 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
             elif axes is None:
                 axes = tuple(range(node.inputs[0].type.ndim))
 
-            # Output dims contain the collapsed dims
-            output_dims = [dims + axes for i, dims in enumerate(input_dims) if i not in axes]
+            if any(input_dims[axis] for axis in axes):
+                raise ValueError(
+                    f"Use of known dimensions as reduced dimensions of op {node.op} not supported."
+                )
+
+            output_dims = [dims for i, dims in enumerate(input_dims) if i not in axes]
             var_dims[node.outputs[0]] = tuple(output_dims)
 
         elif isinstance(node.op, Subtensor):
@@ -198,7 +201,7 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
                     # Dim is kept
                     output_dims.append(value_dims)
                 elif value_dims:
-                    raise NotImplementedError(
+                    raise ValueError(
                         "Partial slicing or indexing of known dimensions not supported."
                     )
                 elif isinstance(idx, slice):
@@ -242,7 +245,7 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
                         non_adv_dims.append(value_dim)
                     elif value_dim:
                         # We are trying to partially slice or index a known dimension
-                        raise NotImplementedError(
+                        raise ValueError(
                             "Partial slicing or advanced integer indexing of known dimensions not supported."
                         )
                     elif isinstance(idx, slice):
@@ -277,16 +280,18 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
     return var_dims
 
 
-def subgraph_dim_connection(
+def subgraph_batch_dim_connection(
     input_var, other_inputs, output_vars
 ) -> list[tuple[tuple[int, ...], ...]]:
-    """Identify how the dims of rv_to_marginalize are consumed by the dims of the output_rvs.
+    """Identify how the batch dims of rv_to_marginalize map to the batch dimensions of the output_rvs.
 
     Raises
     ------
+    ValueError
+        If input batch dimensions are mixed in the graph leading to output_vars.
+
     NotImplementedError
         If variable related to marginalized batch_dims is used in an operation that is not yet supported
-
     """
     var_dims = {input_var: tuple((i,) for i in range(input_var.type.ndim))}
     var_dims = _subgraph_dim_connection(var_dims, [input_var, *other_inputs], output_vars)

pymc_experimental/model/marginal/marginal_model.py

@@ -34,7 +34,7 @@
     find_conditional_dependent_rvs,
     find_conditional_input_rvs,
     is_conditional_dependent,
-    subgraph_dim_connection,
+    subgraph_batch_dim_connection,
 )
 
 ModelRVs = TensorVariable | Sequence[TensorVariable] | str | Sequence[str]
@@ -566,9 +566,15 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
         # 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
-        )
+        try:
+            dependent_rvs_dim_connections = subgraph_batch_dim_connection(
+                rv_to_marginalize, other_direct_rv_ancestors, dependent_rvs
+            )
+        except ValueError as e:
+            # For the perspective of the user this is a NotImplementedError
+            raise NotImplementedError(
+                "The graph between the marginalized and dependent RVs cannot be marginalized"
+            ) from e
 
         if any(
             len(dim) > 1

tests/model/marginal/test_graph_analysis.py

@@ -2,40 +2,46 @@
 import pytest
 
 from pymc.distributions import CustomDist
+from pytensor.tensor.type_other import NoneTypeT
 
-from pymc_experimental.model.marginal.graph_analysis import subgraph_dim_connection
+from pymc_experimental.model.marginal.graph_analysis import subgraph_batch_dim_connection
 
 
-class TestSubgraphDimConnection:
+class TestSubgraphBatchDimConnection:
     def test_dimshuffle(self):
         inp = pt.tensor(shape=(5, 1, 4, 3))
         out1 = pt.matrix_transpose(inp)
         out2 = pt.expand_dims(inp, 1)
         out3 = pt.squeeze(inp)
-        [dims1, dims2, dims3] = subgraph_dim_connection(inp, [], [out1, out2, out3])
+        [dims1, dims2, dims3] = subgraph_batch_dim_connection(inp, [], [out1, out2, out3])
         assert dims1 == ((0,), (1,), (3,), (2,))
         assert dims2 == ((0,), (), (1,), (2,), (3,))
         assert dims3 == ((0,), (2,), (3,))
 
     def test_careduce(self):
         inp = pt.tensor(shape=(4, 3, 2))
-        out = pt.sum(inp, axis=(1,))
-        [dims] = subgraph_dim_connection(inp, [], [out])
-        assert dims == ((0, 1), (2, 1))
+
+        out = pt.sum(inp[:, None], axis=(1,))
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
+        assert dims == ((0,), (1,), (2,))
+
+        invalid_out = pt.sum(inp, axis=(1,))
+        with pytest.raises(ValueError, match="Use of known dimensions"):
+            subgraph_batch_dim_connection(inp, [], [invalid_out])
 
     def test_subtensor(self):
         inp = pt.tensor(shape=(4, 3, 2))
 
         invalid_out = inp[0, :1]
         with pytest.raises(
-            NotImplementedError,
+            ValueError,
             match="Partial slicing or indexing of known dimensions not supported",
         ):
-            subgraph_dim_connection(inp, [], [invalid_out])
+            subgraph_batch_dim_connection(inp, [], [invalid_out])
 
         # If we are selecting dummy / unknown dimensions that's fine
         valid_out = pt.expand_dims(inp, (0, 1))[0, :1]
-        [dims] = subgraph_dim_connection(inp, [], [valid_out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [valid_out])
         assert dims == ((), (0,), (1,), (2,))
 
     def test_advanced_subtensor_value(self):
@@ -44,99 +50,116 @@ def test_advanced_subtensor_value(self):
 
         # Index on an unlabled dim introduced by broadcasting with zeros
         out = intermediate_out[:, [0, 0, 1, 2]]
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((0,), (), (1,), ())
 
         # Indexing that introduces more dimensions
         out = intermediate_out[:, [[0, 0], [1, 2]], :]
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((0,), (), (), (1,), ())
 
         # Special case where advanced dims are moved to the front of the output
         out = intermediate_out[:, [0, 0, 1, 2], :, 0]
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((), (0,), (1,))
 
         # Indexing on a labeled dim fails
         out = intermediate_out[:, :, [0, 0, 1, 2]]
-        with pytest.raises(
-            NotImplementedError, match="Partial slicing or advanced integer indexing"
-        ):
-            subgraph_dim_connection(inp, [], [out])
+        with pytest.raises(ValueError, match="Partial slicing or advanced integer indexing"):
+            subgraph_batch_dim_connection(inp, [], [out])
 
     def test_advanced_subtensor_key(self):
         inp = pt.tensor(shape=(5, 5), dtype=int)
         base = pt.zeros((2, 3, 4))
 
         out = base[inp]
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((0,), (1,), (), ())
 
         out = base[:, :, inp]
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((), (), (0,), (1,))
 
         out = base[1:, 0, inp]
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((), (0,), (1,))
 
         # Special case where advanced dims are moved to the front of the output
         out = base[0, :, inp]
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((0,), (1,), ())
 
         # Mix keys dimensions
         out = base[:, inp, inp.T]
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((), (0, 1), (0, 1))
 
     def test_elemwise(self):
         inp = pt.tensor(shape=(5, 5))
 
         out = inp + inp
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((0,), (1,))
 
         out = inp + inp.T
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == (
             (0, 1),
-            (
-                0,
-                1,
-            ),
+            (0, 1),
         )
 
     def test_blockwise(self):
-        inp = pt.tensor(shape=(5, 4, 3, 2))
-        out = inp @ pt.ones((2, 3))
-        [dims] = subgraph_dim_connection(inp, [], [out])
-        # Every dimension contains information from the core dimensions
-        assert dims == ((0, 2, 3), (1, 2, 3), (2, 3), (2, 3))
+        inp = pt.tensor(shape=(5, 4))
+
+        invalid_out = inp @ pt.ones((4, 3))
+        with pytest.raises(ValueError, match="Use of known dimensions"):
+            subgraph_batch_dim_connection(inp, [], [invalid_out])
+
+        out = (inp[:, :, None, None] + pt.zeros((2, 3))) @ pt.ones((2, 3))
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
+        assert dims == ((0,), (1,), (), ())
 
     def test_random_variable(self):
         inp = pt.tensor(shape=(5, 4, 3))
+
         out1 = pt.random.normal(loc=inp)
-        out2 = pt.random.categorical(p=inp)
-        out3 = pt.random.multivariate_normal(mean=inp, cov=pt.eye(3))
-        [dims1, dims2, dims3] = subgraph_dim_connection(inp, [], [out1, out2, out3])
+        out2 = pt.random.categorical(p=inp[..., None])
+        out3 = pt.random.multivariate_normal(mean=inp[..., None], cov=pt.eye(1))
+        [dims1, dims2, dims3] = subgraph_batch_dim_connection(inp, [], [out1, out2, out3])
         assert dims1 == ((0,), (1,), (2,))
-        assert dims2 == ((0, 2), (1, 2))
-        assert dims3 == ((0, 2), (1, 2), (2,))
+        assert dims2 == ((0,), (1,), (2,))
+        assert dims3 == ((0,), (1,), (2,), ())
+
+        invalid_out = pt.random.categorical(p=inp)
+        with pytest.raises(ValueError, match="Use of known dimensions"):
+            subgraph_batch_dim_connection(inp, [], [invalid_out])
+
+        invalid_out = pt.random.multivariate_normal(mean=inp, cov=pt.eye(3))
+        with pytest.raises(ValueError, match="Use of known dimensions"):
+            subgraph_batch_dim_connection(inp, [], [invalid_out])
 
     def test_symbolic_random_variable(self):
         inp = pt.tensor(shape=(4, 3, 2))
+
+        # Test univariate
         out = CustomDist.dist(
             inp,
             dist=lambda mu, size: pt.random.normal(loc=mu, size=size),
         )
-        [dims] = subgraph_dim_connection(inp, [], [out])
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
         assert dims == ((0,), (1,), (2,))
 
         # Test multivariate
+        def dist(mu, size):
+            if isinstance(size.type, NoneTypeT):
+                size = mu.shape
+            return pt.random.normal(loc=mu[..., None], size=(*size, 2)).sum(-1)
+
         out = CustomDist.dist(
             inp,
-            dist=lambda mu, size: pt.random.normal(loc=mu, size=size).sum(-1),
+            dist=dist,
+            size=(4, 3, 2),
+            ndim_supp=1,
         )
-        [dims] = subgraph_dim_connection(inp, [], [out])
-        assert dims == ((0, 2), (1, 2))
+        [dims] = subgraph_batch_dim_connection(inp, [], [out])
+        assert dims == ((0,), (1,), (2,))