add autograd and jax support

[jcmgray]

Description

This adds support for using autograd and jax. These both allow automatically computing gradients of functions with scalar output. jax additionally compiles the computations for GPU and faster runtime.

Essentially all that was required was just adding the aliases {'jax': 'jax.numpy', 'autograd': 'autograd.numpy'} to backend dispatch mechanism but I've also added explicit tests and jax as a compiled expression backend for numpy arrays (since it seems to have very good performance).

Finally, I've beefed up the backend tests a bit and fixed an issue with tensorflow on travis to make sure those tests are not skipped.

Examples

Compute the gradient of an arbitrary tensor contraction:

import numpy as np
import opt_einsum as oe
from autograd import grad

eq = 'ij,jk,ki->'
shapes = [(2, 3), (3, 4), (4, 2)]


def foo(arrays):
    return oe.contract(eq, *arrays)

# this function now computes the gradient (jacobian) of foo
dfoo = grad(foo)

>>> arrays = [np.random.rand(*s) for s in shapes]
>>> foo(arrays)
array(2.97556246)

>>> dfoo(arrays)
[array([[1.03736352, 1.69989611, 0.99103317],
        [0.61444428, 0.50011247, 0.33026913]]),
 array([[0.31924389, 0.48004359, 1.00855222, 0.55368964],
        [0.3843193 , 0.57097258, 0.74804434, 0.74309046],
        [0.32953833, 0.4940919 , 0.94472643, 0.58736502]]),
 array([[0.96326285, 0.61466171],
        [0.56535028, 0.12002171],
        [0.87289287, 0.53918698],
        [0.74630816, 0.21253244]])]

Compile both the function and gradient using jax:

import jax

jit_foo = jax.jit(foo)
jit_dfoo = jax.jit(jax.grad(foo))

>>> jit_foo(arrays)
DeviceArray(2.9755626, dtype=float32)

>>> jit_dfoo(arrays)
[DeviceArray([[1.03736353, 1.6998961 , 0.9910332 ],
              [0.61444432, 0.50011247, 0.33026916]], dtype=float32),
 DeviceArray([[0.31924388, 0.48004356, 1.00855219, 0.5536896 ],
              [0.38431931, 0.57097256, 0.74804437, 0.74309039],
              [0.32953838, 0.49409193, 0.94472641, 0.58736503]],
             dtype=float32),
 DeviceArray([[0.96326286, 0.61466169],
              [0.56535023, 0.12002171],
              [0.87289286, 0.53918701],
              [0.74630815, 0.21253243]], dtype=float32)]

Or, regardless of gradients, just compile a contraction using jax (that still accepts and gives out numpy arrays):

>>> expr = oe.contract_expression('ij,jk,kl->li', *shapes)
>>> expr(*arrays, backend='jax')
array([[2.4497495, 1.1122954],
       [0.9125154, 0.525813 ]], dtype=float32)

Todos

• Write some docs up (note somewhere that jax by default converts everything to single precision)

Status

• Ready to go

[codecov-io]

Codecov Report

Merging #88 into master will decrease coverage by 0.16%.
The diff coverage is 88.23%.

[dgasmith]

LGTM! I see no issues here, any hindrance to merging?

[jcmgray]

I might just add some docs when I get the chance. Maybe a bullet point on the readme (since automatic differentiation is very nice for optimization etc) and a little section in the 'backends' bit?

[dgasmith]

Sounds good! I am trying to clear some time to get back to this and push a 3.0 out.

[jcmgray]

OK I think all good to go from my end if the docs look good. A 3.0 release sounds great!

[dgasmith]

A few minor comments, feel free to merge if you are ok with them.

This is really awesome to add to the project. Kind of amazing the underlying architecture can handle something like autograd pretty seamlessly.

[dgasmith]

Thats pretty cool that all of this works.

[dgasmith]

It might be good to use the same random seed above and here to show that they produce the same thing?

[jcmgray]

Ah good catch, I thought I had updated them with the same arrays but obviously forgot. Fixed now.

[jcmgray]

There were some new test errors coming from pytest v5.0 that I've also just fixed.

[mattjj]

Woo awesome!

These both allow automatically computing gradients of functions with scalar output.

Actually both Autograd and JAX do a whole lot more than that! Jacobians, Hessians, whatever you like! See for example the JAX Autodiff Cookbook (part 1).

Thanks for adding this. The future of numerical computing is bright :)

[dgasmith]

Thanks for the nice comments! Happy to take a PR that expands on what JAX can do so that we can better highlight the library.