Fix and add tests for Theano printing

Author: brandonwillard

requirements.txt

@@ -1,4 +1,5 @@
 -e ./
+sympy>=1.3
 coveralls
 pydocstyle>=3.0.0
 pytest>=5.0.0

setup.py

@@ -45,7 +45,6 @@ def get_long_description():
         "etuples>=0.3.1",
         "cons>=0.4.0",
         "toolz>=0.9.0",
-        "sympy>=1.3",
         "cachetools",
         "pymc3>=3.6",
         "pymc4 @ git+https://github.com/pymc-devs/pymc4.git@master#egg=pymc4-0.0.1",

symbolic_pymc/theano/printing.py

@@ -10,8 +10,19 @@
 
 from theano import gof
 
-from sympy import Array as SympyArray
-from sympy.printing import latex as sympy_latex
+try:
+    from sympy import Array as SympyArray
+    from sympy.printing import latex as sympy_latex
+
+    def latex_print_array(data):  # pragma: no cover
+        return sympy_latex(SympyArray(data))
+
+
+except ImportError:  # pragma: no cover
+
+    def latex_print_array(data):
+        return data
+
 
 from .opt import FunctionGraph
 from .ops import RandomVariable
@@ -60,16 +71,16 @@ def process_param(self, idx, sform, pstate):
             The printer state.
 
         """
-        return sform
+        return sform  # pragma: no cover
 
     def process(self, output, pstate):
         if output in pstate.memo:
             return pstate.memo[output]
 
         pprinter = pstate.pprinter
-        node = output.owner
+        node = getattr(output, "owner", None)
 
-        if node is None or not isinstance(node.op, RandomVariable):
+        if node is None or not isinstance(node.op, RandomVariable):  # pragma: no cover
             raise TypeError(
                 "Function %s cannot represent a variable that is "
                 "not the result of a RandomVariable operation" % self.name
@@ -78,7 +89,7 @@ def process(self, output, pstate):
         op_name = self.name or getattr(node.op, "print_name", None)
         op_name = op_name or getattr(node.op, "name", None)
 
-        if op_name is None:
+        if op_name is None:  # pragma: no cover
             raise ValueError(f"Could not find a name for {node.op}")
 
         # Allow `Op`s to specify their ascii and LaTeX formats (in a tuple/list
@@ -144,7 +155,7 @@ def process(self, output, pstate):
 
 class GenericSubtensorPrinter(object):
     def process(self, r, pstate):
-        if r.owner is None:
+        if getattr(r, "owner", None) is None:  # pragma: no cover
             raise TypeError("Can only print Subtensor.")
 
         output_latex = getattr(pstate, "latex", False)
@@ -161,13 +172,13 @@ def process(self, r, pstate):
                 if isinstance(entry, slice):
                     s_parts = [""] * 2
                     if entry.start is not None:
-                        s_parts[0] = entry.start
+                        s_parts[0] = pstate.pprinter.process(inputs.pop())
 
                     if entry.stop is not None:
-                        s_parts[1] = entry.stop
+                        s_parts[1] = pstate.pprinter.process(inputs.pop())
 
                     if entry.step is not None:
-                        s_parts.append(entry.stop)
+                        s_parts.append(pstate.pprinter.process(inputs.pop()))
 
                     sidxs.append(":".join(s_parts))
                 else:
@@ -215,16 +226,22 @@ def process(cls, output, pstate):
         using_latex = getattr(pstate, "latex", False)
         # Crude--but effective--means of stopping print-outs for large
         # arrays.
-        constant = isinstance(output, tt.TensorConstant)
+        constant = isinstance(output, (tt.TensorConstant, theano.scalar.basic.ScalarConstant))
         too_large = constant and (output.data.size > cls.max_line_width * cls.max_line_height)
 
         if constant and not too_large:
             # Print constants that aren't too large
             if using_latex and output.ndim > 0:
-                out_name = sympy_latex(SympyArray(output.data))
+                out_name = latex_print_array(output.data)
             else:
                 out_name = str(output.data)
-        elif isinstance(output, tt.TensorVariable) or constant:
+        elif (
+            isinstance(
+                output,
+                (tt.TensorVariable, theano.scalar.basic.Scalar, theano.scalar.basic.ScalarVariable),
+            )
+            or constant
+        ):
             # Process name and shape
 
             # Attempt to get the original variable, in case this is a cloned
@@ -238,7 +255,7 @@ def process(cls, output, pstate):
 
             shape_strings = pstate.preamble_dict.setdefault("shape_strings", OrderedDict())
             shape_strings[output] = shape_info
-        else:
+        else:  # pragma: no cover
             raise TypeError(f"Type {type(output)} not handled by variable printer")
 
         pstate.memo[output] = out_name
@@ -268,7 +285,7 @@ def process_variable_name(cls, output, pstate):
                 _ = [available_names.pop(v.name, None) for v in fgraph.variables]
             setattr(pstate, "available_names", available_names)
 
-        if output.name:
+        if getattr(output, "name", None):
             # Observed an existing name; remove it.
             out_name = output.name
             available_names.pop(out_name, None)
@@ -524,11 +541,18 @@ def __call__(self, *args, latex_env="equation", latex_label=None):
 
 # The order here is important!
 tt_pprint.printers.insert(
-    0, (lambda pstate, r: isinstance(r, tt.Variable), VariableWithShapePrinter)
+    0,
+    (
+        lambda pstate, r: isinstance(r, (theano.scalar.basic.Scalar, tt.Variable)),
+        VariableWithShapePrinter,
+    ),
 )
 tt_pprint.printers.insert(
     0,
-    (lambda pstate, r: r.owner and isinstance(r.owner.op, RandomVariable), RandomVariablePrinter()),
+    (
+        lambda pstate, r: getattr(r, "owner", None) and isinstance(r.owner.op, RandomVariable),
+        RandomVariablePrinter(),
+    ),
 )
 
 
@@ -538,9 +562,9 @@ def process(self, output, pstate):
             return pstate.memo[output]
 
         pprinter = pstate.pprinter
-        node = output.owner
+        node = getattr(output, "owner", None)
 
-        if node is None or not isinstance(node.op, Observed):
+        if node is None or not isinstance(node.op, Observed):  # pragma: no cover
             raise TypeError(f"Node Op is not of type `Observed`: {node.op}")
 
         val = node.inputs[0]

tests/theano/test_printing.py

@@ -1,27 +1,211 @@
+import textwrap
+
 import theano.tensor as tt
 
-from symbolic_pymc.theano.random_variables import NormalRV
-from symbolic_pymc.theano.printing import tt_pprint
+from symbolic_pymc.theano.random_variables import NormalRV, observed
+from symbolic_pymc.theano.printing import tt_pprint, tt_tprint
 
 
 def test_notex_print():
 
     tt_normalrv_noname_expr = tt.scalar("b") * NormalRV(tt.scalar("\\mu"), tt.scalar("\\sigma"))
-    expected = "b in R, \\mu in R, \\sigma in R\na ~ N(\\mu, \\sigma**2) in R\n(b * a)"
-    assert tt_pprint(tt_normalrv_noname_expr) == expected
+    expected = textwrap.dedent(
+        r"""
+    b in R, \mu in R, \sigma in R
+    a ~ N(\mu, \sigma**2) in R
+    (b * a)
+    """
+    )
+    assert tt_pprint(tt_normalrv_noname_expr) == expected.strip()
 
     # Make sure the constant shape is show in values and not symbols.
     tt_normalrv_name_expr = tt.scalar("b") * NormalRV(
         tt.scalar("\\mu"), tt.scalar("\\sigma"), size=[2, 1], name="X"
     )
-    expected = "b in R, \\mu in R, \\sigma in R\nX ~ N(\\mu, \\sigma**2) in R**(2 x 1)\n(b * X)"
-    assert tt_pprint(tt_normalrv_name_expr) == expected
+    expected = textwrap.dedent(
+        r"""
+    b in R, \mu in R, \sigma in R
+    X ~ N(\mu, \sigma**2) in R**(2 x 1)
+    (b * X)
+    """
+    )
+    assert tt_pprint(tt_normalrv_name_expr) == expected.strip()
+
+    tt_2_normalrv_noname_expr = tt.matrix("M") * NormalRV(
+        tt.scalar("\\mu_2"), tt.scalar("\\sigma_2")
+    )
+    tt_2_normalrv_noname_expr *= tt.scalar("b") * NormalRV(
+        tt_2_normalrv_noname_expr, tt.scalar("\\sigma")
+    ) + tt.scalar("c")
+    expected = textwrap.dedent(
+        r"""
+    M in R**(N^M_0 x N^M_1), \mu_2 in R, \sigma_2 in R
+    b in R, \sigma in R, c in R
+    a ~ N(\mu_2, \sigma_2**2) in R, d ~ N((M * a), \sigma**2) in R**(N^d_0 x N^d_1)
+    ((M * a) * ((b * d) + c))
+    """
+    )
+    assert tt_pprint(tt_2_normalrv_noname_expr) == expected.strip()
+
+    expected = textwrap.dedent(
+        r"""
+    b in Z, c in Z, M in R**(N^M_0 x N^M_1)
+    M[b, c]
+    """
+    )
+    # TODO: "c" should be "1".
+    assert (
+        tt_pprint(tt.matrix("M")[tt.iscalar("a"), tt.constant(1, dtype="int")]) == expected.strip()
+    )
+
+    expected = textwrap.dedent(
+        r"""
+    M in R**(N^M_0 x N^M_1)
+    M[1]
+    """
+    )
+    assert tt_pprint(tt.matrix("M")[1]) == expected.strip()
+
+    expected = textwrap.dedent(
+        r"""
+    M in N**(N^M_0)
+    M[2:4:0]
+    """
+    )
+    assert tt_pprint(tt.vector("M", dtype="uint32")[0:4:2]) == expected.strip()
+
+    norm_rv = NormalRV(tt.scalar("\\mu"), tt.scalar("\\sigma"))
+    rv_obs = observed(tt.constant(1.0, dtype=norm_rv.dtype), norm_rv)
+
+    expected = textwrap.dedent(
+        r"""
+    \mu in R, \sigma in R
+    a ~ N(\mu, \sigma**2) in R
+    a = 1.0
+        """
+    )
+    assert tt_pprint(rv_obs) == expected.strip()
+
+
+def test_tex_print():
+
+    tt_normalrv_noname_expr = tt.scalar("b") * NormalRV(tt.scalar("\\mu"), tt.scalar("\\sigma"))
+    expected = textwrap.dedent(
+        r"""
+    \begin{equation}
+      \begin{gathered}
+      b \in \mathbb{R}, \,\mu \in \mathbb{R}, \,\sigma \in \mathbb{R}
+      \\
+      a \sim \operatorname{N}\left(\mu, {\sigma}^{2}\right)\,  \in \mathbb{R}
+      \end{gathered}
+      \\
+      (b \odot a)
+    \end{equation}
+    """
+    )
+    assert tt_tprint(tt_normalrv_noname_expr) == expected.strip()
+
+    tt_normalrv_name_expr = tt.scalar("b") * NormalRV(
+        tt.scalar("\\mu"), tt.scalar("\\sigma"), size=[2, 1], name="X"
+    )
+    expected = textwrap.dedent(
+        r"""
+    \begin{equation}
+      \begin{gathered}
+      b \in \mathbb{R}, \,\mu \in \mathbb{R}, \,\sigma \in \mathbb{R}
+      \\
+      X \sim \operatorname{N}\left(\mu, {\sigma}^{2}\right)\,  \in \mathbb{R}^{2 \times 1}
+      \end{gathered}
+      \\
+      (b \odot X)
+    \end{equation}
+    """
+    )
+    assert tt_tprint(tt_normalrv_name_expr) == expected.strip()
 
     tt_2_normalrv_noname_expr = tt.matrix("M") * NormalRV(
         tt.scalar("\\mu_2"), tt.scalar("\\sigma_2")
     )
     tt_2_normalrv_noname_expr *= tt.scalar("b") * NormalRV(
         tt_2_normalrv_noname_expr, tt.scalar("\\sigma")
     ) + tt.scalar("c")
-    expected = "M in R**(N^M_0 x N^M_1), \\mu_2 in R, \\sigma_2 in R\nb in R, \\sigma in R, c in R\na ~ N(\\mu_2, \\sigma_2**2) in R, d ~ N((M * a), \\sigma**2) in R**(N^d_0 x N^d_1)\n((M * a) * ((b * d) + c))"
-    assert tt_pprint(tt_2_normalrv_noname_expr) == expected
+    expected = textwrap.dedent(
+        r"""
+    \begin{equation}
+      \begin{gathered}
+      M \in \mathbb{R}^{N^{M}_{0} \times N^{M}_{1}}
+      \\
+      \mu_2 \in \mathbb{R}, \,\sigma_2 \in \mathbb{R}
+      \\
+      b \in \mathbb{R}, \,\sigma \in \mathbb{R}, \,c \in \mathbb{R}
+      \\
+      a \sim \operatorname{N}\left(\mu_2, {\sigma_2}^{2}\right)\,  \in \mathbb{R}
+      \\
+      d \sim \operatorname{N}\left((M \odot a), {\sigma}^{2}\right)\,  \in \mathbb{R}^{N^{d}_{0} \times N^{d}_{1}}
+      \end{gathered}
+      \\
+      ((M \odot a) \odot ((b \odot d) + c))
+    \end{equation}
+    """
+    )
+    assert tt_tprint(tt_2_normalrv_noname_expr) == expected.strip()
+
+    expected = textwrap.dedent(
+        r"""
+    \begin{equation}
+      \begin{gathered}
+      b \in \mathbb{Z}, \,c \in \mathbb{Z}, \,M \in \mathbb{R}^{N^{M}_{0} \times N^{M}_{1}}
+      \end{gathered}
+      \\
+      M\left[b, \,c\right]
+    \end{equation}
+    """
+    )
+    # TODO: "c" should be "1".
+    assert (
+        tt_tprint(tt.matrix("M")[tt.iscalar("a"), tt.constant(1, dtype="int")]) == expected.strip()
+    )
+
+    expected = textwrap.dedent(
+        r"""
+    \begin{equation}
+      \begin{gathered}
+      M \in \mathbb{R}^{N^{M}_{0} \times N^{M}_{1}}
+      \end{gathered}
+      \\
+      M\left[1\right]
+    \end{equation}
+    """
+    )
+    assert tt_tprint(tt.matrix("M")[1]) == expected.strip()
+
+    expected = textwrap.dedent(
+        r"""
+    \begin{equation}
+      \begin{gathered}
+      M \in \mathbb{N}^{N^{M}_{0}}
+      \end{gathered}
+      \\
+      M\left[2:4:0\right]
+    \end{equation}
+    """
+    )
+    assert tt_tprint(tt.vector("M", dtype="uint32")[0:4:2]) == expected.strip()
+
+    norm_rv = NormalRV(tt.scalar("\\mu"), tt.scalar("\\sigma"))
+    rv_obs = observed(tt.constant(1.0, dtype=norm_rv.dtype), norm_rv)
+
+    expected = textwrap.dedent(
+        r"""
+    \begin{equation}
+      \begin{gathered}
+      \mu \in \mathbb{R}, \,\sigma \in \mathbb{R}
+      \\
+      a \sim \operatorname{N}\left(\mu, {\sigma}^{2}\right)\,  \in \mathbb{R}
+      \end{gathered}
+      \\
+      a = 1.0
+    \end{equation}
+        """
+    )
+    assert tt_tprint(rv_obs) == expected.strip()