Move tensor_copy rewrite to aesara.tensor.rewriting.basic

Author: brandonwillard

aesara/tensor/rewriting/basic.py

@@ -11,6 +11,7 @@
 from aesara.graph.basic import Constant, Variable
 from aesara.graph.rewriting.basic import (
     NodeRewriter,
+    RemovalNodeRewriter,
     Rewriter,
     copy_stack_trace,
     in2out,
@@ -35,6 +36,7 @@
     join,
     ones_like,
     switch,
+    tensor_copy,
     zeros,
     zeros_like,
 )
@@ -1294,3 +1296,6 @@ def __getattr__(name):
         return fn()
 
     raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
+register_canonicalize(RemovalNodeRewriter(tensor_copy), name="remove_tensor_copy")

aesara/tensor/rewriting/shape.py

@@ -13,7 +13,6 @@
 from aesara.graph.fg import FunctionGraph
 from aesara.graph.rewriting.basic import (
     GraphRewriter,
-    RemovalNodeRewriter,
     check_chain,
     copy_stack_trace,
     node_rewriter,
@@ -27,7 +26,6 @@
     extract_constant,
     get_scalar_constant_value,
     stack,
-    tensor_copy,
 )
 from aesara.tensor.elemwise import DimShuffle, Elemwise
 from aesara.tensor.exceptions import NotScalarConstantError, ShapeError
@@ -972,9 +970,6 @@ def local_reshape_lift(fgraph, node):
         return [e]
 
 
-register_canonicalize(RemovalNodeRewriter(tensor_copy), name="remove_tensor_copy")
-
-
 @register_useless
 @register_canonicalize
 @node_rewriter([SpecifyShape])