Break circular dependency between model_graph.py and mode/core.py with specific lazy imports

Author: ricardoV94

pymc/model/core.py

@@ -52,7 +52,6 @@
 from pymc.logprob.basic import transformed_conditional_logp
 from pymc.logprob.transforms import Transform
 from pymc.logprob.utils import ParameterValueError, replace_rvs_by_values
-from pymc.model_graph import model_to_graphviz
 from pymc.pytensorf import (
     PointFunc,
     SeedSequenceSeed,
@@ -437,6 +436,13 @@ def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None:
         """Exit the context manager."""
         _ = MODEL_MANAGER.active_contexts.pop()
 
+    def _display_(self):
+        import marimo as mo
+
+        from pymc.model_graph import model_to_mermaid
+
+        return mo.mermaid(model_to_mermaid(self))
+
     @staticmethod
     def _validate_name(name):
         if name.endswith(":"):
@@ -1997,6 +2003,8 @@ def to_graphviz(
             # creates the file `schools.pdf`
             schools.to_graphviz().render("schools")
         """
+        from pymc.model_graph import model_to_graphviz
+
         return model_to_graphviz(
             model=self,
             var_names=var_names,

pymc/model_graph.py

@@ -29,8 +29,7 @@
 from pytensor.tensor.shape import Shape
 from pytensor.tensor.variable import TensorVariable
 
-import pymc as pm
-
+from pymc.model.core import modelcontext
 from pymc.util import VarName, get_default_varnames, get_var_name
 
 __all__ = (
@@ -662,7 +661,7 @@ def model_to_networkx(
             stacklevel=2,
         )
 
-    model = pm.modelcontext(model)
+    model = modelcontext(model)
     graph = ModelGraph(model)
     return make_networkx(
         name=model.name,
@@ -777,7 +776,7 @@ def model_to_graphviz(
             stacklevel=2,
         )
 
-    model = pm.modelcontext(model)
+    model = modelcontext(model)
     graph = ModelGraph(model)
     return make_graph(
         model.name,
@@ -904,7 +903,7 @@ def model_to_mermaid(model=None, *, var_names=None, include_dim_lengths: bool =
 
 
     """
-    model = pm.modelcontext(model)
+    model = modelcontext(model)
     graph = ModelGraph(model)
     plates = sorted(graph.get_plates(var_names=var_names), key=lambda plate: hash(plate.dim_info))
     edges = sorted(graph.edges(var_names=var_names))