Fix circular import

Author: fonnesbeck

pymc_extras/model/marginal/distributions.py

@@ -17,7 +17,14 @@
 from pytensor.tensor import TensorVariable
 
 from pymc_extras.distributions import DiscreteMarkovChain
-from pymc_extras.model.marginal.graph_analysis import get_support_axes
+
+
+def get_support_axes(op) -> tuple[tuple[int, ...], ...]:
+    if hasattr(op, "support_axes"):
+        return op.support_axes
+    else:
+        # For vanilla RVs, the support axes are the last ndim_supp
+        return (tuple(range(-op.ndim_supp, 0)),)
 
 
 class MarginalRV(OpFromGraph, MeasurableOp):

pymc_extras/model/marginal/graph_analysis.py

@@ -5,6 +5,7 @@
 
 from pymc import SymbolicRandomVariable
 from pytensor.compile import SharedVariable
+from pytensor.compile.builders import OpFromGraph
 from pytensor.graph import Constant, Variable, ancestors
 from pytensor.graph.basic import io_toposort
 from pytensor.tensor import TensorType, TensorVariable
@@ -16,8 +17,6 @@
 from pytensor.tensor.subtensor import AdvancedSubtensor, Subtensor, get_idx_list
 from pytensor.tensor.type_other import NoneTypeT
 
-from pymc_extras.model.marginal.distributions import MarginalRV
-
 
 def static_shape_ancestors(vars):
     """Identify ancestors Shape Ops of static shapes (therefore constant in a valid graph)."""
@@ -63,7 +62,7 @@ def find_conditional_dependent_rvs(dependable_rv, all_rvs):
 
 
 def get_support_axes(op) -> tuple[tuple[int, ...], ...]:
-    if isinstance(op, MarginalRV):
+    if hasattr(op, "support_axes"):
         return op.support_axes
     else:
         # For vanilla RVs, the support axes are the last ndim_supp
@@ -146,7 +145,7 @@ def _subgraph_batch_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars)
             output_dims = tuple(None if i == "x" else input_dims[i] for i in node.op.new_order)
             var_dims[node.outputs[0]] = output_dims
 
-        elif isinstance(node.op, MarginalRV) or (
+        elif (isinstance(node.op, OpFromGraph) and hasattr(node.op, "support_axes")) or (
             isinstance(node.op, SymbolicRandomVariable) and node.op.extended_signature is None
         ):
             # MarginalRV and SymbolicRandomVariables without signature are a wild-card,
@@ -160,7 +159,7 @@ def _subgraph_batch_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars)
             )
 
             support_axes = iter(get_support_axes(op))
-            if isinstance(op, MarginalRV):
+            if hasattr(op, "support_axes"):
                 # The first output is the marginalized variable for which we don't compute support axes
                 support_axes = itertools.chain(((),), support_axes)
             for i, (out, inner_out) in enumerate(zip(node.outputs, inner_outputs)):