_compgraph3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from IPython.core.debugger import set_trace\n",
    "from IPython.display import Markdown as md\n",
    "import inspect \n",
    "\n",
    "import numpy as _np, numpy as np \n",
    "path_assets = './assets/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the third part of a series on creating a computational graph in Python. Here is [part one](https://tomroth.com.au/compgraph1/) and here is [part two](https://tomroth.com.au/compgraph2/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the last post we created computational graphs for a function, but it was a bit hard to use. We also had these concerns: \n",
    "\n",
    "* we don't want to replace `np.add` with `add_new`, `np.exp` with `exp_new` etc everywhere. That's a pain, especially we have a lot of code to do that for. \n",
    "* currently we have to implement primitives for every `numpy` function we want. Is there a way to get them all?\n",
    "* how do we handle non-differentiable functions? \n",
    "\n",
    "We'll answer these questions here. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Essentially we create a fake `numpy` module called `numpy_autograd` with wrapped versions of all `numpy` functions. This fake `numpy` contains all the functions and objects of the original `numpy`, except some functions (only the differentiable ones) add to a computational graph as they are called.  Then by writing `import numpy_autograd as np`, any functions using numpy functions like `np.add` automatically build a computation graph as they are executed. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Non-differentiable functions "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Autodiff packages like `autograd` have to watch out for non-differentiable functions. Many functions are not differentiable, like `np.asarray`, `np.shape` or `np.argmin`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take `np.floor(x)` as an example. This is a non-differentiable function: its derivative does not exist for integer values of $x$, and the derivative is 0 everywhere else. So this is not something we'd add to the computation graph if we encountered it. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How should we deal with these functions? There are a few approaches. Some packages like `autograd` don't add them to the graph completely. The approach I take here is a bit different: I add them to the computation graph, but I'll add a flag `keepgrad` that indicates if the gradient of this function should be calculated or not. So let's go ahead and modify our `primitive` function from earlier to include this parameter: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def primitive(f, keepgrad=True): \n",
    "    @wraps(f)\n",
    "    def inner(*args, **kwargs):\n",
    "        ## Code to add operation/primitive to computation graph\n",
    "        # We need to separate out the integer/non node case. Sometimes you are adding \n",
    "        # constants to nodes. \n",
    "        def getval(o):      return o.value if type(o) == Node else o\n",
    "        if len(args):       argvals = [getval(o) for o in args]\n",
    "        else:               argvals = args\n",
    "        if len(kwargs):     kwargvals = dict([(k,getval(o)) for k,o in kwargs.items()])\n",
    "        else:               kwargvals =  kwargs\n",
    "         \n",
    "        # get parents \n",
    "        l = list(args) + list(kwargs.values())\n",
    "        parents = [o for o in l if type(o) == Node ]\n",
    "        \n",
    "        value = f(*argvals, **kwargvals)\n",
    "        print(\"add\", \"'\" + f.__name__ + \"'\", \"to graph with value\",value)\n",
    "        return Node(value, f, parents, keepgrad)\n",
    "    return inner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Triage "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A quick note. It can get confusing working with \"original\" `numpy` and \"new\" `numpy`, so note that throughout this post if you see something prefixed with `_np`, that means \"original\" `numpy`. Later I use `anp` to refer to \"new\" `numpy`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Anyway, it's time to create our version of `numpy`. All the attributes of `numpy` are available in `_np.__dict__`. We are going to split the objects in this dict into three categories:   \n",
    "\n",
    "a) differentiable functions (wrap with `primitive, keepgrad=True`  (the default for `primitive`)   \n",
    "b) non-differentiable functions (wrap with `primitive, keepgrad=False`)   \n",
    "c) everything else  (leave unchanged)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We create a function `wrap_namespace` that will copy everything from `_np.__dict__` into a new dictionary, wrapping functions based on the three categories above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as _np\n",
    "def wrap_namespace(old, new):\n",
    "    \"\"\"Performs triage on objects from numpy, copying them from old to new namespace. \n",
    "       old: __dict__ from original numpy\n",
    "       new: dict to copy old into \n",
    "       \"\"\"\n",
    "    # Taken from here: \n",
    "    # https://github.com/mattjj/autodidact/blob/b3b6e0c16863e6c7750b0fc067076c51f34fe271/autograd/numpy/numpy_wrapper.py#L8 \n",
    "    nograd_functions = [\n",
    "        _np.ndim, _np.shape, _np.iscomplexobj, _np.result_type, _np.zeros_like,\n",
    "        _np.ones_like, _np.floor, _np.ceil, _np.round, _np.rint, _np.around,\n",
    "        _np.fix, _np.trunc, _np.all, _np.any, _np.argmax, _np.argmin,\n",
    "        _np.argpartition, _np.argsort, _np.argwhere, _np.nonzero, _np.flatnonzero,\n",
    "        _np.count_nonzero, _np.searchsorted, _np.sign, _np.ndim, _np.shape,\n",
    "        _np.floor_divide, _np.logical_and, _np.logical_or, _np.logical_not,\n",
    "        _np.logical_xor, _np.isfinite, _np.isinf, _np.isnan, _np.isneginf,\n",
    "        _np.isposinf, _np.allclose, _np.isclose, _np.array_equal, _np.array_equiv,\n",
    "        _np.greater, _np.greater_equal, _np.less, _np.less_equal, _np.equal,\n",
    "        _np.not_equal, _np.iscomplexobj, _np.iscomplex, _np.size, _np.isscalar,\n",
    "        _np.isreal, _np.zeros_like, _np.ones_like, _np.result_type\n",
    "    ]\n",
    "    function_types = {_np.ufunc, types.FunctionType, types.BuiltinFunctionType}\n",
    "\n",
    "    for name,obj in old.items(): \n",
    "        if obj in nograd_functions:  \n",
    "            # non-differentiable functions \n",
    "            new[name] = primitive(obj, keepgrad=False)\n",
    "        elif type(obj) in function_types:  # functions with gradients \n",
    "            # differentiable functions\n",
    "            new[name] = primitive(obj)\n",
    "        else: \n",
    "            # just copy over \n",
    "            new[name] = obj\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating our new module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're ready to bring all our code together into a module. Here is an overview of the basic procedure. \n",
    "\n",
    "* make new folder `./numpy_autograd`\n",
    "* put stuff in a file `np_wrapping.py` with definitions of \n",
    "    * `primitive`\n",
    "    * `notrace_primitive`\n",
    "    * `Node` \n",
    "    * `wrap_namespace`\n",
    "* add implementations of the dunder methods to the `Node` class with `setattr`. \n",
    "* make `__init__.py` and at the top put `import * from np_wrapping`\n",
    "* now import using `import numpy_autograd as np` and you are done \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start with the first step. Ccreate a new folder called `numpy_autograd`. Then create a file called `np_wrapper.py`. Now we'll put in all functions we've defined so far here. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start off with imports and `primitive`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as _np\n",
    "import types \n",
    "from functools import wraps\n",
    "\n",
    "def primitive(f, keepgrad=True): \n",
    "    @wraps(f)\n",
    "    def inner(*args, **kwargs):\n",
    "        ## Code to add operation/primitive to computation graph\n",
    "        # We need to separate out the integer/non node case. Sometimes you are adding \n",
    "        # constants to nodes. \n",
    "        def getval(o):      return o.value if type(o) == Node else o\n",
    "        if len(args):       argvals = [getval(o) for o in args]\n",
    "        else:               argvals = args\n",
    "        if len(kwargs):     kwargvals = dict([(k,getval(o)) for k,o in kwargs.items()])\n",
    "        else:               kwargvals =  kwargs\n",
    "         \n",
    "        # get parents \n",
    "        l = list(args) + list(kwargs.values())\n",
    "        parents = [o for o in l if type(o) == Node ]\n",
    "        \n",
    "        value = f(*argvals, **kwargvals)\n",
    "        print(\"add\", \"'\" + f.__name__ + \"'\", \"to graph with value\",value)\n",
    "        return Node(value, f, parents, keepgrad)\n",
    "    return inner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now add the latest iteration of the `Node` class. It's similar to before, except it's modified to incorporate the `keepgrad` parameter. I've also moved the `start_node` function from last time to be a static method of the class, instead of having it float around by itself. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Node:\n",
    "    \"\"\"A node in a computation graph.\"\"\"\n",
    "    def __init__(self, value, fun, parents, keepgrad):\n",
    "        self.parents = parents\n",
    "        self.value = value\n",
    "        self.fun = fun \n",
    "        self.keepgrad = keepgrad\n",
    "        \n",
    "    def __repr__(self): \n",
    "        \"\"\"A (very) basic string representation\"\"\"\n",
    "        if self.value is None: str_val = 'None'\n",
    "        else:                  str_val = str(self.value)\n",
    "        return   \"\\n\" + \"Fun: \" + str(self.fun) +\\\n",
    "                \" Value: \"+ str_val + \\\n",
    "                \" Parents: \" + str(self.parents) \n",
    "    \n",
    "    def start_node(value = None, keepgrad=True): \n",
    "        \"\"\"A function to create an empty node to start off the graph\"\"\"\n",
    "        fun,parents = lambda x: x, []\n",
    "        return Node(value, fun, parents, keepgrad=True)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then the `wrap_namespace` function from earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wrap_namespace(old, new):\n",
    "    \"\"\"Performs triage on objects from numpy, copying them from old to new namespace. \n",
    "       old: __dict__ from original numpy\n",
    "       new: dict to copy old into \n",
    "       \"\"\"\n",
    "    # Taken from here: \n",
    "    # https://github.com/mattjj/autodidact/blob/b3b6e0c16863e6c7750b0fc067076c51f34fe271/autograd/numpy/numpy_wrapper.py#L8 \n",
    "    nograd_functions = [\n",
    "        _np.ndim, _np.shape, _np.iscomplexobj, _np.result_type, _np.zeros_like,\n",
    "        _np.ones_like, _np.floor, _np.ceil, _np.round, _np.rint, _np.around,\n",
    "        _np.fix, _np.trunc, _np.all, _np.any, _np.argmax, _np.argmin,\n",
    "        _np.argpartition, _np.argsort, _np.argwhere, _np.nonzero, _np.flatnonzero,\n",
    "        _np.count_nonzero, _np.searchsorted, _np.sign, _np.ndim, _np.shape,\n",
    "        _np.floor_divide, _np.logical_and, _np.logical_or, _np.logical_not,\n",
    "        _np.logical_xor, _np.isfinite, _np.isinf, _np.isnan, _np.isneginf,\n",
    "        _np.isposinf, _np.allclose, _np.isclose, _np.array_equal, _np.array_equiv,\n",
    "        _np.greater, _np.greater_equal, _np.less, _np.less_equal, _np.equal,\n",
    "        _np.not_equal, _np.iscomplexobj, _np.iscomplex, _np.size, _np.isscalar,\n",
    "        _np.isreal, _np.zeros_like, _np.ones_like, _np.result_type\n",
    "    ]\n",
    "    function_types = {_np.ufunc, types.FunctionType, types.BuiltinFunctionType}\n",
    "\n",
    "    for name,obj in old.items(): \n",
    "        if obj in nograd_functions:  \n",
    "            # non-differentiable functions \n",
    "            new[name] = primitive(obj, keepgrad=False)\n",
    "        elif type(obj) in function_types:  # functions with gradients \n",
    "            # differentiable functions\n",
    "            new[name] = primitive(obj)\n",
    "        else: \n",
    "            # just copy over \n",
    "            new[name] = obj\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we'll call `wrap_namespace()`. We'll hold the new functions in a dict called `anp`, which we init from the current value of `globals()`. Calling this `anp` is to make the next step (operator overloading) a bit clearer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# using globals() here means we can access each np function like np.add: \n",
    "# it means it is available to the global space. \n",
    "anp = globals()\n",
    "wrap_namespace(_np.__dict__, anp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally it's time for operator overloading. Instead of defining these all in `Node`, we use `setattr` to add each method one by one  to the `Node` class.  There are many more functions than before and many more dunder methods to match. There are also properties to deal with, like `np.ndim`, and we use the `property` keyword of Python to handle these. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Definitions taken from here:  \n",
    "## https://github.com/mattjj/autodidact/blob/b3b6e0c16863e6c7750b0fc067076c51f34fe271/autograd/numpy/numpy_boxes.py#L8\n",
    "setattr(Node, 'ndim', property(lambda self: self.value.ndim))\n",
    "setattr(Node, 'size', property(lambda self: self.value.size))\n",
    "setattr(Node, 'dtype',property(lambda self: self.value.dtype))\n",
    "setattr(Node, 'T', property(lambda self: anp['transpose'](self)))\n",
    "setattr(Node, 'shape', property(lambda self: self.value.shape))\n",
    "\n",
    "setattr(Node,'__len__', lambda self, other: len(self._value))\n",
    "setattr(Node,'astype', lambda self,*args,**kwargs: anp['_astype'](self, *args, **kwargs))\n",
    "setattr(Node,'__neg__', lambda self: anp['negative'](self))\n",
    "setattr(Node,'__add__', lambda self, other: anp['add'](     self, other))\n",
    "setattr(Node,'__sub__', lambda self, other: anp['subtract'](self, other))\n",
    "setattr(Node,'__mul__', lambda self, other: anp['multiply'](self, other))\n",
    "setattr(Node,'__pow__', lambda self, other: anp['power'](self, other))\n",
    "setattr(Node,'__div__', lambda self, other: anp['divide'](  self, other))\n",
    "setattr(Node,'__mod__', lambda self, other: anp['mod'](     self, other))\n",
    "setattr(Node,'__truediv__', lambda self, other: anp['true_divide'](self, other))\n",
    "setattr(Node,'__matmul__', lambda self, other: anp['matmul'](self, other))\n",
    "setattr(Node,'__radd__', lambda self, other: anp['add'](     other, self))\n",
    "setattr(Node,'__rsub__', lambda self, other: anp['subtract'](other, self))\n",
    "setattr(Node,'__rmul__', lambda self, other: anp['multiply'](other, self))\n",
    "setattr(Node,'__rpow__', lambda self, other: anp['power'](   other, self))\n",
    "setattr(Node,'__rdiv__', lambda self, other: anp['divide'](  other, self))\n",
    "setattr(Node,'__rmod__', lambda self, other: anp['mod'](     other, self))\n",
    "setattr(Node,'__rtruediv__', lambda self, other: anp['true_divide'](other, self))\n",
    "setattr(Node,'__rmatmul__', lambda self, other: anp['matmul'](other, self))\n",
    "setattr(Node,'__eq__', lambda self, other: anp['equal'](self, other))\n",
    "setattr(Node,'__ne__', lambda self, other: anp['not_equal'](self, other))\n",
    "setattr(Node,'__gt__', lambda self, other: anp['greater'](self, other))\n",
    "setattr(Node,'__ge__', lambda self, other: anp['greater_equal'](self, other))\n",
    "setattr(Node,'__lt__', lambda self, other: anp['less'](self, other))\n",
    "setattr(Node,'__le__', lambda self, other: anp['less_equal'](self, other))\n",
    "setattr(Node,'__abs__', lambda self: anp['abs'](self))\n",
    "setattr(Node,'__hash__', lambda self: id(self))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, make another file `__init__.py` in the  `numpy_autograd` folder and add this to it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#from .np_wrapping import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you're done! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the wrapped numpy "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test it out: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "add 'exp' to graph with value 0.22313016014842982\n",
      "add 'add' to graph with value 1.22313016014843\n",
      "add 'true_divide' to graph with value 0.8175744761936437\n"
     ]
    }
   ],
   "source": [
    "import numpy_autograd as np\n",
    "def logistic(z):  return 1 / (1 + np.exp(-z))\n",
    "a1 = logistic(1.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This works fine. What about something using a non-differentiable function?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "add 'log' to graph with value 1.6094379124341003\n",
      "add 'exp' to graph with value 2.718281828459045\n",
      "add 'multiply' to graph with value 4.374905831402675\n",
      "add 'add' to graph with value 9.374905831402675\n",
      "add 'floor' to graph with value 9.0\n",
      "add 'exp' to graph with value 148.4131591025766\n",
      "add 'add' to graph with value 157.4131591025766\n"
     ]
    }
   ],
   "source": [
    "def f2(x,y): return np.add(np.floor(np.log(x) * np.exp(y) + x*y), np.exp(x))\n",
    "a2 = f2(5,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice! And that's how to create a computation graph. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}