Use type name in Elemwise and CAReduce __str__ implementations

aesara/tensor/elemwise.py

@@ -474,9 +474,9 @@ second dimension
             if self.inplace_pattern:
                 items = list(self.inplace_pattern.items())
                 items.sort()
-                return f"Elemwise{{{self.scalar_op}}}{items}"
+                return f"{type(self).__name__}{{{self.scalar_op}}}{items}"
             else:
-                return "Elemwise{%s}" % (self.scalar_op)
+                return f"{type(self).__name__}{{{self.scalar_op}}}"
         else:
             return self.name
 
@@ -1340,13 +1340,12 @@ class CAReduce(COp):
         self.set_ufunc(self.scalar_op)
 
     def __str__(self):
+        prefix = f"{type(self).__name__}{{{self.scalar_op}}}"
         if self.axis is not None:
-            return "Reduce{{{}}}{{{}}}".format(
-                self.scalar_op,
-                ", ".join(str(x) for x in self.axis),
-            )
+            axes_str = ", ".join(str(x) for x in self.axis)
+            return f"{prefix}{{{axes_str}}}"
         else:
-            return "Reduce{%s}" % self.scalar_op
+            return f"{prefix}"
 
     def perform(self, node, inp, out):
         (input,) = inp
@@ -1750,14 +1749,12 @@ class CAReduceDtype(CAReduce):
         return super(CAReduceDtype, op).make_node(input)
 
     def __str__(self):
-        name = self.__class__.__name__
-        if self.__class__.__name__ == "CAReduceDtype":
-            name = ("ReduceDtype{%s}" % self.scalar_op,)
-        axis = ""
+        prefix = f"{type(self).__name__}{{{self.scalar_op}}}"
         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}}}"
+            return f"{prefix}{{axis=[{axis}], acc_dtype={self.acc_dtype}}}"
+        else:
+            return f"{prefix}{{acc_dtype={self.acc_dtype}}}"
 
 
 def scalar_elemwise(*symbol, nfunc=None, nin=None, nout=None, symbolname=None):

aesara/tensor/math_opt.py

@@ -1578,7 +1578,7 @@ ALL_REDUCE = (
 @local_optimizer(ALL_REDUCE)
 def local_reduce_join(fgraph, node):
     """
-    Reduce{scalar.op}(Join(axis=0, a, b), axis=0) -> Elemwise{scalar.op}(a, b)
+    CAReduce{scalar.op}(Join(axis=0, a, b), axis=0) -> Elemwise{scalar.op}(a, b)
 
     Notes
     -----

tests/tensor/test_elemwise.py

@@ -16,7 +16,7 @@ from aesara.link.basic import PerformLinker
 from aesara.link.c.basic import CLinker, OpWiseCLinker
 from aesara.tensor import as_tensor_variable
 from aesara.tensor.basic import second
-from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise
+from aesara.tensor.elemwise import CAReduce, CAReduceDtype, DimShuffle, Elemwise
 from aesara.tensor.math import all as at_all
 from aesara.tensor.math import any as at_any
 from aesara.tensor.type import (
@@ -622,6 +622,17 @@ class TestCAReduce(unittest_tools.InferShapeTester):
                 warn=0 not in xsh,
             )
 
+    def test_str(self):
+        op = CAReduce(aes.add, axis=None)
+        assert str(op) == "CAReduce{add}"
+        op = CAReduce(aes.add, axis=(1,))
+        assert str(op) == "CAReduce{add}{1}"
+
+        op = CAReduceDtype(aes.add, axis=None, acc_dtype="float64")
+        assert str(op) == "CAReduceDtype{add}{acc_dtype=float64}"
+        op = CAReduceDtype(aes.add, axis=(1,), acc_dtype="float64")
+        assert str(op) == "CAReduceDtype{add}{axis=[1], acc_dtype=float64}"
+
 
 class TestBitOpReduceGrad:
     def setup_method(self):
@@ -722,6 +733,14 @@ class TestElemwise(unittest_tools.InferShapeTester):
     def test_input_dimensions_match_c(self):
         self.check_input_dimensions_match(Mode(linker="c"))
 
+    def test_str(self):
+        op = Elemwise(aes.add, inplace_pattern=None, name=None)
+        assert str(op) == "Elemwise{add}"
+        op = Elemwise(aes.add, inplace_pattern={0: 0}, name=None)
+        assert str(op) == "Elemwise{add}[(0, 0)]"
+        op = Elemwise(aes.add, inplace_pattern=None, name="my_op")
+        assert str(op) == "my_op"
+
 
 def test_not_implemented_elemwise_grad():
     # Regression test for unimplemented gradient in an Elemwise Op.