From 51cb2f49d8f27f696021d534ccef8a6549b7be57 Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Wed, 9 May 2018 14:55:14 -0400 Subject: [PATCH 1/2] Support Numpy-like arrays Several libraries exist that support a broad subset of the Numpy API, but with different underlying data structures: - CuPy for GPU arrays - Sparse arrays - Dask array for parallel and distributed arrays These all support numpy functions like np.sin directly Previously these would not work with autograd because their types and functions were not registered. However, to the extent that they support the numpy API directly (which is increasingly becoming the case) we can just hijack the existing function mappings for Numpy functions. This commit replaces two explicit type checks with more flexible tests that check for "array-like" behavior. If so then they return the corresponding numpy function. --- .travis.yml | 6 +++--- autograd/core.py | 4 ++-- autograd/tracer.py | 6 ++++-- autograd/util.py | 14 ++++++++++++++ tests/test_numpy_like.py | 28 ++++++++++++++++++++++++++++ 5 files changed, 51 insertions(+), 7 deletions(-) create mode 100644 tests/test_numpy_like.py diff --git a/.travis.yml b/.travis.yml index 162163d9a..01108940d 100644 --- a/.travis.yml +++ b/.travis.yml @@ -6,8 +6,8 @@ python: - "2.7" - "3.6" env: - - DEPS="pip nose future numpy scipy" - - DEPS="pip nose future numpy" + - DEPS="pip nose future numpy scipy dask[array]" + - DEPS="pip nose future numpy dask[array]" before_install: - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then wget https://repo.continuum.io/miniconda/Miniconda2-latest-Linux-x86_64.sh -O miniconda.sh; @@ -21,6 +21,6 @@ before_install: install: - conda install --yes python=$TRAVIS_PYTHON_VERSION $DEPS - pip install -v . -script: +script: - cd tests # Run from inside tests directory to make sure Autograd has - nosetests # fully installed. diff --git a/autograd/core.py b/autograd/core.py index 77a041ada..fa148d143 100644 --- a/autograd/core.py +++ b/autograd/core.py @@ -1,7 +1,7 @@ from itertools import count from functools import reduce from .tracer import trace, primitive, toposort, Node, Box, isbox, getval -from .util import func, subval +from .util import func, subval, typeof # -------------------- reverse mode -------------------- @@ -230,7 +230,7 @@ def register(cls, value_type, vspace_maker=None): def vspace(value): try: - return VSpace.mappings[type(value)](value) + return VSpace.mappings[typeof(value)](value) except KeyError: if isbox(value): return vspace(getval(value)) diff --git a/autograd/tracer.py b/autograd/tracer.py index 3dfc7fd84..8dffea589 100644 --- a/autograd/tracer.py +++ b/autograd/tracer.py @@ -1,9 +1,11 @@ import warnings from contextlib import contextmanager from collections import defaultdict -from .util import subvals, toposort +from .util import subvals, toposort, typeof from .wrap_util import wraps +import numpy + def trace(start_node, fun, x): with trace_stack.new_trace() as t: start_box = new_box(x, t, start_node) @@ -115,7 +117,7 @@ def register(cls, value_type): box_type_mappings = Box.type_mappings def new_box(value, trace, node): try: - return box_type_mappings[type(value)](value, trace, node) + return box_type_mappings[typeof(value)](value, trace, node) except KeyError: raise TypeError("Can't differentiate w.r.t. type {}".format(type(value))) diff --git a/autograd/util.py b/autograd/util.py index 4cfd7a369..f20f8e078 100644 --- a/autograd/util.py +++ b/autograd/util.py @@ -38,6 +38,20 @@ def toposort(end_node, parents=operator.attrgetter('parents')): else: child_counts[parent] -= 1 + +def typeof(x): + """ + A Modified type function that returns np.ndarray for any array-like + + This improves portability of autograd to other projects that might support + the numpy API, despite not being exactly numpy. + """ + if all(hasattr(x, attr) for attr in ['__array_ufunc__', 'shape', 'dtype']): + import numpy + return numpy.ndarray + else: + return type(x) + # -------------------- deprecation warnings ----------------------- import warnings diff --git a/tests/test_numpy_like.py b/tests/test_numpy_like.py new file mode 100644 index 000000000..ad023920e --- /dev/null +++ b/tests/test_numpy_like.py @@ -0,0 +1,28 @@ +from __future__ import absolute_import +import warnings + +import autograd.numpy as np +import autograd.numpy.random as npr +from autograd.test_util import check_grads +from autograd import grad +from numpy_utils import combo_check + +from dask.array.utils import assert_eq +import dask.array as da + +npr.seed(1) + +def test_dask(): + x = np.arange(10) + xx = da.arange(10, chunks=(5,)) + + assert_eq(x, xx) + + def f(x): + return np.sin(x).sum() + + f_prime = grad(f) + + assert isinstance(f_prime(xx), type(xx)) + + assert_eq(f_prime(x), f_prime(xx)) From 398ab5132326514b68a6060dedebb097b66bbfac Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Wed, 9 May 2018 17:44:43 -0400 Subject: [PATCH 2/2] Only convert to numpy array if shape and dtype are not present --- autograd/numpy/numpy_vspaces.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/autograd/numpy/numpy_vspaces.py b/autograd/numpy/numpy_vspaces.py index 8eda1b2a2..6bcd096c3 100644 --- a/autograd/numpy/numpy_vspaces.py +++ b/autograd/numpy/numpy_vspaces.py @@ -3,7 +3,8 @@ class ArrayVSpace(VSpace): def __init__(self, value): - value = np.array(value, copy=False) + if not hasattr(value, 'shape') or not hasattr(value, 'dtype'): + value = np.array(value, copy=False) self.shape = value.shape self.dtype = value.dtype