Deprecate str_for_symbolic_dist in favor of str_for_dist

* Display the dists and their inputs used by SymbolicDistributions

Author: ricardoV94

pymc/distributions/censored.py

@@ -30,6 +30,7 @@ class CensoredRV(SymbolicRandomVariable):
     """Censored random variable"""
 
     inline_aeppl = True
+    _print_name = ("Censored", "\\operatorname{Censored}")
 
 
 class Censored(SymbolicDistribution):

pymc/distributions/distribution.py

@@ -52,7 +52,7 @@
     find_size,
     shape_from_dims,
 )
-from pymc.printing import str_for_dist, str_for_symbolic_dist
+from pymc.printing import str_for_dist
 from pymc.util import UNSET
 from pymc.vartypes import string_types
 
@@ -531,9 +531,9 @@ def __new__(
             initval=initval,
         )
         # add in pretty-printing support
-        rv_out.str_repr = types.MethodType(str_for_symbolic_dist, rv_out)
+        rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
         rv_out._repr_latex_ = types.MethodType(
-            functools.partial(str_for_symbolic_dist, formatting="latex"), rv_out
+            functools.partial(str_for_dist, formatting="latex"), rv_out
         )
 
         rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")

pymc/distributions/mixture.py

@@ -45,6 +45,7 @@ class MarginalMixtureRV(SymbolicRandomVariable):
     """A placeholder used to specify a log-likelihood for a mixture sub-graph."""
 
     default_output = 1
+    _print_name = ("MarginalMixture", "\\operatorname{MarginalMixture}")
 
     def update(self, node: Node):
         # Update for the internal mix_indexes RV

pymc/distributions/timeseries.py

@@ -335,6 +335,7 @@ class AutoRegressiveRV(SymbolicRandomVariable):
     default_output = 1
     ar_order: int
     constant_term: bool
+    _print_name = ("AR", "\\operatorname{AR}")
 
     def __init__(self, *args, ar_order, constant_term, **kwargs):
         self.ar_order = ar_order

pymc/printing.py

@@ -20,6 +20,10 @@
 from aesara.tensor.basic import TensorVariable, Variable
 from aesara.tensor.elemwise import DimShuffle
 from aesara.tensor.random.basic import RandomVariable
+from aesara.tensor.random.var import (
+    RandomGeneratorSharedVariable,
+    RandomStateSharedVariable,
+)
 from aesara.tensor.var import TensorConstant
 
 from pymc.model import Model
@@ -31,40 +35,62 @@
 ]
 
 
