Add a way to pass in graph level attributes to graphviz

Author: zaxtax

pymc/model_graph.py

@@ -430,6 +430,9 @@ def edges(
         ]
 
 
+GraphAttrMapping = dict[Any, Any]
+
+
 def make_graph(
     name: str,
     plates: list[Plate],
@@ -439,6 +442,7 @@ def make_graph(
     figsize=None,
     dpi=300,
     node_formatters: NodeTypeFormatterMapping | None = None,
+    graph_attrs: GraphAttrMapping | None = None,
     create_plate_label: PlateLabelFunc = create_plate_label_with_dim_length,
 ):
     """Make graphviz Digraph of PyMC model.
@@ -460,6 +464,9 @@ def make_graph(
     node_formatters = update_node_formatters(node_formatters)
 
     graph = graphviz.Digraph(name)
+    if graph_attrs is not None:
+        graph.attr(**graph_attrs)
+
     for plate in plates:
         if plate.dim_info:
             # must be preceded by 'cluster' to get a box around it
@@ -676,6 +683,7 @@ def model_to_graphviz(
     figsize: tuple[int, int] | None = None,
     dpi: int = 300,
     node_formatters: NodeTypeFormatterMapping | None = None,
+    graph_attrs: GraphAttrMapping | None = None,
     include_dim_lengths: bool = True,
 ):
     """Produce a graphviz Digraph from a PyMC model.
@@ -704,6 +712,10 @@ def model_to_graphviz(
         the size of the saved figure.
     dpi : int, optional
         Dots per inch. It only affects the resolution of the saved figure. The default is 300.
+    graph_attrs : dict, optional
+        A dictionary of top-level layout attributes for graphviz
+        Check out graphviz documentation for more information on available attributes
+        https://graphviz.org/doc/info/attrs.html
     node_formatters : dict, optional
         A dictionary mapping node types to functions that return a dictionary of node attributes.
         Check out graphviz documentation for more information on available
@@ -773,6 +785,7 @@ def model_to_graphviz(
         save=save,
         figsize=figsize,
         dpi=dpi,
+        graph_attrs=graph_attrs,
         node_formatters=node_formatters,
         create_plate_label=create_plate_label_with_dim_length
         if include_dim_lengths