Refactor local_add_specialize and test_local_add_specialize

aesara/tensor/math_opt.py

@@ -2062,13 +2062,15 @@ def local_mul_specialize(fgraph, node):
 @register_specialize
 @local_optimizer([add])
 def local_add_specialize(fgraph, node):
-    def _fill_chain(v):
-        out = fill_chain(v, node.inputs)
-        return out
+    """Remove zeros from ``add``s.
 
+    TODO: This should be a canonicalization, no?
+    """
     # here, we are past the point of canonicalization, so we don't want
     # to put in un-necessary fills.
-    if node.op == add:
+    if node.op != add:
+        return False
+
     new_inputs = []
     for inp in node.inputs:
         try:
@@ -2079,27 +2081,31 @@ def local_add_specialize(fgraph, node):
             continue
         new_inputs.append(inp)
 
-        if len(new_inputs) < len(node.inputs):
-            dtype = node.outputs[0].type.dtype
+    if len(new_inputs) == len(node.inputs):
+        return False
+
+    node_output = node.outputs[0]
+    dtype = node_output.type.dtype
+
     if len(new_inputs) == 0:
         # we got rid of the entire expression!
-                ndim = node.outputs[0].type.ndim
+        ndim = node_output.type.ndim
         # Reuse call to constant for cache()
         cst = constant(np.zeros((1,) * ndim, dtype=dtype))
         assert cst.type.broadcastable == (True,) * ndim
-                return _fill_chain(cst)
+        return fill_chain(cst, node.inputs)
 
     if len(new_inputs) == 1:
-                ret = _fill_chain(new_inputs[0])
+        ret = fill_chain(new_inputs[0], node.inputs)
     else:
-                ret = _fill_chain(add(*new_inputs))
+        ret = fill_chain(add(*new_inputs), node.inputs)
+
     # The dtype should not be changed. It can happen if the input
     # that was forcing upcasting was equal to 0.
     if ret[0].dtype != dtype:
         ret = [cast(ret[0], dtype)]
+
     return ret
-    else:
-        return False
 
 
 mul_canonizer = in2out(

tests/tensor/test_math_opt.py

@@ -77,7 +77,6 @@ from aesara.tensor.math import tan, tanh, true_div, xor
 from aesara.tensor.math_opt import (
     compute_mul,
     is_1pexp,
-    local_add_specialize,
     local_grad_log_erfc_neg,
     local_greedy_distributor,
     mul_canonizer,
@@ -3781,24 +3780,41 @@ class TestLocalSumProdDimshuffle:
     # test_local_sum_divprod_dimshuffle ((a * b) / (c * d))
 
 
-def test_local_add_specialize():
+def test_local_useless_adds():
+    default_mode = get_default_mode()
+
+    # Test for all zeros
+    a = scalar()
+    s = add(aet.zeros_like(a))
+    mode_with_opt = default_mode.including("canonicalization", "local_useless_fill")
+    f = function([a], s, mode=mode_with_opt)
+    assert not any(node.op == add for node in f.maker.fgraph.apply_nodes)
+
     # test of non-zero dimension
     a = vector()
     s = add(aet.zeros_like(a))
-    assert local_add_specialize.transform(None, s.owner)
+    mode_with_opt = default_mode.including("canonicalization", "local_useless_elemwise")
+    f = function([a], s, mode=mode_with_opt)
+    assert not any(node.op == add for node in f.maker.fgraph.apply_nodes)
 
     # test of 0-d
     a = scalar()
     s = add(aet.zeros_like(a))
-    assert local_add_specialize.transform(None, s.owner)
+    mode_with_opt = default_mode.including(
+        "canonicalization", "local_useless_fill", "local_useless_elemwise"
+    )
+    f = function([a], s, mode=mode_with_opt)
+    assert not any(node.op == add for node in f.maker.fgraph.apply_nodes)
 
     # Test when the 0 input is forcing upcasting
     a = aet.constant(0, dtype="int64")
     b = aet.constant(1, dtype="int32")
     s = a + b
-    transformed = local_add_specialize.transform(None, s.owner)
-    assert transformed
-    assert transformed[0].type == s.type
+    mode_with_opt = default_mode.including("canonicalization", "local_add_canonizer")
+    f = function([], s, mode=mode_with_opt)
+    transformed = f.maker.fgraph.outputs[0]
+    assert not any(node.op == add for node in f.maker.fgraph.apply_nodes)
+    assert transformed.type == s.type
 
 
 def test_local_div_to_reciprocal():