Refactor pretty-printing and fix latex underscores, addresses 6508

Author: covertg

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
@@ -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/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,
     )