dispatch special cases for x ** 1 and x ** 2

src/mygrad/tensor_base.py

@@ -5,6 +5,7 @@
 """
 
 from functools import wraps
+from numbers import Number
 from typing import Optional, Set, Type, Union
 
 import numpy as np
@@ -19,6 +20,7 @@
     Negative,
     Positive,
     Power,
+    Square,
     Subtract,
 )
 from mygrad.operation_base import BroadcastableOp, Operation
@@ -672,6 +674,14 @@ def __rmatmul__(self, other):
         return self._op(MatMul, other, self)
 
     def __pow__(self, other):
+        if isinstance(other, Number) or (
+            isinstance(other, np.ndarray) and other.ndim == 0
+        ):
+            if other == 1:
+                return self._op(Positive, self)
+            elif other == 2:
+                return self._op(Square, self)
+
         return self._op(Power, self, other)
 
     def __rpow__(self, other):

tests/tensor_base/test_graph_tracking.py

@@ -16,7 +16,7 @@ def test_op_tracks_graph():
     h = z - f
     assert h.creator.graph == {h.creator} | f.creator.graph
 
-    i = ((h + 3) ** 2) / 5
+    i = ((h + 3) ** 4) / 5
     assert h.creator.graph < i.creator.graph
     assert (
         len(i.creator.graph - h.creator.graph) == 3

tests/tensor_base/test_pow_special_cases.py

@@ -0,0 +1,72 @@
+from functools import partial
+
+import hypothesis.strategies as st
+import numpy as np
+import pytest
+from hypothesis import given
+
+import mygrad as mg
+from mygrad.math.arithmetic.ops import Positive, Square
+
+from ..wrappers.uber import backprop_test_factory, fwdprop_test_factory
+
+
+def custom_pow(x, p, constant=False):
+    out = x ** p
+    if isinstance(out, mg.Tensor):
+        out._constant = constant
+    return out
+
+
+def any_scalar(*args, p):
+    return st.sampled_from([int(p), float(p), np.array(p)])
+
+
+@pytest.mark.parametrize("power, op", [(1, Positive), (2, Square)])
+def test_pow_uses_special_case(power, op):
+    @given(exp=st.sampled_from([int(power), float(power), np.array(power)]))
+    def wrapped_func(exp):
+        out = mg.arange(2) ** exp
+        assert isinstance(out.creator, op)
+
+    wrapped_func()
+
+
+@fwdprop_test_factory(
+    mygrad_func=custom_pow,
+    true_func=custom_pow,
+    num_arrays=1,
+    kwargs={"p": partial(any_scalar, p=1)},
+)
+def test_pow_1_fwd():
+    pass
+
+
+@backprop_test_factory(
+    mygrad_func=custom_pow,
+    true_func=custom_pow,
+    num_arrays=1,
+    kwargs={"p": partial(any_scalar, p=1)},
+)
+def test_pow_1_bkwd():
+    pass
+
+
+@fwdprop_test_factory(
+    mygrad_func=custom_pow,
+    true_func=custom_pow,
+    num_arrays=1,
+    kwargs={"p": partial(any_scalar, p=2)},
+)
+def test_pow_2_fwd():
+    pass
+
+
+@backprop_test_factory(
+    mygrad_func=custom_pow,
+    true_func=custom_pow,
+    num_arrays=1,
+    kwargs={"p": partial(any_scalar, p=2)},
+)
+def test_pow_2_bkwd():
+    pass