From cd57c39ff0da13355009186c87a0699f1550d084 Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Mon, 20 May 2024 18:44:43 -0300
Subject: [PATCH] min max broadcasted autograd and small fix on add broadcasted

---
 norch/__pycache__/tensor.cpython-38.pyc       | Bin 18194 -> 18460 bytes
 .../__pycache__/functions.cpython-38.pyc      | Bin 9830 -> 10085 bytes
 norch/autograd/functions.py                   |  23 ++++---
 norch/tensor.py                               |  20 ++++++
 test.py                                       |  18 ++++--
 tests/test_autograd.py                        |  59 +++++++++++++++---
 tests/test_operations.py                      |  25 ++++++--
 7 files changed, 119 insertions(+), 26 deletions(-)

diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc
index 1705e208cd96ce49ff08d85fed2704db9611c2bc..0e09a81a352e9292f15a031df4e10fab02489361 100644
GIT binary patch
delta 1603
zcmdUvU1$_n6vyxX&dly)-LLFNHrY)zTI)wutRK5Zku++E5~+gLhaz=tT_TJcbR7!P
z*{%|hh##yuEs6^6f`a%}%7Dax(E8v*pQ<AGAoM|@P$>E$_)u)m?947nG;e+B!0b8q
z%;DbO-v7P-(KEPIL3l?f6cGIUGQN56hZ8@AKTF96pno!*bR;#mTtRWU1w*rhp%s*N
zu}>BUEmVC4WiPm{rbG=T49#)AJAGcyh$Xxp;q?g9U)9Z)8X9r0Gy}#ABbX00D&9Ek
z#A=Xdj96tDk(w;2W?)#9G)uL#d}Ksaqcy>qYPODuVQ}W)#W=<a3jf9}%~1Qqm$8!1
z(o|u@E#GP(MLvG$4fiD~IM~jp3GU?A85i%HNlRE5mRxHasx*Q8WtZ7wO-*4<<z3cP
zo%KDaTvr>goNGLN6P3pd%CIid^>>ik87bs3S=x57P^B%i?$Y-iZ_q)f(`ld1sug?b
ztM&2GmH_>nUPiB1MM%qxqz4>95k^p|N-z?WJ{d)V-<V35Gi^!~x;>V83aqEIAMB7f
zC&}&X&HqL2iF<hP?UWt7Hl3)`Yg*<5`Cp1!c?Zc6b&{rA_DDx8_o1D$p3AR653`-L
zD;=Q0u9^O#^L^p_6u;Z$D*m?Xz{;Hv;&8t+%S74X$*!akm3$1~LjVD(LziLbG?tlX
z-<y2}sPQ?-+2=omf^oVq?-pLrruhf(2i=*!2Y=e>1y2<GWxrZ94lK5BEcpiLr5j5=
zvNtYWCGojkUUmVnD7|7k?$gkU<K=DqzjYQzFQ9H2*Vr;z3aDF#e}ZM`PM_EFD_e%!
zBMi4kn1SjHGsu?VMZShGLPj_rX;i#%)QQzFuh_~nwwCLbfu?1!O;lrb%kZ0V&oUgm
zx&;&*1L2Er_}MZNo@GeSGGqmZ)(q`$E;kt@IRzw{cR{l3CN9ZT5U%?hD4grObQ6Sr
zU-=4G>9=ej+9;V@G~s^J)?7E8;?gJ|hQkANFV{J9T5TPT6ZXs8r@%9NX?Ir2my%BP
z^w?L{JO*NPvZvP`Tl)ye+9!K<Lw=#qC*2?0sovR;``x^-?ly4Dj&Cr4i}YPz2|wC@
W_Z@|{S?|#M=~#af19Yzc+kXLy-<hcZ

delta 1352
zcmchWOK1~87{}+E-6Y#I&rO=zr^H8_K4`Vt#<bX%X$wj#SPDh3DN+;`o2GH;K{upO
zsntUX&cQ=HB%lbQq6{c_5pN!}1+}86=s`SqQNfF#;<tGjON$qAhuLrD+hON7-+%s>
z5^!M=Y$H~ync%Z@^}w+w^S5n-rErRMSNFuDA5=rqsG3rci^r)dDMgAoRGiCX=N3IC
z%^Fdx8YtEjjftu>M`8v|)C?i(0!i9Z1V^H3Um$aUBL}W^VsuUj_?9(;A_mD3IcC%h
zB2i?`7>E-}Lh|%2T<D4{<8NKK5WC`vzE|BE(O`~ZrX-xm!(?%@n4UG5o=^tUvk@ld
z#T=+c+7$=>KF^6CkrGJ}CjpVMLZAdhXVlW)NZea?7hn@>@|Lqh*h?#Ru+^O|Hd1k}
zp;xap*uX+X!D5%GNXx-NWoiT?n88e?baYS~7`e0Z1wcFdSTzDiGm<QCaMph+N#Ri=
zcM1*~cSRrlkJ4{cd0)Ep{RKoQMRcsN^qOIq)Cug%=Fjvnu*LEkR^+XwbvZ5er**jI
z4L1%0AtyWkDu$#B!JL&fJi6sF07Eg__~14W;5s{4djp=cAGMS4lCAhk*sN~|Uh@^-
zssL~KmBxDj)qK2Z2B3wFH*M$NoBh<%PjTBlnQhs-<}yCof(ycIW!Es;ZtL#(xLK8w
zMXDKHmyLLxK`B-&A!{y^o!j)7w4hveEQvksI>@^YotAW@U1w6AS=Xg+uWq9W-GgW!
z#Z=%rimoGZbRC`k8a|QV>`x5$8Vq+RgMp4*9|QBI97y5K(lKzZdl);aW*CZMaoDHU
zH*k?XZVN&Qll={Bz&`*IJMV8cp3!U9%kKLtJ?olA=HL8X01tT}&`RMIpX!h){A5!d
zUHrk`r$D^|i*$7IXI);PeFAF_)N^kSjz=@&aPLikNnRCH0M4?7z8EZVWB)8v?32o*
WZtV4vkZxU)8=Ft^NI{kwIQ<LmmsPL;

diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc
index 694d6fd37cb9d953be98737411642d16f71a209e..0b21229cce857c4b4c4e93794269a0aa31a74bcc 100644
GIT binary patch
delta 356
zcmaFn^VE+ol$V!_0SFc!_D(ygu#wM6f$1B=<RFDo#+u0|6#UFO7_t~r7$q69Sb!|%
z6t)!hUZz@hnCJr58uo>ZwJbHPH7p_wwX8L)>5Rb)njDk+6xT6n3QcxbikZAhNr&;)
z<m*Z*j0%(gD9NxfiZKc>aZFZI{>b=lvV@8uKL^ly79eH@;$rGpo+ZM-2sDfhXd)XB
zaxqkC_-5u!=20>Lx$3cs&E!@kIVPYx6;-7J;(@A){6K_1h)4htpfD+N1QFaIf*(kH
x2I67`Ai=;Wz{thO2lRyklMFK(lNd7xkjKR+1o0Wu<N{R-#<!DKs_tP_1^_MpOV9uS

delta 377
zcmaFr_soYcl$V!_0SI~*d#8Pt+sNmnz+}NV`JZwr<G#r!6#T+UggO|q7*ZG|8JZax
z8DQ)MEDITGnM;Ifn2|-ZSb*}(DQqe1y-c<2B|<42HS8&z&CE#pSZkOUGS;$8UZWtz
zQNx<f7|fu_HQ7gT9d{KI7grUtgM)*n@Z>*=G3>Wk3yL!HN+w4r=`co3ZdFoY)R??V
zNrsP!kCBg&gHeo;hf#n@YVsMSkBpxu?^8CkVgq`D8Hk;MxVQ>P)Ii;x#khb8=vpQq
zpE-poo2jS-CIj(B&E#4IdB&Q_a}~727@%%q_A6op8d0PJBs6&@uTZs7)?_PU2MO_k
z2wxx(#hII#7hjf`Q<_?o021_@JW*U+{0ER(3{u3vD8R@CbgTfA0+S5u<a{*?#!r)1
JsO@3Y006aYS4sc?

diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py
index daa0668..e64841a 100644
--- a/norch/autograd/functions.py
+++ b/norch/autograd/functions.py
@@ -226,15 +226,10 @@ class MaxBackward:
             grad_output_shape.insert(self.axis, 1)
             grad_output = gradient.reshape(grad_output_shape)
             grad_output = grad_output + self.input[0].zeros_like()
-            
-            print(self.input[0])
-            max_value = self.input[0].max()
-            print(max_value)
             max_values = self.input[0].max(axis=self.axis, keepdim=True)
-            print('\n\n', max_values, '@@@')      
-            mask = self.input[0].equal(max_value)
+            mask = self.input[0].equal(max_values)
 
-            grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
+            grad_output = (grad_output * mask)
         
         return [grad_output]
 
@@ -256,7 +251,19 @@ class MinBackward:
             grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
 
         else:
-            pass
+            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]
+                
+            # 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)
+            mask = self.input[0].equal(max_values)
+
+            grad_output = (grad_output * mask)
 
         return [grad_output]
 
diff --git a/norch/tensor.py b/norch/tensor.py
index 22cc111..7e60d7f 100644
--- a/norch/tensor.py
+++ b/norch/tensor.py
@@ -270,6 +270,14 @@ class Tensor:
 
         if needs_broadcasting:
             # Call add_broadcasted_tensor if broadcasting is needed
+            if other.ndim == self.ndim - 1:
+                other = other.reshape([1] + other.shape)
+            
+            elif self.ndim == other.ndim - 1:
+                self = self.reshape([1] + self.shape)
+                
+            
+                
             Tensor._C.add_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
             Tensor._C.add_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
 
@@ -361,6 +369,12 @@ class Tensor:
         broadcasted_shape_sub, needs_broadcasting = broadcast_shape(self.shape, other.shape)
 
         if needs_broadcasting:
+            if other.ndim == self.ndim - 1:
+                other = other.reshape([1] + other.shape)
+            
+            elif self.ndim == other.ndim - 1:
+                self = self.reshape([1] + self.shape)
+
             # Call add_broadcasted_tensor if broadcasting is needed
             Tensor._C.sub_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
             Tensor._C.sub_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
@@ -678,6 +692,12 @@ class Tensor:
         broadcasted_shape_add, needs_broadcasting = broadcast_shape(self.shape, other.shape)
         
         if needs_broadcasting:
+            if other.ndim == self.ndim - 1:
+                other = other.reshape([1] + other.shape)
+            
+            elif self.ndim == other.ndim - 1:
+                self = self.reshape([1] + self.shape)
+                
             # Call equal_broadcasted_tensor if broadcasting is needed
             Tensor._C.equal_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
             Tensor._C.equal_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
diff --git a/test.py b/test.py
index 8dae5e0..2306f6a 100644
--- a/test.py
+++ b/test.py
@@ -6,12 +6,20 @@ import norch.optim as optim
 import random
 random.seed(1)
 
-a = norch.Tensor([1, 2, 3])
-b = 3
-for i in range(a.numel):
-    print(a.tensor.contents.data[i])
+torch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True)#.to(self.device)
+torch_tensor2 = norch.Tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]])
+torch_tensor3 = torch_tensor.max(axis=2, keepdim=True)
+print(torch_tensor3)
+print(torch_tensor.equal(torch_tensor3))
 
