fix reduce ops with negative axis

Author: Jencir Lee

micrograd/engine.py

@@ -202,25 +202,23 @@ def sum(self, axis=None):
 
         out = Value(self.data.sum(axis=axis), (self,), 'sum')
 
+        de_neg = lambda x: self.ndim + x if x < 0 else x
         if axis is None:
-            new_shape = self.shape
-        elif isinstance(axis, int):
-            new_shape = [h if j == axis else 1
-                         for j, h in enumerate(self.shape)]
+            expand_axis = tuple(range(self.data.ndim))
+        elif _shape(axis) == ():
+            expand_axis = de_neg(axis)
         else:
-            new_shape = [h if j in axis else 1
-                         for j, h in enumerate(self.shape)]
-        m_new = ones(new_shape)
+            expand_axis = tuple(map(de_neg, axis))
 
-        expand_axis = tuple(range(self.data.ndim)) if axis is None else axis
+        arr_orig_shape = ones(self.shape)
 
         def _forward(**kwds):
             out.data = self.data.sum(axis=axis)
         out._forward = _forward
 
         def _backward():
             self.grad += broadcast_arrays(expand_dims(out.grad, expand_axis),
-                                          m_new)[0]
+                                          arr_orig_shape)[0]
         out._backward = _backward
 
         return out

tests/test_engine.py

@@ -94,6 +94,36 @@ def test_sum_op(self):
         c.backward()
         self.assertTrue(allclose(a.grad, 0))
 
+    def test_sum_op_neg_axis(self):
+
+        a = Value(shape=(2, 2, 3), name='a')
+        b = (a.sum(axis=(0, -1)) - 31).relu()
+        b.forward(a=array([[[1, 2, 3], [4, 5, 6]],
+                           [[7, 8, 9], [10, 11, 12]]]))
+        b.backward()
+        self.assertTrue(allclose(b.data, [0, 17]))
+        self.assertTrue(allclose(a.grad, [[[0, 0, 0], [1, 1, 1]],
+                                          [[0, 0, 0], [1, 1, 1]]]))
+
+        b = (a.sum(axis=-3) - 10).relu()
+        b.forward(a=array([[[1, 2, 3], [4, 5, 6]],
+                           [[7, 8, 9], [10, 11, 12]]]))
+        b.backward()
+        self.assertTrue(allclose(a.grad, [[[0, 0, 1], [1, 1, 1]],
+                                          [[0, 0, 1], [1, 1, 1]]]))
+
+        b = (a.sum() - 77).relu()
+        c = (a.sum() - 79).relu()
+        b.forward(a=array([[[1, 2, 3], [4, 5, 6]],
+                           [[7, 8, 9], [10, 11, 12]]]))
+        b.backward()
+        self.assertTrue(allclose(a.grad, 1))
+
+        c.forward(a=array([[[1, 2, 3], [4, 5, 6]],
+                           [[7, 8, 9], [10, 11, 12]]]))
+        c.backward()
+        self.assertTrue(allclose(a.grad, 0))
+
     def test_mean_op(self):
 
         a = Value(shape=(2, 2, 3), name='a')
@@ -129,6 +159,41 @@ def test_mean_op(self):
         c.backward()
         self.assertTrue(allclose(a.grad, 0))
 
+    def test_mean_op_neg_axis(self):
+
+        a = Value(shape=(2, 2, 3), name='a')
+        self.assertTrue(a.mean()._op == '*')
+
+        b = (a.mean(axis=(0, -1)) - 6).relu()
+        b.forward(a=array([[[1, 2, 3], [4, 5, 6]],
+                           [[7, 8, 9], [10, 11, 12]]]))
+        b.backward()
+        y = 1 / 6
+        self.assertTrue(allclose(b.data, [0, 2]))
+        self.assertTrue(allclose(a.grad, [[[0, 0, 0], [y, y, y]],
+                                          [[0, 0, 0], [y, y, y]]]))
+
+        b = (a.mean(axis=-3) - 5).relu()
+        b.forward(a=array([[[1, 2, 3], [4, 5, 6]],
+                           [[7, 8, 9], [10, 11, 12]]]))
+        b.backward()
+        y = .5
+        self.assertTrue(allclose(a.grad, [[[0, 0, y], [y, y, y]],
+                                          [[0, 0, y], [y, y, y]]]))
+
+        b = (a.mean() - 6).relu()
+        c = (a.mean() - 7).relu()
+        b.forward(a=array([[[1, 2, 3], [4, 5, 6]],
+                           [[7, 8, 9], [10, 11, 12]]]))
+        b.backward()
+        y = 1 / 12
+        self.assertTrue(allclose(a.grad, y))
+
+        c.forward(a=array([[[1, 2, 3], [4, 5, 6]],
+                           [[7, 8, 9], [10, 11, 12]]]))
+        c.backward()
+        self.assertTrue(allclose(a.grad, 0))
+
     def test_tensordot_op(self):
 
         a = Value(empty((2, 3, 4)))

tests/test_vs_torch.py

@@ -166,3 +166,27 @@ def test_reduce_ops(self):
         c2.backward()
         self.assertTrue(allclose(c.data, c2.data))
         self.assertTrue(allclose(a.grad, a2.grad))
+
+    def test_reduce_ops_neg_axis(self):
+
+        a = Value(array([[[1, 2, -2], [2, 1, 0]],
+                         [[-2, 1, 0], [3, 2, 1]]]))
+        b = a.mean(axis=(0, -1)).relu().sum()
+        c = a.mean(axis=(-3, -2)).relu().mean()
+
+        a2 = Tensor([[[1, 2, -2], [2, 1, 0]],
+                     [[-2, 1, 0], [3, 2, 1]]])
+        a2.requires_grad = True
+        b2 = a2.mean(axis=(0, -1)).relu().sum()
+        c2 = a2.mean(axis=(-3, -2)).relu().mean()
+
+        b.backward()
+        b2.backward()
+        self.assertTrue(allclose(b.data, b2.data))
+        self.assertTrue(allclose(a.grad, a2.grad))
+
+        a2.grad = None
+        c.backward()
+        c2.backward()
+        self.assertTrue(allclose(c.data, c2.data))
+        self.assertTrue(allclose(a.grad, a2.grad))