text_plain_output_1.png
import torch

class Var:
    """simple variable container"""
    def __init__(self, value):
        self.v = value
        self.grad = 0
        self.dag = None
    def __add__(self, other):
        """Overloads + operator"""
        if not isinstance(other, Var):
            other = Var(other)
        result = Var(self.v + other.v)
        result.dag = ('+', (self, other))
        return result
    def __radd__(self, other):
        """Reverse addition. Needed when user tries to calculate 2 + a"""
        return self + other
    def __pow__(self, other):
        """Overloads power operator"""
        assert not isinstance(other, Var), 'only support primitives yet'
        result = Var(self.v ** other)
        result.dag = ('^', (self, other))
        return result
    def __mul__(self, other):
        """Overloads multiplication operator"""
        if not isinstance(other, Var):
            other = Var(other)
        result = Var(self.v * other.v)
        result.dag = ('*', (self, other))
        return result
    def __rmul__(self, other):
        return self * other
    def __neg__(self):
        return self * -1
    def __sub__(self, other):
        return self + -other
    def __truediv__(self, other):
        return self * other ** (-1)
    def __rtruediv__(self, other):
        return other * self.v ** (-1)
    def __backward(self, seed):
        """this gets recursively called"""
        if self.dag is None:
            return
        a, b = (None, None)
        op, prevs = self.dag
        if len(prevs) == 2:
            a, b = prevs
        else:
            a = list(prevs)[0]
        if op == '+':
            a.grad += seed
            b.grad += seed
        elif op == '*':
            a.grad += b.v * seed
            b.grad += a.v * seed
        elif op == '^':
            a.grad += b * a.v ** (b - 1) * seed
        a.__backward(a.grad)
        if isinstance(b, Var):
            b.__backward(b.grad)
    def backward(self, seed=1.0):
        """
        do the backward pass and calculate parital derivative of each parent `Var`
        Parameters
        ----------
        seed (float): normally it should be set to 1 because dy/dy = 1.
            the derivatives of the child gets passed to the parents through this variable.
        """
        self.grad = 0
        self.__backward(seed)
        self.grad = seed
    def reset(self):
        """resets the gradient"""
        if self.dag is None:
            return
        a, b = (None, None)
        if len(self.dag[1]) == 2:
            a, b = self.dag[1]
        else:
            a = list(self.dag[1])[0]
        a.grad = 0
        a.reset()
        if b is not None:
            b.grad = 0
            b.reset()
    def __repr__(self):
        return f'[{self.v} | {self.grad}]'

def test_sanity_check(it=None):
    x = Var(-4.0)
    z = 2 * x + 2 + x
    y = z * x
    xmg, ymg = (x, y)
    x = torch.Tensor([-4.0]).double()
    x.requires_grad = True
    z = 2 * x + 2 + x
    y = z * x
    xpt, ypt = (x, y)
    assert ymg.v == ypt.data.item()
    ymg.backward()
    ypt.backward()
    assert xmg.grad == xpt.grad.item()
    a = Var(-3.0)
    a = -a + 20 - 2
    a.backward()
    b = torch.Tensor([-3.0]).double()
    b.requires_grad = True
    b = -b + 20 - 2
    b.backward()
    assert a.v == b.data.item()
    x = Var(2.0)
    y = x ** 3
    z = torch.tensor([2.0], requires_grad=True)
    f = z ** 3
    assert y.v == f.data.item()
    y.backward()
    f.backward()
    assert x.grad == z.grad.item()
test_sanity_check()

def gradTest():
    """test gradient calculations"""
    a = 6.0
    b = 2.2
    c = 3.0
    a1 = Var(a)
    b1 = Var(b)
    z1 = a1 * b1
    z1.backward()
    print(f'SG Gradients of z = ab is, a={a1.grad}, b={b1.grad}')
    a2 = torch.tensor(a, requires_grad=True)
    b2 = torch.tensor(b, requires_grad=True)
    z2 = a2 * b2
    z2.backward()
    print(f'Torch Gradients of z = ab is, a={a2.grad}, b={b2.grad}')
    z1.reset()
    print(a1.grad, b1.grad, len(z1.dag))
    c1 = Var(c)
    d1 = a1 + b1
    z1 = d1 * c1
    z1.backward()
    print(f'SG Gradients of z = (a+b)c = {z1.v} is, a={a1.grad}, b={b1.grad}, c={c1.grad}')
    a2 = torch.tensor(a, requires_grad=True)
    b2 = torch.tensor(b, requires_grad=True)
    c2 = torch.tensor(c, requires_grad=True)
    d2 = a2 + b2
    z2 = d2 * c2
    z2.backward()
    print(f'Torch Gradients of z = (a+b)c = {z2} is, a={a2.grad}, b={b2.grad}, c={c2.grad}')
    a1 = Var(a)
    z1 = a1 + a1
    z1.backward()
    tol = 1e-06
    a2 = torch.tensor([a], requires_grad=True)
    a3 = a2 + a2
    a3.retain_grad()
    a3.backward(gradient=torch.tensor([1.0]))
    assert a1.grad == a2.grad
gradTest()