Enable micrograd to use tanh activation function dynamically

micrograd/engine.py

@@ -1,14 +1,16 @@
+import math
 
 class Value:
     """ stores a single scalar value and its gradient """
 
-    def __init__(self, data, _children=(), _op=''):
+    def __init__(self, data, _children=(), _op='', label=''):
         self.data = data
         self.grad = 0
         # internal variables used for autograd graph construction
         self._backward = lambda: None
         self._prev = set(_children)
         self._op = _op # the op that produced this node, for graphviz / debugging / etc
+        self.label = label # label of this node
 
     def __add__(self, other):
         other = other if isinstance(other, Value) else Value(other)
@@ -50,6 +52,18 @@ def _backward():
         out._backward = _backward
 
         return out
+
+    def tanh(self):
+        x = self.data
+        t = (math.exp(2*x) - 1)/(math.exp(2*x) + 1)
+        out = Value(t, (self, ), 'tanh')
+
+        def _backward():
+            self.grad += (1 - t**2) * out.grad
+        out._backward = _backward
+
+        return out
+
 
     def backward(self):
 

micrograd/nn.py

@@ -12,19 +12,30 @@ def parameters(self):
 
 class Neuron(Module):
 
-    def __init__(self, nin, nonlin=True):
+    def __init__(self, nin, nonlin=True, use_tanh=False):
         self.w = [Value(random.uniform(-1,1)) for _ in range(nin)]
         self.b = Value(0)
         self.nonlin = nonlin
+        self.use_tanh = use_tanh
 
     def __call__(self, x):
         act = sum((wi*xi for wi,xi in zip(self.w, x)), self.b)
+
+        if self.use_tanh:
+            # If specified to use a Tanh activation function
+            return act.tanh() if self.nonlin else act
+
+        # By default uses a relu activation function.
+        # Maybe we can extend this dynamically to pick activation
+        # functions at the caller end.
         return act.relu() if self.nonlin else act
 
     def parameters(self):
         return self.w + [self.b]
 
     def __repr__(self):
+        if self.use_tanh:
+            return f"{'Tanh' if self.nonlin else 'Linear'}Neuron({len(self.w)})"
         return f"{'ReLU' if self.nonlin else 'Linear'}Neuron({len(self.w)})"
 
 class Layer(Module):
@@ -44,9 +55,9 @@ def __repr__(self):
 
 class MLP(Module):
 
-    def __init__(self, nin, nouts):
+    def __init__(self, nin, nouts, use_tanh=False):
         sz = [nin] + nouts
-        self.layers = [Layer(sz[i], sz[i+1], nonlin=i!=len(nouts)-1) for i in range(len(nouts))]
+        self.layers = [Layer(sz[i], sz[i+1], nonlin=i!=len(nouts)-1, use_tanh=use_tanh) for i in range(len(nouts))]
 
     def __call__(self, x):
         for layer in self.layers: