Refactor rewrite

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -1913,20 +1913,29 @@ def local_pow_canonicalize(fgraph, node):
 
     In all cases, the shape of the output is the result of broadcasting the shapes of the inputs.
     """
-
     cst_base = get_underlying_scalar_constant_value(
         node.inputs[0], only_process_constants=True, raise_not_constant=False
     )
-    if cst_base == 1:
-        return [broadcast_arrays(*node.inputs)[0].astype(node.outputs[0].dtype)]
-
     cst_exponent = get_underlying_scalar_constant_value(
         node.inputs[1], only_process_constants=True, raise_not_constant=False
     )
-    if cst_exponent == 0:
-        return [alloc_like(1, node.outputs[0], fgraph)]
-    if cst_exponent == 1:
-        return [alloc_like(node.inputs[0], node.outputs[0], fgraph)]
+
+    new_out = None
+
+    if cst_base == 1:
+        new_out = broadcast_arrays(*node.inputs)[0]
+    elif cst_exponent == 0:
+        new_out = broadcast_arrays(ones_like(node.inputs[0]), node.inputs[1])[0]
+    elif cst_exponent == 1:
+        new_out = broadcast_arrays(*node.inputs)[0]
+
+    if not new_out:
+        return
+
+    if new_out.dtype != node.out.dtype:
+        new_out = cast(new_out, dtype=node.out.dtype)
+
+    return [new_out]
 
 
 @register_specialize