From a57298dfcb5bb3f54f7a76558eab40b20e3828c7 Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Tue, 21 May 2024 02:03:24 -0300
Subject: [PATCH] softmax and small fixes

---
 .../__pycache__/functions.cpython-38.pyc      | Bin 10085 -> 9661 bytes
 norch/autograd/functions.py                   |  34 ++++++-----
 test.py                                       |  10 ++--
 tests/test_autograd.py                        |  56 ++++++++++++++++++
 4 files changed, 80 insertions(+), 20 deletions(-)

diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc
index 0b21229cce857c4b4c4e93794269a0aa31a74bcc..66ce27e5edd74698e2f12ebd7684abb0dc013b74 100644
GIT binary patch
delta 1463
zcmds1OK1~O6rDGpWSTb9>ZF|{&1d?dHvW);i&ES4<F9{djrN0RU~1^Jjrh`viwQ-F
zt`w9@5rsnHLa5?GQ5Rb4qDlo_=(c+W7or>Y-ZxERs})?ha^M~2-Fx1BGw+=HKK8xn
zcE{Xq7l+TY`ku3=vTo082ToFt5iHK+4C3QVlHs#ln#Z|d@Dn^eFfCfdluU=|WcfLg
zXJ+Iq=i$;ql47KYnig-cD#B&n#~DI`Xq>fPMbJ2dV^Iz3CiRWHfvqUhE{mTdY1}qb
zY>8Qt7A&E#X-Srq!AOD(auKnQ68w0&n?WT)#0^GyhbT^tr(n4}2CK5-y@K(S2Z1AH
z#Kl(&4pPtSf((rnpUG`l)B+II+KUe&fl3TC2s;C00W(px69cz!q(lUn$O7$z>)TqP
z$$QSx&B#4aC2J%LZ@dXd4-0hLtM9^quTAa3Kn6#CwM4t@)t`JP$vtTCC&&ZD3Rv>j
zsJ+Nr+`tuRugz_!JVPD<RNf>{U{~O<+K-%vn|S@mYlGK;A+iLa;8pSx?gyKtW5{_9
zU$j1G3JsH2Fds6>0Nl_tvI^dCvpR_Rl}&X9v5p6<@H}}3+x5l<+Z=v8;QZ<TW{&B6
zCL^M01Pju7m{_o&@0`tGs)XLzsP^BwVoOuFZBsbV6pH0QQ%IIn*u2y>MPY7=7FUe+
z6%FYcpLSZXrz$wR==jwu$7uk6KT?(x9p7mS13$u=FcrVJCXJn+iBHn~EE;9d%Ak!w
zI|H3TkO4a?YCGr<zC~o}WC4Z2VT6&Q3aoefLVR*&Y)zbsrzYqTmOBcUs>_*4!yj;(
z(h+Qv5`~oUIv)^oR^?%(x<!p)_AQQ_n_Lo&<rLcoLp5XM8!XrSXR6T?_`iRYXgo=W
zF#k^kni@~8>qCcG$#r?`>tKX6P3Sv2cpYJrK!v-+sEewjnEZ<4r#_?bw63|9-651+
o1(fwTWerSOcT(1Zl&MNb;k~B8-W~O&KE<gNDYD|iAuF{%07N4$g8%>k

delta 1878
zcmeH{OKclu5XX1+<+WRH+BnX}uAL@ng6gUYy>O_M2W~^5gt#<Ln=}fvOYB71#0h-s
zK!LS<D32Zh1xBb;LaI$wL8Kl);=mh0B_!a|vk$!>4n4pHi34z8W)nNH2ogdZS!@29
z`M&?oKK$GHeDudm_)It)lIXa8Yjp10YPhYeP~onTT$UUu2QkTkid>aSGSx~>UX;Z^
zQT4#l9Nh`<{w1h8D@IjnlS;~RhE{^^)XUNw<dmHHf>8{4@+BA{Z%Hoh;=H7KYGt?I
zQ>)4#c~A*^Y8A>Q#Zwv$&8lQkLr<B9td!C-0_euXVUk(S>IdJmN)e%OJ4xmxPkO9<
zMlR9yJpcX{6`^|a>+(}-RucTgB1W>V)a3pn;sB*JaMwWu&heJNzl#?`14|2TadxHf
z<j`wWyE=(ZzmW9Y!R%_`m_57rj4g8gPt=ovrIlHG$sOw7I#KMSlTs~sr|q~E%rCDN
z^1?3WYbtlH=?lwlUSMDJ45ab-Xwr~%xm_{9#7guoT*W<B0#CL^@peqbFWXZvig(-Z
zfwyUWSbLQEzJ(LX48CKn=;Kt0EOr|fRPf%8c|FSu<G4SNfE)O6M>=wv7Oqja{H_cz
zG;tc&_TG<Uu}N6R8?kw~iQ)Ls$OLV_x6P-B3Ey`$egQtf-{M!`LwsZ3iO5;<yx-(I
zOTIK7PE5h4_+sJ}_ySv#gON${eA48b^nKTpQ*aCaNIGy1s~r}6gXPY_$TT^>ZgNkP
zyA6Nnd=tLMXZIf+_YwV-EOk!)f@rb6Q&9=sEuJ(_sS<icod6Cvh7;Tb7}c&ra~_?v
z+g;m6ZFn2CrcbTz>6I4pDY-hedpw<5YM$1(<X#`zM(wD+U(W?zl}mxk_)%A~oHG7%
zl*JIe5RqnS;46+%d55Bj+cA6i;y(v>LL{g`s#R#*+G075zjwt)MmT+d<pj%vEL|)~
zmWNnQvOLT(%JK+F$`Aq0O_nh{)g3*Q;q(+qEmW{Co`2e2TFtu~58*!!e~6-m<ThUG
zjRw`t4+QXL_fTYp*1x0JWICjYnN3=bV7ey{KjXVS{}Z$0@cUb0B$}<w#4KP>vb9Vz
zZRg>(6C?bm#Yt*PVPRQ3&gFe9{~jA3h7E40(R7gdZoneWbK?d^;sTfcN6^gAFT{^k
zAFrSF4o7S<{zS0_kBw&EU@s3=!P8RkC=@*B1P?L6vq>;=!p5(n7QWPXu%*Kcm_gGp
JLllO2;7<qyj?4f6

diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py
index e64841a..f9ad107 100644
--- a/norch/autograd/functions.py
+++ b/norch/autograd/functions.py
@@ -137,13 +137,14 @@ class SumBackward:
             # 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:
-            if not self.keepdim:
-                # Remove dimensions of size 1 from the gradient tensor.
-                input_shape = [s for i, s in enumerate(input_shape) if i != self.axis]
-                
+
+            if self.keepdim:
+                input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:]
+            else:
+                input_shape = input_shape[:self.axis] + input_shape[self.axis+1:]
+
             # Broadcast the gradient to the input shape along the specified axis.
             grad_output_shape = list(input_shape)
-            grad_output_shape.insert(self.axis, 1)
             grad_output = gradient.reshape(grad_output_shape)
             grad_output = grad_output + self.input[0].zeros_like()
         
@@ -216,14 +217,15 @@ class MaxBackward:
             grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
 
         else:
-            
-            if not self.keepdim:
-                # Remove dimensions of size 1 from the gradient tensor.
-                input_shape = [s for i, s in enumerate(input_shape) if i != self.axis]
-                
+
+            if self.keepdim:
+                input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:]
+            else:
+                input_shape = input_shape[:self.axis] + input_shape[self.axis+1:]
+
             # Broadcast the gradient to the input shape along the specified axis.
             grad_output_shape = list(input_shape)
-            grad_output_shape.insert(self.axis, 1)
+                            
             grad_output = gradient.reshape(grad_output_shape)
             grad_output = grad_output + self.input[0].zeros_like()
             max_values = self.input[0].max(axis=self.axis, keepdim=True)
@@ -251,13 +253,13 @@ class MinBackward:
             grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
 
         else:
-            if not self.keepdim:
-                # Remove dimensions of size 1 from the gradient tensor.
-                input_shape = [s for i, s in enumerate(input_shape) if i != self.axis]
-                
+            if self.keepdim:
+                input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:]
+            else:
+                input_shape = input_shape[:self.axis] + input_shape[self.axis+1:]
+
             # Broadcast the gradient to the input shape along the specified axis.
             grad_output_shape = list(input_shape)
-            grad_output_shape.insert(self.axis, 1)
             grad_output = gradient.reshape(grad_output_shape)
             grad_output = grad_output + self.input[0].zeros_like()
             max_values = self.input[0].min(axis=self.axis, keepdim=True)
diff --git a/test.py b/test.py
index 6bd8b2a..d61d329 100644
--- a/test.py
+++ b/test.py
@@ -6,10 +6,12 @@ import norch.optim as optim
 import random
 random.seed(1)
 
-torch_tensor = norch.Tensor([[[2, 2], [-1, -1]], [[1., 2], [3, 3]]], requires_grad=True)#.to(self.device)
-torch_tensor = norch.Tensor([[[2, 2], [2, 2]], [[2, 2.], [2, 2]]])
-b = norch.nn.functional.softmax(torch_tensor)
-print(b)
+torch_tensor = norch.Tensor([[[1, 5], [2, -1]], [[5, 2.], [2, 2]]], requires_grad=True)#.to(self.device)
+b = norch.nn.functional.softmax(torch_tensor, dim=0)
+soma = b.sum()
+print(soma)
+soma.backward()
+print(torch_tensor.grad)
 
 """a = norch.Tensor([[[4.186502456665039]]])
 b = norch.Tensor([[[2.0, 2.0,],[-1.0, -1.0,]],[[1.0, 2.0,],[3.0, 3.0,]]])
diff --git a/tests/test_autograd.py b/tests/test_autograd.py
index eb27544..fb124d6 100644
--- a/tests/test_autograd.py
+++ b/tests/test_autograd.py
@@ -583,6 +583,62 @@ class TestTensorAutograd(unittest.TestCase):
         torch_expected_cos_tensor_grad = torch_cos_tensor.grad
 
         self.assertTrue(utils.compare_torch(torch_result_cos_tensor_grad, torch_expected_cos_tensor_grad))
+
+    def test_sigmoid(self):
+        """
+        Test autograd from sigmoid
+        """
+        norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_sigmoid = norch.sigmoid(norch_tensor)
+        norch_result = norch_sigmoid.sum()
+
+        norch_result.backward()
+        norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
+
+        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_sigmoid = torch.sigmoid(torch_tensor)
+        torch_result = torch_sigmoid.sum()
+
+        torch_result.backward()
+        torch_tensor_grad = torch_tensor.grad
+        
+        self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
+
+    def test_softmax(self):
+        """
+        Test autograd from softmax
+        """
+        norch_tensor = norch.Tensor([[[-5, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_softmax = norch.softmax(norch_tensor, dim=1)
+        norch_result = norch_softmax.sum()
+
+        norch_result.backward()
+        norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
+
+        torch_tensor = torch.tensor([[[-5, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_softmax = torch.softmax(torch_tensor, dim=1)
+        torch_result = torch_softmax.sum()
+
+        torch_result.backward()
+        torch_tensor_grad = torch_tensor.grad
+        
+        self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
+
+        norch_tensor = norch.Tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
+        norch_softmax = norch.softmax(norch_tensor, dim=2)
+        norch_result = norch_softmax.sum()
+
+        norch_result.backward()
+        norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
+
+        torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
+
+        torch_softmax = torch.softmax(torch_tensor, dim=2)
+        torch_result = torch_softmax.sum()
+        torch_result.backward()
+        torch_tensor_grad = torch_tensor.grad
+        self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
+
     
     def test_reshape(self):
         """