Skip to content

Commit 4271195

Browse files
williambdeanericmjl
authored andcommitted
implement _display_ method for model
1 parent f4bdc6c commit 4271195

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

pymc/model/core.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@
5252
from pymc.logprob.basic import transformed_conditional_logp
5353
from pymc.logprob.transforms import Transform
5454
from pymc.logprob.utils import ParameterValueError, replace_rvs_by_values
55-
from pymc.model_graph import model_to_graphviz
55+
from pymc.model_graph import model_to_graphviz, model_to_mermaid
5656
from pymc.pytensorf import (
5757
PointFunc,
5858
SeedSequenceSeed,
@@ -437,6 +437,11 @@ def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None:
437437
"""Exit the context manager."""
438438
_ = MODEL_MANAGER.active_contexts.pop()
439439

440+
def _display_(self):
441+
import marimo as mo
442+
443+
return mo.mermaid(model_to_mermaid(self))
444+
440445
@staticmethod
441446
def _validate_name(name):
442447
if name.endswith(":"):

0 commit comments

Comments
 (0)