Refactor pretty-printing and fix latex underscores

pymc/distributions/distribution.py

@@ -57,7 +57,6 @@
 )
 from pymc.logprob.rewriting import logprob_rewrites_db
 from pymc.model import BlockModelAccess
-from pymc.printing import str_for_dist
 from pymc.pytensorf import collect_default_updates, convert_observed_data
 from pymc.util import UNSET, _add_future_warning_tag
 from pymc.vartypes import string_types
@@ -198,7 +197,7 @@ class SymbolicRandomVariable(OpFromGraph):
     """
 
     _print_name: Tuple[str, str] = ("Unknown", "\\operatorname{Unknown}")
-    """Tuple of (name, latex name) used for for pretty-printing variables of this type"""
+    "Tuple of (name, latex name) used for for pretty-printing variables of this type"
 
     def __init__(self, *args, ndim_supp, **kwargs):
         self.ndim_supp = ndim_supp
@@ -320,9 +319,12 @@ def __new__(
         )
 
         # add in pretty-printing support
-        rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
+        from pymc.printing import str_for_model_var
+
+        rv_out.str_repr = types.MethodType(str_for_model_var, rv_out)
         rv_out._repr_latex_ = types.MethodType(
-            functools.partial(str_for_dist, formatting="latex"), rv_out
+            functools.partial(str_for_model_var, formatting="latex"),
+            rv_out,
         )
 
         rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")

pymc/distributions/multivariate.py

@@ -1116,7 +1116,7 @@ class _LKJCholeskyCovBaseRV(RandomVariable):
     ndim_supp = 1
     ndims_params = [0, 0, 1]
     dtype = "floatX"
-    _print_name = ("_lkjcholeskycovbase", "\\operatorname{_lkjcholeskycovbase}")
+    _print_name = ("_lkjcholeskycovbase", r"\operatorname{\_lkjcholeskycovbase}")
 
     def make_node(self, rng, size, dtype, n, eta, D):
         n = at.as_tensor_variable(n)
@@ -1164,7 +1164,7 @@ def rng_fn(self, rng, n, eta, D, size):
 # be safely resized. Because of this, we add the thin SymbolicRandomVariable wrapper
 class _LKJCholeskyCovRV(SymbolicRandomVariable):
     default_output = 1
-    _print_name = ("_lkjcholeskycov", "\\operatorname{_lkjcholeskycov}")
+    _print_name = ("_lkjcholeskycov", r"\operatorname{\_lkjcholeskycov}")
 
     def update(self, node):
         return {node.inputs[0]: node.outputs[0]}

pymc/model.py

@@ -2015,15 +2015,13 @@ def Deterministic(name, var, model=None, dims=None):
     model.deterministics.append(var)
     model.add_named_variable(var, dims)
 
-    from pymc.printing import str_for_potential_or_deterministic
+    from pymc.printing import str_for_model_var
 
     var.str_repr = types.MethodType(
-        functools.partial(str_for_potential_or_deterministic, dist_name="Deterministic"), var
+        functools.partial(str_for_model_var, dist_name="Deterministic"), var
     )
     var._repr_latex_ = types.MethodType(
-        functools.partial(
-            str_for_potential_or_deterministic, dist_name="Deterministic", formatting="latex"
-        ),
+        functools.partial(str_for_model_var, dist_name="Deterministic", formatting="latex"),
         var,
     )
 
@@ -2047,15 +2045,13 @@ def Potential(name, var, model=None):
     model.potentials.append(var)
     model.add_named_variable(var)
 
-    from pymc.printing import str_for_potential_or_deterministic
+    from pymc.printing import str_for_model_var
 
     var.str_repr = types.MethodType(
-        functools.partial(str_for_potential_or_deterministic, dist_name="Potential"), var
+        functools.partial(str_for_model_var, dist_name="Potential"), var
     )
     var._repr_latex_ = types.MethodType(
-        functools.partial(
-            str_for_potential_or_deterministic, dist_name="Potential", formatting="latex"
-        ),
+        functools.partial(str_for_model_var, dist_name="Potential", formatting="latex"),
         var,
     )