Registered local_add_sub_collector in a new SequenceDB

aesara/tensor/math_opt.py

@@ -22,6 +22,7 @@
     local_optimizer,
 )
 from aesara.graph.opt_utils import get_clients_at_depth
+from aesara.graph.optdb import SequenceDB
 from aesara.misc.safe_asarray import _asarray
 from aesara.tensor.basic import (
     Alloc,
@@ -479,7 +480,6 @@ def local_mul_switch_sink(fgraph, node):
     return False
 
 
-@register_canonicalize
 @local_optimizer([add, sub])
 def local_add_sub_collector(fgraph, node):
     """
@@ -3613,3 +3613,9 @@ def local_reciprocal_1_plus_exp(fgraph, node):
 )
 register_canonicalize(local_sigmoid_logit)
 register_specialize(local_sigmoid_logit)
+
+fastmath_db = SequenceDB()
+fastmath_db.name = "fastmath_db"
+fastmath_db.register(
+    "distributive_terms_collector", local_add_sub_collector, -5, "basic"
+)