aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

batched broadcasted autograd

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-16 21:25:03 -03:00
parent 3627e42469
commit 79de915bb2
2 changed files with 4 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
norch/autograd/__pycache__/functions.cpython-38.pyc
View file

Binary file not shown.

6
norch/autograd/functions.py
View file

@ -83,8 +83,10 @@ class MatmulBackward:
def backward(self, gradient):
x, y = self.input
if x.shape != y.shape: # broadcasted case
return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
if x.ndim != y.ndim: # broadcasted case
print("KKKKKKKKK", x.ndim, y.ndim)
aux = (gradient @ y.transpose(-1,-2))
return [aux.sum(axis=0), x.transpose(-1,-2) @ gradient]
else:
return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]