-def str_for_dist(rv: TensorVariable, formatting: str = "plain", include_params: bool = True) -> str:
-    """Make a human-readable string representation of a RandomVariable in a model, either
+def str_for_dist(
+    dist: TensorVariable, formatting: str = "plain", include_params: bool = True
+) -> str:
+    """Make a human-readable string representation of a Distribution in a model, either
     LaTeX or plain, optionally with distribution parameter values included."""
 
     if include_params:
         # first 3 args are always (rng, size, dtype), rest is relevant for distribution
-        dist_args = [_str_for_input_var(x, formatting=formatting) for x in rv.owner.inputs[3:]]
+        if isinstance(dist.owner.op, RandomVariable):
+            dist_args = [
+                _str_for_input_var(x, formatting=formatting) for x in dist.owner.inputs[3:]
+            ]
+        else:
+            dist_args = [
+                _str_for_input_var(x, formatting=formatting)
+                for x in dist.owner.inputs
+                if not isinstance(x, (RandomStateSharedVariable, RandomGeneratorSharedVariable))
+            ]
+
+    print_name = dist.name
 
-    print_name = rv.name if rv.name is not None else "<unnamed>"
     if "latex" in formatting:
-        print_name = r"\text{" + _latex_escape(print_name) + "}"
-        dist_name = rv.owner.op._print_name[1]
+        if print_name is not None:
+            print_name = r"\text{" + _latex_escape(dist.name) + "}"
+
+        op_name = (
+            dist.owner.op._print_name[1]
+            if hasattr(dist.owner.op, "_print_name")
+            else r"\\operatorname{Unknown}"
+        )
         if include_params:
-            return r"${} \sim {}({})$".format(print_name, dist_name, ",~".join(dist_args))
+            if print_name:
+                return r"${} \sim {}({})$".format(print_name, op_name, ",~".join(dist_args))
+            else:
+                return r"${}({})$".format(op_name, ",~".join(dist_args))
+
         else:
-            return rf"${print_name} \sim {dist_name}$"
+            if print_name:
+                return rf"${print_name} \sim {op_name}$"
+            else:
+                return rf"${op_name}$"
+
     else:  # plain
-        dist_name = rv.owner.op._print_name[0]
+        dist_name = (
+            dist.owner.op._print_name[0] if hasattr(dist.owner.op, "_print_name") else "Unknown"
+        )
         if include_params:
-            return r"{} ~ {}({})".format(print_name, dist_name, ", ".join(dist_args))
+            if print_name:
+                return r"{} ~ {}({})".format(print_name, dist_name, ", ".join(dist_args))
+            else:
+                return r"{}({})".format(dist_name, ", ".join(dist_args))
         else:
-            return rf"{print_name} ~ {dist_name}"
-
-
-def str_for_symbolic_dist(
-    rv: TensorVariable, formatting: str = "plain", include_params: bool = True
-) -> str:
-    """Make a human-readable string representation of a SymbolicDistribution in a model,
-    either LaTeX or plain, optionally with distribution parameter values included."""
-
-    if "latex" in formatting:
-        return rf"$\text{{{rv.name}}} \sim \text{{{rv.owner.op}}}$"
-    else:
-        return rf"{rv.name} ~ {rv.owner.op}"
+            if print_name:
+                return rf"{print_name} ~ {dist_name}"
+            else:
+                return dist_name
 
 
 def str_for_model(model: Model, formatting: str = "plain", include_params: bool = True) -> str:
@@ -145,7 +171,11 @@ def _is_potential_or_determinstic(var: Variable) -> bool:
 
 
 def _str_for_input_rv(var: Variable, formatting: str) -> str:
-    _str = var.name if var.name is not None else "<unnamed>"
+    _str = (
+        var.name
+        if var.name is not None
+        else str_for_dist(var, formatting=formatting, include_params=True)
+    )
     if "latex" in formatting:
         return r"\text{" + _latex_escape(_str) + "}"
     else:

pymc/tests/test_printing.py

@@ -92,7 +92,7 @@ def setup_class(self):
                 r"beta ~ N(0, 10)",
                 r"Z ~ N(f(), f())",
                 r"nb_with_p_n ~ NB(10, nbp)",
-                r"zip ~ MarginalMixtureRV{inline=False}",
+                r"zip ~ MarginalMixture(f(), DiracDelta(0), Pois(5))",
                 r"Y_obs ~ N(mu, sigma)",
                 r"pot ~ Potential(f(beta, alpha))",
             ],
@@ -103,7 +103,7 @@ def setup_class(self):
                 r"beta ~ N",
                 r"Z ~ N",
                 r"nb_with_p_n ~ NB",
-                r"zip ~ MarginalMixtureRV{inline=False}",
+                r"zip ~ MarginalMixture",
                 r"Y_obs ~ N",
                 r"pot ~ Potential",
             ],
@@ -114,7 +114,7 @@ def setup_class(self):
                 r"$\text{beta} \sim \operatorname{N}(0,~10)$",
                 r"$\text{Z} \sim \operatorname{N}(f(),~f())$",
                 r"$\text{nb_with_p_n} \sim \operatorname{NB}(10,~\text{nbp})$",
-                r"$\text{zip} \sim \text{MarginalMixtureRV{inline=False}}$",
+                r"$\text{zip} \sim \operatorname{MarginalMixture}(f(),~\text{\$\operatorname{DiracDelta}(0)\$},~\text{\$\operatorname{Pois}(5)\$})$",
                 r"$\text{Y_obs} \sim \operatorname{N}(\text{mu},~\text{sigma})$",
                 r"$\text{pot} \sim \operatorname{Potential}(f(\text{beta},~\text{alpha}))$",
             ],
@@ -125,7 +125,7 @@ def setup_class(self):
                 r"$\text{beta} \sim \operatorname{N}$",
                 r"$\text{Z} \sim \operatorname{N}$",
                 r"$\text{nb_with_p_n} \sim \operatorname{NB}$",
-                r"$\text{zip} \sim \text{MarginalMixtureRV{inline=False}}$",
+                r"$\text{zip} \sim \operatorname{MarginalMixture}$",
                 r"$\text{Y_obs} \sim \operatorname{N}$",
                 r"$\text{pot} \sim \operatorname{Potential}$",
             ],