Simplify string representation of Elemwise and CAReduce

Author: ricardoV94

pytensor/graph/utils.py

@@ -234,6 +234,7 @@ def __eq__(self, other):
 
                 dct["__eq__"] = __eq__
 
+            # FIXME: This overrides __str__ inheritance when props are provided
             if "__str__" not in dct:
                 if len(props) == 0:
 

pytensor/tensor/elemwise.py

@@ -14,7 +14,7 @@
 from pytensor.link.c.params_type import ParamsType
 from pytensor.misc.frozendict import frozendict
 from pytensor.misc.safe_asarray import _asarray
-from pytensor.printing import FunctionPrinter, Printer, pprint
+from pytensor.printing import Printer, pprint
 from pytensor.scalar import get_scalar_type
 from pytensor.scalar.basic import bool as scalar_bool
 from pytensor.scalar.basic import identity as scalar_identity
@@ -498,15 +498,9 @@ def make_node(self, *inputs):
         return Apply(self, inputs, outputs)
 
     def __str__(self):
-        if self.name is None:
-            if self.inplace_pattern:
-                items = list(self.inplace_pattern.items())
-                items.sort()
-                return f"{type(self).__name__}{{{self.scalar_op}}}{items}"
-            else:
-                return f"{type(self).__name__}{{{self.scalar_op}}}"
-        else:
+        if self.name:
             return self.name
+        return str(self.scalar_op).capitalize()
 
     def R_op(self, inputs, eval_points):
         outs = self(*inputs, return_list=True)
