From 3627e424691b8da84b609e2d821f4fb00647032e Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Thu, 16 May 2024 21:05:30 -0300
Subject: [PATCH] sum axis autograd

---
 .../__pycache__/functions.cpython-38.pyc      | Bin 7539 -> 7975 bytes
 norch/autograd/functions.py                   |  23 ++++++++++++++----
 tests/test_autograd.py                        |  22 +++++++++++++++++
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc
index d48733e29dfa77cff51a2321223754822c1d80b7..2c88c4fb3673489be608f82942d81dc1a1b73f4f 100644
GIT binary patch
delta 1591
zcmah}&2Jl35Z`&b-mL9pYkM6hvD2!}M~Mka(kjp<G$9F1Le(e*MX16Bg{{{Ou@gJY
z+Mrfz9MMP=^+I?>aA+Zs_=*Gvh{S=S-uNdhAr*1zi34XSGaCmIF0AzC&CdLO@4cBf
zZ~hqlYTSIer6r`$_wBuzYYVO4na2aLvpX30S?EW}iYKHUwPcGi4c+qB=JwVeGM=E}
z>Ld&>R71y}-F4%6uzJaIhFx02g}rj!n2?E37M`GG+T>>F3TdAuI%P0CWt}A0_yN3z
zCw=^<;a{Lrdd*d<kGwMn5GwZhh`E-#wG7!kC9k|Hs-miTA`hOLhk}aFc%*F!2OP02
zB8n&Wbq^xSl2}mw&+jM;N?M(N^q9U!f00P57@*y-jYK_IM!V>)1xjuXpXB6x7}KH)
zd`N?Y0wK{1$A!FN4qN^7e~CT+B&45?KPi7RSHO`k9=Q)`*&G?whX{kG<Xq$wEXj8w
zv(_LPvWJL+KCvUZ2ye;j(c5rc24Z6tXS;F;cZP8BXH~LIu>BZ`oRR@Y$p}Yxgrjai
zMs+kNu%To$2Pn{{dU&C-I%Q`IZ`<h1OIuCKsZf`NDdJM(O*A2ghAF<o*gC0xKnfUe
zBTify?KoEB1o5^!+%*!n=x!Y7Q1Va}qzgwW7<jg<dXU%hfxOPXo3M}CYT8!fN>weM
zeowhZH%A9G%7!Oy!g)efH4nCep0*8i)9QaYp1K4zibR`7ppCOE_CG!jXWydXc9B@3
zqIoSzMMKB({Ot=mCLN_HJh3U>)My>RP6A?*L|UzjfgxmKU3HdMcVI6c{9-W7*3h|?
zWw)HIl~&irZ<E1Kjjrg*jlKhE9k|-a@4&g0Qr2E^$4Oxx)6$`T%gfI?4)e--IahQ`
zSU0kzBB^D!t~qwut_MrSoV&QPT*%dfxngBChjuwvt2bfJeZyYMr2~6MG--`Yc)7S%
zDKGktOn3=xXR%Z%lPMED@U!PfKd3RH`H|rf(l%(gypjUVRM}l>0zwcI{<ZIyj31lM
z#ghg(6i!*A1g|l1uk1q_j{2{*KRyj-xrSg<z8}9F9w7`B*m;`d@9}Q9L)K0?*xDVK
zCesJ<VaL4uqP3*Y(jUIy6F-rs6Ia6nzJbLoHxsxge@aYR+$cV#vH#ivI3pu%wZ@j)
zGSZ%q587UW&q*7xE)a6-09wTh{vq$SzYaU{&-P0ADe@wRfbETE4&TU+I%cekWc&IM
z^&(NH<WTat#%8<nR`N28$cM>c_(8U&#;iHQe}9NGN1RFdWon|qxi255UWxfR5q>??
zRd>0F^dS8f0}tf8onv`!CqEuqhb`$SK1~w}7I79mEP7dRmql)*$n6NZwIKHm<ZD9C
lG;-#U6M`HX@>%f`b^A6sakMKMH)CeVjG9em&@@Oh{sy&+S9JgY

delta 1048
zcmY*YOH31C5Z>8sODRy=ZlRPG!&59ONF=;MYb&8jqF~Y(f}(&>kjl#`y<mw4G)B3o
zL*hZAa8Tnx6Xf8*gE7&AN3SFv)Wn2?iHV-knN4|glYjbm=9_Px|BKL{cKi8)0-M6$
z+*ovM+qG^lH$z~f!Td?J_*skP)acWxz}f|s^ew-nsEU;pe3BLxR<LST3s#ZS8QJ(?
zRY9+2ZHt&q88cS1F-*m_5jZGR5A#j*z;==U{IXrvYZw*U2RC7j0e_PB6?~M=U(jpW
z8r_G`u+~7FDOl;PV|`jlC<E$(x}c>Yt*yc+jH+v@LrJSU_cS<^aW$p{w4SYh{B3Qs
z2&iag@6Fsf9G*_j&Lra);PgC2v7J~7*kImsc6e?oHA!E*M|6qSyJMBgu}<nHf8ity
z(cQv#aEQ7b?Un|{;4!-GXoL}Z>o}#?vtlsIv7Q}Q0qKo%2~ret1@&gOU(doc3%M3c
zP0^*YAhcw9XG;r6WuZ{7Ufm%#jZY_^WF?C^B@GEks0nQpyvibQe@t>)HLoddWgKD(
zwlXYW!Q=Ga?Q>+>Ff$!b4o*&7i_>qL*Kt$Wj3DZ~TT+nv+gN!+47;OkGbZt)sB1Sp
zxJ8lTzG5SSwwZ(R<kjKXxDf+AJaIqVVSt}Dy;Zu*eKCtKWbpNSW-Jlb<H=0TutYyS
zSMrV;IuQ|r=|0VtMD#E#?(x~l%8X%y#pFG)m(?;2c*o%pee|~J(ka^Kpkb%saISO+
zR%pF61&68DR{>9{+ZQ!=G4e6xd3xx(FS1AYoI3n{&`QhxCU{Ao{Xt!J^jQ{dlxZV$
zp)9=T_L?@z`dvnMGW(|`C#LbVY#yvr^MPQZn>k)(@rWki#3uehnF+c@REdz+hVq6`
xo(al^qHG(=I-v{<rOOx=A%&yVgOV$D^K<sl-vgCShuvkj*{${*doD}9{Rbt_*W~~J

diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py
index 8ce8625..72cb6a5 100644
--- a/norch/autograd/functions.py
+++ b/norch/autograd/functions.py
@@ -82,7 +82,11 @@ class MatmulBackward:
 
     def backward(self, gradient):
         x, y = self.input
-        return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
+        
+        if x.shape != y.shape: # broadcasted case
+            return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
+        else:
+            return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
         
 """class PowBackward:
     def __init__(self, x, power):
@@ -121,13 +125,22 @@ class LogBackward:
         return [grad_input]
        
 class SumBackward:
-    def __init__(self, x):
+    def __init__(self, x, axis=None):
         self.input = [x]
+        self.axis = axis
 
     def backward(self, gradient):
-        # Since sum reduces a tensor to a scalar, gradient is broadcasted to match the original shape.
-        return [float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()]
-    
+        input_shape = self.input[0].shape
+        if self.axis is None:
+            # If axis is None, sum reduces the tensor to a scalar.
+            grad_output = float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()
+        else:
+            # Broadcast the gradient to the input shape along the specified axis.
+            grad_output_shape = [1 if i in self.axis else dim for i, dim in enumerate(input_shape)]
+            grad_output = gradient.reshape(grad_output_shape)
+            grad_output = grad_output * self.input[0].ones_like()
+        
+        return [grad_output]
 class ReshapeBackward:
     def __init__(self, x):
         self.input = [x]
diff --git a/tests/test_autograd.py b/tests/test_autograd.py
index b22eba6..d40318a 100644
--- a/tests/test_autograd.py
+++ b/tests/test_autograd.py
@@ -34,6 +34,28 @@ class TestTensorAutograd(unittest.TestCase):
         self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
         self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
 
+    def test_sum_axis(self):
+        """
+        Test autograd from sum specifying axis
+        """
+        norch_tensor1 = norch.Tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_tensor2 = norch.Tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        norch_result = (norch_tensor1 + norch_tensor2).sum(axis=0).sum(axis=0).sum()
+        norch_result.backward()
+        norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
+        norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
+
+        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        torch_result = (torch_tensor1 + torch_tensor2).sum(axis=0).sum(axis=0).sum()
+        torch_result.backward()
+        torch_tensor1_grad = torch_tensor1.grad
+        torch_tensor2_grad = torch_tensor2.grad
+
+        self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
+        self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
+
+
     def test_broadcasting_addition_autograd(self):
         """
         Test autograd for broadcasting addition: tensor1 + tensor2