better documentation, removed old code

neuralforge/nn/layers.py

@@ -122,7 +122,6 @@ def forward(self, x):
         out = self.residual_proj(out) # (B, T, D) @ (D, D) -> (B, T, D)
         out = self.residual_dropout(out)
 
-        self.cache = (att, k, v, q)
         return out
 
 # Embedding Layers
@@ -142,7 +141,6 @@ def forward(self, idx):
         # Extracts embedding from row "idx":
         x = self.E[idx._data]
 
-        self.cache = (idx)
         return x
 
 

neuralforge/tensor_operations.py

@@ -816,17 +816,17 @@ def forward(self, a, condition, value):
         # Add new Tensors to "children" and old Tensors to "parents":
         self.parents = (a,)
         a.children.append(z)
-        self.cache = (a)
+        self.cache = (a, condition)
 
         return z 
 
     def backward(self, dz, z):
-        a = self.cache
+        a, condition = self.cache
 
         # Find gradients relative to "a", and pass it downstream:
         if a.requires_grad:
             # Because some activations are just set to a value, this operation is not differentiable.
-            da = dz
+            da = np.where(condition, dz, 0)
 
             a.backward(da, z)
 

tests/test_framework.py

@@ -1,5 +1,3 @@
-import sys
-sys.path.append('..')
 import neuralforge as forge
 import neuralforge.nn as nn
 import unittest
@@ -173,7 +171,7 @@ def forward(self, x):
         hidden_size = 128
         batch_size = 4
         n_heads = 8
-        dropout_p = 0.1
+        dropout_p = 0.2
 
         # Get path to root of repository:
         PATH = '/'.join(os.getcwd().split('/')[:-1])
@@ -189,7 +187,7 @@ def forward(self, x):
 
         # Define loss function and optimizer:
         loss_func = nn.CrossEntropyLoss()
-        optimizer = nn.optim.Adam(model.parameters(), lr=0.005, reg=0)
+        optimizer = nn.optim.Adam(model.parameters(), lr=5e-3, reg=0)
 
         # Training Loop:
         for _ in range(n_iters):
@@ -208,7 +206,6 @@ def forward(self, x):
 
             # Reset the gradients to zero after each training step:
             optimizer.zero_grad()
-            print(loss)
 
         assert loss._data < 1, "Error: Loss is not converging to zero in autograd test."