@@ -1477,23 +1471,17 @@ def clone(
 
         return res
 
-    def __str__(self):
-        prefix = f"{type(self).__name__}{{{self.scalar_op}}}"
-        extra_params = []
-
-        if self.axis is not None:
-            axis = ", ".join(str(x) for x in self.axis)
-            extra_params.append(f"axis=[{axis}]")
-
-        if self.acc_dtype:
-            extra_params.append(f"acc_dtype={self.acc_dtype}")
-
-        extra_params_str = ", ".join(extra_params)
-
-        if extra_params_str:
-            return f"{prefix}{{{extra_params_str}}}"
+    def _axis_str(self):
+        axis = self.axis
+        if axis is None:
+            return "axes=None"
+        elif len(axis) == 1:
+            return f"axis={axis[0]}"
         else:
-            return f"{prefix}"
+            return f"axes={list(axis)}"
+
+    def __str__(self):
+        return f"{type(self).__name__}{{{self.scalar_op}, {self._axis_str()}}}"
 
     def perform(self, node, inp, out):
         (input,) = inp
@@ -1737,21 +1725,17 @@ def construct(symbol):
         symbolname = symbolname or symbol.__name__
 
         if symbolname.endswith("_inplace"):
-            elemwise_name = f"Elemwise{{{symbolname},inplace}}"
-            scalar_op = getattr(scalar, symbolname[: -len("_inplace")])
+            base_symbol_name = symbolname[: -len("_inplace")]
+            scalar_op = getattr(scalar, base_symbol_name)
             inplace_scalar_op = scalar_op.__class__(transfer_type(0))
             rval = Elemwise(
                 inplace_scalar_op,
                 {0: 0},
-                name=elemwise_name,
                 nfunc_spec=(nfunc and (nfunc, nin, nout)),
             )
         else:
-            elemwise_name = f"Elemwise{{{symbolname},no_inplace}}"
             scalar_op = getattr(scalar, symbolname)
-            rval = Elemwise(
-                scalar_op, name=elemwise_name, nfunc_spec=(nfunc and (nfunc, nin, nout))
-            )
+            rval = Elemwise(scalar_op, nfunc_spec=(nfunc and (nfunc, nin, nout)))
 
         if getattr(symbol, "__doc__"):
             rval.__doc__ = symbol.__doc__ + "\n\n    " + rval.__doc__
@@ -1761,8 +1745,6 @@ def construct(symbol):
         rval.__epydoc_asRoutine = symbol
         rval.__module__ = symbol.__module__
 
-        pprint.assign(rval, FunctionPrinter([symbolname.replace("_inplace", "=")]))
-
         return rval
 
     if symbol:

pytensor/tensor/math.py

@@ -583,7 +583,12 @@ def max_and_argmax(a, axis=None, keepdims=False):
     return [out, argout]
 
 
-class NonZeroCAReduce(CAReduce):
+class FixedOpCAReduce(CAReduce):
+    def __str__(self):
+        return f"{type(self).__name__}{{{self._axis_str()}}}"
+
+
+class NonZeroDimsCAReduce(FixedOpCAReduce):
     def _c_all(self, node, name, inames, onames, sub):
         decl, checks, alloc, loop, end = super()._c_all(node, name, inames, onames, sub)
 
@@ -614,7 +619,7 @@ def _c_all(self, node, name, inames, onames, sub):
         return decl, checks, alloc, loop, end
 
 
-class Max(NonZeroCAReduce):
+class Max(NonZeroDimsCAReduce):
     nfunc_spec = ("max", 1, 1)
 
     def __init__(self, axis):
@@ -625,7 +630,7 @@ def clone(self, **kwargs):
         return type(self)(axis=axis)
 
 
-class Min(NonZeroCAReduce):
+class Min(NonZeroDimsCAReduce):
     nfunc_spec = ("min", 1, 1)
 
     def __init__(self, axis):
@@ -1496,7 +1501,7 @@ def complex_from_polar(abs, angle):
     """Return complex-valued tensor from polar coordinate specification."""
 
 
-class Mean(CAReduce):
+class Mean(FixedOpCAReduce):
     __props__ = ("axis",)
     nfunc_spec = ("mean", 1, 1)
 
@@ -2356,7 +2361,7 @@ def outer(x, y):
     return dot(x.dimshuffle(0, "x"), y.dimshuffle("x", 0))
 
 
-class All(CAReduce):
+class All(FixedOpCAReduce):
     """Applies `logical and` to all the values of a tensor along the
     specified axis(es).
 
@@ -2370,12 +2375,6 @@ def __init__(self, axis=None):
     def _output_dtype(self, idtype):
         return "bool"
 
-    def __str__(self):
-        if self.axis is None:
-            return "All"
-        else:
-            return "All{%s}" % ", ".join(map(str, self.axis))
-
     def make_node(self, input):
         input = as_tensor_variable(input)
         if input.dtype != "bool":
@@ -2392,7 +2391,7 @@ def clone(self, **kwargs):
         return type(self)(axis=axis)
 
 
-class Any(CAReduce):
+class Any(FixedOpCAReduce):
     """Applies `bitwise or` to all the values of a tensor along the
     specified axis(es).
 
@@ -2406,12 +2405,6 @@ def __init__(self, axis=None):
     def _output_dtype(self, idtype):
         return "bool"
 
-    def __str__(self):
-        if self.axis is None:
-            return "Any"
-        else:
-            return "Any{%s}" % ", ".join(map(str, self.axis))
-
     def make_node(self, input):
         input = as_tensor_variable(input)
         if input.dtype != "bool":
@@ -2428,7 +2421,7 @@ def clone(self, **kwargs):
         return type(self)(axis=axis)
 
 
-class Sum(CAReduce):
+class Sum(FixedOpCAReduce):
     """
     Sums all the values of a tensor along the specified axis(es).
 
@@ -2449,14 +2442,6 @@ def __init__(self, axis=None, dtype=None, acc_dtype=None):
             upcast_discrete_output=True,
         )
 
-    def __str__(self):
-        name = self.__class__.__name__
-        axis = ""
-        if self.axis is not None:
-            axis = ", ".join(str(x) for x in self.axis)
-            axis = f"axis=[{axis}], "
-        return f"{name}{{{axis}acc_dtype={self.acc_dtype}}}"
-
     def L_op(self, inp, out, grads):
         (x,) = inp
 
@@ -2526,7 +2511,7 @@ def sum(input, axis=None, dtype=None, keepdims=False, acc_dtype=None):
 pprint.assign(Sum, printing.FunctionPrinter(["sum"], ["axis"]))
 
 
-class Prod(CAReduce):
+class Prod(FixedOpCAReduce):
     """
     Multiplies all the values of a tensor along the specified axis(es).
 
@@ -2537,7 +2522,6 @@ class Prod(CAReduce):
     """
 
     __props__ = ("scalar_op", "axis", "dtype", "acc_dtype", "no_zeros_in_input")
-
     nfunc_spec = ("prod", 1, 1)
 
     def __init__(self, axis=None, dtype=None, acc_dtype=None, no_zeros_in_input=False):
@@ -2683,6 +2667,14 @@ def clone(self, **kwargs):
             no_zeros_in_input=no_zeros_in_input,
         )
 
+    def __str__(self):
+        if self.no_zeros_in_input:
+            return f"{super().__str__()[:-1]}, no_zeros_in_input}})"
+        return super().__str__()
+
+    def __repr__(self):
+        return f"{super().__repr__()[:-1]}, no_zeros_in_input={self.no_zeros_in_input})"
+
 
 def prod(
     input,
@@ -2751,7 +2743,7 @@ def c_code_cache_version(self):
 mul_without_zeros = MulWithoutZeros(aes.upcast_out, name="mul_without_zeros")
 
 
-class ProdWithoutZeros(CAReduce):
+class ProdWithoutZeros(FixedOpCAReduce):
     def __init__(self, axis=None, dtype=None, acc_dtype=None):
         super().__init__(
             mul_without_zeros,

pytensor/tensor/rewriting/math.py

@@ -42,7 +42,8 @@
     All,
     Any,
     Dot,
-    NonZeroCAReduce,
+    FixedOpCAReduce,
+    NonZeroDimsCAReduce,
     Prod,
     ProdWithoutZeros,
     Sum,
@@ -1671,7 +1672,8 @@ def local_op_of_op(fgraph, node):
         ProdWithoutZeros,
     ]
     + CAReduce.__subclasses__()
-    + NonZeroCAReduce.__subclasses__()
+    + FixedOpCAReduce.__subclasses__()
+    + NonZeroDimsCAReduce.__subclasses__()
 )
 
 

tests/compile/test_builders.py

@@ -579,9 +579,9 @@ def test_debugprint():
 Inner graphs:
 
 OpFromGraph{inline=False} [id A]
- ← Elemwise{add,no_inplace} [id E]
+ ← Add [id E]
     ├─ *0-<TensorType(float64, (?, ?))> [id F]
-    └─ Elemwise{mul,no_inplace} [id G]
+    └─ Mul [id G]
        ├─ *1-<TensorType(float64, (?, ?))> [id H]
        └─ *2-<TensorType(float64, (?, ?))> [id I]
 """

tests/link/test_utils.py

@@ -156,7 +156,7 @@ def func(*args, op=op):
 
     assert (
         """
-    # Elemwise{add,no_inplace}(Test
+    # Add(Test
     # Op().0, Test
     # Op().1)
     """