Some notes on MyGrad and Autograd

autograd_mygrad_compare.md

A Crude Prototype of Using MyGrad w/ xarray

import xarray as xr
import mygrad as mg

x = xr.DataArray(mg.arange(3.))
y = xr.DataArray(mg.arange(9.).reshape(3, 3))
z = x * y

>>> z
<xarray.DataArray (dim_0: 3, dim_1: 3)>
Tensor([[ 0.,  0.,  0.],
        [ 3.,  4.,  5.],
        [12., 14., 16.]])
Dimensions without coordinates: dim_0, dim_1

>>> z.data.backward()
>>> x.data.grad
>>> array([ 3., 12., 21.])
>>> y.data.grad
array([[0., 0., 0.],
       [1., 1., 1.],
       [2., 2., 2.]])

Getting a Derivative

What it looks like to get a derivative in autograd

# in autograd
import autograd.numpy as anp  # uses autograd's numpy
from autograd import elementwise_grad as egrad

>>> x = anp.array([2., 3.])

# need to define a function to differentiate
>>> def f(x): return anp.prod(x)
>>> egrad(f)(x)
array([3., 2.])

What it looks like to get a derivative in mygrad

# in mygrad
import numpy as np  # works with vanilla numpy
import mygrad as mg  

>>> x = mg.tensor([2., 3.])
>>> y = np.prod(x)
>>> y  # have access to intermediate tensors
Tensor(6.)
>>> y.backward() # trigger backprop from anywhere in "computational graph"
>>> x.grad
array([3., 2.])

Some places where autograd returns incorrect derivatives

Reported in 2018

# in autograd
x = anp.array([2.0, 3.0])
y = anp.array([3.0])


def f(x):
    return anp.einsum("i,i", x, y)

# should be array([3., 3.])
>>> egrad(f)(x)
array([3.])

# in mygrad
x = mg.tensor([2.0, 3.0])
y = np.array([3.0])
np.einsum("i, i", x, y).backward()
>>> x.grad
array([3., 3.])

Reported in 2018

# in autograd
x = anp.array([0., 1., 2.])

def f(x):
    return anp.prod(x)

>>> egrad(f)(x)
array([nan,  0.,  0.])

# in mygrad
x = mg.tensor([0., 1., 2.])
np.prod(x).backward()
>>> x.grad
array([2., 0., 0.])

Some derivatives not supported by autograd

Specifying dtype, where, or out in ufuncs:

>>> x = mg.tensor([1., 2., 3.])
>>> y = mg.zeros_like(x)
>>> np.multiply(x, x, where=[True, False, True], out=y, dtype="float32")
>>> y.backward()
>>> x.grad
array([2., 0., 6.])

In-place operations / rich support for views:

>>> x = mg.arange(16.).reshape(4, 4)
>>> y = +x
>>> bottom_right_corner = y[-2:, -2:]
>>> bottom_right_corner *= -y[:2, :2]  # x top-left corner
>>> y
Tensor([[  0.,   1.,   2.,   3.],
        [  4.,   5.,   6.,   7.],
        [  8.,   9.,  -0., -11.],
        [ 12.,  13., -56., -75.]])
>>> y.backward()
>>> x.grad
array([[ -9., -10.,   1.,   1.],
       [-13., -14.,   1.,   1.],
       [  1.,   1.,   0.,  -1.],
       [  1.,   1.,  -4.,  -5.]])

cumprod support:

# in mygrad
x = mg.tensor([-1., 2., 4.])
y = np.cumprod(x)  # cumprod not implemented in autograd 
y.backward()

>>> x.grad
array([11., -5., -2.])

einsum support:

# in mygrad
x = mg.arange(9.).reshape(3, 3)
y = mg.arange(3.)

# broadcast-reduction in einsum not supported in autograd
z = mg.einsum("...i,i ->", x, y).backward()  
>>> x.grad, y.grad
(array([[0., 1., 2.],
        [0., 1., 2.],
        [0., 1., 2.]]),
 array([ 9., 12., 15.]))

# traces in einsum not supported in autograd
mg.einsum("ii ->", x).backward()
>>> x.grad
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])