-print(a)
+"""print(torch_tensor.shape)
+print('\n\n')
+torch_tensor2 = torch_tensor.max(axis=1)
+print(torch_tensor2.shape)
+print('\n\n')
+c = torch_tensor + torch_tensor2
+print(c)
+"""
 
 """
 
diff --git a/tests/test_autograd.py b/tests/test_autograd.py
index 16e09bd..eb27544 100644
--- a/tests/test_autograd.py
+++ b/tests/test_autograd.py
@@ -115,24 +115,65 @@ class TestTensorAutograd(unittest.TestCase):
         
         self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
 
-
-        norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_tensor = norch.Tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
         norch_max_axis = norch_tensor.max(axis=2)
         norch_result = norch_max_axis.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_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
 
         torch_max_axis, _ = torch_tensor.max(axis=2)
         torch_result = torch_max_axis.sum()
         torch_result.backward()
         torch_tensor_grad = torch_tensor.grad
-        
         self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
 
-        print(norch_tensor_grad, '\n\n', torch_tensor_grad)
+
+    ## evaluate case with some repeated values and axis 2
+
+    #def test_max_axis(self):
+    #    """
+    #    Test autograd from max specifying axis
+    #    """
+    #
+    #    norch_tensor = norch.Tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
+    #    norch_max_axis = norch_tensor.max(axis=2)
+    #    norch_result = norch_max_axis.sum()
+    #
+    #    norch_result.backward()
+    #    norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
+    #
+    #    torch_tensor = torch.tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
+    #
+    #    torch_max_axis, _ = torch_tensor.max(axis=2)
+    #    torch_result = torch_max_axis.sum()
+    #    torch_result.backward()
+    #    torch_tensor_grad = torch_tensor.grad
+    #    self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
+
+
+    #def test_min_axis(self):
+    #    """
+    #    Test autograd from max specifying axis
+    #    """
+    #
+    #    norch_tensor = norch.Tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
+    #    norch_min_axis = norch_tensor.min(axis=2)
+    #    norch_result = norch_min_axis.sum()
+    #
+    #    norch_result.backward()
+    #    norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
+    #
+    #    torch_tensor = torch.tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
+    #
+    #    torch_min_axis, _ = torch_tensor.min(axis=2)
+    #    torch_result = torch_min_axis.sum()
+    #    torch_result.backward()
+    #    torch_tensor_grad = torch_tensor.grad
+    #    self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
+
 
     
     def test_min(self):
