Fixed datatype conversion in KanrenRelationSub

aesara/graph/rewriting/kanren.py

@@ -95,6 +95,10 @@ def transform(self, fgraph, node):
             else:
                 new_outputs = [eval_if_etuple(chosen_res)]
 
+            new_outputs = [
+                _new_out.astype(_inp_expr.dtype)
+                for _inp_expr, _new_out in zip(node.outputs, new_outputs)
+            ]
             return new_outputs
         else:
             return False

aesara/tensor/rewriting/math.py

@@ -9,9 +9,9 @@
 from etuples import etuple
 from kanren import fact, heado, tailo
 from kanren.assoccomm import associative, commutative
+from kanren.constraints import neq
 from kanren.core import lall, lany
 from kanren.graph import mapo
-from kanten.constraints import neq
 from unification import vars as lvars
 
 import aesara.scalar.basic as aes
@@ -3558,13 +3558,11 @@ def distributive_collect(in_lv, out_lv):
     A_lv, op_lv, all_term_lv, all_cdr_lv, cdr_lv, all_flat_lv = lvars(6)
     return lall(
         lany(
-            heado(at.add, all_term_lv),
-            heado(at.sub, all_term_lv),
+            lall(heado(at.add, all_term_lv), eq(cons(at.add, cdr_lv), out_lv)),
+            lall(heado(at.sub, all_term_lv), eq(cons(at.sub, cdr_lv), out_lv)),
         ),
         # Get the flattened `add` arguments
         tailo(all_cdr_lv, all_term_lv),
-        # Add all the arguments and set the output
-        lany(eq(cons(at.add, cdr_lv), out_lv), eq(cons(at.sub, cdr_lv), out_lv)),
         lany(
             lall(
                 lany(

tests/tensor/rewriting/test_math.py

@@ -1958,6 +1958,7 @@ def test_local_elemwise_sub_zeros():
             "ShapeOpt",
             "local_fill_to_alloc",
             "local_elemwise_alloc",
+            "dist_collect_opt",
         )
         .including("local_elemwise_sub_zeros")
     )