@@ -166,23 +207,23 @@ class TestTensorAutograd(unittest.TestCase):
 
         torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
         
-        torch_min = torch_tensor.min(axis=1)
+        torch_min, _ = torch_tensor.min(axis=1)
         torch_result = torch_min.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, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_tensor = norch.Tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
         norch_min = norch_tensor.min(axis=2)
         norch_result = norch_min.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_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
         
-        torch_min = torch_tensor.min(axis=2)
+        torch_min, _ = torch_tensor.min(axis=2)
         torch_result = torch_min.sum()
         torch_result.backward()
         torch_tensor_grad = torch_tensor.grad
diff --git a/tests/test_operations.py b/tests/test_operations.py
index 6a12367..8a684e9 100644
--- a/tests/test_operations.py
+++ b/tests/test_operations.py
@@ -52,6 +52,17 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
+        norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device)  # Shape (1, 2, 3)
+        norch_tensor2 = norch.Tensor([[10, 10], [5, 6]]).to(self.device)  # Shape (3)
+        norch_result = norch_tensor1 + norch_tensor2
+        torch_result = utils.to_torch(norch_result).to(self.device)
+        
+        torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([[[10, 10], [5, 6]]]).to(self.device)  # Shape (3)
+        torch_expected = torch_tensor1 + torch_tensor2
+
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+
         # reversed order broadcasting
         norch_tensor1 = norch.Tensor([[0, 2]]).to(self.device) 
         norch_tensor2 = norch.Tensor([[3, 4], [5, -1]]).to(self.device) 
@@ -274,8 +285,6 @@ class TestTensorOperations(unittest.TestCase):
         torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch.max(torch_tensor)
 
-        print(torch_result, torch_expected)
-
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
     def test_max_axis(self):
@@ -316,8 +325,6 @@ class TestTensorOperations(unittest.TestCase):
         torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch.min(torch_tensor)
 
-        print(torch_result, torch_expected)
-
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
     def test_min_axis(self):
@@ -561,7 +568,17 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
+        norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device)
+        norch_tensor2 = norch.Tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]]).to(self.device)
+        norch_result = norch_tensor1.equal(norch_tensor2)
+        torch_result = utils.to_torch(norch_result).to(self.device)
 
+        torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device)
+        torch_tensor2 = torch.tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]]).to(self.device)
+        torch_expected = (torch_tensor1 == torch_tensor2).float()
+
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+    
 
     def test_zeros_like(self):
         """