[
{
"path": ".github/workflows/check.yml",
"content": "name: Style and package checks\n\non:\n pull_request:\n branches:\n - master\n push:\n branches:\n - master\n workflow_dispatch:\n\nenv:\n PIP_DISABLE_PIP_VERSION_CHECK: \"1\"\n FORCE_COLOR: \"3\"\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}\n cancel-in-progress: true\n\njobs:\n check:\n name: ${{ matrix.env }}\n runs-on: ubuntu-latest\n strategy:\n fail-fast: false\n matrix:\n session:\n # - lint\n - validate-package\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0\n\n - uses: yezz123/setup-uv@ab6be5a42627f19dc36e57b548592a5e52cece4a # v4.1\n\n - name: Run ${{ matrix.env }}\n run: uvx nox -s ${{ matrix.env }}\n"
},
{
"path": ".github/workflows/publish.yml",
"content": "name: Publish\n\non:\n workflow_dispatch:\n release:\n types: [published]\n\nenv:\n PIP_DISABLE_PIP_VERSION_CHECK: '1'\n FORCE_COLOR: '3'\n\njobs:\n build:\n name: Build sdist and wheel\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n name: Checkout repository\n\n - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0\n with:\n python-version: \"3.12\"\n\n - name: Install build tools\n run: |\n pipx run build --outdir dist\n\n - name: Upload wheel and sdist artifacts\n uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2\n with:\n name: artifacts\n path: ./dist/*\n if-no-files-found: error\n\n publish:\n needs: [build]\n name: Upload to PyPI\n runs-on: ubuntu-latest\n environment:\n name: release\n url: https://pypi.org/p/autograd\n permissions:\n id-token: write # mandatory for trusted publishing\n\n steps:\n - name: Download artifacts\n uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0\n with:\n path: dist\n merge-multiple: true\n\n - name: Sanity check artifacts\n run: ls -la dist/\n\n - name: Publish sdist and wheel to PyPI\n uses: pypa/gh-action-pypi-publish@76f52bc884231f62b9a034ebfe128415bbaabdfc # v1.12.4\n with:\n packages-dir: dist/\n"
},
{
"path": ".github/workflows/test.yml",
"content": "name: CI\n\non:\n pull_request:\n branches:\n - master\n push:\n branches:\n - master\n workflow_dispatch:\n schedule:\n - cron: \"0 4 * * *\"\n\nenv:\n PIP_DISABLE_PIP_VERSION_CHECK: \"1\"\n FORCE_COLOR: \"3\"\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}\n cancel-in-progress: true\n\njobs:\n test:\n name: Regular tests / ${{ matrix.platform }} / Python ${{ matrix.python-version }}\n runs-on: ${{ matrix.platform }}\n strategy:\n fail-fast: false\n matrix:\n platform: [ubuntu-latest, ubuntu-22.04-arm, macos-15-intel, macos-latest, windows-latest]\n python-version:\n [\"3.10\", \"3.11\", \"3.12\", \"3.13\", \"3.14\", \"pypy-3.10\"]\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0\n with:\n python-version: ${{ matrix.python-version }}\n allow-prereleases: true\n - uses: yezz123/setup-uv@ab6be5a42627f19dc36e57b548592a5e52cece4a # v4.1\n\n # On PyPy, we skip SciPy because we don't have wheels\n # available, see noxfile.py for more details.\n - name: Run tests\n run: uvx nox -s tests\n\n # In this job, we test against the NumPy nightly wheels hosted on\n # https://anaconda.org/scientific-python-nightly-wheels/numpy\n # on the latest Python version available across platforms, instead of\n # testing all Python versions and implementations on all platforms.\n # We do not test on PyPy.\n #\n # However, \"nox -s nightly-tests\" can be used locally anywhere, on\n # any Python version and implementation on any platform and we leave\n # it to the user to decide what Python version to test against, which\n # might or might not have a corresponding NumPy nightly wheel present.\n nightlies:\n name: Nightly tests / ${{ matrix.platform }} / Python ${{ matrix.python-version }}\n runs-on: ${{ matrix.platform }}\n strategy:\n fail-fast: false\n matrix:\n platform: [ubuntu-latest, ubuntu-22.04-arm, macos-15-intel, macos-latest, windows-latest]\n python-version: [\"3.x\"]\n\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0\n with:\n python-version: ${{ matrix.python-version }}\n allow-prereleases: true\n - uses: yezz123/setup-uv@ab6be5a42627f19dc36e57b548592a5e52cece4a # v4.1\n - name: Run tests against nightly wheels for NumPy and SciPy\n run: uvx nox -s nightly-tests\n"
},
{
"path": ".gitignore",
"content": "__pycache__/\n*.py[cod]\n*$py.class\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\ncoverage.*\n*.cover\n.hypothesis/\nnosetests.xml\n.pytest_cache/\njunit-report.xml\n\n# pyenv\n.python-version\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# mypy\n.mypy_cache/\n\n# OS and IDE config files\n.DS_Store\n.idea/\n\n# project-specific\ndata/\n*.so\n*.c\nscratch/\nexamples/data\n\n.asv/\nasv.conf.json\nbenchmarks/asv.conf.js\n"
},
{
"path": ".pre-commit-config.yaml",
"content": "ci:\n autoupdate_commit_msg: \"chore: update pre-commit hooks\"\n autofix_commit_msg: \"style: pre-commit fixes\"\n\nrepos:\n - repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v6.0.0\n hooks:\n - id: check-added-large-files\n - id: check-case-conflict\n - id: check-merge-conflict\n - id: check-yaml\n exclude: conda_recipe/conda.yaml\n - id: debug-statements\n - id: end-of-file-fixer\n - id: mixed-line-ending\n - id: trailing-whitespace\n\n - repo: https://github.com/asottile/pyupgrade\n rev: v3.21.2\n hooks:\n - id: pyupgrade\n args: [--py310-plus]\n\n - repo: https://github.com/astral-sh/ruff-pre-commit\n rev: \"v0.15.6\"\n hooks:\n - id: ruff\n args: [\"--fix\", \"--show-fixes\"]\n - id: ruff-format\n\n - repo: https://github.com/pre-commit/pygrep-hooks\n rev: v1.10.0\n hooks:\n - id: python-check-blanket-type-ignore\n exclude: ^src/vector/backends/_numba_object.py$\n - id: rst-backticks\n - id: rst-directive-colons\n - id: rst-inline-touching-normal\n"
},
{
"path": "CONTRIBUTING.md",
"content": "# Contributing\n\nUse [Nox](https://nox.thea.codes/en/stable/) to run tests and linting, e.g.,\n\n```shell\npip install nox\n```\n\n`nox` will run all checks in an isolated virtual environment with Autograd and its dependencies, including its optional dependencies, installed.\n\n## Run tests, linting, packaging checks\n\n| Command | Description |\n| ------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| `nox --list` | Lists all available Nox sessions, including selected ones |\n| `nox -s lint` | Runs code style checks with pre-commit and pre-commit hooks as listed in `.pre-commit-config.yaml`. Accepts posargs to pass additional arguments to the linter. |\n| `nox -s tests` | Runs tests with your default Python interpreter. Accepts posargs to pass additional arguments and configuration to `pytest`. |\n| `nox -s nightly-tests` | Similar to `nox -s tests`, except that it runs tests with nightly versions of dependencies (NumPy, SciPy, etc.). |\n| `nox -s validate-package` | Builds a source distribution and a wheel using `pypa/build` and checks the package with `twine` in strict mode. |\n| `nox` | Runs all selected sessions, as listed in `nox.options.sessions` in `noxfile.py`. |\n\nAdditionally, `nox` supports tags to run specific sessions, e.g., `nox --tags tests` runs all sessions tagged with `tests`.\n\nMake sure all tests pass before you push your changes to GitHub.\nGH Actions will run the tests across all supported Python versions.\n\n## Using positional arguments (reformat, upload package, help)\n\nYou can use additional arguments for the tools (`pytest`, `pre-commit`, etc.) called by Nox by\nseparating them from the Nox arguments by a double-hyphen `--`, e.g.,\n\n- `nox -s tests -- tests/test_tuple.py` runs just the tests listed `tests/test_tuple.py`.\n- `nox -s lint -- --fix` runs the linter with the `--fix` flag.\n- and so on.\n"
},
{
"path": "README.md",
"content": "# Autograd [![Checks status][checks-badge]][checks-url] [![Tests status][tests-badge]][tests-url] [![Publish status][publish-badge]][publish-url] [![asv][asv-badge]](#)\n\n[publish-badge]: https://github.com/HIPS/autograd/actions/workflows/publish.yml/badge.svg\n[checks-badge]: https://github.com/HIPS/autograd/actions/workflows/check.yml/badge.svg\n[tests-badge]: https://github.com/HIPS/autograd/actions/workflows/test.yml/badge.svg\n[asv-badge]: http://img.shields.io/badge/benchmarked%20by-asv-green.svg?style=flat\n[publish-url]: https://github.com/HIPS/autograd/actions/workflows/publish.yml\n[checks-url]: https://github.com/HIPS/autograd/actions/workflows/check.yml\n[tests-url]: https://github.com/HIPS/autograd/actions/workflows/test.yml\n\nAutograd can automatically differentiate native Python and Numpy code. It can\nhandle a large subset of Python's features, including loops, ifs, recursion and\nclosures, and it can even take derivatives of derivatives of derivatives. It\nsupports reverse-mode differentiation (a.k.a. backpropagation), which means it\ncan efficiently take gradients of scalar-valued functions with respect to\narray-valued arguments, as well as forward-mode differentiation, and the two can\nbe composed arbitrarily. The main intended application of Autograd is\ngradient-based optimization. For more information, check out the\n[tutorial](docs/tutorial.md) and the [examples directory](examples/).\n\nExample use:\n\n```python\n>>> import autograd.numpy as np # Thinly-wrapped numpy\n>>> from autograd import grad # The only autograd function you may ever need\n>>>\n>>> def tanh(x): # Define a function\n... return (1.0 - np.exp((-2 * x))) / (1.0 + np.exp(-(2 * x)))\n...\n>>> grad_tanh = grad(tanh) # Obtain its gradient function\n>>> grad_tanh(1.0) # Evaluate the gradient at x = 1.0\nnp.float64(0.419974341614026)\n>>> (tanh(1.0001) - tanh(0.9999)) / 0.0002 # Compare to finite differences\nnp.float64(0.41997434264973155)\n```\n\nWe can continue to differentiate as many times as we like, and use numpy's\nvectorization of scalar-valued functions across many different input values:\n\n```python\n>>> from autograd import elementwise_grad as egrad # for functions that vectorize over inputs\n>>> import matplotlib.pyplot as plt\n>>> x = np.linspace(-7, 7, 700)\n>>> plt.plot(x, tanh(x),\n... x, egrad(tanh)(x), # first derivative\n... x, egrad(egrad(tanh))(x), # second derivative\n... x, egrad(egrad(egrad(tanh)))(x), # third derivative\n... x, egrad(egrad(egrad(egrad(tanh))))(x),) # fourth derivative\n>>> plt.show()\n```\n\n![](\"examples/tanh.png\")
\n\nSee the [tanh example file](examples/tanh.py) for the code.\n\n## Documentation\n\nYou can find a tutorial [here.](docs/tutorial.md)\n\n## End-to-end examples\n\n* [Simple neural net](examples/neural_net.py)\n* [Convolutional neural net](examples/convnet.py)\n* [Recurrent neural net](examples/rnn.py)\n* [LSTM](examples/lstm.py)\n* [Neural Turing Machine](https://github.com/DoctorTeeth/diffmem/blob/512aadeefd6dbafc1bdd253a64b6be192a435dc3/ntm/ntm.py)\n* [Backpropagating through a fluid simulation](examples/fluidsim/fluidsim.py)\n\n![](\"examples/fluidsim/animated.gif\")
\n\n* [Variational inference in Bayesian neural network](examples/bayesian_neural_net.py)\n* [Gaussian process regression](examples/gaussian_process.py)\n* [Sampyl, a pure Python MCMC package with HMC and NUTS](https://github.com/mcleonard/sampyl)\n\n## How to install\n\nInstall Autograd using Pip:\n\n```shell\npip install autograd\n```\n\nSome features require SciPy, which you can install separately or as an\noptional dependency along with Autograd:\n\n```shell\npip install \"autograd[scipy]\"\n```\n\n## Authors and maintainers\n\nAutograd was written by [Dougal Maclaurin](https://dougalmaclaurin.com),\n[David Duvenaud](https://www.cs.toronto.edu/~duvenaud/),\n[Matt Johnson](http://people.csail.mit.edu/mattjj/),\n[Jamie Townsend](https://github.com/j-towns)\nand many other contributors. The package is currently being maintained by\n[Agriya Khetarpal](https://github.com/agriyakhetarpal),\n[Fabian Joswig](https://github.com/fjosw) and\n[Jamie Townsend](https://github.com/j-towns).\nPlease feel free to submit any bugs or\nfeature requests. We'd also love to hear about your experiences with Autograd\nin general. Drop us an email!\n\nWe want to thank Jasper Snoek and the rest of the HIPS group (led by Prof. Ryan\nP. Adams) for helpful contributions and advice; Barak Pearlmutter for\nfoundational work on automatic differentiation and for guidance on our\nimplementation; and Analog Devices Inc. (Lyric Labs) and Samsung Advanced Institute\nof Technology for their generous support.\n"
},
{
"path": "autograd/__init__.py",
"content": "from autograd.core import primitive_with_deprecation_warnings as primitive\n\nfrom .builtins import dict, isinstance, list, tuple, type\nfrom .differential_operators import (\n checkpoint,\n deriv,\n elementwise_grad,\n grad,\n grad_and_aux,\n grad_named,\n hessian,\n hessian_tensor_product,\n hessian_vector_product,\n holomorphic_grad,\n jacobian,\n make_ggnvp,\n make_hvp,\n make_jvp,\n make_vjp,\n multigrad_dict,\n tensor_jacobian_product,\n value_and_grad,\n vector_jacobian_product,\n)\n"
},
{
"path": "autograd/builtins.py",
"content": "from .extend import (\n Box,\n SparseObject,\n VSpace,\n defjvp,\n defjvp_argnum,\n defvjp,\n defvjp_argnum,\n notrace_primitive,\n primitive,\n vspace,\n)\nfrom .util import subvals\n\nisinstance_ = isinstance\nisinstance = notrace_primitive(isinstance)\n\ntype_ = type\ntype = notrace_primitive(type)\n\ntuple_, list_, dict_ = tuple, list, dict\n\n\n@primitive\ndef container_take(A, idx):\n return A[idx]\n\n\ndef grad_container_take(ans, A, idx):\n return lambda g: container_untake(g, idx, vspace(A))\n\n\ndefvjp(container_take, grad_container_take)\ndefjvp(container_take, \"same\")\n\n\nclass SequenceBox(Box):\n __slots__ = []\n __getitem__ = container_take\n\n def __len__(self):\n return len(self._value)\n\n def __add__(self, other):\n return sequence_extend_right(self, *other)\n\n def __radd__(self, other):\n return sequence_extend_left(self, *other)\n\n def __contains__(self, elt):\n return elt in self._value\n\n def index(self, elt):\n return self._value.index(elt)\n\n\nSequenceBox.register(tuple_)\nSequenceBox.register(list_)\n\n\nclass DictBox(Box):\n __slots__ = []\n __getitem__ = container_take\n\n def __len__(self):\n return len(self._value)\n\n def __iter__(self):\n return self._value.__iter__()\n\n def __contains__(self, elt):\n return elt in self._value\n\n def items(self):\n return list(self.iteritems())\n\n def keys(self):\n return list(self.iterkeys())\n\n def values(self):\n return list(self.itervalues())\n\n def iteritems(self):\n return ((k, self[k]) for k in self)\n\n def iterkeys(self):\n return iter(self)\n\n def itervalues(self):\n return (self[k] for k in self)\n\n def get(self, k, d=None):\n return self[k] if k in self else d\n\n\nDictBox.register(dict_)\n\n\n@primitive\ndef container_untake(x, idx, vs):\n if isinstance(idx, slice):\n accum = lambda result: [elt_vs._mut_add(a, b) for elt_vs, a, b in zip(vs.shape[idx], result, x)]\n else:\n accum = lambda result: vs.shape[idx]._mut_add(result, x)\n\n def mut_add(A):\n return vs._subval(A, idx, accum(A[idx]))\n\n return SparseObject(vs, mut_add)\n\n\ndefvjp(container_untake, lambda ans, x, idx, _: lambda g: container_take(g, idx))\ndefjvp(container_untake, \"same\")\n\n\n@primitive\ndef sequence_extend_right(seq, *elts):\n return seq + type(seq)(elts)\n\n\ndef grad_sequence_extend_right(argnum, ans, args, kwargs):\n seq, elts = args[0], args[1:]\n return lambda g: g[: len(seq)] if argnum == 0 else g[len(seq) + argnum - 1]\n\n\ndefvjp_argnum(sequence_extend_right, grad_sequence_extend_right)\n\n\n@primitive\ndef sequence_extend_left(seq, *elts):\n return type(seq)(elts) + seq\n\n\ndef grad_sequence_extend_left(argnum, ans, args, kwargs):\n seq, elts = args[0], args[1:]\n return lambda g: g[len(elts) :] if argnum == 0 else g[argnum - 1]\n\n\ndefvjp_argnum(sequence_extend_left, grad_sequence_extend_left)\n\n\n@primitive\ndef make_sequence(seq_type, *args):\n return seq_type(args)\n\n\ndefvjp_argnum(make_sequence, lambda argnum, *args: lambda g: g[argnum - 1])\n\n\ndef fwd_grad_make_sequence(argnum, g, ans, seq_type, *args, **kwargs):\n return container_untake(g, argnum - 1, vspace(ans))\n\n\ndefjvp_argnum(make_sequence, fwd_grad_make_sequence)\n\n\nclass TupleMeta(type(tuple_)):\n def __instancecheck__(self, instance):\n return isinstance(instance, tuple_)\n\n\nclass tuple(tuple_, metaclass=TupleMeta):\n def __new__(cls, xs):\n return make_sequence(tuple_, *xs)\n\n\nclass ListMeta(type_):\n def __instancecheck__(self, instance):\n return isinstance(instance, list_)\n\n\nclass list(list_, metaclass=ListMeta):\n def __new__(cls, xs):\n return make_sequence(list_, *xs)\n\n\nclass DictMeta(type_):\n def __instancecheck__(self, instance):\n return isinstance(instance, dict_)\n\n\nclass dict(dict_, metaclass=DictMeta):\n def __new__(cls, *args, **kwargs):\n result = dict_(*args, **kwargs)\n if result:\n return _make_dict(result.keys(), list(result.values()))\n return result\n\n\n@primitive\ndef _make_dict(keys, vals):\n return dict_(zip(keys, vals))\n\n\ndefvjp(_make_dict, lambda ans, keys, vals: lambda g: list(g[key] for key in keys), argnums=(1,))\n\n\nclass ContainerVSpace(VSpace):\n def __init__(self, value):\n self.shape = value\n self.shape = self._map(vspace)\n\n @property\n def size(self):\n return sum(self._values(self._map(lambda vs: vs.size)))\n\n def zeros(self):\n return self._map(lambda vs: vs.zeros())\n\n def ones(self):\n return self._map(lambda vs: vs.ones())\n\n def randn(self):\n return self._map(lambda vs: vs.randn())\n\n def standard_basis(self):\n zero = self.zeros()\n for i, vs in self._kv_pairs(self.shape):\n for x in vs.standard_basis():\n yield self._subval(zero, i, x)\n\n def _add(self, xs, ys):\n return self._map(lambda vs, x, y: vs._add(x, y), xs, ys)\n\n def _mut_add(self, xs, ys):\n return self._map(lambda vs, x, y: vs._mut_add(x, y), xs, ys)\n\n def _scalar_mul(self, xs, a):\n return self._map(lambda vs, x: vs._scalar_mul(x, a), xs)\n\n def _inner_prod(self, xs, ys):\n return sum(self._values(self._map(lambda vs, x, y: vs._inner_prod(x, y), xs, ys)))\n\n def _covector(self, xs):\n return self._map(lambda vs, x: vs._covector(x), xs)\n\n\nclass SequenceVSpace(ContainerVSpace):\n def _values(self, x):\n return x\n\n def _kv_pairs(self, x):\n return enumerate(x)\n\n def _map(self, f, *args):\n return self.seq_type(map(f, self.shape, *args))\n\n def _subval(self, xs, idx, x):\n return self.seq_type(subvals(xs, [(idx, x)]))\n\n\nclass ListVSpace(SequenceVSpace):\n seq_type = list_\n\n\nclass TupleVSpace(SequenceVSpace):\n seq_type = tuple_\n\n\nclass DictVSpace(ContainerVSpace):\n def _values(self, x):\n return x.values()\n\n def _kv_pairs(self, x):\n return x.items()\n\n def _map(self, f, *args):\n return {k: f(vs, *[x[k] for x in args]) for k, vs in self.shape.items()}\n\n def _subval(self, xs, idx, x):\n d = dict(xs.items())\n d[idx] = x\n return d\n\n\nListVSpace.register(list_)\nTupleVSpace.register(tuple_)\nDictVSpace.register(dict_)\n\n\nclass NamedTupleVSpace(SequenceVSpace):\n def _map(self, f, *args):\n return self.seq_type(*map(f, self.shape, *args))\n\n def _subval(self, xs, idx, x):\n return self.seq_type(*subvals(xs, [(idx, x)]))\n"
},
{
"path": "autograd/core.py",
"content": "from functools import reduce\nfrom itertools import count\n\nfrom .tracer import Box, Node, getval, isbox, primitive, toposort, trace\nfrom .util import func, subval\n\n# -------------------- reverse mode --------------------\n\n\ndef make_vjp(fun, x):\n start_node = VJPNode.new_root()\n end_value, end_node = trace(start_node, fun, x)\n if end_node is None:\n\n def vjp(g):\n return vspace(x).zeros()\n else:\n\n def vjp(g):\n return backward_pass(g, end_node)\n\n return vjp, end_value\n\n\ndef backward_pass(g, end_node):\n outgrads = {end_node: (g, False)}\n for node in toposort(end_node):\n outgrad = outgrads.pop(node)\n ingrads = node.vjp(outgrad[0])\n for parent, ingrad in zip(node.parents, ingrads):\n outgrads[parent] = add_outgrads(outgrads.get(parent), ingrad)\n return outgrad[0]\n\n\nclass VJPNode(Node):\n __slots__ = [\"parents\", \"vjp\"]\n\n def __init__(self, value, fun, args, kwargs, parent_argnums, parents):\n self.parents = parents\n try:\n vjpmaker = primitive_vjps[fun]\n except KeyError:\n fun_name = getattr(fun, \"__name__\", fun)\n raise NotImplementedError(f\"VJP of {fun_name} wrt argnums {parent_argnums} not defined\")\n self.vjp = vjpmaker(parent_argnums, value, args, kwargs)\n\n def initialize_root(self):\n self.parents = []\n self.vjp = lambda g: ()\n\n\nprimitive_vjps = {}\n\n\ndef defvjp_argnums(fun, vjpmaker):\n primitive_vjps[fun] = vjpmaker\n\n\ndef defvjp_argnum(fun, vjpmaker):\n def vjp_argnums(argnums, *args):\n vjps = [vjpmaker(argnum, *args) for argnum in argnums]\n return lambda g: (vjp(g) for vjp in vjps)\n\n defvjp_argnums(fun, vjp_argnums)\n\n\ndef defvjp(fun, *vjpmakers, **kwargs):\n argnums = kwargs.get(\"argnums\", count())\n vjps_dict = {\n argnum: translate_vjp(vjpmaker, fun, argnum) for argnum, vjpmaker in zip(argnums, vjpmakers)\n }\n\n def vjp_argnums(argnums, ans, args, kwargs):\n L = len(argnums)\n # These first two cases are just optimizations\n if L == 1:\n argnum = argnums[0]\n try:\n vjpfun = vjps_dict[argnum]\n except KeyError:\n raise NotImplementedError(f\"VJP of {fun.__name__} wrt argnum 0 not defined\")\n vjp = vjpfun(ans, *args, **kwargs)\n return lambda g: (vjp(g),)\n elif L == 2:\n argnum_0, argnum_1 = argnums\n try:\n vjp_0_fun = vjps_dict[argnum_0]\n vjp_1_fun = vjps_dict[argnum_1]\n except KeyError:\n raise NotImplementedError(f\"VJP of {fun.__name__} wrt argnums 0, 1 not defined\")\n vjp_0 = vjp_0_fun(ans, *args, **kwargs)\n vjp_1 = vjp_1_fun(ans, *args, **kwargs)\n return lambda g: (vjp_0(g), vjp_1(g))\n else:\n vjps = [vjps_dict[argnum](ans, *args, **kwargs) for argnum in argnums]\n return lambda g: (vjp(g) for vjp in vjps)\n\n defvjp_argnums(fun, vjp_argnums)\n\n\ndef translate_vjp(vjpfun, fun, argnum):\n if vjpfun is None:\n return lambda ans, *args, **kwargs: lambda g: vspace(args[argnum]).zeros()\n elif callable(vjpfun):\n return vjpfun\n else:\n raise Exception(f\"Bad VJP '{vjpfun}' for '{fun.__name__}'\")\n\n\n# -------------------- forward mode --------------------\n\n\ndef make_jvp(fun, x):\n def jvp(g):\n start_node = JVPNode.new_root(g)\n end_value, end_node = trace(start_node, fun, x)\n if end_node is None:\n return end_value, vspace(end_value).zeros()\n else:\n return end_value, end_node.g\n\n return jvp\n\n\nclass JVPNode(Node):\n __slots__ = [\"g\"]\n\n def __init__(self, value, fun, args, kwargs, parent_argnums, parents):\n parent_gs = [parent.g for parent in parents]\n try:\n jvpmaker = primitive_jvps[fun]\n except KeyError:\n name = getattr(fun, \"__name__\", fun)\n raise NotImplementedError(f\"JVP of {name} wrt argnums {parent_argnums} not defined\")\n self.g = jvpmaker(parent_argnums, parent_gs, value, args, kwargs)\n\n def initialize_root(self, g):\n self.g = g\n\n\nprimitive_jvps = {}\n\n\ndef defjvp_argnums(fun, jvpmaker):\n primitive_jvps[fun] = jvpmaker\n\n\ndef defjvp_argnum(fun, jvpmaker):\n def jvp_argnums(argnums, gs, ans, args, kwargs):\n return sum_outgrads(jvpmaker(argnum, g, ans, args, kwargs) for argnum, g in zip(argnums, gs))\n\n defjvp_argnums(fun, jvp_argnums)\n\n\ndef defjvp(fun, *jvpfuns, **kwargs):\n argnums = kwargs.get(\"argnums\", count())\n jvps_dict = {argnum: translate_jvp(jvpfun, fun, argnum) for argnum, jvpfun in zip(argnums, jvpfuns)}\n\n def jvp_argnums(argnums, gs, ans, args, kwargs):\n return sum_outgrads(jvps_dict[argnum](g, ans, *args, **kwargs) for argnum, g in zip(argnums, gs))\n\n defjvp_argnums(fun, jvp_argnums)\n\n\ndef translate_jvp(jvpfun, fun, argnum):\n if jvpfun is None:\n return lambda g, ans, *a, **k: vspace(ans).zeros()\n elif jvpfun == \"same\":\n return lambda g, ans, *args, **kwargs: fun(*subval(args, argnum, g), **kwargs)\n elif callable(jvpfun):\n return jvpfun\n else:\n raise Exception(f\"Bad JVP '{jvpfun}' for '{fun.__name__}'\")\n\n\ndef def_linear(fun):\n \"\"\"Flags that a function is linear wrt all args\"\"\"\n defjvp_argnum(fun, lambda argnum, g, ans, args, kwargs: fun(*subval(args, argnum, g), **kwargs))\n\n\n# -------------------- vector behavior --------------------\n\n\ndef add_outgrads(prev_g_flagged, g):\n sparse = type(g) in sparse_object_types\n if prev_g_flagged:\n vs = vspace(g)\n prev_g, mutable = prev_g_flagged\n if mutable:\n if sparse:\n return sparse_add(vs, prev_g, g), True\n else:\n return vs.mut_add(prev_g, g), True\n else:\n if sparse:\n prev_g_mutable = vs.mut_add(None, prev_g)\n return sparse_add(vs, prev_g_mutable, g), True\n else:\n return vs.add(prev_g, g), True\n else:\n if sparse:\n return sparse_add(vspace(g), None, g), True\n else:\n return g, False\n\n\ndef sum_outgrads(gs):\n return reduce(add_outgrads, gs, None)[0]\n\n\n@primitive\ndef sparse_add(vs, x_prev, x_new):\n x_prev = x_prev if x_prev is not None else vs.zeros()\n return x_new.mut_add(x_prev)\n\n\nclass VSpace:\n __slots__ = []\n mappings = {}\n iscomplex = False\n\n def __init__(self, value):\n pass\n\n def zeros(self):\n assert False, repr(self)\n\n def ones(self):\n assert False, repr(self)\n\n def standard_basis(self):\n assert False, repr(self)\n\n def randn(self):\n assert False, repr(self)\n\n @primitive\n def mut_add(self, x_prev, x_new):\n x_prev = x_prev if x_prev is not None else self.zeros()\n return self._mut_add(x_prev, x_new)\n\n @primitive\n def add(self, x_prev, x_new):\n return self._add(x_prev, x_new)\n\n @primitive\n def scalar_mul(self, x, a):\n return self._scalar_mul(x, a)\n\n @primitive\n def inner_prod(self, x, y):\n return self._inner_prod(x, y)\n\n @primitive\n def covector(self, x):\n return self._covector(x)\n\n def _add(self, x, y):\n return x + y\n\n def _mut_add(self, x, y):\n x += y\n return x\n\n def _scalar_mul(self, x, a):\n return x * a\n\n def _inner_prod(self, x, y):\n assert False\n\n def _covector(self, x):\n return x\n\n def __eq__(self, other):\n return type(self) == type(other) and self.__dict__ == other.__dict__\n\n def __repr__(self):\n return f\"{type(self).__name__}_{self.__dict__}\"\n\n @classmethod\n def register(cls, value_type, vspace_maker=None):\n if vspace_maker:\n VSpace.mappings[value_type] = vspace_maker\n else:\n VSpace.mappings[value_type] = cls\n\n\ndef vspace(value):\n try:\n return VSpace.mappings[type(value)](value)\n except KeyError:\n if isbox(value):\n return vspace(getval(value))\n else:\n raise TypeError(\n \"Can't find vector space for value {} of type {}. Valid types are {}\".format(\n value, type(value), VSpace.mappings.keys()\n )\n )\n\n\nclass SparseBox(Box):\n __slots__ = []\n\n\nclass SparseObject:\n __slots__ = [\"vs\", \"mut_add\"]\n\n def __init__(self, vs, mut_add):\n self.vs = vs\n self.mut_add = mut_add\n\n\nVSpace.register(SparseObject, lambda x: x.vs)\nSparseBox.register(SparseObject)\nsparse_object_types = {SparseObject, SparseBox}\n\n# -------------------- core reverse mode grads --------------------\n\nidentity_vjp = lambda argnums, *args: lambda g: g\ndefvjp(sparse_add, None, identity_vjp, identity_vjp)\ndefvjp(func(VSpace.add), None, identity_vjp, identity_vjp)\ndefvjp(func(VSpace.mut_add), None, identity_vjp, identity_vjp)\ndefvjp(\n func(VSpace.inner_prod),\n None,\n lambda ans, vs, x, y: lambda g: vs.covector(vs.scalar_mul(y, g)),\n lambda ans, vs, x, y: lambda g: vs.covector(vs.scalar_mul(x, g)),\n)\ndefvjp(func(VSpace.covector), None, lambda ans, vs, x: lambda g: vs.covector(g))\ndefvjp(\n func(VSpace.scalar_mul),\n None,\n lambda ans, vs, x, a: lambda g: vs.covector(vs.scalar_mul(vs.covector(g), a)),\n lambda ans, vs, x, a: lambda g: vs.inner_prod(g, vs.covector(x)),\n)\n\n# -------------------- core forward mode grads --------------------\n\nidentity_jvp = lambda g, *args, **kwargs: g\ndefjvp(sparse_add, None, identity_jvp, identity_jvp)\ndefjvp(func(VSpace.mut_add), None, identity_jvp, identity_jvp)\ndefjvp(func(VSpace.add), None, identity_jvp, identity_jvp)\ndefjvp(func(VSpace.scalar_mul), None, \"same\", \"same\")\ndefjvp(func(VSpace.inner_prod), None, \"same\", \"same\")\ndefjvp(func(VSpace.covector), None, \"same\")\n\n# -------------------- deprecation warnings -----------------------\n\nimport warnings\n\ndeprecated_defvjp_message = \"\"\"\nThe {} method is deprecated. See the update guide and tutorial:\nhttps://github.com/HIPS/autograd/blob/master/docs/updateguide.md\nhttps://github.com/HIPS/autograd/blob/master/docs/tutorial.md\"\"\"\n\n\ndef deprecated_defvjp(primitive_fun):\n deprecation_msg = deprecated_defvjp_message.format(\"defvjp\")\n vjpfuns = {}\n\n def defvjp_unstaged(vjpmaker, argnum=0):\n warnings.warn(deprecation_msg)\n\n def staged_vjpmaker(ans, *args, **kwargs):\n def vjp(g):\n vs, gvs = vspace(args[argnum]), vspace(g)\n return vjpmaker(g, ans, vs, gvs, *args, **kwargs)\n\n return vjp\n\n vjpfuns[argnum] = staged_vjpmaker\n argnums, vjpmakers = zip(*[(argnum, vjpfuns[argnum]) for argnum in sorted(vjpfuns.keys())])\n defvjp(primitive_fun, *vjpmakers, argnums=argnums)\n\n return defvjp_unstaged\n\n\ndef deprecated_defvjp_is_zero(primitive_fun):\n deprecation_msg = deprecated_defvjp_message.format(\"defvjp_is_zero\")\n zero_vjps = [set()]\n\n def defvjp_is_zero(argnums=(0,)):\n warnings.warn(deprecation_msg)\n zero_vjps[0] |= set(argnums)\n nones = [None] * len(zero_vjps[0])\n defvjp(primitive_fun, *nones, argnums=sorted(zero_vjps[0]))\n\n return defvjp_is_zero\n\n\ndef deprecated_defgrad(primitive_fun):\n deprecation_msg = deprecated_defvjp_message.format(\"defgrad\")\n gradfuns = {}\n\n def defgrad(gradfun, argnum=0):\n warnings.warn(deprecation_msg)\n gradfuns[argnum] = gradfun\n argnums, vjpmakers = zip(*[(argnum, gradfuns[argnum]) for argnum in sorted(gradfuns.keys())])\n defvjp(primitive_fun, *vjpmakers, argnums=argnums)\n\n return defgrad\n\n\nprimitive_ = primitive\n\n\ndef primitive_with_deprecation_warnings(f_raw):\n f_wrapped = primitive_(f_raw)\n f_wrapped.defvjp = deprecated_defvjp(f_wrapped)\n f_wrapped.defvjp_is_zero = deprecated_defvjp_is_zero(f_wrapped)\n f_wrapped.defgrad = deprecated_defgrad(f_wrapped)\n return f_wrapped\n\n\nprimitive = primitive_with_deprecation_warnings\n"
},
{
"path": "autograd/differential_operators.py",
"content": "\"\"\"Convenience functions built on top of `make_vjp`.\"\"\"\n\nfrom collections import OrderedDict\n\ntry:\n from inspect import getfullargspec as _getargspec # Python 3\nexcept ImportError:\n from inspect import getargspec as _getargspec # Python 2\nimport warnings\n\nimport autograd.numpy as np\n\nfrom .builtins import tuple as atuple\nfrom .core import make_jvp as _make_jvp\nfrom .core import make_vjp as _make_vjp\nfrom .extend import defvjp_argnum, primitive, vspace\nfrom .wrap_util import unary_to_nary\n\nmake_vjp = unary_to_nary(_make_vjp)\nmake_jvp = unary_to_nary(_make_jvp)\n\n\n@unary_to_nary\ndef grad(fun, x):\n \"\"\"\n Returns a function which computes the gradient of `fun` with respect to\n positional argument number `argnum`. The returned function takes the same\n arguments as `fun`, but returns the gradient instead. The function `fun`\n should be scalar-valued. The gradient has the same type as the argument.\"\"\"\n vjp, ans = _make_vjp(fun, x)\n if not vspace(ans).size == 1:\n raise TypeError(\n \"Grad only applies to real scalar-output functions. \"\n \"Try jacobian, elementwise_grad or holomorphic_grad.\"\n )\n return vjp(vspace(ans).ones())\n\n\n@unary_to_nary\ndef elementwise_grad(fun, x):\n \"\"\"\n Returns a function that computes the sum of each column of the Jacobian of\n `fun`, in one pass. If the Jacobian is diagonal, then this is the diagonal\n of the Jacobian.\n \"\"\"\n vjp, ans = _make_vjp(fun, x)\n if vspace(ans).iscomplex:\n raise TypeError(\"Elementwise_grad only applies to real-output functions.\")\n return vjp(vspace(ans).ones())\n\n\n@unary_to_nary\ndef deriv(fun, x):\n return _make_jvp(fun, x)(vspace(x).ones())[1]\n\n\n@unary_to_nary\ndef jacobian(fun, x):\n \"\"\"\n Returns a function which computes the Jacobian of `fun` with respect to\n positional argument number `argnum`, which must be a scalar or array. Unlike\n `grad` it is not restricted to scalar-output functions, but also it cannot\n take derivatives with respect to some argument types (like lists or dicts).\n If the input to `fun` has shape (in1, in2, ...) and the output has shape\n (out1, out2, ...) then the Jacobian has shape (out1, out2, ..., in1, in2, ...).\n \"\"\"\n vjp, ans = _make_vjp(fun, x)\n ans_vspace = vspace(ans)\n jacobian_shape = ans_vspace.shape + vspace(x).shape\n grads = map(vjp, ans_vspace.standard_basis())\n return np.reshape(np.stack(grads), jacobian_shape)\n\n\n@unary_to_nary\ndef holomorphic_grad(fun, x):\n if not vspace(x).iscomplex:\n warnings.warn(\"Input to holomorphic_grad is not complex\")\n return grad(lambda x: np.real(fun(x)))(x)\n\n\ndef grad_named(fun, argname):\n \"\"\"Takes gradients with respect to a named argument.\n Doesn't work on *args or **kwargs.\"\"\"\n arg_index = _getargspec(fun).args.index(argname)\n return grad(fun, arg_index)\n\n\n@unary_to_nary\ndef hessian(fun, x):\n \"Returns a function that computes the exact Hessian.\"\n return jacobian(jacobian(fun))(x)\n\n\n@unary_to_nary\ndef make_hvp(fun, x):\n \"\"\"Builds a function for evaluating the Hessian-vector product at a point,\n which may be useful when evaluating many Hessian-vector products at the same\n point while caching the results of the forward pass.\"\"\"\n return _make_vjp(grad(fun), x)\n\n\ndef hessian_tensor_product(fun, argnum=0):\n \"\"\"Builds a function that returns the exact Hessian-tensor product.\n The returned function has arguments (*args, tensor, **kwargs), and for\n vectors takes roughly 4x as long to evaluate as the original function.\"\"\"\n fun_grad = grad(fun, argnum)\n\n def vector_dot_grad(*args, **kwargs):\n args, vector = args[:-1], args[-1]\n return np.tensordot(fun_grad(*args, **kwargs), vector, np.ndim(vector))\n\n return grad(vector_dot_grad, argnum)\n\n\nhessian_vector_product = hessian_tensor_product\n\n\ndef tensor_jacobian_product(fun, argnum=0):\n \"\"\"Builds a function that returns the exact tensor-Jacobian product, that\n is the Jacobian matrix left-multiplied by tensor. The returned function\n has arguments (*args, tensor, **kwargs).\"\"\"\n\n def vector_dot_fun(*args, **kwargs):\n args, vector = args[:-1], args[-1]\n return np.tensordot(vector, fun(*args, **kwargs), axes=np.ndim(vector))\n\n return jacobian(vector_dot_fun, argnum)\n\n\nvector_jacobian_product = tensor_jacobian_product\n\n\n@unary_to_nary\ndef make_jvp_reversemode(fun, x):\n \"\"\"Builds a function for evaluating the Jacobian-vector product at a\n point. Roughly 1.5x more FLOPs than forward-mode, plus memory requirements\n that scale with the number of primitives applied in the evaluation of f, as\n well as other overheads. See j-towns.github.io/2017/06/12/A-new-trick.html.\"\"\"\n vjp, y = _make_vjp(fun, x)\n vjp_vjp, _ = _make_vjp(vjp, vspace(y).zeros())\n return vjp_vjp # vjp_vjp is just jvp by linearity\n\n\n# TODO(mattjj): update this function using make_jvp and const_graph\ndef make_ggnvp(f, g=lambda x: 1.0 / 2 * np.sum(x**2, axis=-1), f_argnum=0):\n \"\"\"Builds a function for evaluating generalized-Gauss-Newton-vector products\n at a point. Slightly more expensive than mixed-mode.\"\"\"\n\n @unary_to_nary\n def _make_ggnvp(f, x):\n f_vjp, f_x = _make_vjp(f, x)\n g_hvp, grad_g_x = _make_vjp(grad(g), f_x)\n f_jvp, _ = _make_vjp(f_vjp, vspace(grad_g_x).zeros())\n\n def ggnvp(v):\n return f_vjp(g_hvp(f_jvp(v)))\n\n return ggnvp\n\n return _make_ggnvp(f, f_argnum)\n\n\n@unary_to_nary\ndef value_and_grad(fun, x):\n \"\"\"Returns a function that returns both value and gradient. Suitable for use\n in scipy.optimize\"\"\"\n vjp, ans = _make_vjp(fun, x)\n if not vspace(ans).size == 1:\n raise TypeError(\n \"value_and_grad only applies to real scalar-output \"\n \"functions. Try jacobian, elementwise_grad or \"\n \"holomorphic_grad.\"\n )\n return ans, vjp(vspace(ans).ones())\n\n\n@unary_to_nary\ndef grad_and_aux(fun, x):\n \"\"\"Builds a function that returns the gradient of the first output and the\n (unmodified) second output of a function that returns two outputs.\"\"\"\n vjp, (ans, aux) = _make_vjp(lambda x: atuple(fun(x)), x)\n return vjp((vspace(ans).ones(), vspace(aux).zeros())), aux\n\n\ndef multigrad_dict(fun):\n \"Takes gradients wrt all arguments simultaneously,\"\n \"returns a dict mapping 'argname' to 'gradval'\"\n\n import funcsigs\n\n sig = funcsigs.signature(fun)\n\n def select(preds, lst):\n idx = lambda item: next((i for i, pred in enumerate(preds) if pred(item)), len(preds))\n results = [[] for _ in preds] + [[]]\n for item in lst:\n results[idx(item)].append(item)\n return results\n\n is_var_pos = lambda name: sig.parameters[name].kind == sig.parameters[name].VAR_POSITIONAL\n is_var_kwd = lambda name: sig.parameters[name].kind == sig.parameters[name].VAR_KEYWORD\n var_pos, var_kwd, argnames = select([is_var_pos, is_var_kwd], sig.parameters)\n\n todict = lambda dct: {key: dct[key] for key in dct}\n\n def apply_defaults(arguments):\n defaults = {\n name: param.default for name, param in sig.parameters.items() if param.default is not param.empty\n }\n return OrderedDict(\n (name, arguments[name] if name in arguments else defaults[name]) for name in sig.parameters\n )\n\n def gradfun(*args, **kwargs):\n bindings = sig.bind(*args, **kwargs)\n\n args = lambda dct: tuple(dct[var_pos[0]]) if var_pos else ()\n kwargs = lambda dct: todict(dct[var_kwd[0]]) if var_kwd else {}\n others = lambda dct: tuple(dct[argname] for argname in argnames if argname not in var_kwd + var_pos)\n\n newfun = lambda dct: fun(*(others(dct) + args(dct)), **kwargs(dct))\n\n argdict = apply_defaults(bindings.arguments)\n grad_dict = grad(newfun)(dict(argdict))\n return OrderedDict((argname, grad_dict[argname]) for argname in argdict)\n\n return gradfun\n\n\ndef checkpoint(fun):\n \"\"\"Returns a checkpointed version of `fun`, where intermediate values\n computed during the forward pass of `fun` are discarded and then recomputed\n for the backward pass. Useful to save memory, effectively trading off time\n and memory. See e.g. arxiv.org/abs/1604.06174.\n \"\"\"\n\n def wrapped_grad(argnum, ans, args, kwargs):\n return make_vjp(fun, argnum)(*args, **kwargs)[0]\n\n wrapped = primitive(fun)\n defvjp_argnum(wrapped, wrapped_grad)\n return wrapped\n"
},
{
"path": "autograd/extend.py",
"content": "# Exposes API for extending autograd\nfrom .core import (\n JVPNode,\n SparseObject,\n VJPNode,\n VSpace,\n def_linear,\n defjvp,\n defjvp_argnum,\n defjvp_argnums,\n defvjp,\n defvjp_argnum,\n defvjp_argnums,\n vspace,\n)\nfrom .tracer import Box, notrace_primitive, primitive, register_notrace\n"
},
{
"path": "autograd/misc/__init__.py",
"content": "from .flatten import flatten\nfrom .tracers import const_graph\n"
},
{
"path": "autograd/misc/fixed_points.py",
"content": "from autograd import make_vjp\nfrom autograd.builtins import tuple\nfrom autograd.extend import defvjp, primitive, vspace\n\n\n@primitive\ndef fixed_point(f, a, x0, distance, tol):\n _f = f(a)\n x, x_prev = _f(x0), x0\n while distance(x, x_prev) > tol:\n x, x_prev = _f(x), x\n return x\n\n\ndef fixed_point_vjp(ans, f, a, x0, distance, tol):\n def rev_iter(params):\n a, x_star, x_star_bar = params\n vjp_x, _ = make_vjp(f(a))(x_star)\n vs = vspace(x_star)\n return lambda g: vs.add(vjp_x(g), x_star_bar)\n\n vjp_a, _ = make_vjp(lambda x, y: f(x)(y))(a, ans)\n return lambda g: vjp_a(fixed_point(rev_iter, tuple((a, ans, g)), vspace(x0).zeros(), distance, tol))\n\n\ndefvjp(fixed_point, None, fixed_point_vjp, None)\n"
},
{
"path": "autograd/misc/flatten.py",
"content": "\"\"\"\nHandy functions for flattening nested containers containing numpy\narrays. The main purpose is to make examples and optimizers simpler.\n\"\"\"\n\nimport autograd.numpy as np\nfrom autograd import make_vjp\nfrom autograd.builtins import type\n\n\ndef flatten(value):\n \"\"\"Flattens any nesting of tuples, lists, or dicts, with numpy arrays or\n scalars inside. Returns 1D numpy array and an unflatten function.\n Doesn't preserve mixed numeric types (e.g. floats and ints). Assumes dict\n keys are sortable.\"\"\"\n unflatten, flat_value = make_vjp(_flatten)(value)\n return flat_value, unflatten\n\n\ndef _flatten(value):\n t = type(value)\n if t in (list, tuple):\n return _concatenate(map(_flatten, value))\n elif t is dict:\n return _concatenate(_flatten(value[k]) for k in sorted(value))\n else:\n return np.ravel(value)\n\n\ndef _concatenate(lst):\n lst = list(lst)\n return np.concatenate(lst) if lst else np.array([])\n\n\ndef flatten_func(func, example):\n _ex, unflatten = flatten(example)\n _func = lambda _x, *args: flatten(func(unflatten(_x), *args))[0]\n return _func, unflatten, _ex\n"
},
{
"path": "autograd/misc/optimizers.py",
"content": "\"\"\"Some standard gradient-based stochastic optimizers.\n\nThese are just standard routines that don't make any use of autograd,\nthough you could take gradients of these functions too if you want\nto do meta-optimization.\n\nThese routines can optimize functions whose inputs are structured\nobjects, such as dicts of numpy arrays.\"\"\"\n\nimport autograd.numpy as np\nfrom autograd.misc import flatten\nfrom autograd.wrap_util import wraps\n\n\ndef unflatten_optimizer(optimize):\n \"\"\"Takes an optimizer that operates on flat 1D numpy arrays and returns a\n wrapped version that handles trees of nested containers (lists/tuples/dicts)\n with arrays/scalars at the leaves.\"\"\"\n\n @wraps(optimize)\n def _optimize(grad, x0, callback=None, *args, **kwargs):\n _x0, unflatten = flatten(x0)\n _grad = lambda x, i: flatten(grad(unflatten(x), i))[0]\n if callback:\n _callback = lambda x, i, g: callback(unflatten(x), i, unflatten(g))\n else:\n _callback = None\n return unflatten(optimize(_grad, _x0, _callback, *args, **kwargs))\n\n return _optimize\n\n\n@unflatten_optimizer\ndef sgd(grad, x, callback=None, num_iters=200, step_size=0.1, mass=0.9):\n \"\"\"Stochastic gradient descent with momentum.\n grad() must have signature grad(x, i), where i is the iteration number.\"\"\"\n velocity = np.zeros(len(x))\n for i in range(num_iters):\n g = grad(x, i)\n if callback:\n callback(x, i, g)\n velocity = mass * velocity - (1.0 - mass) * g\n x = x + step_size * velocity\n return x\n\n\n@unflatten_optimizer\ndef rmsprop(grad, x, callback=None, num_iters=100, step_size=0.1, gamma=0.9, eps=10**-8):\n \"\"\"Root mean squared prop: See Adagrad paper for details.\"\"\"\n avg_sq_grad = np.ones(len(x))\n for i in range(num_iters):\n g = grad(x, i)\n if callback:\n callback(x, i, g)\n avg_sq_grad = avg_sq_grad * gamma + g**2 * (1 - gamma)\n x = x - step_size * g / (np.sqrt(avg_sq_grad) + eps)\n return x\n\n\n@unflatten_optimizer\ndef adam(grad, x, callback=None, num_iters=100, step_size=0.001, b1=0.9, b2=0.999, eps=10**-8):\n \"\"\"Adam as described in http://arxiv.org/pdf/1412.6980.pdf.\n It's basically RMSprop with momentum and some correction terms.\"\"\"\n m = np.zeros(len(x))\n v = np.zeros(len(x))\n for i in range(num_iters):\n g = grad(x, i)\n if callback:\n callback(x, i, g)\n m = (1 - b1) * g + b1 * m # First moment estimate.\n v = (1 - b2) * (g**2) + b2 * v # Second moment estimate.\n mhat = m / (1 - b1 ** (i + 1)) # Bias correction.\n vhat = v / (1 - b2 ** (i + 1))\n x = x - step_size * mhat / (np.sqrt(vhat) + eps)\n return x\n"
},
{
"path": "autograd/misc/tracers.py",
"content": "from functools import partial\nfrom itertools import repeat\n\nfrom autograd.tracer import Node, trace\nfrom autograd.util import subvals, toposort\nfrom autograd.wrap_util import wraps\n\n\nclass ConstGraphNode(Node):\n __slots__ = [\"parents\", \"partial_fun\"]\n\n def __init__(self, value, fun, args, kwargs, parent_argnums, parents):\n args = subvals(args, zip(parent_argnums, repeat(None)))\n\n def partial_fun(partial_args):\n return fun(*subvals(args, zip(parent_argnums, partial_args)), **kwargs)\n\n self.parents = parents\n self.partial_fun = partial_fun\n\n def initialize_root(self):\n self.parents = []\n\n\ndef const_graph_unary(fun):\n graph = []\n _fun = [fun] # Allow fun to be freed, since it may have bound args\n\n def maybe_cached_fun(x):\n if graph:\n _graph = graph[0]\n vals = {_graph[0]: x}\n for node in _graph[1:]:\n vals[node] = node.partial_fun([vals[p] for p in node.parents])\n return vals[node]\n else:\n start_node = ConstGraphNode.new_root()\n end_value, end_node = trace(start_node, _fun.pop(), x)\n if end_node is None:\n raise Exception(\"Output is independent of input\")\n graph.append(list(toposort(end_node))[::-1])\n return end_value\n\n return maybe_cached_fun\n\n\ndef const_graph(fun, *args, **kwargs):\n partial_fun = partial(fun, *args, **kwargs)\n unary_fun = lambda args: partial_fun(*args)\n maybe_cached_unary_fun = const_graph_unary(unary_fun)\n\n @wraps(fun)\n def _fun(*args):\n return maybe_cached_unary_fun(args)\n\n return _fun\n\n\nclass FullGraphNode(Node):\n __slots__ = [\"value\", \"recipe\"]\n\n def __init__(self, value, fun, args, kwargs, parent_argnums, parents):\n self.value = value\n self.recipe = (fun, args, kwargs, zip(parent_argnums, parents))\n\n def initialize_root(self):\n self.value = None\n self.recipe = (lambda x: x, (), {}, [])\n\n\ndef full_graph(fun, *args, **kwargs):\n unary_fun = lambda args: fun(*args, **kwargs)\n start_node = FullGraphNode.new_root()\n end_value, end_node = trace(start_node, unary_fun, args)\n return end_node\n"
},
{
"path": "autograd/numpy/__init__.py",
"content": "from . import fft, linalg, numpy_boxes, numpy_jvps, numpy_vjps, numpy_vspaces, random\nfrom .numpy_wrapper import *\nfrom .numpy_wrapper import numpy_version as __version__\n"
},
{
"path": "autograd/numpy/fft.py",
"content": "import numpy.fft as ffto\n\nfrom autograd.extend import defvjp, primitive, vspace\n\nfrom . import numpy_wrapper as anp\nfrom .numpy_vjps import match_complex\nfrom .numpy_wrapper import wrap_namespace\n\nwrap_namespace(ffto.__dict__, globals())\n\n\n# TODO: make fft gradient work for a repeated axis,\n# e.g. by replacing fftn with repeated calls to 1d fft along each axis\ndef fft_grad(get_args, fft_fun, ans, x, *args, **kwargs):\n axes, s, norm = get_args(x, *args, **kwargs)\n check_no_repeated_axes(axes)\n vs = vspace(x)\n return lambda g: match_complex(x, truncate_pad(fft_fun(g, *args, **kwargs), vs.shape))\n\n\ndefvjp(fft, lambda *args, **kwargs: fft_grad(get_fft_args, fft, *args, **kwargs))\ndefvjp(ifft, lambda *args, **kwargs: fft_grad(get_fft_args, ifft, *args, **kwargs))\n\ndefvjp(fft2, lambda *args, **kwargs: fft_grad(get_fft_args, fft2, *args, **kwargs))\ndefvjp(ifft2, lambda *args, **kwargs: fft_grad(get_fft_args, ifft2, *args, **kwargs))\n\ndefvjp(fftn, lambda *args, **kwargs: fft_grad(get_fft_args, fftn, *args, **kwargs))\ndefvjp(ifftn, lambda *args, **kwargs: fft_grad(get_fft_args, ifftn, *args, **kwargs))\n\n\ndef rfft_grad(get_args, irfft_fun, ans, x, *args, **kwargs):\n axes, s, norm = get_args(x, *args, **kwargs)\n vs = vspace(x)\n gvs = vspace(ans)\n check_no_repeated_axes(axes)\n if s is None:\n s = [vs.shape[i] for i in axes]\n check_even_shape(s)\n\n # s is the full fft shape\n # gs is the compressed shape\n gs = list(s)\n gs[-1] = gs[-1] // 2 + 1\n fac = make_rfft_factors(axes, gvs.shape, gs, s, norm)\n\n def vjp(g):\n g = anp.conj(g / fac)\n r = match_complex(x, truncate_pad((irfft_fun(g, *args, **kwargs)), vs.shape))\n return r\n\n return vjp\n\n\ndef irfft_grad(get_args, rfft_fun, ans, x, *args, **kwargs):\n axes, gs, norm = get_args(x, *args, **kwargs)\n vs = vspace(x)\n gvs = vspace(ans)\n check_no_repeated_axes(axes)\n if gs is None:\n gs = [gvs.shape[i] for i in axes]\n check_even_shape(gs)\n\n # gs is the full fft shape\n # s is the compressed shape\n s = list(gs)\n s[-1] = s[-1] // 2 + 1\n\n def vjp(g):\n r = match_complex(x, truncate_pad((rfft_fun(g, *args, **kwargs)), vs.shape))\n fac = make_rfft_factors(axes, vs.shape, s, gs, norm)\n r = anp.conj(r) * fac\n return r\n\n return vjp\n\n\ndefvjp(rfft, lambda *args, **kwargs: rfft_grad(get_fft_args, irfft, *args, **kwargs))\n\ndefvjp(irfft, lambda *args, **kwargs: irfft_grad(get_fft_args, rfft, *args, **kwargs))\n\ndefvjp(rfft2, lambda *args, **kwargs: rfft_grad(get_fft2_args, irfft2, *args, **kwargs))\n\ndefvjp(irfft2, lambda *args, **kwargs: irfft_grad(get_fft2_args, rfft2, *args, **kwargs))\n\ndefvjp(rfftn, lambda *args, **kwargs: rfft_grad(get_fftn_args, irfftn, *args, **kwargs))\n\ndefvjp(irfftn, lambda *args, **kwargs: irfft_grad(get_fftn_args, rfftn, *args, **kwargs))\n\ndefvjp(\n fftshift, lambda ans, x, axes=None: lambda g: match_complex(x, anp.conj(ifftshift(anp.conj(g), axes)))\n)\ndefvjp(\n ifftshift, lambda ans, x, axes=None: lambda g: match_complex(x, anp.conj(fftshift(anp.conj(g), axes)))\n)\n\n\n@primitive\ndef truncate_pad(x, shape):\n # truncate/pad x to have the appropriate shape\n slices = [slice(n) for n in shape]\n pads = tuple(\n zip(anp.zeros(len(shape), dtype=int), anp.maximum(0, anp.array(shape) - anp.array(x.shape)))\n )\n return anp.pad(x, pads, \"constant\")[tuple(slices)]\n\n\ndefvjp(truncate_pad, lambda ans, x, shape: lambda g: match_complex(x, truncate_pad(g, vspace(x).shape)))\n\n\n## TODO: could be made less stringent, to fail only when repeated axis has different values of s\ndef check_no_repeated_axes(axes):\n axes_set = set(axes)\n if len(axes) != len(axes_set):\n raise NotImplementedError(\"FFT gradient for repeated axes not implemented.\")\n\n\ndef check_even_shape(shape):\n if shape[-1] % 2 != 0:\n raise NotImplementedError(\"Real FFT gradient for odd lengthed last axes is not implemented.\")\n\n\ndef get_fft_args(a, d=None, axis=-1, norm=None, *args, **kwargs):\n axes = [axis]\n if d is not None:\n d = [d]\n return axes, d, norm\n\n\ndef get_fft2_args(a, s=None, axes=(-2, -1), norm=None, *args, **kwargs):\n return axes, s, norm\n\n\ndef get_fftn_args(a, s=None, axes=None, norm=None, *args, **kwargs):\n if axes is None:\n axes = list(range(a.ndim))\n return axes, s, norm\n\n\ndef make_rfft_factors(axes, resshape, facshape, normshape, norm):\n \"\"\"make the compression factors and compute the normalization\n for irfft and rfft.\n \"\"\"\n N = 1.0\n for n in normshape:\n N = N * n\n\n # inplace modification is fine because we produce a constant\n # which doesn't go into autograd.\n # For same reason could have used numpy rather than anp.\n # but we already imported anp, so use it instead.\n fac = anp.zeros(resshape)\n fac[...] = 2\n index = [slice(None)] * len(resshape)\n if facshape[-1] <= resshape[axes[-1]]:\n index[axes[-1]] = (0, facshape[-1] - 1)\n else:\n index[axes[-1]] = (0,)\n fac[tuple(index)] = 1\n if norm is None:\n fac /= N\n return fac\n"
},
{
"path": "autograd/numpy/linalg.py",
"content": "from functools import partial\n\nimport numpy.linalg as npla\n\nfrom autograd.extend import defjvp, defvjp\n\nfrom . import numpy_wrapper as anp\nfrom .numpy_wrapper import wrap_namespace\n\nwrap_namespace(npla.__dict__, globals())\n\n# Some formulas are from\n# \"An extended collection of matrix derivative results\n# for forward and reverse mode algorithmic differentiation\"\n# by Mike Giles\n# https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf\n\n\n# transpose by swapping last two dimensions\ndef T(x):\n return anp.swapaxes(x, -1, -2)\n\n\n_dot = partial(anp.einsum, \"...ij,...jk->...ik\")\n\n# batched diag\n_diag = lambda a: anp.eye(a.shape[-1]) * a\n\n\n# batched diagonal, similar to matrix_diag in tensorflow\ndef _matrix_diag(a):\n reps = anp.array(a.shape)\n reps[:-1] = 1\n reps[-1] = a.shape[-1]\n newshape = list(a.shape) + [a.shape[-1]]\n return _diag(anp.tile(a, reps).reshape(newshape))\n\n\n# add two dimensions to the end of x\ndef add2d(x):\n return anp.reshape(x, anp.shape(x) + (1, 1))\n\n\ndefvjp(det, lambda ans, x: lambda g: add2d(g) * add2d(ans) * T(inv(x)))\ndefvjp(slogdet, lambda ans, x: lambda g: add2d(g[1]) * T(inv(x)))\n\n\ndef grad_inv(ans, x):\n return lambda g: -_dot(_dot(T(ans), g), T(ans))\n\n\ndefvjp(inv, grad_inv)\n\n\ndef grad_pinv(ans, x):\n # https://mathoverflow.net/questions/25778/analytical-formula-for-numerical-derivative-of-the-matrix-pseudo-inverse\n return lambda g: T(\n -_dot(_dot(ans, T(g)), ans)\n + _dot(_dot(_dot(ans, T(ans)), g), anp.eye(x.shape[-2]) - _dot(x, ans))\n + _dot(_dot(_dot(anp.eye(ans.shape[-2]) - _dot(ans, x), g), T(ans)), ans)\n )\n\n\ndefvjp(pinv, grad_pinv)\n\n\ndef grad_solve(argnum, ans, a, b):\n updim = lambda x: x if x.ndim == a.ndim else x[..., None]\n if argnum == 0:\n return lambda g: -_dot(updim(solve(T(a), g)), T(updim(ans)))\n else:\n return lambda g: solve(T(a), g)\n\n\ndefvjp(solve, partial(grad_solve, 0), partial(grad_solve, 1))\n\n\ndef norm_vjp(ans, x, ord=None, axis=None):\n def check_implemented():\n matrix_norm = (x.ndim == 2 and axis is None) or isinstance(axis, tuple)\n\n if matrix_norm:\n if not (ord is None or ord == \"fro\" or ord == \"nuc\"):\n raise NotImplementedError(f\"Gradient of matrix norm not implemented for ord={ord}\")\n elif not (ord is None or ord > 1):\n raise NotImplementedError(f\"Gradient of norm not implemented for ord={ord}\")\n\n if axis is None:\n expand = lambda a: a\n elif isinstance(axis, tuple):\n row_axis, col_axis = axis\n if row_axis > col_axis:\n row_axis = row_axis - 1\n expand = lambda a: anp.expand_dims(anp.expand_dims(a, row_axis), col_axis)\n else:\n expand = lambda a: anp.expand_dims(a, axis=axis)\n\n if ord == \"nuc\":\n if axis is None:\n roll = lambda a: a\n unroll = lambda a: a\n else:\n row_axis, col_axis = axis\n if row_axis > col_axis:\n row_axis = row_axis - 1\n # Roll matrix axes to the back\n roll = lambda a: anp.rollaxis(anp.rollaxis(a, col_axis, a.ndim), row_axis, a.ndim - 1)\n # Roll matrix axes to their original position\n unroll = lambda a: anp.rollaxis(anp.rollaxis(a, a.ndim - 2, row_axis), a.ndim - 1, col_axis)\n\n check_implemented()\n\n def vjp(g):\n if ord in (None, 2, \"fro\"):\n return expand(g / ans) * anp.conj(x)\n elif ord == \"nuc\":\n x_rolled = roll(x)\n u, s, vt = svd(x_rolled, full_matrices=False)\n uvt_rolled = _dot(u, vt)\n # Roll the matrix axes back to their correct positions\n uvt = unroll(uvt_rolled)\n g = expand(g)\n return g * anp.conj(uvt)\n else:\n # see https://en.wikipedia.org/wiki/Norm_(mathematics)#p-norm\n return expand(g / ans ** (ord - 1)) * anp.conj(x) * anp.abs(x) ** (ord - 2)\n\n return vjp\n\n\ndefvjp(norm, norm_vjp)\n\n\ndef norm_jvp(g, ans, x, ord=None, axis=None):\n def check_implemented():\n matrix_norm = (x.ndim == 2 and axis is None) or isinstance(axis, tuple)\n\n if matrix_norm:\n if not (ord is None or ord == \"fro\" or ord == \"nuc\"):\n raise NotImplementedError(f\"Gradient of matrix norm not implemented for ord={ord}\")\n elif not (ord is None or ord > 1):\n raise NotImplementedError(f\"Gradient of norm not implemented for ord={ord}\")\n\n if axis is None:\n contract = lambda a: anp.sum(a)\n else:\n contract = partial(anp.sum, axis=axis)\n\n if ord == \"nuc\":\n if axis is None:\n roll = lambda a: a\n unroll = lambda a: a\n else:\n row_axis, col_axis = axis\n if row_axis > col_axis:\n row_axis = row_axis - 1\n # Roll matrix axes to the back\n roll = lambda a: anp.rollaxis(anp.rollaxis(a, col_axis, a.ndim), row_axis, a.ndim - 1)\n # Roll matrix axes to their original position\n unroll = lambda a: anp.rollaxis(anp.rollaxis(a, a.ndim - 2, row_axis), a.ndim - 1, col_axis)\n\n check_implemented()\n if ord in (None, 2, \"fro\"):\n return contract(g * anp.conj(x)) / ans\n elif ord == \"nuc\":\n x_rolled = roll(x)\n u, s, vt = svd(x_rolled, full_matrices=False)\n uvt_rolled = _dot(u, vt)\n # Roll the matrix axes back to their correct positions\n uvt = unroll(uvt_rolled)\n return contract(g * anp.conj(uvt))\n else:\n # see https://en.wikipedia.org/wiki/Norm_(mathematics)#p-norm\n return contract(g * anp.conj(x) * anp.abs(x) ** (ord - 2)) / ans ** (ord - 1)\n\n\ndefjvp(norm, norm_jvp)\n\n\ndef grad_eigh(ans, x, UPLO=\"L\"):\n \"\"\"Gradient for eigenvalues and vectors of a symmetric matrix.\"\"\"\n N = x.shape[-1]\n w, v = ans # Eigenvalues, eigenvectors.\n vc = anp.conj(v)\n\n def vjp(g):\n wg, vg = g # Gradient w.r.t. eigenvalues, eigenvectors.\n w_repeated = anp.repeat(w[..., anp.newaxis], N, axis=-1)\n\n # Eigenvalue part\n vjp_temp = _dot(vc * wg[..., anp.newaxis, :], T(v))\n\n # Add eigenvector part only if non-zero backward signal is present.\n # This can avoid NaN results for degenerate cases if the function depends\n # on the eigenvalues only.\n if anp.any(vg):\n off_diag = anp.ones((N, N)) - anp.eye(N)\n F = off_diag / (T(w_repeated) - w_repeated + anp.eye(N))\n vjp_temp += _dot(_dot(vc, F * _dot(T(v), vg)), T(v))\n\n # eigh always uses only the lower or the upper part of the matrix\n # we also have to make sure broadcasting works\n reps = anp.array(x.shape)\n reps[-2:] = 1\n\n if UPLO == \"L\":\n tri = anp.tile(anp.tril(anp.ones(N), -1), reps)\n elif UPLO == \"U\":\n tri = anp.tile(anp.triu(anp.ones(N), 1), reps)\n\n return anp.real(vjp_temp) * anp.eye(vjp_temp.shape[-1]) + (vjp_temp + anp.conj(T(vjp_temp))) * tri\n\n return vjp\n\n\ndefvjp(eigh, grad_eigh)\n\n\n# https://arxiv.org/pdf/1701.00392.pdf Eq(4.77)\n# Note the formula from Sec3.1 in https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf is incomplete\ndef grad_eig(ans, x):\n \"\"\"Gradient of a general square (complex valued) matrix\"\"\"\n e, u = ans # eigenvalues as 1d array, eigenvectors in columns\n n = e.shape[-1]\n\n def vjp(g):\n ge, gu = g\n ge = _matrix_diag(ge)\n f = 1 / (e[..., anp.newaxis, :] - e[..., :, anp.newaxis] + 1.0e-20)\n f -= _diag(f)\n ut = anp.swapaxes(u, -1, -2)\n r1 = f * _dot(ut, gu)\n r2 = -f * (_dot(_dot(ut, anp.conj(u)), anp.real(_dot(ut, gu)) * anp.eye(n)))\n r = _dot(_dot(inv(ut), ge + r1 + r2), ut)\n if not anp.iscomplexobj(x):\n r = anp.real(r)\n # the derivative is still complex for real input (imaginary delta is allowed), real output\n # but the derivative should be real in real input case when imaginary delta is forbidden\n return r\n\n return vjp\n\n\ndefvjp(eig, grad_eig)\n\n\ndef grad_cholesky(L, A):\n # Based on Iain Murray's note http://arxiv.org/abs/1602.07527\n # scipy's dtrtrs wrapper, solve_triangular, doesn't broadcast along leading\n # dimensions, so we just call a generic LU solve instead of directly using\n # backsubstitution (also, we factor twice...)\n solve_trans = lambda a, b: solve(T(a), b)\n phi = lambda X: anp.tril(X) / (1.0 + anp.eye(X.shape[-1]))\n\n def conjugate_solve(L, X):\n # X -> L^{-T} X L^{-1}\n return solve_trans(L, T(solve_trans(L, T(X))))\n\n def vjp(g):\n S = conjugate_solve(L, phi(anp.einsum(\"...ki,...kj->...ij\", L, g)))\n return (S + T(S)) / 2.0\n\n return vjp\n\n\ndefvjp(cholesky, grad_cholesky)\n\n\n# https://j-towns.github.io/papers/svd-derivative.pdf\n# https://arxiv.org/abs/1909.02659\ndef grad_svd(usv_, a, full_matrices=True, compute_uv=True):\n def vjp(g):\n usv = usv_\n\n if not compute_uv:\n s = usv\n\n # Need U and V so do the whole svd anyway...\n usv = svd(a, full_matrices=False)\n u = usv[0]\n v = anp.conj(T(usv[2]))\n\n return _dot(anp.conj(u) * g[..., anp.newaxis, :], T(v))\n\n elif full_matrices:\n raise NotImplementedError(\"Gradient of svd not implemented for full_matrices=True\")\n\n else:\n u = usv[0]\n s = usv[1]\n v = anp.conj(T(usv[2]))\n\n m, n = a.shape[-2:]\n\n k = anp.min((m, n))\n # broadcastable identity array with shape (1, 1, ..., 1, k, k)\n i = anp.reshape(anp.eye(k), anp.concatenate((anp.ones(a.ndim - 2, dtype=int), (k, k))))\n\n f = 1 / (s[..., anp.newaxis, :] ** 2 - s[..., :, anp.newaxis] ** 2 + i)\n\n gu = g[0]\n gs = g[1]\n gv = anp.conj(T(g[2]))\n\n utgu = _dot(T(u), gu)\n vtgv = _dot(T(v), gv)\n t1 = (f * (utgu - anp.conj(T(utgu)))) * s[..., anp.newaxis, :]\n t1 = t1 + i * gs[..., :, anp.newaxis]\n t1 = t1 + s[..., :, anp.newaxis] * (f * (vtgv - anp.conj(T(vtgv))))\n\n if anp.iscomplexobj(u):\n t1 = t1 + 1j * anp.imag(_diag(utgu)) / s[..., anp.newaxis, :]\n\n t1 = _dot(_dot(anp.conj(u), t1), T(v))\n\n if m < n:\n i_minus_vvt = anp.reshape(\n anp.eye(n), anp.concatenate((anp.ones(a.ndim - 2, dtype=int), (n, n)))\n ) - _dot(v, anp.conj(T(v)))\n t1 = t1 + anp.conj(_dot(_dot(u / s[..., anp.newaxis, :], T(gv)), i_minus_vvt))\n\n return t1\n\n elif m == n:\n return t1\n\n elif m > n:\n i_minus_uut = anp.reshape(\n anp.eye(m), anp.concatenate((anp.ones(a.ndim - 2, dtype=int), (m, m)))\n ) - _dot(u, anp.conj(T(u)))\n t1 = t1 + T(_dot(_dot(v / s[..., anp.newaxis, :], T(gu)), i_minus_uut))\n\n return t1\n\n return vjp\n\n\ndefvjp(svd, grad_svd)\n"
},
{
"path": "autograd/numpy/numpy_boxes.py",
"content": "import numpy as np\n\nfrom autograd.builtins import SequenceBox\nfrom autograd.extend import Box, primitive\n\nfrom . import numpy_wrapper as anp\n\nBox.__array_priority__ = 90.0\n\n\nclass ArrayBox(Box):\n __slots__ = []\n __array_priority__ = 100.0\n\n @primitive\n def __getitem__(A, idx):\n return A[idx]\n\n # Constants w.r.t float data just pass though\n shape = property(lambda self: self._value.shape)\n ndim = property(lambda self: self._value.ndim)\n size = property(lambda self: self._value.size)\n dtype = property(lambda self: self._value.dtype)\n T = property(lambda self: anp.transpose(self))\n\n def __array_namespace__(self, *, api_version: str | None = None):\n return anp\n\n def __len__(self):\n return len(self._value)\n\n def astype(self, *args, **kwargs):\n return anp._astype(self, *args, **kwargs)\n\n def __neg__(self):\n return anp.negative(self)\n\n def __add__(self, other):\n return anp.add(self, other)\n\n def __sub__(self, other):\n return anp.subtract(self, other)\n\n def __mul__(self, other):\n return anp.multiply(self, other)\n\n def __pow__(self, other):\n return anp.power(self, other)\n\n def __div__(self, other):\n return anp.divide(self, other)\n\n def __mod__(self, other):\n return anp.mod(self, other)\n\n def __truediv__(self, other):\n return anp.true_divide(self, other)\n\n def __matmul__(self, other):\n return anp.matmul(self, other)\n\n def __radd__(self, other):\n return anp.add(other, self)\n\n def __rsub__(self, other):\n return anp.subtract(other, self)\n\n def __rmul__(self, other):\n return anp.multiply(other, self)\n\n def __rpow__(self, other):\n return anp.power(other, self)\n\n def __rdiv__(self, other):\n return anp.divide(other, self)\n\n def __rmod__(self, other):\n return anp.mod(other, self)\n\n def __rtruediv__(self, other):\n return anp.true_divide(other, self)\n\n def __rmatmul__(self, other):\n return anp.matmul(other, self)\n\n def __eq__(self, other):\n return anp.equal(self, other)\n\n def __ne__(self, other):\n return anp.not_equal(self, other)\n\n def __gt__(self, other):\n return anp.greater(self, other)\n\n def __ge__(self, other):\n return anp.greater_equal(self, other)\n\n def __lt__(self, other):\n return anp.less(self, other)\n\n def __le__(self, other):\n return anp.less_equal(self, other)\n\n def __abs__(self):\n return anp.abs(self)\n\n def __hash__(self):\n return id(self)\n\n\nArrayBox.register(np.ndarray)\nfor type_ in [\n float,\n np.longdouble,\n np.float64,\n np.float32,\n np.float16,\n complex,\n np.clongdouble,\n np.complex64,\n np.complex128,\n]:\n ArrayBox.register(type_)\n\n# These numpy.ndarray methods are just refs to an equivalent numpy function\nnondiff_methods = [\n \"all\",\n \"any\",\n \"argmax\",\n \"argmin\",\n \"argpartition\",\n \"argsort\",\n \"nonzero\",\n \"searchsorted\",\n \"round\",\n]\ndiff_methods = [\n \"clip\",\n \"compress\",\n \"cumprod\",\n \"cumsum\",\n \"diagonal\",\n \"max\",\n \"mean\",\n \"min\",\n \"prod\",\n \"ptp\",\n \"ravel\",\n \"repeat\",\n \"reshape\",\n \"squeeze\",\n \"std\",\n \"sum\",\n \"swapaxes\",\n \"take\",\n \"trace\",\n \"transpose\",\n \"var\",\n]\nfor method_name in nondiff_methods + diff_methods:\n setattr(ArrayBox, method_name, anp.__dict__[method_name])\n\n# Flatten has no function, only a method.\nsetattr(ArrayBox, \"flatten\", anp.__dict__[\"ravel\"])\n\nif np.lib.NumpyVersion(np.__version__) >= \"2.0.0\":\n SequenceBox.register(np.linalg._linalg.EigResult)\n SequenceBox.register(np.linalg._linalg.EighResult)\n SequenceBox.register(np.linalg._linalg.QRResult)\n SequenceBox.register(np.linalg._linalg.SlogdetResult)\n SequenceBox.register(np.linalg._linalg.SVDResult)\nelif np.__version__ >= \"1.25\":\n SequenceBox.register(np.linalg.linalg.EigResult)\n SequenceBox.register(np.linalg.linalg.EighResult)\n SequenceBox.register(np.linalg.linalg.QRResult)\n SequenceBox.register(np.linalg.linalg.SlogdetResult)\n SequenceBox.register(np.linalg.linalg.SVDResult)\n"
},
{
"path": "autograd/numpy/numpy_jvps.py",
"content": "import numpy as onp\n\nfrom autograd.extend import JVPNode, def_linear, defjvp, defjvp_argnum, register_notrace, vspace\n\nfrom ..util import func\nfrom . import numpy_wrapper as anp\nfrom .numpy_boxes import ArrayBox\nfrom .numpy_vjps import (\n balanced_eq,\n dot_adjoint_0,\n dot_adjoint_1,\n match_complex,\n nograd_functions,\n replace_zero,\n tensordot_adjoint_0,\n tensordot_adjoint_1,\n untake,\n)\n\nfor fun in nograd_functions:\n register_notrace(JVPNode, fun)\n\ndefjvp(func(ArrayBox.__getitem__), \"same\")\ndefjvp(untake, \"same\")\n\ndefjvp_argnum(anp.array_from_args, lambda argnum, g, ans, args, kwargs: untake(g, argnum - 2, vspace(ans)))\ndefjvp(\n anp._array_from_scalar_or_array,\n None,\n None,\n lambda g, ans, args, kwargs, _: anp._array_from_scalar_or_array(args, kwargs, g),\n)\n\n# ----- Functions that are constant w.r.t. continuous inputs -----\ndefjvp(anp.nan_to_num, lambda g, ans, x: anp.where(anp.isfinite(x), g, 0.0))\n\n# ----- Binary ufuncs (linear) -----\ndef_linear(anp.multiply)\n\n# ----- Binary ufuncs -----\ndefjvp(anp.add, lambda g, ans, x, y: broadcast(g, ans), lambda g, ans, x, y: broadcast(g, ans))\ndefjvp(anp.subtract, lambda g, ans, x, y: broadcast(g, ans), lambda g, ans, x, y: broadcast(-g, ans))\ndefjvp(anp.divide, \"same\", lambda g, ans, x, y: -g * x / y**2)\ndefjvp(\n anp.maximum,\n lambda g, ans, x, y: g * balanced_eq(x, ans, y),\n lambda g, ans, x, y: g * balanced_eq(y, ans, x),\n)\ndefjvp(\n anp.minimum,\n lambda g, ans, x, y: g * balanced_eq(x, ans, y),\n lambda g, ans, x, y: g * balanced_eq(y, ans, x),\n)\ndefjvp(\n anp.fmax,\n lambda g, ans, x, y: g * balanced_eq(x, ans, y),\n lambda g, ans, x, y: g * balanced_eq(y, ans, x),\n)\ndefjvp(\n anp.fmin,\n lambda g, ans, x, y: g * balanced_eq(x, ans, y),\n lambda g, ans, x, y: g * balanced_eq(y, ans, x),\n)\ndefjvp(anp.logaddexp, lambda g, ans, x, y: g * anp.exp(x - ans), lambda g, ans, x, y: g * anp.exp(y - ans))\ndefjvp(anp.logaddexp2, lambda g, ans, x, y: g * 2 ** (x - ans), lambda g, ans, x, y: g * 2 ** (y - ans))\ndefjvp(anp.true_divide, \"same\", lambda g, ans, x, y: -g * x / y**2)\ndefjvp(anp.mod, lambda g, ans, x, y: broadcast(g, ans), lambda g, ans, x, y: -g * anp.floor(x / y))\ndefjvp(anp.remainder, lambda g, ans, x, y: broadcast(g, ans), lambda g, ans, x, y: -g * anp.floor(x / y))\ndefjvp(\n anp.power,\n lambda g, ans, x, y: g * y * x ** anp.where(y, y - 1, 1.0),\n lambda g, ans, x, y: g * anp.log(replace_zero(x, 1.0)) * ans,\n)\ndefjvp(anp.arctan2, lambda g, ans, x, y: g * y / (x**2 + y**2), lambda g, ans, x, y: g * -x / (x**2 + y**2))\n\n# ----- Simple grads (linear) -----\ndefjvp(anp.negative, \"same\")\ndefjvp(anp.rad2deg, \"same\")\ndefjvp(anp.degrees, \"same\")\ndefjvp(anp.deg2rad, \"same\")\ndefjvp(anp.radians, \"same\")\ndefjvp(anp.reshape, \"same\")\ndefjvp(anp.roll, \"same\")\ndefjvp(anp.array_split, \"same\")\ndefjvp(anp.split, \"same\")\ndefjvp(anp.vsplit, \"same\")\ndefjvp(anp.hsplit, \"same\")\ndefjvp(anp.dsplit, \"same\")\ndefjvp(anp.ravel, \"same\")\ndefjvp(anp.expand_dims, \"same\")\ndefjvp(anp.squeeze, \"same\")\ndefjvp(anp.diag, \"same\")\ndefjvp(anp.diagonal, \"same\")\ndefjvp(anp.make_diagonal, \"same\")\ndefjvp(anp.flipud, \"same\")\ndefjvp(anp.fliplr, \"same\")\ndefjvp(anp.rot90, \"same\")\ndefjvp(anp.trace, \"same\")\ndefjvp(anp.full, \"same\", argnums=(1,))\ndefjvp(anp.triu, \"same\")\ndefjvp(anp.tril, \"same\")\ndefjvp(anp.swapaxes, \"same\")\ndefjvp(anp.rollaxis, \"same\")\ndefjvp(anp.moveaxis, \"same\")\ndefjvp(anp.broadcast_to, \"same\")\ndef_linear(anp.cross)\n\n# ----- Simple grads -----\nnp_abs_jvp = lambda g, ans, x: anp.real(g * replace_zero(anp.conj(x), 0.0)) / replace_zero(ans, 1.0)\ndefjvp(anp.abs, np_abs_jvp)\ndefjvp(anp.absolute, np_abs_jvp)\ndefjvp(anp.fabs, lambda g, ans, x: anp.sign(x) * g) # fabs doesn't take complex numbers.\ndefjvp(anp.reciprocal, lambda g, ans, x: -g / x**2)\ndefjvp(anp.exp, lambda g, ans, x: ans * g)\ndefjvp(anp.exp2, lambda g, ans, x: ans * anp.log(2) * g)\ndefjvp(anp.expm1, lambda g, ans, x: (ans + 1) * g)\ndefjvp(anp.log, lambda g, ans, x: g / x)\ndefjvp(anp.log2, lambda g, ans, x: g / x / anp.log(2))\ndefjvp(anp.log10, lambda g, ans, x: g / x / anp.log(10))\ndefjvp(anp.log1p, lambda g, ans, x: g / (x + 1))\ndefjvp(anp.sin, lambda g, ans, x: g * anp.cos(x))\ndefjvp(anp.cos, lambda g, ans, x: -g * anp.sin(x))\ndefjvp(anp.tan, lambda g, ans, x: g / anp.cos(x) ** 2)\ndefjvp(anp.arcsin, lambda g, ans, x: g / anp.sqrt(1 - x**2))\ndefjvp(anp.arccos, lambda g, ans, x: -g / anp.sqrt(1 - x**2))\ndefjvp(anp.arctan, lambda g, ans, x: g / (1 + x**2))\ndefjvp(anp.sinh, lambda g, ans, x: g * anp.cosh(x))\ndefjvp(anp.cosh, lambda g, ans, x: g * anp.sinh(x))\ndefjvp(anp.tanh, lambda g, ans, x: g / anp.cosh(x) ** 2)\ndefjvp(anp.arcsinh, lambda g, ans, x: g / anp.sqrt(x**2 + 1))\ndefjvp(anp.arccosh, lambda g, ans, x: g / anp.sqrt(x**2 - 1))\ndefjvp(anp.arctanh, lambda g, ans, x: g / (1 - x**2))\ndefjvp(anp.square, lambda g, ans, x: g * 2 * x)\ndefjvp(anp.sqrt, lambda g, ans, x: g * 0.5 * x**-0.5)\ndefjvp(\n anp.sinc,\n lambda g, ans, x: g * (anp.cos(anp.pi * x) * anp.pi * x - anp.sin(anp.pi * x)) / (anp.pi * x**2),\n)\ndefjvp(anp.clip, lambda g, ans, x, a_min, a_max: g * anp.logical_and(ans != a_min, ans != a_max))\ndefjvp(anp.real_if_close, lambda g, ans, x: match_complex(ans, g))\ndefjvp(anp.real, lambda g, ans, x: anp.real(g))\ndefjvp(anp.imag, lambda g, ans, x: match_complex(ans, -1j * g))\nnp_conj_jvp = lambda g, ans, x: anp.conj(g)\ndefjvp(anp.conj, np_conj_jvp)\ndefjvp(anp.conjugate, np_conj_jvp)\ndefjvp(anp.angle, lambda g, ans, x: match_complex(ans, g * anp.conj(x * 1j) / anp.abs(x) ** 2))\ndefjvp(\n anp.where,\n None,\n lambda g, ans, c, x=None, y=None: anp.where(c, g, anp.zeros(anp.shape(g))),\n lambda g, ans, c, x=None, y=None: anp.where(c, anp.zeros(g.shape), g),\n)\n\n# ----- Trickier grads -----\ndefjvp(anp.kron, \"same\", \"same\")\ndefjvp(anp.diff, \"same\")\ndefjvp(anp.gradient, \"same\")\ndefjvp(anp.repeat, \"same\")\ndefjvp(anp.tile, \"same\")\ndefjvp(anp.transpose, \"same\")\ndefjvp(anp.sum, \"same\")\ndefjvp(anp.mean, \"same\")\ndefjvp(\n anp.prod, lambda g, ans, x, axis=None, keepdims=False: ans * anp.sum(g / x, axis=axis, keepdims=keepdims)\n)\ndefjvp(\n anp.linspace,\n lambda g, ans, start, stop, *args, **kwargs: anp.linspace(g, 0, *args, **kwargs),\n lambda g, ans, start, stop, *args, **kwargs: anp.linspace(0, g, *args, **kwargs),\n)\n\n\ndef forward_grad_np_var(g, ans, x, axis=None, ddof=0, keepdims=False):\n if axis is None:\n num_reps = anp.size(g)\n elif isinstance(axis, int):\n num_reps = anp.shape(g)[axis]\n elif isinstance(axis, tuple):\n num_reps = anp.prod(anp.array(np.shape(g))[list(axis)])\n\n x_minus_mean = anp.conj(x - anp.mean(x, axis=axis, keepdims=True))\n return 2.0 * anp.sum(anp.real(g * x_minus_mean), axis=axis, keepdims=keepdims) / (num_reps - ddof)\n\n\ndefjvp(anp.var, forward_grad_np_var)\n\n\ndef forward_grad_np_std(g, ans, x, axis=None, ddof=0, keepdims=False):\n if axis is None:\n num_reps = anp.size(g)\n elif isinstance(axis, int):\n num_reps = anp.shape(g)[axis]\n elif isinstance(axis, tuple):\n num_reps = anp.prod(anp.array(anp.shape(g))[list(axis)])\n\n if num_reps <= 1:\n return anp.zeros_like(ans)\n x_minus_mean = anp.conj(x - anp.mean(x, axis=axis, keepdims=True))\n return anp.sum(anp.real(g * x_minus_mean), axis=axis, keepdims=keepdims) / ((num_reps - ddof) * ans)\n\n\ndefjvp(anp.std, forward_grad_np_std)\n\n\ndef fwd_grad_chooser(g, ans, x, axis=None, keepdims=False):\n if anp.isscalar(x):\n return g\n if not keepdims:\n if isinstance(axis, int):\n ans = anp.expand_dims(ans, axis)\n elif isinstance(axis, tuple):\n for ax in sorted(axis):\n ans = anp.expand_dims(ans, ax)\n chosen_locations = x == ans\n return anp.sum((g * chosen_locations), axis=axis, keepdims=keepdims) / anp.sum(\n chosen_locations, axis=axis, keepdims=keepdims\n )\n\n\ndefjvp(anp.max, fwd_grad_chooser)\ndefjvp(anp.min, fwd_grad_chooser)\ndefjvp(anp.amax, fwd_grad_chooser)\ndefjvp(anp.amin, fwd_grad_chooser)\n\ndefjvp(anp.cumsum, \"same\")\n\ndef_linear(anp.inner)\ndef_linear(anp.matmul)\ndef_linear(anp.dot)\ndef_linear(anp.tensordot)\ndef_linear(anp.outer)\n\ndef_linear(dot_adjoint_0)\ndef_linear(dot_adjoint_1)\n\ndef_linear(tensordot_adjoint_0)\ndef_linear(tensordot_adjoint_1)\n\n\ndef fwd_grad_concatenate_args(argnum, g, ans, axis_args, kwargs):\n result = []\n for i in range(1, len(axis_args)):\n if i == argnum:\n result.append(g)\n else:\n result.append(anp.zeros_like(axis_args[i]))\n return anp.concatenate_args(axis_args[0], *result)\n\n\ndefjvp_argnum(anp.concatenate_args, fwd_grad_concatenate_args)\n\n\ndef fwd_grad_sort(g, ans, x, axis=-1, kind=\"quicksort\", order=None):\n sort_perm = anp.argsort(x, axis, kind, order)\n return g[sort_perm]\n\n\ndefjvp(anp.sort, fwd_grad_sort)\nif onp.lib.NumpyVersion(onp.__version__) < \"2.0.0\":\n defjvp(anp.msort, lambda g, ans, x: fwd_grad_sort(g, ans, x, axis=0))\n\n\ndef fwd_grad_partition(g, ans, x, kth, axis=-1, kind=\"introselect\", order=None):\n partition_perm = anp.argpartition(x, kth, axis, kind, order)\n return g[partition_perm]\n\n\ndefjvp(anp.partition, fwd_grad_partition)\n\n\ndef atleast_jvpmaker(fun):\n def jvp(g, ans, *arys):\n if len(arys) > 1:\n raise NotImplementedError(\"Can't handle multiple arguments yet.\")\n return fun(g)\n\n return jvp\n\n\ndefjvp(anp.atleast_1d, atleast_jvpmaker(anp.atleast_1d))\ndefjvp(anp.atleast_2d, atleast_jvpmaker(anp.atleast_2d))\ndefjvp(anp.atleast_3d, atleast_jvpmaker(anp.atleast_3d))\n\ndef_linear(anp.einsum)\n\n\n# TODO(mattjj): can we call np.broadcast_to or a related function instead?\ndef broadcast(x, target):\n target_shape, target_ndim, target_dtype, target_iscomplex = anp.metadata(target)\n while anp.ndim(x) < target_ndim:\n x = anp.expand_dims(x, 0)\n for axis, size in enumerate(anp.shape(x)):\n if size == 1:\n x = anp.repeat(x, target_shape[axis], axis=axis)\n if target_iscomplex and not anp.iscomplexobj(x):\n x = x + 0j # TODO(mattjj): this might promote the dtype\n return x\n\n\ndefjvp(anp.pad, lambda g, ans, array, width, mode, **kwargs: anp.pad(g, width, mode))\n"
},
{
"path": "autograd/numpy/numpy_vjps.py",
"content": "from functools import partial\n\nimport numpy as onp\n\nfrom autograd.extend import SparseObject, VJPNode, defvjp, defvjp_argnum, primitive, register_notrace, vspace\n\nfrom ..util import func\nfrom . import numpy_wrapper as anp\nfrom .numpy_boxes import ArrayBox\n\n# ----- Non-differentiable functions -----\n\nnograd_functions = [\n anp.floor,\n anp.ceil,\n anp.round,\n anp.rint,\n anp.around,\n anp.fix,\n anp.trunc,\n anp.all,\n anp.any,\n anp.argmax,\n anp.argmin,\n anp.argpartition,\n anp.argsort,\n anp.argwhere,\n anp.nonzero,\n anp.flatnonzero,\n anp.count_nonzero,\n anp.searchsorted,\n anp.sign,\n anp.ndim,\n anp.shape,\n anp.floor_divide,\n anp.logical_and,\n anp.logical_or,\n anp.logical_not,\n anp.logical_xor,\n anp.isfinite,\n anp.isinf,\n anp.isnan,\n anp.isneginf,\n anp.isposinf,\n anp.allclose,\n anp.isclose,\n anp.array_equal,\n anp.array_equiv,\n anp.greater,\n anp.greater_equal,\n anp.less,\n anp.less_equal,\n anp.equal,\n anp.not_equal,\n anp.iscomplexobj,\n anp.iscomplex,\n anp.size,\n anp.isscalar,\n anp.isreal,\n anp.zeros_like,\n anp.ones_like,\n anp.empty_like,\n anp.full_like,\n anp.result_type,\n]\n\nfor fun in nograd_functions:\n register_notrace(VJPNode, fun)\n\n# ----- Functions that are constant w.r.t. continuous inputs -----\n\ndefvjp(anp.nan_to_num, lambda ans, x: lambda g: anp.where(anp.isfinite(x), g, 0.0))\n\n# ----- Binary ufuncs -----\n\ndefvjp(\n anp.add, lambda ans, x, y: unbroadcast_f(x, lambda g: g), lambda ans, x, y: unbroadcast_f(y, lambda g: g)\n)\ndefvjp(\n anp.multiply,\n lambda ans, x, y: unbroadcast_f(x, lambda g: y * g),\n lambda ans, x, y: unbroadcast_f(y, lambda g: x * g),\n)\ndefvjp(\n anp.subtract,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g),\n lambda ans, x, y: unbroadcast_f(y, lambda g: -g),\n)\ndefvjp(\n anp.divide,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g / y),\n lambda ans, x, y: unbroadcast_f(y, lambda g: -g * x / y**2),\n)\ndefvjp(\n anp.maximum,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * balanced_eq(x, ans, y)),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * balanced_eq(y, ans, x)),\n)\ndefvjp(\n anp.minimum,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * balanced_eq(x, ans, y)),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * balanced_eq(y, ans, x)),\n)\ndefvjp(\n anp.fmax,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * balanced_eq(x, ans, y)),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * balanced_eq(y, ans, x)),\n)\ndefvjp(\n anp.fmin,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * balanced_eq(x, ans, y)),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * balanced_eq(y, ans, x)),\n)\ndefvjp(\n anp.logaddexp,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * anp.exp(x - ans)),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * anp.exp(y - ans)),\n)\ndefvjp(\n anp.logaddexp2,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * 2 ** (x - ans)),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * 2 ** (y - ans)),\n)\ndefvjp(\n anp.true_divide,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g / y),\n lambda ans, x, y: unbroadcast_f(y, lambda g: -g * x / y**2),\n)\ndefvjp(\n anp.mod,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g),\n lambda ans, x, y: unbroadcast_f(y, lambda g: -g * anp.floor(x / y)),\n)\ndefvjp(\n anp.remainder,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g),\n lambda ans, x, y: unbroadcast_f(y, lambda g: -g * anp.floor(x / y)),\n)\ndefvjp(\n anp.power,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * y * x ** anp.where(y, y - 1, 1.0)),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * anp.log(replace_zero(x, 1.0)) * ans),\n)\ndefvjp(\n anp.arctan2,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * y / (x**2 + y**2)),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * -x / (x**2 + y**2)),\n)\ndefvjp(\n anp.hypot,\n lambda ans, x, y: unbroadcast_f(x, lambda g: g * x / ans),\n lambda ans, x, y: unbroadcast_f(y, lambda g: g * y / ans),\n)\n\n# ----- Simple grads -----\n\ndefvjp(anp.negative, lambda ans, x: lambda g: -g)\nnp_abs_vjp = lambda ans, x: lambda g: g * replace_zero(anp.conj(x), 0.0) / replace_zero(ans, 1.0)\ndefvjp(anp.abs, np_abs_vjp)\ndefvjp(anp.absolute, np_abs_vjp)\ndefvjp(anp.fabs, lambda ans, x: lambda g: anp.sign(x) * g) # fabs doesn't take complex numbers.\ndefvjp(anp.reciprocal, lambda ans, x: lambda g: -g / x**2)\ndefvjp(anp.exp, lambda ans, x: lambda g: ans * g)\ndefvjp(anp.exp2, lambda ans, x: lambda g: ans * anp.log(2) * g)\ndefvjp(anp.expm1, lambda ans, x: lambda g: (ans + 1) * g)\ndefvjp(anp.log, lambda ans, x: lambda g: g / x)\ndefvjp(anp.log2, lambda ans, x: lambda g: g / x / anp.log(2))\ndefvjp(anp.log10, lambda ans, x: lambda g: g / x / anp.log(10))\ndefvjp(anp.log1p, lambda ans, x: lambda g: g / (x + 1))\ndefvjp(anp.sin, lambda ans, x: lambda g: g * anp.cos(x))\ndefvjp(anp.cos, lambda ans, x: lambda g: -g * anp.sin(x))\ndefvjp(anp.tan, lambda ans, x: lambda g: g / anp.cos(x) ** 2)\ndefvjp(anp.arcsin, lambda ans, x: lambda g: g / anp.sqrt(1 - x**2))\ndefvjp(anp.arccos, lambda ans, x: lambda g: -g / anp.sqrt(1 - x**2))\ndefvjp(anp.arctan, lambda ans, x: lambda g: g / (1 + x**2))\ndefvjp(anp.sinh, lambda ans, x: lambda g: g * anp.cosh(x))\ndefvjp(anp.cosh, lambda ans, x: lambda g: g * anp.sinh(x))\ndefvjp(anp.tanh, lambda ans, x: lambda g: g / anp.cosh(x) ** 2)\ndefvjp(anp.arcsinh, lambda ans, x: lambda g: g / anp.sqrt(x**2 + 1))\ndefvjp(anp.arccosh, lambda ans, x: lambda g: g / anp.sqrt(x**2 - 1))\ndefvjp(anp.arctanh, lambda ans, x: lambda g: g / (1 - x**2))\ndefvjp(anp.rad2deg, lambda ans, x: lambda g: g / anp.pi * 180.0)\ndefvjp(anp.degrees, lambda ans, x: lambda g: g / anp.pi * 180.0)\ndefvjp(anp.deg2rad, lambda ans, x: lambda g: g * anp.pi / 180.0)\ndefvjp(anp.radians, lambda ans, x: lambda g: g * anp.pi / 180.0)\ndefvjp(anp.square, lambda ans, x: lambda g: g * 2 * x)\ndefvjp(anp.sqrt, lambda ans, x: lambda g: g * 0.5 * x**-0.5)\ndefvjp(\n anp.sinc,\n lambda ans, x: lambda g: g * (anp.cos(anp.pi * x) * anp.pi * x - anp.sin(anp.pi * x)) / (anp.pi * x**2),\n)\ndefvjp(anp.reshape, lambda ans, x, shape, order=None: lambda g: anp.reshape(g, anp.shape(x), order=order))\ndefvjp(anp.roll, lambda ans, x, shift, axis=None: lambda g: anp.roll(g, -shift, axis=axis))\ndefvjp(anp.array_split, lambda ans, ary, idxs, axis=0: lambda g: anp.concatenate(g, axis=axis))\ndefvjp(anp.split, lambda ans, ary, idxs, axis=0: lambda g: anp.concatenate(g, axis=axis))\ndefvjp(anp.vsplit, lambda ans, ary, idxs: lambda g: anp.concatenate(g, axis=0))\ndefvjp(anp.hsplit, lambda ans, ary, idxs: lambda g: anp.concatenate(g, axis=1))\ndefvjp(anp.dsplit, lambda ans, ary, idxs: lambda g: anp.concatenate(g, axis=2))\ndefvjp(anp.ravel, lambda ans, x, order=None: lambda g: anp.reshape(g, anp.shape(x), order=order))\ndefvjp(anp.expand_dims, lambda ans, x, axis: lambda g: anp.reshape(g, anp.shape(x)))\ndefvjp(anp.squeeze, lambda ans, x, axis=None: lambda g: anp.reshape(g, anp.shape(x)))\ndefvjp(anp.diag, lambda ans, x, k=0: lambda g: anp.diag(g, k))\ndefvjp(anp.flipud, lambda ans, x,: lambda g: anp.flipud(g))\ndefvjp(anp.fliplr, lambda ans, x,: lambda g: anp.fliplr(g))\ndefvjp(anp.rot90, lambda ans, x, k=1: lambda g: anp.rot90(g, -k))\ndefvjp(\n anp.trace,\n lambda ans, x, offset=0: (\n lambda g: anp.einsum(\"ij,...->ij...\", anp.eye(x.shape[0], x.shape[1], k=offset), g)\n ),\n)\ndefvjp(anp.full, lambda ans, shape, fill_value, dtype=None: lambda g: anp.sum(g), argnums=(1,))\ndefvjp(anp.triu, lambda ans, x, k=0: lambda g: anp.triu(g, k=k))\ndefvjp(anp.tril, lambda ans, x, k=0: lambda g: anp.tril(g, k=k))\ndefvjp(anp.clip, lambda ans, x, a_min, a_max: lambda g: g * anp.logical_and(ans != a_min, ans != a_max))\ndefvjp(anp.swapaxes, lambda ans, x, axis1, axis2: lambda g: anp.swapaxes(g, axis2, axis1))\ndefvjp(anp.moveaxis, lambda ans, a, source, destination: lambda g: anp.moveaxis(g, destination, source))\ndefvjp(anp.real_if_close, lambda ans, x: lambda g: match_complex(x, g))\ndefvjp(anp.real, lambda ans, x: lambda g: match_complex(x, g))\ndefvjp(anp.imag, lambda ans, x: lambda g: match_complex(x, -1j * g))\nnp_conj_vjp = lambda ans, x: lambda g: anp.conj(g)\ndefvjp(anp.conj, np_conj_vjp)\ndefvjp(anp.conjugate, np_conj_vjp)\ndefvjp(anp.angle, lambda ans, x: lambda g: match_complex(x, g * anp.conj(x * 1j) / anp.abs(x) ** 2))\ndefvjp(\n anp.where,\n None,\n lambda ans, c, x=None, y=None: lambda g: anp.where(c, g, anp.zeros(g.shape)),\n lambda ans, c, x=None, y=None: lambda g: anp.where(c, anp.zeros(g.shape), g),\n)\ndefvjp(\n anp.cross,\n lambda ans, a, b, axisa=-1, axisb=-1, axisc=-1, axis=None: (\n lambda g: anp.cross(b, g, axisb, axisc, axisa, axis)\n ),\n lambda ans, a, b, axisa=-1, axisb=-1, axisc=-1, axis=None: (\n lambda g: anp.cross(g, a, axisc, axisa, axisb, axis)\n ),\n)\ndefvjp(\n anp.linspace,\n lambda ans, start, stop, num: lambda g: anp.dot(anp.linspace(1.0, 0.0, num), g),\n lambda ans, start, stop, num: lambda g: anp.dot(anp.linspace(0.0, 1.0, num), g),\n)\n\ndefvjp(\n anp._astype,\n lambda ans, A, dtype, order=\"K\", casting=\"unsafe\", subok=True, copy=True: (\n lambda g: anp._astype(g, A.dtype)\n ),\n)\n\n\n# ----- Trickier grads -----\ndef grad_rollaxis(ans, a, axis, start=0):\n if axis < 0:\n raise NotImplementedError(\n \"Gradient of rollaxis not implemented for axis < 0. Please use moveaxis instead.\"\n )\n elif start < 0:\n raise NotImplementedError(\n \"Gradient of rollaxis not implemented for start < 0. Please use moveaxis instead.\"\n )\n return lambda g: anp.rollaxis(g, start - 1, axis) if start > axis else anp.rollaxis(g, start, axis + 1)\n\n\ndefvjp(anp.rollaxis, grad_rollaxis)\n\n\ndef grad_diff(ans, a, n=1, axis=-1):\n nd = anp.ndim(a)\n ans_shape = anp.shape(ans)\n sl1 = [slice(None)] * nd\n sl1[axis] = slice(None, 1)\n\n sl2 = [slice(None)] * nd\n sl2[axis] = slice(-1, None)\n\n def undiff(g):\n if g.shape[axis] > 0:\n return anp.concatenate((-g[tuple(sl1)], -anp.diff(g, axis=axis), g[tuple(sl2)]), axis=axis)\n shape = list(ans_shape)\n shape[axis] = 1\n return anp.zeros(shape)\n\n def helper(g, n):\n if n == 0:\n return g\n return helper(undiff(g), n - 1)\n\n return lambda g: helper(g, n)\n\n\ndefvjp(anp.diff, grad_diff)\n\n\ndef grad_gradient(ans, x, *vargs, **kwargs):\n axis = kwargs.pop(\"axis\", None)\n if vargs or kwargs:\n raise NotImplementedError(\"The only optional argument currently supported for np.gradient is axis.\")\n if axis is None:\n axis = range(x.ndim)\n elif type(axis) is int:\n axis = [axis]\n else:\n axis = list(axis)\n\n x_dtype = x.dtype\n x_shape = x.shape\n nd = x.ndim\n\n def vjp(g):\n if anp.ndim(g) == nd:\n # add axis if gradient was along one axis only\n g = g[anp.newaxis]\n\n # accumulate gradient\n out = anp.zeros(x_shape, dtype=x_dtype)\n\n for i, a in enumerate(axis):\n # swap gradient axis to the front\n g_swap = anp.swapaxes(g[i], 0, a)[:, anp.newaxis]\n\n out_axis = anp.concatenate(\n (\n -g_swap[0] - 0.5 * g_swap[1],\n g_swap[0] - 0.5 * g_swap[2],\n (-1.0) * anp.gradient(g_swap, axis=0)[2:-2, 0],\n 0.5 * g_swap[-3] - g_swap[-1],\n 0.5 * g_swap[-2] + g_swap[-1],\n ),\n axis=0,\n )\n\n out = out + anp.swapaxes(out_axis, 0, a)\n\n return out\n\n return vjp\n\n\ndefvjp(anp.gradient, grad_gradient)\n\n\ndef grad_repeat(ans, x, repeats, axis=None):\n shape = anp.shape(x)\n\n def vjp(g):\n if axis is None: # If axis is none, np.repeat() repeats the flattened array.\n expanded = anp.reshape(g, (anp.prod(shape),) + (repeats,))\n return anp.reshape(anp.sum(expanded, axis=1, keepdims=False), shape)\n else:\n if shape[axis] == 1: # For this common case, the logic is simple.\n return anp.sum(g, axis=axis, keepdims=True)\n else:\n expanded = anp.reshape(g, shape[0 : axis + 1] + (repeats,) + shape[axis + 1 :])\n return anp.sum(expanded, axis=axis + 1, keepdims=False)\n\n return vjp\n\n\ndefvjp(anp.repeat, grad_repeat)\n\n\ndef grad_tile(ans, x, reps):\n reps = [reps] if anp.isscalar(reps) else reps\n x_shape = anp.shape(x)\n\n def vjp(g):\n for axis, rep in enumerate(reps):\n g = sum(anp.split(g, rep, axis))\n return anp.reshape(g, x_shape)\n\n return vjp\n\n\ndefvjp(anp.tile, grad_tile)\n\n\ndef grad_kron(argnum, ans, orig_A, orig_B):\n # kron has different promotion rules than dot. the reshapes are necessary if\n # and only if (1) orig_B is 1D or (2) orig_A and/or orig_B are 0D\n orig_A_shape = anp.shape(orig_A)\n orig_B_shape = anp.shape(orig_B)\n\n def vjp(G):\n A, B = anp.atleast_2d(orig_A), anp.atleast_2d(orig_B)\n shape = list(A.shape + B.shape)\n n = anp.ndim(A)\n shape[n - 1], shape[n] = shape[n], shape[n - 1]\n reshaped_G = anp.swapaxes(anp.reshape(G, shape), n - 1, n)\n if argnum == 0:\n return match_complex(\n orig_A, anp.reshape(anp.tensordot(reshaped_G, B, axes=anp.ndim(B)), orig_A_shape)\n )\n else:\n return match_complex(\n orig_B, anp.reshape(anp.tensordot(A, reshaped_G, axes=anp.ndim(A)), orig_B_shape)\n )\n\n return vjp\n\n\ndefvjp(anp.kron, partial(grad_kron, 0), partial(grad_kron, 1))\n\n\ndef grad_transpose(ans, x, axes=None):\n if axes is not None:\n axes = anp.argsort(axes)\n return lambda g: anp.transpose(g, axes)\n\n\ndefvjp(anp.transpose, grad_transpose)\n\n\ndef repeat_to_match_shape(g, shape, dtype, axis, keepdims):\n \"\"\"Returns the array g repeated along axis to fit vector space vs.\n Also returns the number of repetitions of the array.\"\"\"\n if shape == ():\n return g, 1\n axis = list(axis) if isinstance(axis, tuple) else axis\n new_shape = onp.array(shape)\n new_shape[axis] = 1\n num_reps = onp.prod(onp.array(shape)[axis])\n # Can't use broadcast_to because of numpy bug: https://github.com/numpy/numpy/issues/9165\n # return anp.broadcast_to(anp.reshape(g, new_shape), shape), num_reps\n return anp.reshape(g, new_shape) + onp.zeros(shape, dtype=dtype), num_reps\n\n\ndef grad_broadcast_to(ans, x, new_shape):\n old_shape = anp.shape(x)\n assert anp.shape(ans) == new_shape\n assert len(old_shape) == len(new_shape), \"Can't handle extra leading dims\"\n broadcast_axes = tuple(\n onp.where(onp.logical_and(onp.array(old_shape) == 1, onp.array(new_shape) > 1))[0]\n )\n return lambda g: anp.sum(g, axis=broadcast_axes, keepdims=True)\n\n\ndefvjp(anp.broadcast_to, grad_broadcast_to)\n\n\ndef grad_np_sum(ans, x, axis=None, keepdims=False, dtype=None):\n shape, dtype = anp.shape(x), anp.result_type(x)\n return lambda g: repeat_to_match_shape(g, shape, dtype, axis, keepdims)[0]\n\n\ndefvjp(anp.sum, grad_np_sum)\n\n\ndef grad_np_mean(ans, x, axis=None, keepdims=False):\n shape, dtype = anp.shape(x), anp.result_type(x)\n\n def vjp(g):\n g_repeated, num_reps = repeat_to_match_shape(g, shape, dtype, axis, keepdims)\n return g_repeated / num_reps\n\n return vjp\n\n\ndefvjp(anp.mean, grad_np_mean)\n\n\ndef grad_np_prod(ans, x, axis=None, keepdims=False): # TODO: Support tuples of axes.\n shape, dtype = anp.shape(x), anp.result_type(x)\n\n def vjp(g):\n g_repeated, _ = repeat_to_match_shape(g * ans, shape, dtype, axis, keepdims)\n return g_repeated / x\n\n return vjp\n\n\ndefvjp(anp.prod, grad_np_prod)\n\n\ndef grad_np_var(ans, x, axis=None, ddof=0, keepdims=False):\n shape, _, dtype, iscomplex = anp.metadata(x)\n\n def vjp(g):\n if iscomplex:\n g = g + 0j\n g_repeated, num_reps = repeat_to_match_shape(g, shape, dtype, axis, keepdims)\n x_minus_mean = anp.conj(x - anp.mean(x, axis=axis, keepdims=True))\n return 2.0 * g_repeated * x_minus_mean / (num_reps - ddof)\n\n return vjp\n\n\ndefvjp(anp.var, grad_np_var)\n\n\ndef grad_np_std(ans, x, axis=None, ddof=0, keepdims=False):\n shape, _, dtype, iscomplex = anp.metadata(x)\n\n def vjp(g):\n if iscomplex:\n g = g + 0j\n g_repeated, num_reps = repeat_to_match_shape(\n g, shape, dtype, axis, keepdims\n ) # Avoid division by zero.\n if num_reps <= 1:\n return g_repeated * 0.0\n else:\n g_repeated, num_reps = repeat_to_match_shape(g / ans, shape, dtype, axis, keepdims)\n x_minus_mean = anp.conj(x - anp.mean(x, axis=axis, keepdims=True))\n return g_repeated * x_minus_mean / (num_reps - ddof)\n\n return vjp\n\n\ndefvjp(anp.std, grad_np_std)\n\n\ndef grad_chooser(ans, x, axis=None, keepdims=None):\n shape, dtype = anp.shape(x), anp.result_type(x)\n\n def vjp(g):\n \"\"\"Builds gradient of functions that choose a single item, such as min or max.\"\"\"\n g_repeated, _ = repeat_to_match_shape(g, shape, dtype, axis, keepdims)\n argmax_locations = x == repeat_to_match_shape(ans, shape, dtype, axis, keepdims)[0]\n return g_repeated * argmax_locations / onp.sum(argmax_locations, axis=axis, keepdims=True)\n\n return vjp\n\n\ndefvjp(anp.max, grad_chooser)\ndefvjp(anp.min, grad_chooser)\ndefvjp(anp.amax, grad_chooser)\ndefvjp(anp.amin, grad_chooser)\n\n\ndef reverse_axis(x, axis):\n x = x.swapaxes(axis, 0)\n x = x[::-1, ...]\n return x.swapaxes(0, axis)\n\n\ndef grad_np_cumsum(ans, x, axis=None):\n def vjp(g):\n if axis:\n return reverse_axis(anp.cumsum(reverse_axis(g, axis), axis), axis)\n else:\n return anp.reshape(anp.cumsum(g[::-1], axis)[::-1], x.shape)\n\n return vjp\n\n\ndefvjp(anp.cumsum, grad_np_cumsum)\n\n\ndef grad_inner(argnum, ans, A, B):\n A_ndim, B_ndim = anp.ndim(A), anp.ndim(B)\n if A_ndim == 0 or B_ndim == 0:\n axes = ([], [])\n else:\n axes = ([A_ndim - 1], [B_ndim - 1])\n if argnum == 0:\n return lambda G: tensordot_adjoint_0(B, G, axes, A_ndim, B_ndim)\n elif argnum == 1:\n return lambda G: tensordot_adjoint_1(A, G, axes, A_ndim, B_ndim)\n\n\ndefvjp(anp.inner, partial(grad_inner, 0), partial(grad_inner, 1))\n\n\ndef matmul_adjoint_0(B, G, A_meta, B_ndim):\n if anp.ndim(G) == 0: # A_ndim == B_ndim == 1\n return unbroadcast(G * B, A_meta)\n _, A_ndim, _, _ = A_meta\n if A_ndim == 1:\n G = anp.expand_dims(G, anp.ndim(G) - 1)\n if B_ndim == 1: # The result we need is an outer product\n B = anp.expand_dims(B, 0)\n G = anp.expand_dims(G, anp.ndim(G))\n else: # We need to swap the last two axes of B\n B = anp.swapaxes(B, B_ndim - 2, B_ndim - 1)\n result = anp.matmul(G, B)\n return unbroadcast(result, A_meta)\n\n\ndef matmul_adjoint_1(A, G, A_ndim, B_meta):\n if anp.ndim(G) == 0: # A_ndim == B_ndim == 1\n return unbroadcast(G * A, B_meta)\n _, B_ndim, _, _ = B_meta\n B_is_vec = B_ndim == 1\n if B_is_vec:\n G = anp.expand_dims(G, anp.ndim(G))\n if A_ndim == 1: # The result we need is an outer product\n A = anp.expand_dims(A, 1)\n G = anp.expand_dims(G, anp.ndim(G) - 1)\n else: # We need to swap the last two axes of A\n A = anp.swapaxes(A, A_ndim - 2, A_ndim - 1)\n result = anp.matmul(A, G)\n if B_is_vec:\n result = anp.squeeze(result, anp.ndim(G) - 1)\n return unbroadcast(result, B_meta)\n\n\ndef matmul_vjp_0(ans, A, B):\n A_meta = anp.metadata(A)\n B_ndim = anp.ndim(B)\n return lambda g: matmul_adjoint_0(B, g, A_meta, B_ndim)\n\n\ndef matmul_vjp_1(ans, A, B):\n A_ndim = anp.ndim(A)\n B_meta = anp.metadata(B)\n return lambda g: matmul_adjoint_1(A, g, A_ndim, B_meta)\n\n\ndefvjp(anp.matmul, matmul_vjp_0, matmul_vjp_1)\n\n\n@primitive\ndef dot_adjoint_0(B, G, A_meta, B_meta):\n _, A_ndim, A_dtype, _ = A_meta\n _, B_ndim, _, _ = B_meta\n if B_ndim == 0 or B_ndim == 1 or A_ndim == 0:\n contract_num = max(0, B_ndim - (A_ndim != 0))\n out = onp.tensordot(G, B, contract_num)\n else:\n out = onp.tensordot(G, onp.swapaxes(B, -1, -2), B_ndim - 1)\n return onp.asarray(out, dtype=A_dtype)\n\n\n@primitive\ndef dot_adjoint_1(A, G, A_meta, B_meta):\n _, A_ndim, _, _ = A_meta\n _, B_ndim, B_dtype, _ = B_meta\n needs_transpose = B_ndim > 1 and A_ndim != 0\n swap = (lambda x: onp.swapaxes(x, -1, -2)) if needs_transpose else (lambda x: x)\n if A_ndim == 0 or A_ndim == 1 or B_ndim == 0:\n contract_num = max(0, A_ndim - (B_ndim != 0))\n out = swap(onp.tensordot(G, A, contract_num))\n else:\n out = swap(onp.tensordot(G, A, [range(-A_ndim - B_ndim + 2, -B_ndim + 1), range(A_ndim - 1)]))\n return onp.asarray(out, dtype=B_dtype)\n\n\ndef dot_vjp_0(ans, A, B):\n A_meta, B_meta = anp.metadata(A), anp.metadata(B)\n return lambda g: match_complex(A, dot_adjoint_0(B, g, A_meta, B_meta))\n\n\ndef dot_vjp_1(ans, A, B):\n A_meta, B_meta = anp.metadata(A), anp.metadata(B)\n return lambda g: match_complex(B, dot_adjoint_1(A, g, A_meta, B_meta))\n\n\ndefvjp(anp.dot, dot_vjp_0, dot_vjp_1)\n\ndefvjp(\n dot_adjoint_0,\n lambda ans, B, g, An, Bn: lambda A: match_complex(B, dot_adjoint_1(A, g, An, Bn)),\n lambda ans, B, g, An, Bn: lambda A: match_complex(g, anp.dot(A, B)),\n)\n\ndefvjp(\n dot_adjoint_1,\n lambda ans, A, g, An, Bn: lambda B: match_complex(A, dot_adjoint_0(B, g, An, Bn)),\n lambda ans, A, g, An, Bn: lambda B: match_complex(g, anp.dot(A, B)),\n)\n\n\n@primitive\ndef tensordot_adjoint_0(B, G, axes, A_ndim, B_ndim):\n # The adjoint of the operator\n # A |--> np.tensordot(A, B, axes)\n if B_ndim == 0:\n return G * B\n\n G_axes = onp.arange(onp.ndim(G))\n if type(axes) is int:\n axes = max(axes, 0)\n B_axes = onp.arange(B_ndim)\n return onp.tensordot(G, B, [G_axes[A_ndim - axes :], B_axes[axes:]])\n else:\n axes0 = [axes[0]] if type(axes[0]) is int else axes[0]\n axes1 = [axes[1]] if type(axes[1]) is int else axes[1]\n axes = [axes0, axes1]\n A_axes = onp.arange(A_ndim)\n B_axes = onp.arange(B_ndim)\n summed_axes = [\n onp.asarray(axes[0], dtype=\"int64\") % A_ndim,\n onp.asarray(axes[1], dtype=\"int64\") % B_ndim,\n ]\n other_axes = [onp.delete(A_axes, summed_axes[0]), onp.delete(B_axes, summed_axes[1])]\n out = onp.tensordot(G, B, [G_axes[len(other_axes[0]) :], other_axes[1]])\n perm = onp.argsort(onp.concatenate((other_axes[0], summed_axes[0][onp.argsort(summed_axes[1])])))\n return onp.transpose(out, perm)\n\n\n@primitive\ndef tensordot_adjoint_1(A, G, axes, A_ndim, B_ndim):\n # The adjoint of the operator\n # B |--> np.tensordot(A, B, axes)\n if A_ndim == 0:\n return G * A\n\n G_axes = onp.arange(onp.ndim(G))\n if type(axes) is int:\n axes = max(axes, 0)\n A_axes = onp.arange(A_ndim)\n return onp.tensordot(A, G, [A_axes[: A_ndim - axes], G_axes[: A_ndim - axes]])\n else:\n axes0 = [axes[0]] if type(axes[0]) is int else axes[0]\n axes1 = [axes[1]] if type(axes[1]) is int else axes[1]\n axes = [axes0, axes1]\n A_axes = onp.arange(A_ndim)\n B_axes = onp.arange(B_ndim)\n summed_axes = [\n onp.asarray(axes[0], dtype=\"int64\") % A_ndim,\n onp.asarray(axes[1], dtype=\"int64\") % B_ndim,\n ]\n other_axes = [onp.delete(A_axes, summed_axes[0]), onp.delete(B_axes, summed_axes[1])]\n out = onp.tensordot(A, G, [other_axes[0], G_axes[: len(other_axes[0])]])\n perm = onp.argsort(onp.concatenate((summed_axes[1][onp.argsort(summed_axes[0])], other_axes[1])))\n return onp.transpose(out, perm)\n\n\ndef tensordot_vjp_0(ans, A, B, axes=2):\n A_ndim, B_ndim = anp.ndim(A), anp.ndim(B)\n return lambda G: match_complex(A, tensordot_adjoint_0(B, G, axes, A_ndim, B_ndim))\n\n\ndef tensordot_vjp_1(ans, A, B, axes=2):\n A_ndim, B_ndim = anp.ndim(A), anp.ndim(B)\n return lambda G: match_complex(B, tensordot_adjoint_1(A, G, axes, A_ndim, B_ndim))\n\n\ndefvjp(anp.tensordot, tensordot_vjp_0, tensordot_vjp_1)\ndefvjp(\n tensordot_adjoint_0,\n lambda ans, B, G, axes, An, Bn: lambda A: match_complex(B, tensordot_adjoint_1(A, G, axes, An, Bn)),\n lambda ans, B, G, axes, An, Bn: lambda A: match_complex(G, anp.tensordot(A, B, axes)),\n)\ndefvjp(\n tensordot_adjoint_1,\n lambda ans, A, G, axes, An, Bn: lambda B: match_complex(A, tensordot_adjoint_0(B, G, axes, An, Bn)),\n lambda ans, A, G, axes, An, Bn: lambda B: match_complex(G, anp.tensordot(A, B, axes)),\n)\ndefvjp(\n anp.outer,\n lambda ans, a, b: lambda g: match_complex(a, anp.dot(g, b.T)),\n lambda ans, a, b: lambda g: match_complex(b, anp.dot(a.T, g)),\n)\n\n\ndef grad_concatenate_args(argnum, ans, axis_args, kwargs):\n axis, args = axis_args[0], axis_args[1:]\n sizes = [anp.shape(a)[axis] for a in args[:argnum]]\n start = sum(sizes[:-1])\n idxs = [slice(None)] * ans.ndim\n idxs[axis] = slice(start, start + sizes[-1])\n return lambda g: g[tuple(idxs)]\n\n\ndefvjp_argnum(anp.concatenate_args, grad_concatenate_args)\n\n\ndef wrapped_reshape(x, *args, **kwargs):\n # The reshape method can be called like A.reshape((5,4)) or A.reshape(5,4).\n # The reshape function doesn't support both ways, so we have to wrap it.\n if isinstance(args[0], int):\n return anp.reshape(x, args, **kwargs)\n else:\n return anp.reshape(x, *args, **kwargs)\n\n\nsetattr(ArrayBox, \"reshape\", wrapped_reshape)\n\n\ndef grad_sort(ans, x, axis=-1, kind=\"quicksort\", order=None):\n # TODO: Cast input with np.asanyarray()\n if len(x.shape) > 1:\n raise NotImplementedError(\"Gradient of sort not implemented for multi-dimensional arrays.\")\n sort_perm = anp.argsort(x, axis, kind, order)\n return lambda g: unpermuter(g, sort_perm)\n\n\ndefvjp(anp.sort, grad_sort)\nif onp.lib.NumpyVersion(onp.__version__) < \"2.0.0\":\n defvjp(anp.msort, grad_sort) # Until multi-D is allowed, these are the same.\n\n\ndef grad_partition(ans, x, kth, axis=-1, kind=\"introselect\", order=None):\n # TODO: Cast input with np.asanyarray()\n if len(x.shape) > 1:\n raise NotImplementedError(\"Gradient of partition not implemented for multi-dimensional arrays.\")\n partition_perm = anp.argpartition(x, kth, axis, kind, order)\n return lambda g: unpermuter(g, partition_perm)\n\n\ndefvjp(anp.partition, grad_partition)\n\n\ndef unpermuter(g, permutation):\n unsort = anp.zeros(len(permutation), dtype=int)\n unsort[permutation] = list(range(len(permutation)))\n return g[unsort]\n\n\ndef grad_reshape_list(ans, *arys):\n if len(arys) > 1:\n raise NotImplementedError(\"Can't handle multiple arguments yet.\")\n return lambda g: anp.reshape(g, anp.shape(arys[0]))\n\n\ndefvjp(anp.atleast_1d, grad_reshape_list)\ndefvjp(anp.atleast_2d, grad_reshape_list)\ndefvjp(anp.atleast_3d, grad_reshape_list)\n\n\ndef grad_einsum(argnum, ans, operands_, kwargs):\n result_meta = anp.metadata(operands_[argnum])\n\n def vjp(g):\n operands = operands_\n if isinstance(operands[0], str): # using \"ijk\" convention.\n in_subs, out_subs, _ = anp.parse_einsum_input(*operands)\n string, operands = operands[0], operands[1:]\n\n in_subs_list = in_subs.split(\",\")\n op_num = argnum - 1\n subs_wrt = in_subs_list[op_num]\n rest_of_ops = operands[:op_num] + operands[op_num + 1 :]\n rest_of_subs = in_subs_list[:op_num] + in_subs_list[op_num + 1 :]\n\n # subscripts that only appear in subs_wrt (and not in other subscript lists\n # or in the output) are implicitly being summed out, as if contracted\n # against a tensor of ones. we make that tensor of ones explicit to handle\n # the necessary vjp broadcasting inside einsum.\n other_named_subs = set(\"\".join([out_subs] + rest_of_subs))\n naked_summed = [(i, sub) for i, sub in enumerate(subs_wrt) if sub not in other_named_subs]\n if naked_summed:\n naked_summed_dims, ones_subs = zip(*naked_summed)\n ones_subs = \"\".join(ones_subs)\n ones = onp.ones(onp.array(operands[op_num].shape)[list(naked_summed_dims)])\n new_input_subs = \",\".join([out_subs, ones_subs] + rest_of_subs)\n new_operands = (g, ones) + rest_of_ops\n else:\n new_input_subs = \",\".join([out_subs] + rest_of_subs)\n new_operands = (g,) + rest_of_ops\n\n new_subscripts = new_input_subs + \"->\" + subs_wrt\n return unbroadcast(anp.einsum(new_subscripts, *new_operands), result_meta)\n else: # using (op0, sublist0, op1, sublist1, ..., sublistout) convention\n if len(operands) % 2 == 0:\n raise NotImplementedError(\"Need sublistout argument\")\n operands = list(operands)\n rest_of_ops = (\n [operands[-1]] + operands[:argnum] + operands[(argnum + 2) : -1] + [operands[argnum + 1]]\n )\n return unbroadcast_einsum(anp.einsum(g, *rest_of_ops), result_meta, operands[argnum + 1])\n\n return vjp\n\n\ndefvjp_argnum(anp.einsum, grad_einsum)\n\ndefvjp(\n anp.diagonal,\n lambda ans, A, offset=0, axis1=0, axis2=1: lambda g: anp.make_diagonal(g, offset, axis1, axis2),\n)\ndefvjp(\n anp.make_diagonal,\n lambda ans, D, offset=0, axis1=0, axis2=1: lambda g: anp.diagonal(g, offset, axis1, axis2),\n)\n\n\ndef match_complex(target, x):\n target_iscomplex = anp.iscomplexobj(target)\n x_iscomplex = anp.iscomplexobj(x)\n if x_iscomplex and not target_iscomplex:\n return anp.real(x)\n elif not x_iscomplex and target_iscomplex:\n return x + 0j\n else:\n return x\n\n\ndef unbroadcast(x, target_meta, broadcast_idx=0):\n target_shape, target_ndim, dtype, target_iscomplex = target_meta\n while anp.ndim(x) > target_ndim:\n x = anp.sum(x, axis=broadcast_idx)\n for axis, size in enumerate(target_shape):\n if size == 1:\n x = anp.sum(x, axis=axis, keepdims=True)\n if anp.iscomplexobj(x) and not target_iscomplex:\n x = anp.real(x)\n return x\n\n\ndef unbroadcast_f(target, f):\n target_meta = anp.metadata(target)\n return lambda g: unbroadcast(f(g), target_meta)\n\n\ndef unbroadcast_einsum(x, target_meta, subscript):\n if Ellipsis not in subscript:\n return x\n elif subscript[0] == Ellipsis:\n return unbroadcast(x, target_meta, 0)\n elif subscript[-1] == Ellipsis:\n return unbroadcast(x, target_meta, -1)\n else:\n return unbroadcast(x, target_meta, subscript.index(Ellipsis))\n\n\ndef balanced_eq(x, z, y):\n return (x == z) / (1.0 + (x == y))\n\n\ndef replace_zero(x, val):\n return anp.where(x, x, val)\n\n\n# ----- extra functions used internally -----\n\n\ndef array_from_args_gradmaker(argnum, ans, args, kwargs):\n return lambda g: g[argnum - 2]\n\n\ndefvjp_argnum(anp.array_from_args, array_from_args_gradmaker)\n\n\ndef array_from_scalar_or_array_gradmaker(ans, array_args, array_kwargs, scarray):\n ndmin = array_kwargs.get(\"ndmin\", 0)\n scarray_ndim = anp.ndim(scarray)\n if ndmin > scarray_ndim:\n return lambda g: anp.squeeze(g, axis=tuple(range(ndmin - scarray_ndim)))\n else:\n return lambda g: g\n\n\ndefvjp(anp._array_from_scalar_or_array, array_from_scalar_or_array_gradmaker, argnums=(2, 3))\n\n\n@primitive\ndef untake(x, idx, vs):\n if isinstance(idx, list) and (len(idx) == 0 or not isinstance(idx[0], slice)):\n idx = onp.array(idx, dtype=\"int64\")\n\n def mut_add(A):\n onp.add.at(A, idx, x)\n return A\n\n return SparseObject(vs, mut_add)\n\n\ndefvjp(func(ArrayBox.__getitem__), lambda ans, A, idx: lambda g: untake(g, idx, vspace(A)))\ndefvjp(untake, lambda ans, x, idx, _: lambda g: g[idx])\n\n\ndef _unpad(array, width):\n if anp.isscalar(width):\n width = [[width, width]]\n elif anp.shape(width) == (1,):\n width = [anp.concatenate((width, width))]\n elif anp.shape(width) == (2,):\n width = [width]\n if anp.shape(width)[0] == 1:\n width = anp.repeat(width, anp.ndim(array), 0)\n idxs = tuple(slice(l, -u or None) for l, u in width)\n return array[idxs]\n\n\ndef pad_vjp(ans, array, pad_width, mode, **kwargs):\n assert mode == \"constant\", \"Only constant mode padding is supported.\"\n return lambda g: _unpad(g, pad_width)\n\n\ndefvjp(anp.pad, pad_vjp)\n"
},
{
"path": "autograd/numpy/numpy_vspaces.py",
"content": "import numpy as np\n\nfrom autograd.builtins import NamedTupleVSpace\nfrom autograd.extend import VSpace\n\n\nclass ArrayVSpace(VSpace):\n def __init__(self, value):\n value = np.asarray(value)\n self.shape = value.shape\n self.dtype = value.dtype\n\n @property\n def size(self):\n return np.prod(self.shape)\n\n @property\n def ndim(self):\n return len(self.shape)\n\n def zeros(self):\n return np.zeros(self.shape, dtype=self.dtype)\n\n def ones(self):\n return np.ones(self.shape, dtype=self.dtype)\n\n def standard_basis(self):\n for idxs in np.ndindex(*self.shape):\n vect = np.zeros(self.shape, dtype=self.dtype)\n vect[idxs] = 1\n yield vect\n\n def randn(self):\n return np.array(np.random.randn(*self.shape)).astype(self.dtype)\n\n def _inner_prod(self, x, y):\n return np.dot(np.ravel(x), np.ravel(y))\n\n\nclass ComplexArrayVSpace(ArrayVSpace):\n iscomplex = True\n\n @property\n def size(self):\n return np.prod(self.shape) * 2\n\n def ones(self):\n return np.ones(self.shape, dtype=self.dtype) + 1.0j * np.ones(self.shape, dtype=self.dtype)\n\n def standard_basis(self):\n for idxs in np.ndindex(*self.shape):\n for v in [1.0, 1.0j]:\n vect = np.zeros(self.shape, dtype=self.dtype)\n vect[idxs] = v\n yield vect\n\n def randn(self):\n return np.array(np.random.randn(*self.shape)).astype(self.dtype) + 1.0j * np.array(\n np.random.randn(*self.shape)\n ).astype(self.dtype)\n\n def _inner_prod(self, x, y):\n return np.real(np.dot(np.conj(np.ravel(x)), np.ravel(y)))\n\n def _covector(self, x):\n return np.conj(x)\n\n\nVSpace.register(np.ndarray, lambda x: ComplexArrayVSpace(x) if np.iscomplexobj(x) else ArrayVSpace(x))\n\nfor type_ in [float, np.longdouble, np.float64, np.float32, np.float16]:\n ArrayVSpace.register(type_)\n\nfor type_ in [complex, np.clongdouble, np.complex64, np.complex128]:\n ComplexArrayVSpace.register(type_)\n\n\nif np.lib.NumpyVersion(np.__version__) >= \"2.0.0\":\n\n class EigResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg._linalg.EigResult\n\n class EighResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg._linalg.EighResult\n\n class QRResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg._linalg.QRResult\n\n class SlogdetResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg._linalg.SlogdetResult\n\n class SVDResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg._linalg.SVDResult\n\n EigResultVSpace.register(np.linalg._linalg.EigResult)\n EighResultVSpace.register(np.linalg._linalg.EighResult)\n QRResultVSpace.register(np.linalg._linalg.QRResult)\n SlogdetResultVSpace.register(np.linalg._linalg.SlogdetResult)\n SVDResultVSpace.register(np.linalg._linalg.SVDResult)\nelif np.__version__ >= \"1.25\":\n\n class EigResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg.linalg.EigResult\n\n class EighResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg.linalg.EighResult\n\n class QRResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg.linalg.QRResult\n\n class SlogdetResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg.linalg.SlogdetResult\n\n class SVDResultVSpace(NamedTupleVSpace):\n seq_type = np.linalg.linalg.SVDResult\n\n EigResultVSpace.register(np.linalg.linalg.EigResult)\n EighResultVSpace.register(np.linalg.linalg.EighResult)\n QRResultVSpace.register(np.linalg.linalg.QRResult)\n SlogdetResultVSpace.register(np.linalg.linalg.SlogdetResult)\n SVDResultVSpace.register(np.linalg.linalg.SVDResult)\n"
},
{
"path": "autograd/numpy/numpy_wrapper.py",
"content": "import warnings\n\nimport numpy as _np\n\nimport autograd.builtins as builtins\nfrom autograd.extend import notrace_primitive, primitive\n\nif _np.lib.NumpyVersion(_np.__version__) >= \"2.0.0\":\n from numpy._core.einsumfunc import _parse_einsum_input\nelse:\n from numpy.core.einsumfunc import _parse_einsum_input\n\nnumpy_version = _np.__version__\n\nnotrace_functions = [_np.ndim, _np.shape, _np.iscomplexobj, _np.result_type]\n\n\ndef wrap_intdtype(cls):\n class IntdtypeSubclass(cls):\n __new__ = notrace_primitive(cls.__new__)\n\n return IntdtypeSubclass\n\n\ndef wrap_namespace(old, new):\n unchanged_types = {float, int, type(None), type}\n int_types = {_np.int8, _np.int16, _np.int32, _np.int64, _np.integer}\n for name, obj in old.items():\n if obj in notrace_functions:\n new[name] = notrace_primitive(obj)\n elif callable(obj) and type(obj) is not type:\n new[name] = primitive(obj)\n elif type(obj) is type and obj in int_types:\n new[name] = wrap_intdtype(obj)\n elif type(obj) in unchanged_types:\n new[name] = obj\n\n\nwrap_namespace(_np.__dict__, globals())\n\n# ----- Special treatment of list-input functions -----\n\n\n@primitive\ndef concatenate_args(axis, *args):\n return _np.concatenate(args, axis).view(ndarray)\n\n\nconcatenate = lambda arr_list, axis=0: concatenate_args(axis, *arr_list)\nvstack = row_stack = lambda tup: concatenate([atleast_2d(_m) for _m in tup], axis=0)\n\n\ndef hstack(tup):\n arrs = [atleast_1d(_m) for _m in tup]\n if arrs[0].ndim == 1:\n return concatenate(arrs, 0)\n return concatenate(arrs, 1)\n\n\ndef column_stack(tup):\n arrays = []\n for v in tup:\n arr = array(v)\n if arr.ndim < 2:\n arr = array(arr, ndmin=2).T\n arrays.append(arr)\n return concatenate(arrays, 1)\n\n\ndef array(A, *args, **kwargs):\n t = builtins.type(A)\n if t in (list, tuple):\n return array_from_args(args, kwargs, *map(array, A))\n else:\n return _array_from_scalar_or_array(args, kwargs, A)\n\n\ndef wrap_if_boxes_inside(raw_array, slow_op_name=None):\n if raw_array.dtype is _np.dtype(\"O\"):\n if slow_op_name:\n warnings.warn(f\"{slow_op_name} is slow for array inputs. np.concatenate() is faster.\")\n return array_from_args((), {}, *raw_array.ravel()).reshape(raw_array.shape)\n else:\n return raw_array\n\n\n@primitive\ndef _array_from_scalar_or_array(array_args, array_kwargs, scalar):\n return _np.array(scalar, *array_args, **array_kwargs)\n\n\n@primitive\ndef array_from_args(array_args, array_kwargs, *args):\n return _np.array(args, *array_args, **array_kwargs)\n\n\ndef select(condlist, choicelist, default=0):\n raw_array = _np.select(list(condlist), list(choicelist), default=default)\n return array(list(raw_array.ravel())).reshape(raw_array.shape)\n\n\ndef stack(arrays, axis=0):\n # this code is basically copied from numpy/core/shape_base.py's stack\n # we need it here because we want to re-implement stack in terms of the\n # primitives defined in this file\n\n arrays = [array(arr) for arr in arrays]\n if not arrays:\n raise ValueError(\"need at least one array to stack\")\n\n shapes = {arr.shape for arr in arrays}\n if len(shapes) != 1:\n raise ValueError(\"all input arrays must have the same shape\")\n\n result_ndim = arrays[0].ndim + 1\n if not -result_ndim <= axis < result_ndim:\n raise IndexError(\"axis {0} out of bounds [-{1}, {1})\".format(axis, result_ndim))\n if axis < 0:\n axis += result_ndim\n\n sl = (slice(None),) * axis + (None,)\n return concatenate([arr[sl] for arr in arrays], axis=axis)\n\n\ndef append(arr, values, axis=None):\n # this code is basically copied from numpy/lib/function_base.py's append\n arr = array(arr)\n if axis is None:\n if ndim(arr) != 1:\n arr = ravel(arr)\n values = ravel(array(values))\n axis = ndim(arr) - 1\n return concatenate((arr, values), axis=axis)\n\n\n# ----- Enable functions called using [] ----\n\n\nclass r_class:\n def __getitem__(self, args):\n raw_array = _np.r_[args]\n return wrap_if_boxes_inside(raw_array, slow_op_name=\"r_\")\n\n\nr_ = r_class()\n\n\nclass c_class:\n def __getitem__(self, args):\n raw_array = _np.c_[args]\n return wrap_if_boxes_inside(raw_array, slow_op_name=\"c_\")\n\n\nc_ = c_class()\n\n\n# ----- misc -----\n@primitive\ndef make_diagonal(D, offset=0, axis1=0, axis2=1):\n # Numpy doesn't offer a complement to np.diagonal: a function to create new\n # diagonal arrays with extra dimensions. We need such a function for the\n # gradient of np.diagonal and it's also quite handy to have. So here it is.\n if not (offset == 0 and axis1 == -1 and axis2 == -2):\n raise NotImplementedError(\"Currently make_diagonal only supports offset=0, axis1=-1, axis2=-2\")\n\n # We use a trick: calling np.diagonal returns a view on the original array,\n # so we can modify it in-place. (only valid for numpy version >= 1.10.)\n new_array = _np.zeros(D.shape + (D.shape[-1],))\n new_array_diag = _np.diagonal(new_array, offset=0, axis1=-1, axis2=-2)\n new_array_diag.flags.writeable = True\n new_array_diag[:] = D\n return new_array\n\n\n@notrace_primitive\ndef metadata(A):\n return _np.shape(A), _np.ndim(A), _np.result_type(A), _np.iscomplexobj(A)\n\n\n@notrace_primitive\ndef parse_einsum_input(*args):\n return _parse_einsum_input(args)\n\n\nif _np.lib.NumpyVersion(_np.__version__) >= \"2.0.0\":\n # Wrapped above\n _astype = astype\nelse:\n\n @primitive\n def _astype(A, dtype, order=\"K\", casting=\"unsafe\", subok=True, copy=True):\n return A.astype(dtype, order, casting, subok, copy)\n"
},
{
"path": "autograd/numpy/random.py",
"content": "import numpy.random as npr\n\nfrom .numpy_wrapper import wrap_namespace\n\nwrap_namespace(npr.__dict__, globals())\n"
},
{
"path": "autograd/scipy/__init__.py",
"content": "from . import integrate, signal, special, stats\n"
},
{
"path": "autograd/scipy/integrate.py",
"content": "import scipy.integrate\n\nimport autograd.numpy as np\nfrom autograd import make_vjp\nfrom autograd.builtins import tuple\nfrom autograd.extend import defvjp_argnums, primitive\nfrom autograd.misc import flatten\n\nodeint = primitive(scipy.integrate.odeint)\n\n\ndef grad_odeint(yt, func, y0, t, func_args, **kwargs):\n # Extended from \"Scalable Inference of Ordinary Differential\n # Equation Models of Biochemical Processes\", Sec. 2.4.2\n # Fabian Froehlich, Carolin Loos, Jan Hasenauer, 2017\n # https://arxiv.org/abs/1711.08079\n\n T, D = np.shape(yt)\n flat_args, unflatten = flatten(func_args)\n\n def flat_func(y, t, flat_args):\n return func(y, t, *unflatten(flat_args))\n\n def unpack(x):\n # y, vjp_y, vjp_t, vjp_args\n return x[0:D], x[D : 2 * D], x[2 * D], x[2 * D + 1 :]\n\n def augmented_dynamics(augmented_state, t, flat_args):\n # Orginal system augmented with vjp_y, vjp_t and vjp_args.\n y, vjp_y, _, _ = unpack(augmented_state)\n vjp_all, dy_dt = make_vjp(flat_func, argnum=(0, 1, 2))(y, t, flat_args)\n vjp_y, vjp_t, vjp_args = vjp_all(-vjp_y)\n return np.hstack((dy_dt, vjp_y, vjp_t, vjp_args))\n\n def vjp_all(g):\n vjp_y = g[-1, :]\n vjp_t0 = 0\n time_vjp_list = []\n vjp_args = np.zeros(np.size(flat_args))\n\n for i in range(T - 1, 0, -1):\n # Compute effect of moving measurement time.\n vjp_cur_t = np.dot(func(yt[i, :], t[i], *func_args), g[i, :])\n time_vjp_list.append(vjp_cur_t)\n vjp_t0 = vjp_t0 - vjp_cur_t\n\n # Run augmented system backwards to the previous observation.\n aug_y0 = np.hstack((yt[i, :], vjp_y, vjp_t0, vjp_args))\n aug_ans = odeint(\n augmented_dynamics, aug_y0, np.array([t[i], t[i - 1]]), tuple((flat_args,)), **kwargs\n )\n _, vjp_y, vjp_t0, vjp_args = unpack(aug_ans[1])\n\n # Add gradient from current output.\n vjp_y = vjp_y + g[i - 1, :]\n\n time_vjp_list.append(vjp_t0)\n vjp_times = np.hstack(time_vjp_list)[::-1]\n\n return None, vjp_y, vjp_times, unflatten(vjp_args)\n\n return vjp_all\n\n\ndef argnums_unpack(all_vjp_builder):\n # A generic autograd helper function. Takes a function that\n # builds vjps for all arguments, and wraps it to return only required vjps.\n def build_selected_vjps(argnums, ans, combined_args, kwargs):\n vjp_func = all_vjp_builder(ans, *combined_args, **kwargs)\n\n def chosen_vjps(g): # Returns whichever vjps were asked for.\n all_vjps = vjp_func(g)\n return [all_vjps[argnum] for argnum in argnums]\n\n return chosen_vjps\n\n return build_selected_vjps\n\n\ndefvjp_argnums(odeint, argnums_unpack(grad_odeint))\n"
},
{
"path": "autograd/scipy/linalg.py",
"content": "from functools import partial\n\nimport scipy.linalg\n\nimport autograd.numpy as anp\nfrom autograd.extend import defjvp, defjvp_argnums, defvjp, defvjp_argnums\nfrom autograd.numpy.numpy_wrapper import wrap_namespace\n\nwrap_namespace(scipy.linalg.__dict__, globals()) # populates module namespace\n\n\ndef _vjp_sqrtm(ans, A, disp=True, blocksize=64):\n assert disp, \"sqrtm vjp not implemented for disp=False\"\n ans_transp = anp.transpose(ans)\n\n def vjp(g):\n return anp.real(solve_sylvester(ans_transp, ans_transp, g))\n\n return vjp\n\n\ndefvjp(sqrtm, _vjp_sqrtm)\n\n\ndef _flip(a, trans):\n if anp.iscomplexobj(a):\n return \"H\" if trans in (\"N\", 0) else \"N\"\n else:\n return \"T\" if trans in (\"N\", 0) else \"N\"\n\n\ndef grad_solve_triangular(ans, a, b, trans=0, lower=False, **kwargs):\n tri = anp.tril if (lower ^ (_flip(a, trans) == \"N\")) else anp.triu\n transpose = lambda x: x if _flip(a, trans) != \"N\" else x.T\n al2d = lambda x: x if x.ndim > 1 else x[..., None]\n\n def vjp(g):\n v = al2d(solve_triangular(a, g, trans=_flip(a, trans), lower=lower))\n return -transpose(tri(anp.dot(v, al2d(ans).T)))\n\n return vjp\n\n\ndefvjp(\n solve_triangular,\n grad_solve_triangular,\n lambda ans, a, b, trans=0, lower=False, **kwargs: (\n lambda g: solve_triangular(a, g, trans=_flip(a, trans), lower=lower)\n ),\n)\n\n\ndef grad_solve_banded(argnum, ans, l_and_u, a, b):\n updim = lambda x: x if x.ndim == a.ndim else x[..., None]\n\n def transpose_banded(l_and_u, a):\n # Compute the transpose of a banded matrix.\n # The transpose is itself a banded matrix.\n\n num_rows = a.shape[0]\n\n shifts = anp.arange(-l_and_u[1], l_and_u[0] + 1)\n\n T_a = anp.roll(a[:1, :], shifts[0])\n for rr in range(1, num_rows):\n T_a = anp.vstack([T_a, anp.flipud(anp.roll(a[rr : rr + 1, :], shifts[rr]))])\n T_a = anp.flipud(T_a)\n\n T_l_and_u = anp.flip(l_and_u)\n\n return T_l_and_u, T_a\n\n def banded_dot(l_and_u, uu, vv):\n # Compute tensor product of vectors uu and vv.\n # Tensor product elements are resticted to the bands specified by l_and_u.\n\n # TODO: replace the brute-force ravel() by smarter dimension handeling of uu and vv\n\n # main diagonal\n banded_uv = anp.ravel(uu) * anp.ravel(vv)\n\n # stack below the sub-diagonals\n for rr in range(1, l_and_u[0] + 1):\n banded_uv_rr = anp.hstack([anp.ravel(uu)[rr:] * anp.ravel(vv)[:-rr], anp.zeros(rr)])\n banded_uv = anp.vstack([banded_uv, banded_uv_rr])\n\n # stack above the sup-diagonals\n for rr in range(1, l_and_u[1] + 1):\n banded_uv_rr = anp.hstack([anp.zeros(rr), anp.ravel(uu)[:-rr] * anp.ravel(vv)[rr:]])\n banded_uv = anp.vstack([banded_uv_rr, banded_uv])\n\n return banded_uv\n\n T_l_and_u, T_a = transpose_banded(l_and_u, a)\n\n if argnum == 1:\n return lambda g: (\n -banded_dot(l_and_u, updim(solve_banded(T_l_and_u, T_a, g)), anp.transpose(updim(ans)))\n )\n elif argnum == 2:\n return lambda g: solve_banded(T_l_and_u, T_a, g)\n\n\ndefvjp(solve_banded, partial(grad_solve_banded, 1), partial(grad_solve_banded, 2), argnums=[1, 2])\n\n\ndef _jvp_sqrtm(dA, ans, A, disp=True, blocksize=64):\n assert disp, \"sqrtm jvp not implemented for disp=False\"\n return solve_sylvester(ans, ans, dA)\n\n\ndefjvp(sqrtm, _jvp_sqrtm)\n\n\ndef _jvp_sylvester(argnums, dms, ans, args, _):\n a, b, q = args\n if 0 in argnums:\n da = dms[0]\n db = dms[1] if 1 in argnums else 0\n else:\n da = 0\n db = dms[0] if 1 in argnums else 0\n dq = dms[-1] if 2 in argnums else 0\n rhs = dq - anp.dot(da, ans) - anp.dot(ans, db)\n return solve_sylvester(a, b, rhs)\n\n\ndefjvp_argnums(solve_sylvester, _jvp_sylvester)\n\n\ndef _vjp_sylvester(argnums, ans, args, _):\n a, b, q = args\n\n def vjp(g):\n vjps = []\n q_vjp = solve_sylvester(anp.transpose(a), anp.transpose(b), g)\n if 0 in argnums:\n vjps.append(-anp.dot(q_vjp, anp.transpose(ans)))\n if 1 in argnums:\n vjps.append(-anp.dot(anp.transpose(ans), q_vjp))\n if 2 in argnums:\n vjps.append(q_vjp)\n return tuple(vjps)\n\n return vjp\n\n\ndefvjp_argnums(solve_sylvester, _vjp_sylvester)\n"
},
{
"path": "autograd/scipy/signal.py",
"content": "from functools import partial\n\nimport numpy as npo # original numpy\nfrom numpy.lib.stride_tricks import as_strided\n\nimport autograd.numpy as np\nfrom autograd.extend import defvjp, primitive\n\n\n@primitive\ndef convolve(A, B, axes=None, dot_axes=[(), ()], mode=\"full\"):\n assert mode in [\"valid\", \"full\"], f\"Mode {mode} not yet implemented\"\n if axes is None:\n axes = [list(range(A.ndim)), list(range(A.ndim))]\n wrong_order = any([B.shape[ax_B] < A.shape[ax_A] for ax_A, ax_B in zip(*axes)])\n if wrong_order:\n if mode == \"valid\" and not all([B.shape[ax_B] <= A.shape[ax_A] for ax_A, ax_B in zip(*axes)]):\n raise Exception(\"One array must be larger than the other along all convolved dimensions\")\n elif mode != \"full\" or B.size <= A.size: # Tie breaker\n i1 = B.ndim - len(dot_axes[1]) - len(axes[1]) # B ignore\n i2 = i1 + A.ndim - len(dot_axes[0]) - len(axes[0]) # A ignore\n i3 = i2 + len(axes[0])\n ignore_B = list(range(i1))\n ignore_A = list(range(i1, i2))\n conv = list(range(i2, i3))\n return convolve(B, A, axes=axes[::-1], dot_axes=dot_axes[::-1], mode=mode).transpose(\n ignore_A + ignore_B + conv\n )\n\n if mode == \"full\":\n B = pad_to_full(B, A, axes[::-1])\n B_view_shape = list(B.shape)\n B_view_strides = list(B.strides)\n flipped_idxs = [slice(None)] * A.ndim\n for ax_A, ax_B in zip(*axes):\n B_view_shape.append(abs(B.shape[ax_B] - A.shape[ax_A]) + 1)\n B_view_strides.append(B.strides[ax_B])\n B_view_shape[ax_B] = A.shape[ax_A]\n flipped_idxs[ax_A] = slice(None, None, -1)\n B_view = as_strided(B, B_view_shape, B_view_strides)\n A_view = A[tuple(flipped_idxs)]\n all_axes = [list(axes[i]) + list(dot_axes[i]) for i in [0, 1]]\n return einsum_tensordot(A_view, B_view, all_axes)\n\n\ndef einsum_tensordot(A, B, axes, reverse=False):\n # Does tensor dot product using einsum, which shouldn't require a copy.\n A_axnums = list(range(A.ndim))\n B_axnums = list(range(A.ndim, A.ndim + B.ndim))\n sum_axnum = A.ndim + B.ndim\n for i_sum, (i_A, i_B) in enumerate(zip(*axes)):\n A_axnums[i_A] = sum_axnum + i_sum\n B_axnums[i_B] = sum_axnum + i_sum\n return npo.einsum(A, A_axnums, B, B_axnums)\n\n\ndef pad_to_full(A, B, axes):\n A_pad = [(0, 0)] * A.ndim\n for ax_A, ax_B in zip(*axes):\n A_pad[ax_A] = (B.shape[ax_B] - 1,) * 2\n return npo.pad(A, A_pad, mode=\"constant\")\n\n\ndef parse_axes(A_shape, B_shape, conv_axes, dot_axes, mode):\n A_ndim, B_ndim = len(A_shape), len(B_shape)\n if conv_axes is None:\n conv_axes = (\n tuple(range(A_ndim)),\n tuple(range(A_ndim)),\n )\n axes = {\n \"A\": {\n \"conv\": tuple(conv_axes[0]),\n \"dot\": tuple(dot_axes[0]),\n \"ignore\": tuple(i for i in range(A_ndim) if i not in conv_axes[0] and i not in dot_axes[0]),\n },\n \"B\": {\n \"conv\": tuple(conv_axes[1]),\n \"dot\": tuple(dot_axes[1]),\n \"ignore\": tuple(i for i in range(B_ndim) if i not in conv_axes[1] and i not in dot_axes[1]),\n },\n }\n assert len(axes[\"A\"][\"dot\"]) == len(axes[\"B\"][\"dot\"])\n assert len(axes[\"A\"][\"conv\"]) == len(axes[\"B\"][\"conv\"])\n i1 = len(axes[\"A\"][\"ignore\"])\n i2 = i1 + len(axes[\"B\"][\"ignore\"])\n i3 = i2 + len(axes[\"A\"][\"conv\"])\n axes[\"out\"] = {\n \"ignore_A\": tuple(range(i1)),\n \"ignore_B\": tuple(range(i1, i2)),\n \"conv\": tuple(range(i2, i3)),\n }\n conv_shape = (\n compute_conv_size(A_shape[i], B_shape[j], mode) for i, j in zip(axes[\"A\"][\"conv\"], axes[\"B\"][\"conv\"])\n )\n shapes = {\n \"A\": {s: tuple(A_shape[i] for i in ax) for s, ax in axes[\"A\"].items()},\n \"B\": {s: tuple(B_shape[i] for i in ax) for s, ax in axes[\"B\"].items()},\n }\n shapes[\"out\"] = {\n \"ignore_A\": shapes[\"A\"][\"ignore\"],\n \"ignore_B\": shapes[\"B\"][\"ignore\"],\n \"conv\": conv_shape,\n }\n return axes, shapes\n\n\ndef compute_conv_size(A_size, B_size, mode):\n if mode == \"full\":\n return A_size + B_size - 1\n elif mode == \"same\":\n return A_size\n elif mode == \"valid\":\n return abs(A_size - B_size) + 1\n else:\n raise Exception(f\"Mode {mode} not recognized\")\n\n\ndef flipped_idxs(ndim, axes):\n new_idxs = [slice(None)] * ndim\n for ax in axes:\n new_idxs[ax] = slice(None, None, -1)\n return tuple(new_idxs)\n\n\ndef grad_convolve(argnum, ans, A, B, axes=None, dot_axes=[(), ()], mode=\"full\"):\n assert mode in [\"valid\", \"full\"], f\"Grad for mode {mode} not yet implemented\"\n axes, shapes = parse_axes(A.shape, B.shape, axes, dot_axes, mode)\n if argnum == 0:\n X, Y = A, B\n _X_, _Y_ = \"A\", \"B\"\n ignore_Y = \"ignore_B\"\n elif argnum == 1:\n X, Y = B, A\n _X_, _Y_ = \"B\", \"A\"\n ignore_Y = \"ignore_A\"\n else:\n raise NotImplementedError(f\"Can't take grad of convolve w.r.t. arg {argnum}\")\n\n if mode == \"full\":\n new_mode = \"valid\"\n else:\n if any([x_size > y_size for x_size, y_size in zip(shapes[_X_][\"conv\"], shapes[_Y_][\"conv\"])]):\n new_mode = \"full\"\n else:\n new_mode = \"valid\"\n\n def vjp(g):\n result = convolve(\n g,\n Y[flipped_idxs(Y.ndim, axes[_Y_][\"conv\"])],\n axes=[axes[\"out\"][\"conv\"], axes[_Y_][\"conv\"]],\n dot_axes=[axes[\"out\"][ignore_Y], axes[_Y_][\"ignore\"]],\n mode=new_mode,\n )\n new_order = npo.argsort(axes[_X_][\"ignore\"] + axes[_X_][\"dot\"] + axes[_X_][\"conv\"])\n return np.transpose(result, new_order)\n\n return vjp\n\n\ndefvjp(convolve, partial(grad_convolve, 0), partial(grad_convolve, 1))\n"
},
{
"path": "autograd/scipy/special.py",
"content": "import scipy.special\n\nimport autograd.numpy as np\nfrom autograd.extend import defjvp, defvjp, primitive\nfrom autograd.numpy.numpy_vjps import repeat_to_match_shape, unbroadcast_f\n\n### Beta function ###\nbeta = primitive(scipy.special.beta)\nbetainc = primitive(scipy.special.betainc)\nbetaln = primitive(scipy.special.betaln)\n\ndefvjp(\n beta,\n lambda ans, a, b: unbroadcast_f(a, lambda g: g * ans * (psi(a) - psi(a + b))),\n lambda ans, a, b: unbroadcast_f(b, lambda g: g * ans * (psi(b) - psi(a + b))),\n)\ndefvjp(\n betainc,\n lambda ans, a, b, x: unbroadcast_f(\n x, lambda g: g * np.power(x, a - 1) * np.power(1 - x, b - 1) / beta(a, b)\n ),\n argnums=[2],\n)\ndefvjp(\n betaln,\n lambda ans, a, b: unbroadcast_f(a, lambda g: g * (psi(a) - psi(a + b))),\n lambda ans, a, b: unbroadcast_f(b, lambda g: g * (psi(b) - psi(a + b))),\n)\n\n### Gamma functions ###\npolygamma = primitive(scipy.special.polygamma)\npsi = primitive(scipy.special.psi) # psi(x) is just polygamma(0, x)\ndigamma = primitive(scipy.special.digamma) # digamma is another name for psi.\ngamma = primitive(scipy.special.gamma)\ngammaln = primitive(scipy.special.gammaln)\ngammainc = primitive(scipy.special.gammainc)\ngammaincc = primitive(scipy.special.gammaincc)\ngammasgn = primitive(scipy.special.gammasgn)\nrgamma = primitive(scipy.special.rgamma)\nmultigammaln = primitive(scipy.special.multigammaln)\n\ndefvjp(gammasgn, None)\ndefvjp(polygamma, None, lambda ans, n, x: lambda g: g * polygamma(n + 1, x))\ndefvjp(psi, lambda ans, x: lambda g: g * polygamma(1, x))\ndefvjp(digamma, lambda ans, x: lambda g: g * polygamma(1, x))\ndefvjp(gamma, lambda ans, x: lambda g: g * ans * psi(x))\ndefvjp(gammaln, lambda ans, x: lambda g: g * psi(x))\ndefvjp(rgamma, lambda ans, x: lambda g: g * psi(x) / -gamma(x))\ndefvjp(\n multigammaln,\n lambda ans, a, d: lambda g: g * np.sum(digamma(np.expand_dims(a, -1) - np.arange(d) / 2.0), -1),\n None,\n)\n\n\ndef make_gammainc_vjp_arg1(sign):\n def gammainc_vjp_arg1(ans, a, x):\n coeffs = sign * np.exp(-x) * np.power(x, a - 1) / gamma(a)\n return unbroadcast_f(x, lambda g: g * coeffs)\n\n return gammainc_vjp_arg1\n\n\ndefvjp(gammainc, make_gammainc_vjp_arg1(1), argnums=[1])\ndefvjp(gammaincc, make_gammainc_vjp_arg1(-1), argnums=[1])\n\n### Bessel functions ###\n\nj0 = primitive(scipy.special.j0)\ny0 = primitive(scipy.special.y0)\nj1 = primitive(scipy.special.j1)\ny1 = primitive(scipy.special.y1)\njn = primitive(scipy.special.jn)\nyn = primitive(scipy.special.yn)\n\ndefvjp(j0, lambda ans, x: lambda g: -g * j1(x))\ndefvjp(y0, lambda ans, x: lambda g: -g * y1(x))\ndefvjp(j1, lambda ans, x: lambda g: g * (j0(x) - jn(2, x)) / 2.0)\ndefvjp(y1, lambda ans, x: lambda g: g * (y0(x) - yn(2, x)) / 2.0)\ndefvjp(jn, None, lambda ans, n, x: lambda g: g * (jn(n - 1, x) - jn(n + 1, x)) / 2.0)\ndefvjp(yn, None, lambda ans, n, x: lambda g: g * (yn(n - 1, x) - yn(n + 1, x)) / 2.0)\n\n\n### Faster versions of common Bessel functions ###\ni0 = primitive(scipy.special.i0)\ni1 = primitive(scipy.special.i1)\niv = primitive(scipy.special.iv)\nive = primitive(scipy.special.ive)\n\ndefvjp(i0, lambda ans, x: lambda g: g * i1(x))\ndefvjp(i1, lambda ans, x: lambda g: g * (i0(x) + iv(2, x)) / 2.0)\ndefvjp(iv, None, lambda ans, n, x: lambda g: g * (iv(n - 1, x) + iv(n + 1, x)) / 2.0)\ndefvjp(ive, None, lambda ans, n, x: lambda g: g * (ans * (n / x - np.sign(x)) + ive(n + 1, x)))\n\n### Error Function ###\ninv_root_pi = 0.56418958354775627928\nerf = primitive(scipy.special.erf)\nerfc = primitive(scipy.special.erfc)\n\ndefvjp(erf, lambda ans, x: lambda g: 2.0 * g * inv_root_pi * np.exp(-(x**2)))\ndefvjp(erfc, lambda ans, x: lambda g: -2.0 * g * inv_root_pi * np.exp(-(x**2)))\n\n\n### Inverse error function ###\nroot_pi = 1.7724538509055159\nerfinv = primitive(scipy.special.erfinv)\nerfcinv = primitive(scipy.special.erfcinv)\n\ndefvjp(erfinv, lambda ans, x: lambda g: g * root_pi / 2 * np.exp(erfinv(x) ** 2))\ndefvjp(erfcinv, lambda ans, x: lambda g: -g * root_pi / 2 * np.exp(erfcinv(x) ** 2))\n\n### Logit and Expit ###\nlogit = primitive(scipy.special.logit)\nexpit = primitive(scipy.special.expit)\n\ndefvjp(logit, lambda ans, x: lambda g: g / (x * (1 - x)))\ndefvjp(expit, lambda ans, x: lambda g: g * ans * (1 - ans))\n\n### logsumexp ###\nlogsumexp = primitive(scipy.special.logsumexp)\n\n\ndef make_grad_logsumexp(ans, x, axis=None, b=1.0, keepdims=False):\n shape, dtype = np.shape(x), np.result_type(x)\n\n def vjp(g):\n g_repeated, _ = repeat_to_match_shape(g, shape, dtype, axis, keepdims)\n ans_repeated, _ = repeat_to_match_shape(ans, shape, dtype, axis, keepdims)\n return g_repeated * b * np.exp(x - ans_repeated)\n\n return vjp\n\n\ndefvjp(logsumexp, make_grad_logsumexp)\n\n\ndef fwd_grad_logsumexp(g, ans, x, axis=None, b=1.0, keepdims=False):\n if not keepdims:\n if isinstance(axis, int):\n ans = np.expand_dims(ans, axis)\n elif isinstance(axis, tuple):\n for ax in sorted(axis):\n ans = np.expand_dims(ans, ax)\n return np.sum(g * b * np.exp(x - ans), axis=axis, keepdims=keepdims)\n\n\ndefjvp(logsumexp, fwd_grad_logsumexp)\n"
},
{
"path": "autograd/scipy/stats/__init__.py",
"content": "from . import beta, chi2, gamma, norm, poisson, t\n\n# Try block needed in case the user has an\n# old version of scipy without multivariate normal.\ntry:\n from . import multivariate_normal\nexcept AttributeError:\n pass\n\ntry:\n from . import dirichlet\nexcept AttributeError:\n pass\n"
},
{
"path": "autograd/scipy/stats/beta.py",
"content": "import scipy.stats\n\nimport autograd.numpy as np\nfrom autograd.extend import defvjp, primitive\nfrom autograd.numpy.numpy_vjps import unbroadcast_f\nfrom autograd.scipy.special import beta, psi\n\ncdf = primitive(scipy.stats.beta.cdf)\nlogpdf = primitive(scipy.stats.beta.logpdf)\npdf = primitive(scipy.stats.beta.pdf)\n\n\ndef grad_beta_logpdf_arg0(x, a, b):\n return (1 + a * (x - 1) + x * (b - 2)) / (x * (x - 1))\n\n\ndef grad_beta_logpdf_arg1(x, a, b):\n return np.log(x) - psi(a) + psi(a + b)\n\n\ndef grad_beta_logpdf_arg2(x, a, b):\n return np.log1p(-x) - psi(b) + psi(a + b)\n\n\ndefvjp(\n cdf,\n lambda ans, x, a, b: unbroadcast_f(\n x, lambda g: g * np.power(x, a - 1) * np.power(1 - x, b - 1) / beta(a, b)\n ),\n argnums=[0],\n)\ndefvjp(\n logpdf,\n lambda ans, x, a, b: unbroadcast_f(x, lambda g: g * grad_beta_logpdf_arg0(x, a, b)),\n lambda ans, x, a, b: unbroadcast_f(a, lambda g: g * grad_beta_logpdf_arg1(x, a, b)),\n lambda ans, x, a, b: unbroadcast_f(b, lambda g: g * grad_beta_logpdf_arg2(x, a, b)),\n)\ndefvjp(\n pdf,\n lambda ans, x, a, b: unbroadcast_f(x, lambda g: g * ans * grad_beta_logpdf_arg0(x, a, b)),\n lambda ans, x, a, b: unbroadcast_f(a, lambda g: g * ans * grad_beta_logpdf_arg1(x, a, b)),\n lambda ans, x, a, b: unbroadcast_f(b, lambda g: g * ans * grad_beta_logpdf_arg2(x, a, b)),\n)\n"
},
{
"path": "autograd/scipy/stats/chi2.py",
"content": "import scipy.stats\n\nimport autograd.numpy as np\nfrom autograd.extend import defvjp, primitive\nfrom autograd.numpy.numpy_vjps import unbroadcast_f\nfrom autograd.scipy.special import gamma\n\ncdf = primitive(scipy.stats.chi2.cdf)\nlogpdf = primitive(scipy.stats.chi2.logpdf)\npdf = primitive(scipy.stats.chi2.pdf)\n\n\ndef grad_chi2_logpdf(x, df):\n return np.where(df % 1 == 0, (df - x - 2) / (2 * x), 0)\n\n\ndefvjp(\n cdf,\n lambda ans, x, df: unbroadcast_f(\n x, lambda g: g * np.power(2.0, -df / 2) * np.exp(-x / 2) * np.power(x, df / 2 - 1) / gamma(df / 2)\n ),\n argnums=[0],\n)\ndefvjp(logpdf, lambda ans, x, df: unbroadcast_f(x, lambda g: g * grad_chi2_logpdf(x, df)), argnums=[0])\ndefvjp(pdf, lambda ans, x, df: unbroadcast_f(x, lambda g: g * ans * grad_chi2_logpdf(x, df)), argnums=[0])\n"
},
{
"path": "autograd/scipy/stats/dirichlet.py",
"content": "import scipy.stats\n\nimport autograd.numpy as np\nfrom autograd.extend import defvjp, primitive\nfrom autograd.scipy.special import digamma\n\nrvs = primitive(scipy.stats.dirichlet.rvs)\npdf = primitive(scipy.stats.dirichlet.pdf)\nlogpdf = primitive(scipy.stats.dirichlet.logpdf)\n\ndefvjp(\n logpdf,\n lambda ans, x, alpha: lambda g: g * (alpha - 1) / x,\n lambda ans, x, alpha: lambda g: g * (digamma(np.sum(alpha)) - digamma(alpha) + np.log(x)),\n)\n\n# Same as log pdf, but multiplied by the pdf (ans).\ndefvjp(\n pdf,\n lambda ans, x, alpha: lambda g: g * ans * (alpha - 1) / x,\n lambda ans, x, alpha: lambda g: g * ans * (digamma(np.sum(alpha)) - digamma(alpha) + np.log(x)),\n)\n"
},
{
"path": "autograd/scipy/stats/gamma.py",
"content": "import scipy.stats\n\nimport autograd.numpy as np\nfrom autograd.extend import defvjp, primitive\nfrom autograd.numpy.numpy_vjps import unbroadcast_f\nfrom autograd.scipy.special import gamma, psi\n\ncdf = primitive(scipy.stats.gamma.cdf)\nlogpdf = primitive(scipy.stats.gamma.logpdf)\npdf = primitive(scipy.stats.gamma.pdf)\n\n\ndef grad_gamma_logpdf_arg0(x, a):\n return (a - x - 1) / x\n\n\ndef grad_gamma_logpdf_arg1(x, a):\n return np.log(x) - psi(a)\n\n\ndefvjp(\n cdf,\n lambda ans, x, a: unbroadcast_f(x, lambda g: g * np.exp(-x) * np.power(x, a - 1) / gamma(a)),\n argnums=[0],\n)\ndefvjp(\n logpdf,\n lambda ans, x, a: unbroadcast_f(x, lambda g: g * grad_gamma_logpdf_arg0(x, a)),\n lambda ans, x, a: unbroadcast_f(a, lambda g: g * grad_gamma_logpdf_arg1(x, a)),\n)\ndefvjp(\n pdf,\n lambda ans, x, a: unbroadcast_f(x, lambda g: g * ans * grad_gamma_logpdf_arg0(x, a)),\n lambda ans, x, a: unbroadcast_f(a, lambda g: g * ans * grad_gamma_logpdf_arg1(x, a)),\n)\n"
},
{
"path": "autograd/scipy/stats/multivariate_normal.py",
"content": "import scipy.stats\n\nimport autograd.numpy as np\nfrom autograd.extend import defvjp, primitive\nfrom autograd.numpy.numpy_vjps import unbroadcast_f\n\npdf = primitive(scipy.stats.multivariate_normal.pdf)\nlogpdf = primitive(scipy.stats.multivariate_normal.logpdf)\nentropy = primitive(scipy.stats.multivariate_normal.entropy)\n\n# With thanks to Eric Bresch.\n# Some formulas are from\n# \"An extended collection of matrix derivative results\n# for forward and reverse mode algorithmic differentiation\"\n# by Mike Giles\n# https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf\n\n\ndef generalized_outer_product(x):\n if np.ndim(x) == 1:\n return np.outer(x, x)\n return np.matmul(x, np.swapaxes(x, -1, -2))\n\n\ndef covgrad(x, mean, cov, allow_singular=False):\n if allow_singular:\n raise NotImplementedError(\n \"The multivariate normal pdf is not differentiable w.r.t. a singular covariance matix\"\n )\n J = np.linalg.inv(cov)\n solved = np.matmul(J, np.expand_dims(x - mean, -1))\n return 1.0 / 2 * (generalized_outer_product(solved) - J)\n\n\ndef solve(allow_singular):\n if allow_singular:\n return lambda A, x: np.dot(np.linalg.pinv(A), x)\n else:\n return np.linalg.solve\n\n\ndefvjp(\n logpdf,\n lambda ans, x, mean, cov, allow_singular=False: unbroadcast_f(\n x, lambda g: -np.expand_dims(np.atleast_1d(g), 1) * solve(allow_singular)(cov, (x - mean).T).T\n ),\n lambda ans, x, mean, cov, allow_singular=False: unbroadcast_f(\n mean, lambda g: np.expand_dims(np.atleast_1d(g), 1) * solve(allow_singular)(cov, (x - mean).T).T\n ),\n lambda ans, x, mean, cov, allow_singular=False: unbroadcast_f(\n cov, lambda g: np.reshape(g, np.shape(g) + (1, 1)) * covgrad(x, mean, cov, allow_singular)\n ),\n)\n\n# Same as log pdf, but multiplied by the pdf (ans).\ndefvjp(\n pdf,\n lambda ans, x, mean, cov, allow_singular=False: unbroadcast_f(\n x, lambda g: -np.expand_dims(np.atleast_1d(ans * g), 1) * solve(allow_singular)(cov, (x - mean).T).T\n ),\n lambda ans, x, mean, cov, allow_singular=False: unbroadcast_f(\n mean,\n lambda g: np.expand_dims(np.atleast_1d(ans * g), 1) * solve(allow_singular)(cov, (x - mean).T).T,\n ),\n lambda ans, x, mean, cov, allow_singular=False: unbroadcast_f(\n cov, lambda g: np.reshape(ans * g, np.shape(g) + (1, 1)) * covgrad(x, mean, cov, allow_singular)\n ),\n)\n\ndefvjp(entropy, None, lambda ans, mean, cov: unbroadcast_f(cov, lambda g: 0.5 * g * np.linalg.inv(cov).T))\n"
},
{
"path": "autograd/scipy/stats/norm.py",
"content": "\"\"\"Gradients of the normal distribution.\"\"\"\n\nimport scipy.stats\n\nimport autograd.numpy as anp\nfrom autograd.extend import defvjp, primitive\nfrom autograd.numpy.numpy_vjps import unbroadcast_f\n\npdf = primitive(scipy.stats.norm.pdf)\ncdf = primitive(scipy.stats.norm.cdf)\nsf = primitive(scipy.stats.norm.sf)\nlogpdf = primitive(scipy.stats.norm.logpdf)\nlogcdf = primitive(scipy.stats.norm.logcdf)\nlogsf = primitive(scipy.stats.norm.logsf)\n\ndefvjp(\n pdf,\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: -g * ans * (x - loc) / scale**2),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(loc, lambda g: g * ans * (x - loc) / scale**2),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n scale, lambda g: g * ans * (((x - loc) / scale) ** 2 - 1.0) / scale\n ),\n)\n\ndefvjp(\n cdf,\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: g * pdf(x, loc, scale)),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(loc, lambda g: -g * pdf(x, loc, scale)),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n scale, lambda g: -g * pdf(x, loc, scale) * (x - loc) / scale\n ),\n)\n\ndefvjp(\n logpdf,\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: -g * (x - loc) / scale**2),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(loc, lambda g: g * (x - loc) / scale**2),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n scale, lambda g: g * (-1.0 / scale + (x - loc) ** 2 / scale**3)\n ),\n)\n\ndefvjp(\n logcdf,\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n x, lambda g: g * anp.exp(logpdf(x, loc, scale) - logcdf(x, loc, scale))\n ),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n loc, lambda g: -g * anp.exp(logpdf(x, loc, scale) - logcdf(x, loc, scale))\n ),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n scale, lambda g: -g * anp.exp(logpdf(x, loc, scale) - logcdf(x, loc, scale)) * (x - loc) / scale\n ),\n)\n\ndefvjp(\n logsf,\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n x, lambda g: -g * anp.exp(logpdf(x, loc, scale) - logsf(x, loc, scale))\n ),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n loc, lambda g: g * anp.exp(logpdf(x, loc, scale) - logsf(x, loc, scale))\n ),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n scale, lambda g: g * anp.exp(logpdf(x, loc, scale) - logsf(x, loc, scale)) * (x - loc) / scale\n ),\n)\n\ndefvjp(\n sf,\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: -g * pdf(x, loc, scale)),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(loc, lambda g: g * pdf(x, loc, scale)),\n lambda ans, x, loc=0.0, scale=1.0: unbroadcast_f(\n scale, lambda g: g * pdf(x, loc, scale) * (x - loc) / scale\n ),\n)\n"
},
{
"path": "autograd/scipy/stats/poisson.py",
"content": "import scipy.stats\n\nimport autograd.numpy as np\nfrom autograd.extend import defvjp, primitive\nfrom autograd.numpy.numpy_vjps import unbroadcast_f\n\ncdf = primitive(scipy.stats.poisson.cdf)\nlogpmf = primitive(scipy.stats.poisson.logpmf)\npmf = primitive(scipy.stats.poisson.pmf)\n\n\ndef grad_poisson_logpmf(k, mu):\n return np.where(k % 1 == 0, k / mu - 1, 0)\n\n\ndefvjp(cdf, lambda ans, k, mu: unbroadcast_f(mu, lambda g: g * -pmf(np.floor(k), mu)), argnums=[1])\ndefvjp(logpmf, lambda ans, k, mu: unbroadcast_f(mu, lambda g: g * grad_poisson_logpmf(k, mu)), argnums=[1])\ndefvjp(\n pmf, lambda ans, k, mu: unbroadcast_f(mu, lambda g: g * ans * grad_poisson_logpmf(k, mu)), argnums=[1]\n)\n"
},
{
"path": "autograd/scipy/stats/t.py",
"content": "\"\"\"Gradients of the univariate t distribution.\"\"\"\n\nimport scipy.stats\n\nimport autograd.numpy as np\nfrom autograd.extend import defvjp, primitive\nfrom autograd.numpy.numpy_vjps import unbroadcast_f\nfrom autograd.scipy.special import psi\n\npdf = primitive(scipy.stats.t.pdf)\ncdf = primitive(scipy.stats.t.cdf)\nlogpdf = primitive(scipy.stats.t.logpdf)\nlogcdf = primitive(scipy.stats.t.logcdf)\n\n\ndef grad_tlogpdf_diff(diff, df):\n return -diff * (1.0 + df) / (diff**2 + df)\n\n\ndef grad_tlogpdf_x(x, df, loc, scale):\n return grad_tlogpdf_diff((x - loc) / scale, df) / scale\n\n\ndef grad_tlogpdf_loc(x, df, loc, scale):\n return -grad_tlogpdf_diff((x - loc) / scale, df) / scale\n\n\ndef grad_tlogpdf_scale(x, df, loc, scale):\n diff = x - loc\n return -(df * (scale**2 - diff**2)) / (scale * (df * scale**2 + diff**2))\n\n\ndef grad_tlogpdf_df(x, df, loc, scale):\n y = (x - loc) / scale\n return 0.5 * (\n (y**2 * (df + 1)) / (df * (y**2 + df))\n - np.log(y**2 / df + 1)\n - 1.0 / df\n - psi(df / 2.0)\n + psi((df + 1) / 2.0)\n )\n\n\ndefvjp(\n pdf,\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n x, lambda g: g * ans * grad_tlogpdf_x(x, df, loc, scale)\n ),\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n df, lambda g: g * ans * grad_tlogpdf_df(x, df, loc, scale)\n ),\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n loc, lambda g: g * ans * grad_tlogpdf_loc(x, df, loc, scale)\n ),\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n scale, lambda g: g * ans * grad_tlogpdf_scale(x, df, loc, scale)\n ),\n)\n\ndefvjp(\n cdf,\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: g * pdf(x, df, loc, scale)),\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(loc, lambda g: -g * pdf(x, df, loc, scale)),\n argnums=(0, 2),\n)\n\ndefvjp(\n logpdf,\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(x, lambda g: g * grad_tlogpdf_x(x, df, loc, scale)),\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n df, lambda g: g * grad_tlogpdf_df(x, df, loc, scale)\n ),\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n loc, lambda g: g * grad_tlogpdf_loc(x, df, loc, scale)\n ),\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n scale, lambda g: g * grad_tlogpdf_scale(x, df, loc, scale)\n ),\n)\n\ndefvjp(\n logcdf,\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n x, lambda g: g * np.exp(logpdf(x, df, loc, scale) - logcdf(x, df, loc, scale))\n ),\n lambda ans, x, df, loc=0.0, scale=1.0: unbroadcast_f(\n loc, lambda g: -g * np.exp(logpdf(x, df, loc, scale) - logcdf(x, df, loc, scale))\n ),\n argnums=(0, 2),\n)\n"
},
{
"path": "autograd/test_util.py",
"content": "from itertools import product\n\nfrom .core import make_jvp, make_vjp, vspace\nfrom .wrap_util import get_name, unary_to_nary\n\nTOL = 1e-6\nRTOL = 1e-6\n\n\ndef scalar_close(a, b):\n return abs(a - b) < TOL or abs(a - b) / abs(a + b) < RTOL\n\n\nEPS = 1e-6\n\n\ndef make_numerical_jvp(f, x):\n y = f(x)\n x_vs, y_vs = vspace(x), vspace(y)\n\n def jvp(v):\n # (f(x + v*eps/2) - f(x - v*eps/2)) / eps\n f_x_plus = f(x_vs.add(x, x_vs.scalar_mul(v, EPS / 2)))\n f_x_minus = f(x_vs.add(x, x_vs.scalar_mul(v, -EPS / 2)))\n neg_f_x_minus = y_vs.scalar_mul(f_x_minus, -1.0)\n return y_vs.scalar_mul(y_vs.add(f_x_plus, neg_f_x_minus), 1.0 / EPS)\n\n return jvp\n\n\ndef check_vjp(f, x):\n vjp, y = make_vjp(f, x)\n jvp = make_numerical_jvp(f, x)\n x_vs, y_vs = vspace(x), vspace(y)\n x_v, y_v = x_vs.randn(), y_vs.randn()\n\n vjp_y = x_vs.covector(vjp(y_vs.covector(y_v)))\n assert vspace(vjp_y) == x_vs\n vjv_exact = x_vs.inner_prod(x_v, vjp_y)\n vjv_numeric = y_vs.inner_prod(y_v, jvp(x_v))\n assert scalar_close(vjv_numeric, vjv_exact), (\n \"Derivative (VJP) check of {} failed with arg {}:\\nanalytic: {}\\nnumeric: {}\".format(\n get_name(f), x, vjv_exact, vjv_numeric\n )\n )\n\n\ndef check_jvp(f, x):\n jvp = make_jvp(f, x)\n jvp_numeric = make_numerical_jvp(f, x)\n x_v = vspace(x).randn()\n check_equivalent(jvp(x_v)[1], jvp_numeric(x_v))\n\n\ndef check_equivalent(x, y):\n x_vs, y_vs = vspace(x), vspace(y)\n assert x_vs == y_vs, f\"VSpace mismatch:\\nx: {x_vs}\\ny: {y_vs}\"\n v = x_vs.randn()\n assert scalar_close(x_vs.inner_prod(x, v), x_vs.inner_prod(y, v)), f\"Value mismatch:\\nx: {x}\\ny: {y}\"\n\n\n@unary_to_nary\ndef check_grads(f, x, modes=[\"fwd\", \"rev\"], order=2):\n assert all(m in [\"fwd\", \"rev\"] for m in modes)\n if \"fwd\" in modes:\n check_jvp(f, x)\n if order > 1:\n grad_f = lambda x, v: make_jvp(f, x)(v)[1]\n grad_f.__name__ = f\"jvp_{get_name(f)}\"\n v = vspace(x).randn()\n check_grads(grad_f, (0, 1), modes, order=order - 1)(x, v)\n if \"rev\" in modes:\n check_vjp(f, x)\n if order > 1:\n grad_f = lambda x, v: make_vjp(f, x)[0](v)\n grad_f.__name__ = f\"vjp_{get_name(f)}\"\n v = vspace(f(x)).randn()\n check_grads(grad_f, (0, 1), modes, order=order - 1)(x, v)\n\n\ndef combo_check(fun, *args, **kwargs):\n # Tests all combinations of args and kwargs given.\n _check_grads = lambda f: check_grads(f, *args, **kwargs)\n\n def _combo_check(*args, **kwargs):\n kwarg_key_vals = [[(k, x) for x in xs] for k, xs in kwargs.items()]\n for _args in product(*args):\n for _kwargs in product(*kwarg_key_vals):\n _check_grads(fun)(*_args, **dict(_kwargs))\n\n return _combo_check\n"
},
{
"path": "autograd/tracer.py",
"content": "import warnings\nfrom collections import defaultdict\nfrom contextlib import contextmanager\n\nfrom .util import subvals, toposort\nfrom .wrap_util import wraps\n\n\ndef trace(start_node, fun, x):\n with trace_stack.new_trace() as t:\n start_box = new_box(x, t, start_node)\n end_box = fun(start_box)\n if isbox(end_box) and end_box._trace == start_box._trace:\n return end_box._value, end_box._node\n else:\n warnings.warn(\"Output seems independent of input.\")\n return end_box, None\n\n\nclass Node:\n __slots__ = []\n\n def __init__(self, value, fun, args, kwargs, parent_argnums, parents):\n assert False\n\n def initialize_root(self, *args, **kwargs):\n assert False\n\n @classmethod\n def new_root(cls, *args, **kwargs):\n root = cls.__new__(cls)\n root.initialize_root(*args, **kwargs)\n return root\n\n\ndef primitive(f_raw):\n \"\"\"\n Wraps a function so that its gradient can be specified and its invocation\n can be recorded. For examples, see the docs.\"\"\"\n\n @wraps(f_raw)\n def f_wrapped(*args, **kwargs):\n boxed_args, trace, node_constructor = find_top_boxed_args(args)\n if boxed_args:\n argvals = subvals(args, [(argnum, box._value) for argnum, box in boxed_args])\n if f_wrapped in notrace_primitives[node_constructor]:\n return f_wrapped(*argvals, **kwargs)\n parents = tuple(box._node for _, box in boxed_args)\n argnums = tuple(argnum for argnum, _ in boxed_args)\n ans = f_wrapped(*argvals, **kwargs)\n node = node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)\n return new_box(ans, trace, node)\n else:\n return f_raw(*args, **kwargs)\n\n f_wrapped.fun = f_raw\n f_wrapped._is_autograd_primitive = True\n return f_wrapped\n\n\nnotrace_primitives = defaultdict(set)\n\n\ndef register_notrace(trace_type, primitive_fun):\n notrace_primitives[trace_type].add(primitive_fun)\n\n\ndef notrace_primitive(f_raw):\n @wraps(f_raw)\n def f_wrapped(*args, **kwargs):\n argvals = map(getval, args)\n return f_raw(*argvals, **kwargs)\n\n f_wrapped._is_primitive = True\n return f_wrapped\n\n\ndef find_top_boxed_args(args):\n top_trace = -1\n top_boxes = []\n top_node_type = None\n for argnum, arg in enumerate(args):\n if isbox(arg):\n trace = arg._trace\n if trace > top_trace:\n top_boxes = [(argnum, arg)]\n top_trace = trace\n top_node_type = type(arg._node)\n elif trace == top_trace:\n top_boxes.append((argnum, arg))\n return top_boxes, top_trace, top_node_type\n\n\nclass TraceStack:\n def __init__(self):\n self.top = -1\n\n @contextmanager\n def new_trace(self):\n self.top += 1\n yield self.top\n self.top -= 1\n\n\ntrace_stack = TraceStack()\n\n\nclass Box:\n type_mappings = {}\n types = set()\n\n __slots__ = [\"_value\", \"_trace\", \"_node\"]\n\n def __init__(self, value, trace, node):\n self._value = value\n self._node = node\n self._trace = trace\n\n def __bool__(self):\n return bool(self._value)\n\n __nonzero__ = __bool__\n\n def __str__(self):\n return f\"Autograd {type(self).__name__} with value {str(self._value)}\"\n\n @classmethod\n def register(cls, value_type):\n Box.types.add(cls)\n Box.type_mappings[value_type] = cls\n Box.type_mappings[cls] = cls\n\n\nbox_type_mappings = Box.type_mappings\n\n\ndef new_box(value, trace, node):\n try:\n return box_type_mappings[type(value)](value, trace, node)\n except KeyError:\n raise TypeError(f\"Can't differentiate w.r.t. type {type(value)}\")\n\n\nbox_types = Box.types\nisbox = lambda x: type(x) in box_types # almost 3X faster than isinstance(x, Box)\ngetval = lambda x: getval(x._value) if isbox(x) else x\n"
},
{
"path": "autograd/util.py",
"content": "import operator\n\n\ndef subvals(x, ivs):\n x_ = list(x)\n for i, v in ivs:\n x_[i] = v\n return tuple(x_)\n\n\ndef subval(x, i, v):\n x_ = list(x)\n x_[i] = v\n return tuple(x_)\n\n\ndef func(f):\n return f\n\n\ndef toposort(end_node, parents=operator.attrgetter(\"parents\")):\n child_counts = {}\n stack = [end_node]\n while stack:\n node = stack.pop()\n if node in child_counts:\n child_counts[node] += 1\n else:\n child_counts[node] = 1\n stack.extend(parents(node))\n\n childless_nodes = [end_node]\n while childless_nodes:\n node = childless_nodes.pop()\n yield node\n for parent in parents(node):\n if child_counts[parent] == 1:\n childless_nodes.append(parent)\n else:\n child_counts[parent] -= 1\n\n\n# -------------------- deprecation warnings -----------------------\n\nimport warnings\n\ndeprecation_msg = \"\"\"\nThe quick_grad_check function is deprecated. See the update guide:\nhttps://github.com/HIPS/autograd/blob/master/docs/updateguide.md\"\"\"\n\n\ndef quick_grad_check(\n fun, arg0, extra_args=(), kwargs={}, verbose=True, eps=1e-4, rtol=1e-4, atol=1e-6, rs=None\n):\n warnings.warn(deprecation_msg)\n from autograd.test_util import check_grads\n\n fun_ = lambda arg0: fun(arg0, *extra_args, **kwargs)\n check_grads(fun_, modes=[\"rev\"], order=1)(arg0)\n"
},
{
"path": "autograd/wrap_util.py",
"content": "from .util import subvals\n\n\ndef unary_to_nary(unary_operator):\n @wraps(unary_operator)\n def nary_operator(fun, argnum=0, *nary_op_args, **nary_op_kwargs):\n assert type(argnum) in (int, tuple, list), argnum\n\n @wrap_nary_f(fun, unary_operator, argnum)\n def nary_f(*args, **kwargs):\n @wraps(fun)\n def unary_f(x):\n if isinstance(argnum, int):\n subargs = subvals(args, [(argnum, x)])\n else:\n subargs = subvals(args, zip(argnum, x))\n return fun(*subargs, **kwargs)\n\n if isinstance(argnum, int):\n x = args[argnum]\n else:\n x = tuple(args[i] for i in argnum)\n return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)\n\n return nary_f\n\n return nary_operator\n\n\ndef wraps(fun, namestr=\"{fun}\", docstr=\"{doc}\", **kwargs):\n def _wraps(f):\n try:\n f.__name__ = namestr.format(fun=get_name(fun), **kwargs)\n f.__doc__ = docstr.format(fun=get_name(fun), doc=get_doc(fun), **kwargs)\n except BaseException:\n pass\n return f\n\n return _wraps\n\n\ndef wrap_nary_f(fun, op, argnum):\n namestr = \"{op}_of_{fun}_wrt_argnum_{argnum}\"\n docstr = \"\"\"\\\n {op} of function {fun} with respect to argument number {argnum}. Takes the\n same arguments as {fun} but returns the {op}.\n \"\"\"\n return wraps(fun, namestr, docstr, op=get_name(op), argnum=argnum)\n\n\nget_name = lambda f: getattr(f, \"__name__\", \"[unknown name]\")\nget_doc = lambda f: getattr(f, \"__doc__\", \"\")\n"
},
{
"path": "benchmarks/__init__.py",
"content": ""
},
{
"path": "benchmarks/asv.conf.json.sample",
"content": "{\n \"version\": 1,\n \"project\": \"autograd\",\n \"project_url\": \"http://github.com/hips/autograd\",\n \"branches\": [\"master\"],\n \"dvcs\": \"git\",\n \"environment_type\": \"virtualenv\",\n \"install_timeout\": 600,\n \"repo\" : \"..\",\n \"benchmark_dir\" : \".\",\n \"env_dir\" : \"../.asv/env\",\n \"results_dir\" : \"../.asv/results\",\n \"html_dir\" : \"../.asv/html\",\n}\n"
},
{
"path": "benchmarks/bench_core.py",
"content": "import numpy as onp\n\nimport autograd.numpy as np\nfrom autograd import grad\n\ntry:\n from autograd.core import VJPNode, backward_pass, vspace\n from autograd.tracer import new_box, trace\n\n MASTER_BRANCH = False\nexcept ImportError:\n from autograd.core import backward_pass, forward_pass, new_progenitor, vspace\n\n MASTER_BRANCH = True\n\n\n## SHORT FUNCTION\ndef f_short(x):\n return x**2\n\n\ndef time_short_fun():\n f_short(2.0)\n\n\ndef time_short_forward_pass():\n if MASTER_BRANCH:\n forward_pass(f_short, (2.0,), {})\n else:\n start_node = VJPNode.new_root()\n trace(start_node, f_short, x)\n\n\ndef time_short_backward_pass():\n if MASTER_BRANCH:\n backward_pass(1.0, short_end_node, short_start_node)\n else:\n backward_pass(1.0, short_end_node)\n\n\ndef time_short_grad():\n grad(f_short)(2.0)\n\n\n## LONG FUNCTION\ndef f_long(x):\n for i in range(50):\n x = np.sin(x)\n return x\n\n\ndef time_long_fun():\n f_long(2.0)\n\n\ndef time_long_forward_pass():\n if MASTER_BRANCH:\n forward_pass(f_long, (2.0,), {})\n else:\n start_node = VJPNode.new_root()\n trace(start_node, f_long, x)\n\n\ndef time_long_backward_pass():\n if MASTER_BRANCH:\n backward_pass(1.0, long_end_node, long_start_node)\n else:\n backward_pass(1.0, long_end_node)\n\n\ndef time_long_grad():\n grad(f_long)(2.0)\n\n\n## 'PEARLMUTTER TEST' FUNCTION\ndef fan_out_fan_in(x):\n for i in range(10**4):\n x = (x + x) / 2.0\n return np.sum(x)\n\n\ndef time_fan_out_fan_in_fun():\n fan_out_fan_in(2.0)\n\n\ndef time_fan_out_fan_in_forward_pass():\n if MASTER_BRANCH:\n forward_pass(fan_out_fan_in, (2.0,), {})\n else:\n start_node = VJPNode.new_root()\n trace(start_node, fan_out_fan_in, x)\n\n\ndef time_fan_out_fan_in_backward_pass():\n if MASTER_BRANCH:\n backward_pass(1.0, fan_end_node, fan_start_node)\n else:\n backward_pass(1.0, fan_end_node)\n\n\ndef time_fan_out_fan_in_grad():\n grad(fan_out_fan_in)(2.0)\n\n\n## UNIT BENCHMARKS\ndef time_vspace_float():\n vspace(1.0)\n\n\nA = np.array([[1.0, 2.0, 3.0]])\n\n\ndef time_vspace_array():\n vspace(A)\n\n\ndef time_new_box_float():\n new_box(1.0, 0, start_node)\n\n\ndef time_new_box_array():\n new_box(A, 0, start_node)\n\n\ndef time_exp_call():\n onp.exp(2.0)\n\n\ndef time_exp_primitive_call_unboxed():\n np.exp(2.0)\n\n\ndef time_exp_primitive_call_boxed():\n if MASTER_BRANCH:\n np.exp(progenitor)\n else:\n np.exp(start_box)\n\n\ndef time_no_autograd_control():\n # Test whether the benchmarking machine is running slowly independent of autograd\n A = np.random.randn(200, 200)\n np.dot(A, A)\n\n\nif MASTER_BRANCH:\n short_start_node, short_end_node = forward_pass(f_short, (2.0,), {})\n long_start_node, long_end_node = forward_pass(f_long, (2.0,), {})\n fan_start_node, fan_end_node = forward_pass(fan_out_fan_in, (2.0,), {})\n progenitor = new_progenitor(2.0)\nelse:\n x = 2.0\n start_node = VJPNode.new_root()\n start_box = new_box(x, 0, start_node)\n _, short_end_node = trace(VJPNode.new_root(), f_short, x)\n _, long_end_node = trace(VJPNode.new_root(), f_long, x)\n _, fan_end_node = trace(VJPNode.new_root(), fan_out_fan_in, x)\n"
},
{
"path": "benchmarks/bench_mem.py",
"content": "import autograd.numpy as np\nfrom autograd import grad\n\n\ndef peakmem_needless_nodes():\n N, M = 1000, 100\n\n def fun(x):\n for i in range(M):\n x = x + 1\n return np.sum(x)\n\n grad(fun)(np.zeros((N, N)))\n"
},
{
"path": "benchmarks/bench_numpy_vjps.py",
"content": "import autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import make_vjp\n\ndot_0 = lambda a, b, g: make_vjp(np.dot, argnum=0)(a, b)[0](g)\ndot_1 = lambda a, b, g: make_vjp(np.dot, argnum=1)(a, b)[0](g)\n\ndot_0_0 = lambda a, b, g: make_vjp(dot_0, argnum=0)(a, b, g)[0](a)\ndot_0_1 = lambda a, b, g: make_vjp(dot_0, argnum=1)(a, b, g)[0](a)\ndot_0_2 = lambda a, b, g: make_vjp(dot_0, argnum=2)(a, b, g)[0](a)\n\ndot_1_0 = lambda a, b, g: make_vjp(dot_1, argnum=0)(a, b, g)[0](b)\ndot_1_1 = lambda a, b, g: make_vjp(dot_1, argnum=1)(a, b, g)[0](b)\ndot_1_2 = lambda a, b, g: make_vjp(dot_1, argnum=2)(a, b, g)[0](b)\n\na = npr.randn(2, 3, 4, 5)\nb = npr.randn(2, 3, 5, 4)\ng = npr.randn(2, 3, 4, 2, 3, 4)\n\n\ndef time_dot_0():\n dot_0(a, b, g)\n\n\ndef time_dot_1():\n dot_1(a, b, g)\n\n\ndef time_dot_0_0():\n dot_0_0(a, b, g)\n\n\ndef time_dot_0_1():\n dot_0_1(a, b, g)\n\n\ndef time_dot_0_2():\n dot_0_2(a, b, g)\n\n\ndef time_dot_1_0():\n dot_1_0(a, b, g)\n\n\ndef time_dot_1_1():\n dot_1_1(a, b, g)\n\n\ndef time_dot_1_2():\n dot_1_2(a, b, g)\n\n\ntensordot_0 = lambda A, B, G: make_vjp(np.tensordot, argnum=0)(A, B, 2)[0](G)\ntensordot_1 = lambda A, B, G: make_vjp(np.tensordot, argnum=1)(A, B, 2)[0](G)\n\ntensordot_0_0 = lambda A, B, G: make_vjp(tensordot_0, argnum=0)(A, B, G)[0](A)\ntensordot_0_1 = lambda A, B, G: make_vjp(tensordot_0, argnum=1)(A, B, G)[0](A)\ntensordot_0_2 = lambda A, B, G: make_vjp(tensordot_0, argnum=2)(A, B, G)[0](A)\n\ntensordot_1_0 = lambda A, B, G: make_vjp(tensordot_1, argnum=0)(A, B, G)[0](B)\ntensordot_1_1 = lambda A, B, G: make_vjp(tensordot_1, argnum=1)(A, B, G)[0](B)\ntensordot_1_2 = lambda A, B, G: make_vjp(tensordot_1, argnum=2)(A, B, G)[0](B)\n\nA = npr.randn(2, 3, 5, 4)\nB = npr.randn(5, 4, 2, 3)\nG = npr.randn(2, 3, 2, 3)\n\n\ndef time_tensordot_0():\n tensordot_0(A, B, G)\n\n\ndef time_tensordot_1():\n tensordot_1(A, B, G)\n\n\ndef time_tensordot_0_0():\n tensordot_0_0(A, B, G)\n\n\ndef time_tensordot_0_1():\n tensordot_0_1(A, B, G)\n\n\ndef time_tensordot_0_2():\n tensordot_0_2(A, B, G)\n\n\ndef time_tensordot_1_0():\n tensordot_1_0(A, B, G)\n\n\ndef time_tensordot_1_1():\n tensordot_1_1(A, B, G)\n\n\ndef time_tensordot_1_2():\n tensordot_1_2(A, B, G)\n"
},
{
"path": "benchmarks/bench_rnn.py",
"content": "# Write the benchmarking functions here.\n# See \"Writing benchmarks\" in the asv docs for more information.\n# http://asv.readthedocs.io/en/latest/writing_benchmarks.html\nimport autograd.numpy as np\nfrom autograd import grad\n\n\nclass RNNSuite:\n \"\"\"\n Checking speed on a vanilla RNN.\n \"\"\"\n\n # NOTE: this is run each time we run a benchmark.\n # Might want to switch to setup_cache, which has to return an object which is loaded and unpacked in setup().\n def setup(self):\n self.batch_size = 16\n self.dtype = \"float32\"\n self.D = 2**10\n self.x = 0.01 * np.random.randn(self.batch_size, self.D).astype(self.dtype)\n self.W1 = 0.01 * np.random.randn(self.D, self.D).astype(self.dtype)\n self.b1 = 0.01 * np.random.randn(self.D).astype(self.dtype)\n self.Wout = 0.01 * np.random.randn(self.D, 1).astype(self.dtype)\n self.bout = 0.01 * np.random.randn(1).astype(self.dtype)\n self.l = (np.random.rand(self.batch_size, 1) > 0.5).astype(self.dtype)\n self.n = 50\n\n def autograd_rnn(params, x, label, n):\n W, b, Wout, bout = params\n h1 = x\n for i in range(n):\n h1 = np.tanh(np.dot(h1, W) + b)\n logit = np.dot(h1, Wout) + bout\n loss = -np.sum(label * logit - (logit + np.log(1 + np.exp(-logit))))\n return loss\n\n self.fn = autograd_rnn\n self.grad_fn = grad(self.fn)\n\n def rnn_grad(self):\n self.grad_fn((self.W1, self.b1, self.Wout, self.bout), self.x, self.l, self.n)\n\n def time_rnn_grad(self):\n self.rnn_grad()\n\n def peakmem_rnn_grad(self):\n self.rnn_grad()\n\n def time_manual_rnn_grad(self):\n self.manual_rnn_grad()\n\n def peakmem_manual_rnn_grad(self):\n self.manual_rnn_grad()\n\n def manual_rnn_grad(self):\n def repeat_to_match_shape(g, A, axis=None):\n gout = np.empty_like(A)\n if np.ndim(gout) == 0:\n gout = g\n else:\n gout = np.ones_like(A) * g\n return gout\n\n def sum_to_match_shape(sum_this, to_match_this):\n sum_this = np.sum(sum_this, axis=tuple(range(0, np.ndim(sum_this) - np.ndim(to_match_this))))\n for axis, size in enumerate(np.shape(to_match_this)):\n if size == 1:\n sum_this = np.sum(sum_this, axis=axis, keepdims=True)\n return sum_this\n\n def grad_dot_A(g, A, B):\n ga = np.dot(g, B.T)\n ga = np.reshape(ga, np.shape(A))\n return ga\n\n def grad_dot_B(g, A, B):\n gb = np.dot(A.T, g)\n gb = np.reshape(gb, np.shape(B))\n return gb\n\n def _rnn_grad(x, W, b, Wout, bout, label, n):\n h1__1_stack, h1__1 = [], None\n h1__0_stack, h1__0 = [], None\n out_stack, out = [], None\n h1_stack = []\n h1 = x\n _for1 = list(range(n))\n\n for i in _for1:\n h1__1_stack.append(h1__1)\n h1__1 = np.dot(h1, W)\n h1__0_stack.append(h1__0)\n h1__0 = h1__1 + b\n h1_stack.append(h1)\n h1 = np.tanh(h1__0)\n out__0 = np.dot(h1, Wout)\n out = out__0 + bout\n loss__2 = label * out\n loss__7 = -out\n loss__6 = np.exp(loss__7)\n loss__5 = 1 + loss__6\n loss__4 = np.log(loss__5)\n loss__3 = out + loss__4\n loss__1 = loss__2 - loss__3\n\n # Begin Backward Pass\n g_loss = 1\n g_h1__0 = 0\n g_h1__1 = 0\n g_b = 0\n g_W = 0\n\n # Reverse of: loss = -loss__0\n g_loss__0 = -g_loss\n\n # Reverse of: loss__0 = np.sum(loss__1)\n g_loss__1 = repeat_to_match_shape(g_loss__0, loss__1)\n\n # Reverse of: loss__1 = loss__2 - loss__3\n g_loss__2 = sum_to_match_shape(g_loss__1, loss__2)\n g_loss__3 = sum_to_match_shape(-g_loss__1, loss__3)\n\n # Reverse of: loss__3 = out + loss__4\n g_out = sum_to_match_shape(g_loss__3, out)\n g_loss__4 = sum_to_match_shape(g_loss__3, loss__4)\n\n # Reverse of: loss__4 = np.log(loss__5)\n g_loss__5 = g_loss__4 / loss__5\n\n # Reverse of: loss__5 = 1 + loss__6\n g_loss__6 = sum_to_match_shape(g_loss__5, loss__6)\n\n # Reverse of: loss__6 = np.exp(loss__7)\n g_loss__7 = g_loss__6 * np.exp(loss__7)\n\n # Reverse of: loss__7 = -out\n g_out += -g_loss__7\n g_out += sum_to_match_shape(g_loss__2 * label, out)\n\n # Reverse of: out = out__0 + bout\n g_out__0 = sum_to_match_shape(g_out, out__0)\n g_bout = sum_to_match_shape(g_out, bout)\n\n # Reverse of: out__0 = np.dot(h1, Wout)\n g_h1 = grad_dot_A(g_out__0, h1, Wout)\n g_Wout = grad_dot_B(g_out__0, h1, Wout)\n _for1 = reversed(_for1)\n for i in _for1:\n h1 = h1_stack.pop()\n tmp_g0 = g_h1 / np.cosh(h1__0) ** 2.0\n g_h1 = 0\n g_h1__0 += tmp_g0\n h1__0 = h1__0_stack.pop()\n tmp_g1 = sum_to_match_shape(g_h1__0, h1__1)\n tmp_g2 = sum_to_match_shape(g_h1__0, b)\n g_h1__0 = 0\n g_h1__1 += tmp_g1\n g_b += tmp_g2\n h1__1 = h1__1_stack.pop()\n tmp_g3 = grad_dot_A(g_h1__1, h1, W)\n tmp_g4 = grad_dot_B(g_h1__1, h1, W)\n g_h1__1 = 0\n g_h1 += tmp_g3\n g_W += tmp_g4\n return g_W, g_b, g_Wout, g_bout\n\n _rnn_grad(self.x, self.W1, self.b1, self.Wout, self.bout, self.l, self.n)\n pass\n"
},
{
"path": "benchmarks/bench_util.py",
"content": "import autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\n\ntry:\n from autograd.misc.flatten import flatten\nexcept ImportError:\n from autograd.util import flatten\n\n\ndef time_flatten():\n val = {\n \"k\": npr.random((4, 4)),\n \"k2\": npr.random((3, 3)),\n \"k3\": 3.0,\n \"k4\": [1.0, 4.0, 7.0, 9.0],\n \"k5\": np.array([4.0, 5.0, 6.0]),\n \"k6\": np.array([[7.0, 8.0], [9.0, 10.0]]),\n }\n\n vect, unflatten = flatten(val)\n val_recovered = unflatten(vect)\n vect_2, _ = flatten(val_recovered)\n\n\n# def time_vspace_flatten():\n# val = {'k': npr.random((4, 4)),\n# 'k2': npr.random((3, 3)),\n# 'k3': 3.0,\n# 'k4': [1.0, 4.0, 7.0, 9.0],\n# 'k5': np.array([4., 5., 6.]),\n# 'k6': np.array([[7., 8.], [9., 10.]])}\n\n# vspace_flatten(val)\n\n\ndef time_grad_flatten():\n val = {\n \"k\": npr.random((4, 4)),\n \"k2\": npr.random((3, 3)),\n \"k3\": 3.0,\n \"k4\": [1.0, 4.0, 7.0, 9.0],\n \"k5\": np.array([4.0, 5.0, 6.0]),\n \"k6\": np.array([[7.0, 8.0], [9.0, 10.0]]),\n }\n\n vect, unflatten = flatten(val)\n\n def fun(vec):\n v = unflatten(vec)\n return np.sum(v[\"k5\"]) + np.sum(v[\"k6\"])\n\n grad(fun)(vect)\n"
},
{
"path": "conda_recipe/conda.yaml",
"content": "package:\n name: autograd\n # there are ways to derive version from other sources; for now, it's hard-coded\n version: 1.1.1\n\nsource:\n {% if not environ.get('BINSTAR_PLATFORM', None) %}\n git_url: ../\n {% else %}\n # we're building on binstar, we already have the repo; treat as local path\n path: ../\n {% endif %}\n\nrequirements:\n build:\n - python\n - hatch\n - hatchling\n - future\n - numpy >=1.9\n\n run:\n - python\n - future\n - numpy >=1.9\n\nbuild:\n script: pip install . --no-deps\n\ntest:\n # Python imports\n imports:\n - autograd\n - autograd.numpy\n\nabout:\n home: https://github.com/HIPS/autograd\n license: MIT\n summary: 'Efficiently computes derivatives of numpy code.'\n"
},
{
"path": "docs/tutorial.md",
"content": "# Autograd tutorial\n\n## Motivation\n\nImagine you want to test out a new machine learning model for your data. This\nusually means coming up with some loss function to capture how well your model\nfits the data and optimizing that loss with respect to the model parameters. If\nthere are many model parameters (neural nets can have millions) then you need\ngradients. You then have two options: derive and code them up yourself, or\nimplement your model using the syntactic and semantic constraints of a system\nlike [Theano](http://deeplearning.net/software/theano/) or\n[TensorFlow](https://github.com/tensorflow/tensorflow).\n\nWe want to provide a third way: just write down the loss function using a\nstandard numerical library like Numpy, and Autograd will give you its gradient.\n\n## How to use Autograd\n\nAutograd's `grad` function takes in a function, and gives you a function that computes its derivative.\nYour function must have a scalar-valued output (i.e. a float).\nThis covers the common case when you want to use gradients to optimize something.\n\nAutograd works on ordinary Python and Numpy code containing all the usual control structures, including `while` loops, `if` statements, and closures. Here's a simple example of using an open-ended loop to compute the sine function:\n\n```python\nimport autograd.numpy as np # Thinly-wrapped version of Numpy\nfrom autograd import grad\n\ndef taylor_sine(x): # Taylor approximation to sine function\n ans = currterm = x\n i = 0\n while np.abs(currterm) > 0.001:\n currterm = -currterm * x**2 / ((2 * i + 3) * (2 * i + 2))\n ans = ans + currterm\n i += 1\n return ans\n\ngrad_sine = grad(taylor_sine)\nprint \"Gradient of sin(pi) is\", grad_sine(np.pi)\n```\n\n## Complete example: logistic regression\n\nA common use case for automatic differentiation is to train a probabilistic model.\nHere we present a very simple (but complete) example of specifying and training\na logistic regression model for binary classification:\n\n```python\nimport autograd.numpy as np\nfrom autograd import grad\n\ndef sigmoid(x):\n return 0.5 * (np.tanh(x / 2.) + 1)\n\ndef logistic_predictions(weights, inputs):\n # Outputs probability of a label being true according to logistic model.\n return sigmoid(np.dot(inputs, weights))\n\ndef training_loss(weights):\n # Training loss is the negative log-likelihood of the training labels.\n preds = logistic_predictions(weights, inputs)\n label_probabilities = preds * targets + (1 - preds) * (1 - targets)\n return -np.sum(np.log(label_probabilities))\n\n# Build a toy dataset.\ninputs = np.array([[0.52, 1.12, 0.77],\n [0.88, -1.08, 0.15],\n [0.52, 0.06, -1.30],\n [0.74, -2.49, 1.39]])\ntargets = np.array([True, True, False, True])\n\n# Define a function that returns gradients of training loss using Autograd.\ntraining_gradient_fun = grad(training_loss)\n\n# Optimize weights using gradient descent.\nweights = np.array([0.0, 0.0, 0.0])\nprint(\"Initial loss:\", training_loss(weights))\nfor i in range(100):\n weights -= training_gradient_fun(weights) * 0.01\n\nprint(\"Trained loss:\", training_loss(weights))\n```\n\nPython syntax is pretty good for specifying probabilistic models. The biggest\nwin is that it becomes a lot easier to modify a model and rapidly iterate.\n\nFor more complex examples, see our [examples directory](../examples/), which includes:\n* [a simple neural net](../examples/neural_net.py)\n* [a convolutional neural net](../examples/convnet.py)\n* [a recurrent neural net](../examples/rnn.py)\n* [a long short-term memory (LSTM)](../examples/lstm.py)\n* [backpropagating through a fluid simulation](../examples/fluidsim/fluidsim.py)\n\n\n## What's going on under the hood?\n\nTo compute the gradient, Autograd first has to record every transformation that was applied to the input as it was turned into the output of your function.\nTo do this, Autograd wraps functions (using the function `primitive`) so that when they're called, they add themselves to a list of operations performed.\nAutograd's core has a table mapping these wrapped primitives to their corresponding gradient functions (or, more precisely, their vector-Jacobian product functions).\nTo flag the variables we're taking the gradient with respect to, we wrap them using the `Box` class.\nYou should never have to think about the `Box` class, but you might notice it when printing out debugging info.\n\nAfter the function is evaluated, Autograd has a graph specifying all operations that were performed on the inputs with respect to which we want to differentiate.\nThis is the computational graph of the function evaluation.\nTo compute the derivative, we simply apply the rules of differentiation to each node in the graph.\n\n### Reverse mode differentiation\n\nGiven a function made up of several nested function calls, there are several ways to compute its derivative.\n\nFor example, given L(x) = F(G(H(x))), the chain rule says that its gradient is dL/dx = dF/dG * dG/dH * dH/dx. If we evaluate this product from right-to-left: (dF/dG * (dG/dH * dH/dx)), the same order as the computations themselves were performed, this is called forward-mode differentiation.\nIf we evaluate this product from left-to-right: ((dF/dG * dG/dH) * dH/dx), the reverse order as the computations themselves were performed, this is called reverse-mode differentiation.\n\nCompared to finite differences or forward-mode, reverse-mode differentiation is by far the more practical method for differentiating functions that take in a large vector and output a single number.\nIn the machine learning community, reverse-mode differentiation is known as 'backpropagation', since the gradients propagate backwards through the function.\nIt's particularly nice since you don't need to instantiate the intermediate Jacobian matrices explicitly, and instead only rely on applying a sequence of matrix-free vector-Jacobian product functions (VJPs).\nBecause Autograd supports higher derivatives as well, Hessian-vector products (a form of second-derivative) are also available and efficient to compute.\n\n### How can you support ifs, while loops and recursion?\n\nSome autodiff packages (such as [TensorFlow](https://github.com/tensorflow/tensorflow)) work by having you specify a graph of the computation that your function performs, including all the control flow (such as if and for loops), and then turn that graph into another one that computes gradients.\nThis has some benefits (such as allowing compile-time optimizations), but it requires you to express control flow in a limited mini-language that those packages know how to handle. (For example, the `tf.while` and `tf.cond` operations in TensorFlow.)\n\nIn contrast, Autograd doesn't have to know about any ifs, branches, loops or recursion that were used to decide which operations were called. To compute the gradient of a particular input, one only needs to know which continuous transforms were applied to that particular input, not which other transforms might have been applied.\nSince Autograd keeps track of the relevant operations on each function call separately, it's not a problem that all the Python control flow operations are invisible to Autograd. In fact, it greatly simplifies the implementation.\n\n\n## What can Autograd differentiate?\n\nThe main constraint is that any function that operates on a `Box` is marked as `primitive`, and has its gradient implemented.\nThis is taken care of for most functions in the Numpy library, and it's easy to write your own gradients.\n\nThe input can be a scalar, complex number, vector, tuple, a tuple of vectors, a tuple of tuples, etc.\n\nWhen using the `grad` function, the output must be a scalar, but the functions `elementwise_grad` and `jacobian` allow gradients of vectors.\n\n\n## Supported and unsupported parts of numpy/scipy\n\nNumpy has [a lot of features](http://docs.scipy.org/doc/numpy/reference/). We've done our best to support most of them. So far, we've implemented gradients for:\n* most of the [mathematical operations](../autograd/numpy/numpy_vjps.py)\n* most of the [array and matrix manipulation routines](../autograd/numpy/numpy_vjps.py)\n* some [linear algebra](../autograd/numpy/linalg.py) functions\n* most of the [fast fourier transform](../autograd/numpy/fft.py) routines\n* full support for complex numbers\n* [N-dimensional convolutions](../autograd/scipy/signal.py)\n* Some scipy routines, including [`scipy.stats.norm`](../autograd/scipy/stats/norm.py)\n\nSome things remain to be implemented. For example, we support indexing (`x = A[i, j, :]`) but not assignment (`A[i,j] = x`) in arrays that are being differentiated with respect to.\nAssignment is hard to support because it requires keeping copies of the overwritten data, and so even when you write code that looks like it's performing assignment, the system would have to be making copies behind the scenes, often defeating the purpose of in-place operations.\n\nSimilarly, we don't support the syntax `A.dot(B)`; use the equivalent `np.dot(A, B)` instead.\nThe reason we don't support the first way is that subclassing `ndarray` raises a host of issues.\nAs another consequence of not subclassing `ndarray`, some subclass checks can break, like `isinstance(x, np.ndarray)` can return `False`.\nHowever, those `isinstance` checks will work if you instead use Autograd's provided one, writing `from autograd.builtins import isinstance`.\n\nIn-place modification of arrays not being differentiated with respect to (for example, `A[i] = x` or `A += B`) won't raise an error, but be careful.\nIt's easy to accidentally change something without Autograd knowing about it.\nThis can be a problem because Autograd keeps references to variables used in the forward pass if they will be needed on the reverse pass.\nMaking copies would be too slow.\n\nLists and dicts can be used freely - like control flow, Autograd usually doesn't even need to know about them.\nThe exception is passing in a list to a primitive function, such as `autograd.numpy.sum`.\nThis requires special care, since the list contents need to be examined for boxes.\nWe do support passing lists to `autograd.numpy.array` and `autograd.numpy.concatenate`, but in other cases, you may need to explicitly construct an array using `autograd.numpy.array` before passing a list or tuple argument into a primitive.\nAn alternative is to use the `list`, `dict`, and `tuple` classes in `autograd.builtins`, which should work just like the Python builtins while also ensuring boxes don't get hidden inside those containers.\nRemember, these issues typically only come up when you're passing a `list` or `tuple` to a primitive function; when passing around lists or tuples in your own (non-primitive) functions, you can put boxed values inside lists, tuples, or dicts without having to worry about it.\n\n#### TL;DR: Do use\n* [Most](../autograd/numpy/numpy_vjps.py) of numpy's functions\n* [Most](../autograd/numpy/numpy_boxes.py) numpy.ndarray methods\n* [Some](../autograd/scipy/) scipy functions\n* Indexing and slicing of arrays `x = A[3, :, 2:4]`\n* Explicit array creation from lists `A = np.array([x, y])`\n\n#### Don't use\n* Assignment to arrays `A[0,0] = x`\n* Implicit casting of lists to arrays `A = np.sum([x, y])`, use `A = np.sum(np.array([x, y]))` instead.\n* `A.dot(B)` notation (use `np.dot(A, B)` instead)\n* In-place operations (such as `a += b`, use `a = a + b` instead)\n* Some isinstance checks, like `isinstance(x, np.ndarray)` or `isinstance(x, tuple)`, without first doing `from autograd.builtins import isinstance, tuple`.\n\nLuckily, it's easy to check gradients numerically if you're worried that something's wrong.\n\n## Extend Autograd by defining your own primitives\n\nWhat if Autograd doesn't support a function you need to take the gradient of?\nThis can happen if your code depends on external library calls or C code.\nIt can sometimes even be a good idea to provide the gradient of a pure Python function for speed or numerical stability.\n\nFor example, let's add the gradient of a numerically stable version of `log(sum(exp(x)))`.\nThis function is included in `scipy.special` and already supported, but let's make our own version.\n\nNext, we define our function using standard Python, using `@primitive` as a decorator:\n\n```python\nimport autograd.numpy as np\nfrom autograd.extend import primitive, defvjp\n\n@primitive\ndef logsumexp(x):\n \"\"\"Numerically stable log(sum(exp(x)))\"\"\"\n max_x = np.max(x)\n return max_x + np.log(np.sum(np.exp(x - max_x)))\n```\n\n`@primitive` tells Autograd not to look inside the function, but instead to treat it as a black box whose gradient can be specified later.\nFunctions with this decorator can contain anything that Python knows how to execute, including calls to other languages.\n\nNext, we write a function that specifies the gradient of the primitive `logsumexp`:\n\n```python\ndef logsumexp_vjp(ans, x):\n x_shape = x.shape\n return lambda g: np.full(x_shape, g) * np.exp(x - np.full(x_shape, ans))\n```\n\n`logsumexp_vjp` returns a vector-Jacobian product (VJP) operator, which is a function that right-multiplies its argument `g` by the Jacobian matrix of `logsumexp` (without explicitly forming the matrix's coefficients).\n`g` will be the gradient of the final objective with respect to `ans` (the output of `logsumexp`).\nThe calculation can depend on both the input (`x`) and the output (`ans`) of the original function.\nIf you want to be able to take higher-order derivatives, then the code inside the VJP function must be itself differentiable by Autograd, which usually just means you write it in terms of other primitives which themselves have VJPs (like Numpy functions).\n\nThe final step is to tell Autograd about `logsumexp`'s vector-Jacobian product function:\n```python\ndefvjp(logsumexp, logsumexp_vjp)\n```\n\nNow we can use `logsumexp` anywhere, including inside of a larger function that we want to differentiate:\n\n```python\nfrom autograd import grad\n\ndef example_func(y):\n z = y**2\n lse = logsumexp(z)\n return np.sum(lse)\n\ngrad_of_example = grad(example_func)\nprint \"Gradient: \", grad_of_example(np.array([1.5, 6.7, 1e-10])\n```\n\nThis example can be found as a Python script [here](../examples/define_gradient.py).\n\n## Complex numbers\n\nAutograd supports complex arrays and scalars using a convention described as follows.\nConsider a complex-to-complex function, `f`,\nexpressed in terms of real-to-real components, `u` and `v`:\n\n```python\ndef f(z):\n x, y = real(z), imag(z)\n return u(x, y) + v(x, y) * 1j\n```\n\nWe define `grad` of `f` as\n\n```python\ndef grad_f(z):\n x, y = real(z), imag(z)\n return grad(u, 0)(x, y) - i * grad(u, 1)(x, y)\n```\n\n(The second argument of `grad` specifies which argument we're differentiating with respect to.)\nSo we throw out v, the imaginary part of f, entirely.\n\nOur convention covers three important cases:\n * If `f` is holomorphic, we get the usual complex derivative\n (since `grad(u, 0) == grad(v, 1)` and `grad(u, 1) == - grad(v, 0)`).\n * If `f` is a real-valued loss function of a complex parameter, `x`,\n we get a result that we can use in a gradient-based optimizer,\n by taking steps in the direction of the complex conjugate of `grad(f)(x)`.\n * If `f` is a real-to-real function that happens to use complex primitives internally,\n some of which must necessarily be non-holomorphic\n (maybe you use FFTs to implement convolutions for example)\n then we get the same result that a purely real implementation would have given.\n\nOur convention doesn't handle the case where `f` is a non-holomorphic function\nand you're interested in all of du/dx, du/dy, dv/dx and dv/dy.\nBut then the answer would have to contain four real values\nand there would be no way to express it as a single complex number.\n\nWe define primitive vector-Jacobian products of complex functions like this\n\n```python\ndef f_vjp(g, z):\n z_x, z_y = real(z), imag(z)\n g_x, g_y = real(g), imag(g)\n return ( g_x * grad(u, 0)(x, y)\n - i * g_x * grad(u, 1)(x, y)\n - g_y * grad(v, 0)(x, y)\n + i * g_y * grad(v, 1)(x, y))\n```\n\nFor holomorphic primitives, this is just the regular complex derivative multiplied by `g`,\nso most simple math primitives don't need to be changed from their real implementations.\nFor non-holomorphic primitives, it preserves all four real partial derivatives as if we\nwere treating complex numbers as real 2-tuples\n(though it throws a couple of negative signs in there).\nChapter 4 of [Dougal's PhD thesis](https://dougalmaclaurin.com/phd-thesis.pdf)\ngoes into a bit more detail about how we define the primitive vector-Jacobian products.\n\n## Autograd Lecture\nFor more information on automatic differentiation, autograd's implementation, and advanced automatic differentiation techniques, see a [talk by Matt at the Deep Learning Summer School, Montreal 2017](https://videolectures.net/videos/deeplearning2017_johnson_automatic_differentiation/).\n\n## Support\n\nAutograd was written by\n[Dougal Maclaurin](https://dougalmaclaurin.com),\n[David Duvenaud](http://mlg.eng.cam.ac.uk/duvenaud/), and\n[Matthew Johnson](http://www.mit.edu/~mattjj/)\nand we're actively developing it. Please\nfeel free to submit any bugs or feature requests. We'd also love to hear about\nyour experiences with Autograd in general. Drop us an email!\n"
},
{
"path": "docs/updateguide.md",
"content": "# Autograd v1.2 update guide\n\nAutograd v1.2 changed the interface for defining custom vector-Jacobian\nproducts (VJPs). Luckily the change only affects users writing custom VJPs, and\nshould only require minor updates to the custom VJP code.\n\nThis guide is meant to explain why we made these changes (and others) in\nAutograd v1.2, and to summarize everything you need to know to update your\ncustom VJP code.\n\n- [Reasoning for the changes](#reasoning-for-the-changes)\n- [New defvjp interface](#new-defvjp-interface)\n- [Gradient checking](#gradient-checking)\n\n## Reasoning for the changes\n\nHere are some of the most important reasons for this update:\n1. To allow us to make Autograd faster and more memory efficient, we staged the\n VJP functions to allow more garbage collection and eliminated almost all of\n the vspace metadata checks.\n1. Forward-mode now comes built-in with `make_jvp`.\n1. There's now a clear extension API in `autograd.extend`, so you can write\n custom VJPs or wrap your own numerical libraries.\n1. Autograd is now backend-independent, making it easy to wrap other numerical\n libraries.\n1. Autograd's tracing functionality is now parameterized and easily reusable,\n and we added some new tracers for\n [computation graph visualization](https://github.com/hips/autograd/blob/master/examples/dot_graph.py)\n and\n [pure-Python constant folding](https://github.com/hips/autograd/blob/master/autograd/misc/tracers.py).\n1. More exhaustive, fast reverse- and forward-mode checking with `autograd.test_util.check_grads`.\n1. Expensive VJPs can share work across arguments using `defvjp_argnums`.\n1. These changes enabled some internal cleanups, and more features to come!\n\n## New defvjp interface\nFirst, here's an example of the old way to write custom primitives and VJPs:\n```python\nimport autograd.numpy as np\nfrom autograd import primitive\n\n@primitive\ndef func(x, y, z):\n assert z != 0\n return x * y**2\n\nfunc.defvjp(lambda g, ans, vs, gvs, x, y, z: g * y**2)\nfunc.defvjp(lambda g, ans, vs, gvs, x, y, z: 2 * g * x * y, argnum=1)\nfunc.defvjp_is_zero(argnums=[2])\n```\n\nHere's the new way to write custom VJPs for that same primitive:\n```python\nimport autograd.numpy as np\nfrom autograd.extend import primitive, defvjp # defvjp is now a function\n\n# primitives look the same as before\n@primitive\ndef func(x, y, z):\n assert z != 0\n return x * y**2\n\n# but we call defvjp differently\ndefvjp(func,\n lambda ans, x, y, z: lambda g: g * y**2,\n lambda ans, x, y, z: lambda g: 2 * g * x * y,\n None)\n```\n\nHere's a list of the `defvjp` changes illustrated in that example:\n1. `defvjp` is a function, rather than a method on the `primitive` class. (Actually, `primitive` is now just a function, and no longer a class.) As a result, `func.defvjp(...)` became `defvjp(func, ...)`.\n1. VJPs are staged, so that instead of writing `lambda g, ans, vs, gvs, *args: ...` we write `lambda ans, *args: lambda g: ...`. This change enables a lot of automatic garbage collection. In the above example, if we were differentiating only with respect to `x` argument of `func`, because the VJP for `func` with respect to argument index 0 doesn't need the values of `x` or `z` from the forward pass, those values aren't stored and can instead be immediately garbage-collected.\n1. There are no more `vs` and `gvs` arguments. These usually weren't used, and computing vspace metadata for every intermediate value proved to contribute significant overhead for some programs. Autograd now avoids computing vspace metadata unless necessary.\n1. `defvjp` lets you define VJPs with respect to multiple arguments at once, and the argnum(s) involved are often implicit.\n\nHere's another example, this time showing how to define VJPs with respect to\nspecific argnums, leaving the others undefined.\n```python\n# OLD way to leave some VJPs undefined\nfunc.defvjp(lambda g, ans, vs, gvs, x, y, z, w: ..., argnum=2)\nfunc.defvjp(lambda g, ans, vs, gvs, x, y, z, w: ..., argnum=3)\n\n# NEW way to leave some VJPs undefined\ndefvjp(func,\n lambda ans, x, y, z, w: lambda g: ...,\n lambda ans, x, y, z, w: lambda g: ...,\n argnums=[2, 3])\n```\n\n## Gradient checking\nHere's how to do gradient checking, whether on a composite function or on your\nprimitive with a custom VJP:\n\n```python\nfrom autograd.test_util import check_grads\n\n# check reverse-mode to second order\ncheck_grads(my_func, modes=['rev'], order=2)(*args_for_my_func)\n```\n"
},
{
"path": "examples/README.md",
"content": "# Autograd examples\n\n## Usage instructions\n\nSome of the examples require additional dependencies beyond Autograd and its\ncore dependencies. These are set up under the `examples` dependency group. To\ninstall them, navigate to the root directory of where you cloned Autograd and\nrun\n```sh\npip install --group examples\n```\nfrom the command line. Note that dependency groups are a recent feature so you\nmay need to upgrade `pip` with\n```sh\npip install --upgrade pip\n```\n\nHaving installed the additional dependencies, you may navigate to the `examples`\nsubdirectory and run any of the Python scripts. For example:\n```sh\npython3 tanh.py\n```\nSome of the examples print to the terminal and others open pop-up windows for\nplots.\n"
},
{
"path": "examples/__init__.py",
"content": ""
},
{
"path": "examples/bayesian_neural_net.py",
"content": "import matplotlib.pyplot as plt\nfrom black_box_svi import black_box_variational_inference\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd.misc.optimizers import adam\n\n\ndef make_nn_funs(layer_sizes, L2_reg, noise_variance, nonlinearity=np.tanh):\n \"\"\"These functions implement a standard multi-layer perceptron,\n vectorized over both training examples and weight samples.\"\"\"\n shapes = list(zip(layer_sizes[:-1], layer_sizes[1:]))\n num_weights = sum((m + 1) * n for m, n in shapes)\n\n def unpack_layers(weights):\n num_weight_sets = len(weights)\n for m, n in shapes:\n yield (\n weights[:, : m * n].reshape((num_weight_sets, m, n)),\n weights[:, m * n : m * n + n].reshape((num_weight_sets, 1, n)),\n )\n weights = weights[:, (m + 1) * n :]\n\n def predictions(weights, inputs):\n \"\"\"weights is shape (num_weight_samples x num_weights)\n inputs is shape (num_datapoints x D)\"\"\"\n inputs = np.expand_dims(inputs, 0)\n for W, b in unpack_layers(weights):\n outputs = np.einsum(\"mnd,mdo->mno\", inputs, W) + b\n inputs = nonlinearity(outputs)\n return outputs\n\n def logprob(weights, inputs, targets):\n log_prior = -L2_reg * np.sum(weights**2, axis=1)\n preds = predictions(weights, inputs)\n log_lik = -np.sum((preds - targets) ** 2, axis=1)[:, 0] / noise_variance\n return log_prior + log_lik\n\n return num_weights, predictions, logprob\n\n\ndef build_toy_dataset(n_data=40, noise_std=0.1):\n D = 1\n rs = npr.RandomState(0)\n inputs = np.concatenate([np.linspace(0, 2, num=n_data // 2), np.linspace(6, 8, num=n_data // 2)])\n targets = np.cos(inputs) + rs.randn(n_data) * noise_std\n inputs = (inputs - 4.0) / 4.0\n inputs = inputs.reshape((len(inputs), D))\n targets = targets.reshape((len(targets), D))\n return inputs, targets\n\n\nif __name__ == \"__main__\":\n # Specify inference problem by its unnormalized log-posterior.\n rbf = lambda x: np.exp(-(x**2))\n relu = lambda x: np.maximum(x, 0.0)\n num_weights, predictions, logprob = make_nn_funs(\n layer_sizes=[1, 20, 20, 1], L2_reg=0.1, noise_variance=0.01, nonlinearity=rbf\n )\n\n inputs, targets = build_toy_dataset()\n log_posterior = lambda weights, t: logprob(weights, inputs, targets)\n\n # Build variational objective.\n objective, gradient, unpack_params = black_box_variational_inference(\n log_posterior, num_weights, num_samples=20\n )\n\n # Set up figure.\n fig = plt.figure(figsize=(12, 8), facecolor=\"white\")\n ax = fig.add_subplot(111, frameon=False)\n plt.ion()\n plt.show(block=False)\n\n def callback(params, t, g):\n print(f\"Iteration {t} lower bound {-objective(params, t)}\")\n\n # Sample functions from posterior.\n rs = npr.RandomState(0)\n mean, log_std = unpack_params(params)\n # rs = npr.RandomState(0)\n sample_weights = rs.randn(10, num_weights) * np.exp(log_std) + mean\n plot_inputs = np.linspace(-8, 8, num=400)\n outputs = predictions(sample_weights, np.expand_dims(plot_inputs, 1))\n\n # Plot data and functions.\n plt.cla()\n ax.plot(inputs.ravel(), targets.ravel(), \"bx\")\n ax.plot(plot_inputs, outputs[:, :, 0].T)\n ax.set_ylim([-2, 3])\n plt.draw()\n plt.pause(1.0 / 60.0)\n\n # Initialize variational parameters\n rs = npr.RandomState(0)\n init_mean = rs.randn(num_weights)\n init_log_std = -5 * np.ones(num_weights)\n init_var_params = np.concatenate([init_mean, init_log_std])\n\n print(\"Optimizing variational parameters...\")\n variational_params = adam(gradient, init_var_params, step_size=0.1, num_iters=1000, callback=callback)\n"
},
{
"path": "examples/bayesian_optimization.py",
"content": "\"\"\"This Bayesian optimization demo using gradient-based optimization\nto find the next query point.\"\"\"\n\nimport matplotlib.pyplot as plt\nfrom gaussian_process import make_gp_funs, rbf_covariance\nfrom scipy.optimize import minimize\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import value_and_grad\nfrom autograd.scipy.stats import norm\n\n\ndef probability_of_improvement(mean, std, max_so_far):\n return norm.cdf(max_so_far, mean, std)\n\n\ndef expected_new_max(mean, std, max_so_far):\n return (\n max_so_far\n - (mean - max_so_far) * norm.cdf(mean, max_so_far, std)\n + std * norm.pdf(mean, max_so_far, std)\n )\n\n\ndef init_covariance_params(num_params):\n return np.zeros(num_params)\n\n\ndef defaultmax(x, default=-np.inf):\n if x.size == 0:\n return default\n return np.max(x)\n\n\ndef bayesian_optimize(func, domain_min, domain_max, num_iters=20, callback=None):\n D = len(domain_min)\n\n num_params, predict, log_marginal_likelihood = make_gp_funs(rbf_covariance, num_cov_params=D + 1)\n\n model_params = init_covariance_params(num_params)\n\n def optimize_gp_params(init_params, X, y):\n log_hyperprior = lambda params: np.sum(norm.logpdf(params, 0.0, 100.0))\n objective = lambda params: -log_marginal_likelihood(params, X, y) - log_hyperprior(params)\n return minimize(value_and_grad(objective), init_params, jac=True, method=\"CG\").x\n\n def choose_next_point(domain_min, domain_max, acquisition_function, num_tries=15, rs=npr.RandomState(0)):\n \"\"\"Uses gradient-based optimization to find next query point.\"\"\"\n init_points = rs.rand(num_tries, D) * (domain_max - domain_min) + domain_min\n\n grad_obj = value_and_grad(lambda x: -acquisition_function(x))\n\n def optimize_point(init_point):\n print(\".\", end=\"\")\n result = minimize(\n grad_obj,\n x0=init_point,\n jac=True,\n method=\"L-BFGS-B\",\n options={\"maxiter\": 10},\n bounds=list(zip(domain_min, domain_max)),\n )\n return result.x, acquisition_function(result.x)\n\n optimzed_points, optimized_values = list(zip(*list(map(optimize_point, init_points))))\n print()\n best_ix = np.argmax(optimized_values)\n return np.atleast_2d(optimzed_points[best_ix])\n\n # Start by evaluating once in the middle of the domain.\n X = np.zeros((0, D))\n y = np.zeros(0)\n X = np.concatenate((X, np.reshape((domain_max - domain_min) / 2.0, (D, 1))))\n y = np.concatenate((y, np.reshape(np.array(func(X)), (1,))))\n\n for i in range(num_iters):\n if i > 1:\n print(\"Optimizing model parameters...\")\n model_params = optimize_gp_params(model_params, X, y)\n\n print(\"Choosing where to look next\", end=\"\")\n\n def predict_func(xstar):\n mean, cov = predict(model_params, X, y, xstar)\n return mean, np.sqrt(np.diag(cov))\n\n def acquisition_function(xstar):\n xstar = np.atleast_2d(xstar) # To work around a bug in scipy.minimize\n mean, std = predict_func(xstar)\n return expected_new_max(mean, std, defaultmax(y))\n\n next_point = choose_next_point(domain_min, domain_max, acquisition_function)\n\n print(\"Evaluating expensive function...\")\n new_value = func(next_point)\n\n X = np.concatenate((X, next_point))\n y = np.concatenate((y, np.reshape(np.array(new_value), (1,))))\n\n if callback:\n callback(X, y, predict_func, acquisition_function, next_point, new_value)\n\n best_ix = np.argmax(y)\n return X[best_ix, :], y[best_ix]\n\n\nif __name__ == \"__main__\":\n\n def example_function(x):\n return np.sum(x * np.sin(10.0 * x) + x) - 1\n\n domain_min = np.array([0.0])\n domain_max = np.array([1.1])\n\n # Set up figure.\n fig = plt.figure(figsize=(12, 8), facecolor=\"white\")\n ax = fig.add_subplot(111, frameon=False)\n plt.show(block=False)\n\n def callback(X, y, predict_func, acquisition_function, next_point, new_value):\n plt.cla()\n\n # Show posterior marginals.\n plot_xs = np.reshape(np.linspace(domain_min, domain_max, 300), (300, 1))\n pred_mean, pred_std = predict_func(plot_xs)\n ax.plot(plot_xs, pred_mean, \"b\")\n ax.fill(\n np.concatenate([plot_xs, plot_xs[::-1]]),\n np.concatenate([pred_mean - 1.96 * pred_std, (pred_mean + 1.96 * pred_std)[::-1]]),\n alpha=0.15,\n fc=\"Blue\",\n ec=\"None\",\n )\n\n ax.plot(X, y, \"kx\")\n ax.plot(next_point, new_value, \"ro\")\n\n alphas = acquisition_function(plot_xs)\n ax.plot(plot_xs, alphas, \"r\")\n ax.set_ylim([-1.5, 1.5])\n ax.set_xticks([])\n ax.set_yticks([])\n plt.draw()\n plt.pause(1)\n\n best_x, best_y = bayesian_optimize(example_function, domain_min, domain_max, callback=callback)\n"
},
{
"path": "examples/black_box_svi.py",
"content": "import matplotlib.pyplot as plt\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nimport autograd.scipy.stats.multivariate_normal as mvn\nimport autograd.scipy.stats.norm as norm\nfrom autograd import grad\nfrom autograd.misc.optimizers import adam\n\n\ndef black_box_variational_inference(logprob, D, num_samples):\n \"\"\"Implements http://arxiv.org/abs/1401.0118, and uses the\n local reparameterization trick from http://arxiv.org/abs/1506.02557\"\"\"\n\n def unpack_params(params):\n # Variational dist is a diagonal Gaussian.\n mean, log_std = params[:D], params[D:]\n return mean, log_std\n\n def gaussian_entropy(log_std):\n return 0.5 * D * (1.0 + np.log(2 * np.pi)) + np.sum(log_std)\n\n rs = npr.RandomState(0)\n\n def variational_objective(params, t):\n \"\"\"Provides a stochastic estimate of the variational lower bound.\"\"\"\n mean, log_std = unpack_params(params)\n samples = rs.randn(num_samples, D) * np.exp(log_std) + mean\n lower_bound = gaussian_entropy(log_std) + np.mean(logprob(samples, t))\n return -lower_bound\n\n gradient = grad(variational_objective)\n\n return variational_objective, gradient, unpack_params\n\n\nif __name__ == \"__main__\":\n # Specify an inference problem by its unnormalized log-density.\n D = 2\n\n def log_density(x, t):\n mu, log_sigma = x[:, 0], x[:, 1]\n sigma_density = norm.logpdf(log_sigma, 0, 1.35)\n mu_density = norm.logpdf(mu, 0, np.exp(log_sigma))\n return sigma_density + mu_density\n\n # Build variational objective.\n objective, gradient, unpack_params = black_box_variational_inference(log_density, D, num_samples=2000)\n\n # Set up plotting code\n def plot_isocontours(ax, func, xlimits=[-2, 2], ylimits=[-4, 2], numticks=101):\n x = np.linspace(*xlimits, num=numticks)\n y = np.linspace(*ylimits, num=numticks)\n X, Y = np.meshgrid(x, y)\n zs = func(np.concatenate([np.atleast_2d(X.ravel()), np.atleast_2d(Y.ravel())]).T)\n Z = zs.reshape(X.shape)\n plt.contour(X, Y, Z)\n ax.set_yticks([])\n ax.set_xticks([])\n\n # Set up figure.\n fig = plt.figure(figsize=(8, 8), facecolor=\"white\")\n ax = fig.add_subplot(111, frameon=False)\n plt.ion()\n plt.show(block=False)\n\n def callback(params, t, g):\n print(f\"Iteration {t} lower bound {-objective(params, t)}\")\n\n plt.cla()\n target_distribution = lambda x: np.exp(log_density(x, t))\n plot_isocontours(ax, target_distribution)\n\n mean, log_std = unpack_params(params)\n variational_contour = lambda x: mvn.pdf(x, mean, np.diag(np.exp(2 * log_std)))\n plot_isocontours(ax, variational_contour)\n plt.draw()\n plt.pause(1.0 / 30.0)\n\n print(\"Optimizing variational parameters...\")\n init_mean = -1 * np.ones(D)\n init_log_std = -5 * np.ones(D)\n init_var_params = np.concatenate([init_mean, init_log_std])\n variational_params = adam(gradient, init_var_params, step_size=0.1, num_iters=2000, callback=callback)\n"
},
{
"path": "examples/convnet.py",
"content": "\"\"\"Convolutional neural net on MNIST, modeled on 'LeNet-5',\nhttp://yann.lecun.com/exdb/publis/pdf/lecun-98.pdf\"\"\"\n\nimport data_mnist\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nimport autograd.scipy.signal\nfrom autograd import grad\n\nconvolve = autograd.scipy.signal.convolve\n\n\nclass WeightsParser:\n \"\"\"A helper class to index into a parameter vector.\"\"\"\n\n def __init__(self):\n self.idxs_and_shapes = {}\n self.N = 0\n\n def add_weights(self, name, shape):\n start = self.N\n self.N += np.prod(shape)\n self.idxs_and_shapes[name] = (slice(start, self.N), shape)\n\n def get(self, vect, name):\n idxs, shape = self.idxs_and_shapes[name]\n return np.reshape(vect[idxs], shape)\n\n\ndef make_batches(N_total, N_batch):\n start = 0\n batches = []\n while start < N_total:\n batches.append(slice(start, start + N_batch))\n start += N_batch\n return batches\n\n\ndef logsumexp(X, axis, keepdims=False):\n max_X = np.max(X)\n return max_X + np.log(np.sum(np.exp(X - max_X), axis=axis, keepdims=keepdims))\n\n\ndef make_nn_funs(input_shape, layer_specs, L2_reg):\n parser = WeightsParser()\n cur_shape = input_shape\n for layer in layer_specs:\n N_weights, cur_shape = layer.build_weights_dict(cur_shape)\n parser.add_weights(layer, (N_weights,))\n\n def predictions(W_vect, inputs):\n \"\"\"Outputs normalized log-probabilities.\n shape of inputs : [data, color, y, x]\"\"\"\n cur_units = inputs\n for layer in layer_specs:\n cur_weights = parser.get(W_vect, layer)\n cur_units = layer.forward_pass(cur_units, cur_weights)\n return cur_units\n\n def loss(W_vect, X, T):\n log_prior = -L2_reg * np.dot(W_vect, W_vect)\n log_lik = np.sum(predictions(W_vect, X) * T)\n return -log_prior - log_lik\n\n def frac_err(W_vect, X, T):\n return np.mean(np.argmax(T, axis=1) != np.argmax(pred_fun(W_vect, X), axis=1))\n\n return parser.N, predictions, loss, frac_err\n\n\nclass conv_layer:\n def __init__(self, kernel_shape, num_filters):\n self.kernel_shape = kernel_shape\n self.num_filters = num_filters\n\n def forward_pass(self, inputs, param_vector):\n # Input dimensions: [data, color_in, y, x]\n # Params dimensions: [color_in, color_out, y, x]\n # Output dimensions: [data, color_out, y, x]\n params = self.parser.get(param_vector, \"params\")\n biases = self.parser.get(param_vector, \"biases\")\n conv = convolve(inputs, params, axes=([2, 3], [2, 3]), dot_axes=([1], [0]), mode=\"valid\")\n return conv + biases\n\n def build_weights_dict(self, input_shape):\n # Input shape : [color, y, x] (don't need to know number of data yet)\n self.parser = WeightsParser()\n self.parser.add_weights(\"params\", (input_shape[0], self.num_filters) + self.kernel_shape)\n self.parser.add_weights(\"biases\", (1, self.num_filters, 1, 1))\n output_shape = (self.num_filters,) + self.conv_output_shape(input_shape[1:], self.kernel_shape)\n return self.parser.N, output_shape\n\n def conv_output_shape(self, A, B):\n return (A[0] - B[0] + 1, A[1] - B[1] + 1)\n\n\nclass maxpool_layer:\n def __init__(self, pool_shape):\n self.pool_shape = pool_shape\n\n def build_weights_dict(self, input_shape):\n # input_shape dimensions: [color, y, x]\n output_shape = list(input_shape)\n for i in [0, 1]:\n assert input_shape[i + 1] % self.pool_shape[i] == 0, \"maxpool shape should tile input exactly\"\n output_shape[i + 1] = input_shape[i + 1] / self.pool_shape[i]\n return 0, output_shape\n\n def forward_pass(self, inputs, param_vector):\n new_shape = inputs.shape[:2]\n for i in [0, 1]:\n pool_width = self.pool_shape[i]\n img_width = inputs.shape[i + 2]\n new_shape += (img_width // pool_width, pool_width)\n result = inputs.reshape(new_shape)\n return np.max(np.max(result, axis=3), axis=4)\n\n\nclass full_layer:\n def __init__(self, size):\n self.size = size\n\n def build_weights_dict(self, input_shape):\n # Input shape is anything (all flattened)\n input_size = np.prod(input_shape, dtype=int)\n self.parser = WeightsParser()\n self.parser.add_weights(\"params\", (input_size, self.size))\n self.parser.add_weights(\"biases\", (self.size,))\n return self.parser.N, (self.size,)\n\n def forward_pass(self, inputs, param_vector):\n params = self.parser.get(param_vector, \"params\")\n biases = self.parser.get(param_vector, \"biases\")\n if inputs.ndim > 2:\n inputs = inputs.reshape((inputs.shape[0], np.prod(inputs.shape[1:])))\n return self.nonlinearity(np.dot(inputs[:, :], params) + biases)\n\n\nclass tanh_layer(full_layer):\n def nonlinearity(self, x):\n return np.tanh(x)\n\n\nclass softmax_layer(full_layer):\n def nonlinearity(self, x):\n return x - logsumexp(x, axis=1, keepdims=True)\n\n\nif __name__ == \"__main__\":\n # Network parameters\n L2_reg = 1.0\n input_shape = (1, 28, 28)\n layer_specs = [\n conv_layer((5, 5), 6),\n maxpool_layer((2, 2)),\n conv_layer((5, 5), 16),\n maxpool_layer((2, 2)),\n tanh_layer(120),\n tanh_layer(84),\n softmax_layer(10),\n ]\n\n # Training parameters\n param_scale = 0.1\n learning_rate = 1e-3\n momentum = 0.9\n batch_size = 256\n num_epochs = 50\n\n # Load and process MNIST data\n print(\"Loading training data...\")\n add_color_channel = lambda x: x.reshape((x.shape[0], 1, x.shape[1], x.shape[2]))\n one_hot = lambda x, K: np.array(x[:, None] == np.arange(K)[None, :], dtype=int)\n train_images, train_labels, test_images, test_labels = data_mnist.mnist()\n train_images = add_color_channel(train_images) / 255.0\n test_images = add_color_channel(test_images) / 255.0\n train_labels = one_hot(train_labels, 10)\n test_labels = one_hot(test_labels, 10)\n N_data = train_images.shape[0]\n\n # Make neural net functions\n N_weights, pred_fun, loss_fun, frac_err = make_nn_funs(input_shape, layer_specs, L2_reg)\n loss_grad = grad(loss_fun)\n\n # Initialize weights\n rs = npr.RandomState()\n W = rs.randn(N_weights) * param_scale\n\n # Check the gradients numerically, just to be safe\n # quick_grad_check(loss_fun, W, (train_images[:50], train_labels[:50]))\n\n print(\" Epoch | Train err | Test error \")\n\n def print_perf(epoch, W):\n test_perf = frac_err(W, test_images, test_labels)\n train_perf = frac_err(W, train_images, train_labels)\n print(f\"{epoch:15}|{train_perf:15}|{test_perf:15}\")\n\n # Train with sgd\n batch_idxs = make_batches(N_data, batch_size)\n cur_dir = np.zeros(N_weights)\n\n for epoch in range(num_epochs):\n print_perf(epoch, W)\n for idxs in batch_idxs:\n grad_W = loss_grad(W, train_images[idxs], train_labels[idxs])\n cur_dir = momentum * cur_dir + (1.0 - momentum) * grad_W\n W -= learning_rate * cur_dir\n"
},
{
"path": "examples/data.py",
"content": "import data_mnist\nimport matplotlib.image\nimport matplotlib.pyplot as plt\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\n\n\ndef load_mnist():\n partial_flatten = lambda x: np.reshape(x, (x.shape[0], np.prod(x.shape[1:])))\n one_hot = lambda x, k: np.array(x[:, None] == np.arange(k)[None, :], dtype=int)\n train_images, train_labels, test_images, test_labels = data_mnist.mnist()\n train_images = partial_flatten(train_images) / 255.0\n test_images = partial_flatten(test_images) / 255.0\n train_labels = one_hot(train_labels, 10)\n test_labels = one_hot(test_labels, 10)\n N_data = train_images.shape[0]\n\n return N_data, train_images, train_labels, test_images, test_labels\n\n\ndef plot_images(\n images,\n ax,\n ims_per_row=5,\n padding=5,\n digit_dimensions=(28, 28),\n cmap=matplotlib.cm.binary,\n vmin=None,\n vmax=None,\n):\n \"\"\"Images should be a (N_images x pixels) matrix.\"\"\"\n N_images = images.shape[0]\n N_rows = (N_images - 1) // ims_per_row + 1\n pad_value = np.min(images.ravel())\n concat_images = np.full(\n (\n (digit_dimensions[0] + padding) * N_rows + padding,\n (digit_dimensions[1] + padding) * ims_per_row + padding,\n ),\n pad_value,\n )\n for i in range(N_images):\n cur_image = np.reshape(images[i, :], digit_dimensions)\n row_ix = i // ims_per_row\n col_ix = i % ims_per_row\n row_start = padding + (padding + digit_dimensions[0]) * row_ix\n col_start = padding + (padding + digit_dimensions[1]) * col_ix\n concat_images[\n row_start : row_start + digit_dimensions[0], col_start : col_start + digit_dimensions[1]\n ] = cur_image\n cax = ax.matshow(concat_images, cmap=cmap, vmin=vmin, vmax=vmax)\n plt.xticks(np.array([]))\n plt.yticks(np.array([]))\n return cax\n\n\ndef save_images(images, filename, **kwargs):\n fig = plt.figure(1)\n fig.clf()\n ax = fig.add_subplot(111)\n plot_images(images, ax, **kwargs)\n fig.patch.set_visible(False)\n ax.patch.set_visible(False)\n plt.savefig(filename)\n\n\ndef make_pinwheel(radial_std, tangential_std, num_classes, num_per_class, rate, rs=npr.RandomState(0)):\n \"\"\"Based on code by Ryan P. Adams.\"\"\"\n rads = np.linspace(0, 2 * np.pi, num_classes, endpoint=False)\n\n features = rs.randn(num_classes * num_per_class, 2) * np.array([radial_std, tangential_std])\n features[:, 0] += 1\n labels = np.repeat(np.arange(num_classes), num_per_class)\n\n angles = rads[labels] + rate * np.exp(features[:, 0])\n rotations = np.stack([np.cos(angles), -np.sin(angles), np.sin(angles), np.cos(angles)])\n rotations = np.reshape(rotations.T, (-1, 2, 2))\n\n return np.einsum(\"ti,tij->tj\", features, rotations)\n"
},
{
"path": "examples/data_mnist.py",
"content": "import array\nimport gzip\nimport os\nimport struct\nfrom urllib.request import urlretrieve\n\nimport numpy as np\n\n\ndef download(url, filename):\n if not os.path.exists(\"data\"):\n os.makedirs(\"data\")\n out_file = os.path.join(\"data\", filename)\n if not os.path.isfile(out_file):\n urlretrieve(url, out_file)\n\n\ndef mnist():\n base_url = \"https://storage.googleapis.com/cvdf-datasets/mnist/\"\n\n def parse_labels(filename):\n with gzip.open(filename, \"rb\") as fh:\n magic, num_data = struct.unpack(\">II\", fh.read(8))\n return np.array(array.array(\"B\", fh.read()), dtype=np.uint8)\n\n def parse_images(filename):\n with gzip.open(filename, \"rb\") as fh:\n magic, num_data, rows, cols = struct.unpack(\">IIII\", fh.read(16))\n return np.array(array.array(\"B\", fh.read()), dtype=np.uint8).reshape(num_data, rows, cols)\n\n for filename in [\n \"train-images-idx3-ubyte.gz\",\n \"train-labels-idx1-ubyte.gz\",\n \"t10k-images-idx3-ubyte.gz\",\n \"t10k-labels-idx1-ubyte.gz\",\n ]:\n download(base_url + filename, filename)\n\n train_images = parse_images(\"data/train-images-idx3-ubyte.gz\")\n train_labels = parse_labels(\"data/train-labels-idx1-ubyte.gz\")\n test_images = parse_images(\"data/t10k-images-idx3-ubyte.gz\")\n test_labels = parse_labels(\"data/t10k-labels-idx1-ubyte.gz\")\n\n return train_images, train_labels, test_images, test_labels\n"
},
{
"path": "examples/deep_gaussian_process.py",
"content": "import matplotlib.pyplot as plt\nfrom gaussian_process import make_gp_funs, rbf_covariance\nfrom scipy.optimize import minimize\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import value_and_grad\n\n\ndef build_step_function_dataset(D=1, n_data=40, noise_std=0.1):\n rs = npr.RandomState(0)\n inputs = np.linspace(-2, 2, num=n_data)\n targets = np.sign(inputs) + rs.randn(n_data) * noise_std\n inputs = inputs.reshape((len(inputs), D))\n return inputs, targets\n\n\ndef build_deep_gp(input_dimension, hidden_dimension, covariance_function):\n # GP going from input to hidden\n num_params_layer1, predict_layer1, log_marginal_likelihood_layer1 = make_gp_funs(\n covariance_function, num_cov_params=input_dimension + 1\n )\n\n # GP going from hidden to output\n num_params_layer2, predict_layer2, log_marginal_likelihood_layer2 = make_gp_funs(\n covariance_function, num_cov_params=hidden_dimension + 1\n )\n\n num_hidden_params = hidden_dimension * n_data\n total_num_params = num_params_layer1 + num_params_layer2 + num_hidden_params\n\n def unpack_all_params(all_params):\n layer1_params = all_params[:num_params_layer1]\n layer2_params = all_params[num_params_layer1 : num_params_layer1 + num_params_layer2]\n hiddens = all_params[num_params_layer1 + num_params_layer2 :]\n return layer1_params, layer2_params, hiddens\n\n def combined_predict_fun(all_params, X, y, xs):\n layer1_params, layer2_params, hiddens = unpack_all_params(all_params)\n h_star_mean, h_star_cov = predict_layer1(layer1_params, X, hiddens, xs)\n y_star_mean, y_star_cov = predict_layer2(\n layer2_params, np.atleast_2d(hiddens).T, y, np.atleast_2d(h_star_mean).T\n )\n return y_star_mean, y_star_cov\n\n def log_marginal_likelihood(all_params):\n layer1_params, layer2_params, h = unpack_all_params(all_params)\n return log_marginal_likelihood_layer1(layer1_params, X, h) + log_marginal_likelihood_layer2(\n layer2_params, np.atleast_2d(h).T, y\n )\n\n predict_layer_funcs = [predict_layer1, predict_layer2]\n\n return (\n total_num_params,\n log_marginal_likelihood,\n combined_predict_fun,\n unpack_all_params,\n predict_layer_funcs,\n )\n\n\nif __name__ == \"__main__\":\n n_data = 20\n input_dimension = 1\n hidden_dimension = 1\n X, y = build_step_function_dataset(D=input_dimension, n_data=n_data)\n\n (\n total_num_params,\n log_marginal_likelihood,\n combined_predict_fun,\n unpack_all_params,\n predict_layer_funcs,\n ) = build_deep_gp(input_dimension, hidden_dimension, rbf_covariance)\n\n # Set up figure.\n fig = plt.figure(figsize=(12, 8), facecolor=\"white\")\n ax_end_to_end = fig.add_subplot(311, frameon=False)\n ax_x_to_h = fig.add_subplot(312, frameon=False)\n ax_h_to_y = fig.add_subplot(313, frameon=False)\n plt.show(block=False)\n\n def plot_gp(ax, X, y, pred_mean, pred_cov, plot_xs):\n ax.cla()\n marg_std = np.sqrt(np.diag(pred_cov))\n ax.plot(plot_xs, pred_mean, \"b\")\n ax.fill(\n np.concatenate([plot_xs, plot_xs[::-1]]),\n np.concatenate([pred_mean - 1.96 * marg_std, (pred_mean + 1.96 * marg_std)[::-1]]),\n alpha=0.15,\n fc=\"Blue\",\n ec=\"None\",\n )\n\n # Show samples from posterior.\n rs = npr.RandomState(0)\n sampled_funcs = rs.multivariate_normal(pred_mean, pred_cov, size=10)\n ax.plot(plot_xs, sampled_funcs.T)\n ax.plot(X, y, \"kx\")\n ax.set_ylim([-1.5, 1.5])\n ax.set_xticks([])\n ax.set_yticks([])\n\n def callback(params):\n print(f\"Log marginal likelihood {log_marginal_likelihood(params)}\")\n\n # Show posterior marginals.\n plot_xs = np.reshape(np.linspace(-5, 5, 300), (300, 1))\n pred_mean, pred_cov = combined_predict_fun(params, X, y, plot_xs)\n plot_gp(ax_end_to_end, X, y, pred_mean, pred_cov, plot_xs)\n ax_end_to_end.set_title(\"X to y\")\n\n layer1_params, layer2_params, hiddens = unpack_all_params(params)\n h_star_mean, h_star_cov = predict_layer_funcs[0](layer1_params, X, hiddens, plot_xs)\n y_star_mean, y_star_cov = predict_layer_funcs[0](layer2_params, np.atleast_2d(hiddens).T, y, plot_xs)\n\n plot_gp(ax_x_to_h, X, hiddens, h_star_mean, h_star_cov, plot_xs)\n ax_x_to_h.set_title(\"X to hiddens\")\n\n plot_gp(ax_h_to_y, np.atleast_2d(hiddens).T, y, y_star_mean, y_star_cov, plot_xs)\n ax_h_to_y.set_title(\"hiddens to y\")\n\n plt.draw()\n plt.pause(1.0 / 60.0)\n\n # Initialize covariance parameters and hiddens.\n rs = npr.RandomState(0)\n init_params = 0.1 * rs.randn(total_num_params)\n\n print(\"Optimizing covariance parameters...\")\n objective = lambda params: -log_marginal_likelihood(params)\n cov_params = minimize(value_and_grad(objective), init_params, jac=True, method=\"CG\", callback=callback)\n plt.pause(10.0)\n"
},
{
"path": "examples/define_gradient.py",
"content": "\"\"\"This example shows how to define the gradient of your own functions.\nThis can be useful for speed, numerical stability, or in cases where\nyour code depends on external library calls.\"\"\"\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.extend import defvjp, primitive\nfrom autograd.test_util import check_grads\n\n\n# @primitive tells Autograd not to look inside this function, but instead\n# to treat it as a black box, whose gradient might be specified later.\n# Functions with this decorator can contain anything that Python knows\n# how to execute, and you can do things like in-place operations on arrays.\n@primitive\ndef logsumexp(x):\n \"\"\"Numerically stable log(sum(exp(x))), also defined in scipy.special\"\"\"\n max_x = np.max(x)\n return max_x + np.log(np.sum(np.exp(x - max_x)))\n\n\n# Next, we write a function that specifies the gradient with a closure.\n# The reason for the closure is so that the gradient can depend\n# on both the input to the original function (x), and the output of the\n# original function (ans).\n\n\ndef logsumexp_vjp(ans, x):\n # If you want to be able to take higher-order derivatives, then all the\n # code inside this function must be itself differentiable by Autograd.\n # This closure multiplies g with the Jacobian of logsumexp (d_ans/d_x).\n # Because Autograd uses reverse-mode differentiation, g contains\n # the gradient of the objective w.r.t. ans, the output of logsumexp.\n # This returned VJP function doesn't close over `x`, so Python can\n # garbage-collect `x` if there are no references to it elsewhere.\n x_shape = x.shape\n return lambda g: np.full(x_shape, g) * np.exp(x - np.full(x_shape, ans))\n\n\n# Now we tell Autograd that logsumexmp has a gradient-making function.\ndefvjp(logsumexp, logsumexp_vjp)\n\nif __name__ == \"__main__\":\n # Now we can use logsumexp() inside a larger function that we want\n # to differentiate.\n def example_func(y):\n z = y**2\n lse = logsumexp(z)\n return np.sum(lse)\n\n grad_of_example = grad(example_func)\n print(\"Gradient: \\n\", grad_of_example(npr.randn(10)))\n\n # Check the gradients numerically, just to be safe.\n check_grads(example_func, modes=[\"rev\"])(npr.randn(10))\n"
},
{
"path": "examples/dot_graph.py",
"content": "\"\"\"Generates a graphviz DOT file of an evaluation trace.\nUsage (need the dot binary, from the graphviz package, www.graphviz.org):\n\npython2 dot_graph.py | dot -Tpdf -o graph.pdf\n\"\"\"\n\nimport autograd.numpy as np\nfrom autograd.tracer import Node, trace\n\n\nclass GraphNode(Node):\n # Records the full graph (could having this in tracer.py)\n def __init__(self, value, fun, args, kwargs, parent_argnums, parents):\n self.fun_name = fun.__name__\n self.args = args\n self.parents = dict(zip(parent_argnums, parents))\n self.isroot = False\n\n def initialize_root(self, x):\n self.isroot = True\n\n def __repr__(self):\n return f\"node_{id(self)}\"\n\n\ndef trace_graph(f, x):\n start_node = GraphNode.new_root(x)\n _, node = trace(start_node, f, x)\n return node\n\n\ndot_edge = \"{} -> {} [color=gray30];\\n\".format\ndot_function_node = '{} [label=\"{}\", shape=box, color=lightblue, style=filled];\\n'.format\ndot_variable_node = '{} [label=\"{}\", color=orange, style=filled];\\n'.format\ndot_graph = \"digraph G {{{}}}\".format\n\n\ndef graph_to_dotfile(graph):\n visited = set()\n\n def node_to_fragment(node):\n visited.add(node)\n if node.isroot:\n return dot_variable_node(node, \"input\")\n fragment = dot_function_node(node, node.fun_name)\n for argnum, arg in enumerate(node.args):\n if argnum in node.parents:\n parent = node.parents[argnum]\n fragment += dot_edge(parent, node)\n if parent not in visited:\n fragment += node_to_fragment(parent)\n else:\n argnode = f\"{node}_arg_{argnum}\"\n fragment += dot_edge(argnode, node)\n fragment += dot_variable_node(argnode, arg)\n\n return fragment\n\n dot_body = node_to_fragment(graph)\n dot_body += dot_variable_node(\"output\", \"output\")\n dot_body += dot_edge(graph, \"output\")\n return dot_graph(dot_body)\n\n\nif __name__ == \"__main__\":\n\n def fun(x):\n y = np.sin(x)\n return (y + np.exp(x) - 0.5) * y\n\n print(graph_to_dotfile(trace_graph(fun, 1.0)))\n"
},
{
"path": "examples/fixed_points.py",
"content": "import autograd.numpy as np\nfrom autograd import grad\nfrom autograd.misc.fixed_points import fixed_point\n\n\ndef newton_sqrt_iter(a):\n return lambda x: 0.5 * (x + a / x)\n\n\ndef grad_descent_sqrt_iter(a):\n return lambda x: x - 0.05 * (x**2 - a)\n\n\ndef sqrt(a, guess=10.0):\n # return fixed_point(newton_sqrt_iter, a, guess, distance, 1e-4)\n return fixed_point(grad_descent_sqrt_iter, a, guess, distance, 1e-4)\n\n\ndef distance(x, y):\n return np.abs(x - y)\n\n\nprint(np.sqrt(2.0))\nprint(sqrt(2.0))\nprint()\nprint(grad(np.sqrt)(2.0))\nprint(grad(sqrt)(2.0))\nprint()\nprint(grad(grad(np.sqrt))(2.0))\nprint(grad(grad(sqrt))(2.0))\nprint()\n"
},
{
"path": "examples/fluidsim/fluidsim.py",
"content": "import os\n\nimport matplotlib\nimport matplotlib.pyplot as plt\nfrom matplotlib.pyplot import imread\nfrom scipy.optimize import minimize\n\nimport autograd.numpy as np\nfrom autograd import value_and_grad\n\n# Fluid simulation code based on\n# \"Real-Time Fluid Dynamics for Games\" by Jos Stam\n# https://www.josstam.com/_files/ugd/cf1fd6_9989229efbd34a26ba5ccd913721a2ac.pdf\n\n\ndef project(vx, vy):\n \"\"\"Project the velocity field to be approximately mass-conserving,\n using a few iterations of Gauss-Seidel.\"\"\"\n p = np.zeros(vx.shape)\n h = 1.0 / vx.shape[0]\n div = (\n -0.5\n * h\n * (\n np.roll(vx, -1, axis=0)\n - np.roll(vx, 1, axis=0)\n + np.roll(vy, -1, axis=1)\n - np.roll(vy, 1, axis=1)\n )\n )\n\n for k in range(10):\n p = (\n div\n + np.roll(p, 1, axis=0)\n + np.roll(p, -1, axis=0)\n + np.roll(p, 1, axis=1)\n + np.roll(p, -1, axis=1)\n ) / 4.0\n\n vx -= 0.5 * (np.roll(p, -1, axis=0) - np.roll(p, 1, axis=0)) / h\n vy -= 0.5 * (np.roll(p, -1, axis=1) - np.roll(p, 1, axis=1)) / h\n return vx, vy\n\n\ndef advect(f, vx, vy):\n \"\"\"Move field f according to x and y velocities (u and v)\n using an implicit Euler integrator.\"\"\"\n rows, cols = f.shape\n cell_ys, cell_xs = np.meshgrid(np.arange(rows), np.arange(cols))\n center_xs = (cell_xs - vx).ravel()\n center_ys = (cell_ys - vy).ravel()\n\n # Compute indices of source cells.\n left_ix = np.floor(center_xs).astype(int)\n top_ix = np.floor(center_ys).astype(int)\n rw = center_xs - left_ix # Relative weight of right-hand cells.\n bw = center_ys - top_ix # Relative weight of bottom cells.\n left_ix = np.mod(left_ix, rows) # Wrap around edges of simulation.\n right_ix = np.mod(left_ix + 1, rows)\n top_ix = np.mod(top_ix, cols)\n bot_ix = np.mod(top_ix + 1, cols)\n\n # A linearly-weighted sum of the 4 surrounding cells.\n flat_f = (1 - rw) * ((1 - bw) * f[left_ix, top_ix] + bw * f[left_ix, bot_ix]) + rw * (\n (1 - bw) * f[right_ix, top_ix] + bw * f[right_ix, bot_ix]\n )\n return np.reshape(flat_f, (rows, cols))\n\n\ndef simulate(vx, vy, smoke, num_time_steps, ax=None, render=False):\n print(\"Running simulation...\")\n for t in range(num_time_steps):\n if ax:\n plot_matrix(ax, smoke, t, render)\n vx_updated = advect(vx, vx, vy)\n vy_updated = advect(vy, vx, vy)\n vx, vy = project(vx_updated, vy_updated)\n smoke = advect(smoke, vx, vy)\n if ax:\n plot_matrix(ax, smoke, num_time_steps, render)\n return smoke\n\n\ndef plot_matrix(ax, mat, t, render=False):\n plt.cla()\n ax.matshow(mat)\n ax.set_xticks([])\n ax.set_yticks([])\n plt.draw()\n if render:\n matplotlib.image.imsave(f\"step{t:03d}.png\", mat)\n plt.pause(0.001)\n\n\nif __name__ == \"__main__\":\n simulation_timesteps = 100\n basepath = os.path.dirname(__file__)\n\n print(\"Loading initial and target states...\")\n init_smoke = imread(os.path.join(basepath, \"init_smoke.png\"))[:, :, 0]\n # target = imread('peace.png')[::2,::2,3]\n target = imread(os.path.join(basepath, \"skull.png\"))[::2, ::2]\n rows, cols = target.shape\n\n init_dx_and_dy = np.zeros((2, rows, cols)).ravel()\n\n def distance_from_target_image(smoke):\n return np.mean((target - smoke) ** 2)\n\n def convert_param_vector_to_matrices(params):\n vx = np.reshape(params[: (rows * cols)], (rows, cols))\n vy = np.reshape(params[(rows * cols) :], (rows, cols))\n return vx, vy\n\n def objective(params):\n init_vx, init_vy = convert_param_vector_to_matrices(params)\n final_smoke = simulate(init_vx, init_vy, init_smoke, simulation_timesteps)\n return distance_from_target_image(final_smoke)\n\n # Specify gradient of objective function using autograd.\n objective_with_grad = value_and_grad(objective)\n\n fig = plt.figure(figsize=(8, 8))\n ax = fig.add_subplot(111, frameon=False)\n\n def callback(params):\n init_vx, init_vy = convert_param_vector_to_matrices(params)\n simulate(init_vx, init_vy, init_smoke, simulation_timesteps, ax)\n\n print(\"Optimizing initial conditions...\")\n result = minimize(\n objective_with_grad,\n init_dx_and_dy,\n jac=True,\n method=\"CG\",\n options={\"maxiter\": 25, \"disp\": True},\n callback=callback,\n )\n\n print(\"Rendering optimized flow...\")\n init_vx, init_vy = convert_param_vector_to_matrices(result.x)\n simulate(init_vx, init_vy, init_smoke, simulation_timesteps, ax, render=True)\n\n print(\"Converting frames to an animated GIF...\")\n os.system(\"convert -delay 5 -loop 0 step*.png -delay 250 step100.png surprise.gif\") # Using imagemagick.\n os.system(\"rm step*.png\")\n"
},
{
"path": "examples/fluidsim/wing.py",
"content": "import os\n\nimport matplotlib.pyplot as plt\nfrom scipy.optimize import minimize\n\nimport autograd.numpy as np\nfrom autograd import value_and_grad\n\nrows, cols = 40, 60\n\n# Fluid simulation code based on\n# \"Real-Time Fluid Dynamics for Games\" by Jos Stam\n# http://www.intpowertechcorp.com/GDC03.pdf\n\n\ndef occlude(f, occlusion):\n return f * (1 - occlusion)\n\n\ndef project(vx, vy, occlusion):\n \"\"\"Project the velocity field to be approximately mass-conserving,\n using a few iterations of Gauss-Seidel.\"\"\"\n p = np.zeros(vx.shape)\n div = -0.5 * (\n np.roll(vx, -1, axis=1) - np.roll(vx, 1, axis=1) + np.roll(vy, -1, axis=0) - np.roll(vy, 1, axis=0)\n )\n div = make_continuous(div, occlusion)\n\n for k in range(50):\n p = (\n div\n + np.roll(p, 1, axis=1)\n + np.roll(p, -1, axis=1)\n + np.roll(p, 1, axis=0)\n + np.roll(p, -1, axis=0)\n ) / 4.0\n p = make_continuous(p, occlusion)\n\n vx = vx - 0.5 * (np.roll(p, -1, axis=1) - np.roll(p, 1, axis=1))\n vy = vy - 0.5 * (np.roll(p, -1, axis=0) - np.roll(p, 1, axis=0))\n\n vx = occlude(vx, occlusion)\n vy = occlude(vy, occlusion)\n return vx, vy\n\n\ndef advect(f, vx, vy):\n \"\"\"Move field f according to x and y velocities (u and v)\n using an implicit Euler integrator.\"\"\"\n rows, cols = f.shape\n cell_xs, cell_ys = np.meshgrid(np.arange(cols), np.arange(rows))\n center_xs = (cell_xs - vx).ravel()\n center_ys = (cell_ys - vy).ravel()\n\n # Compute indices of source cells.\n left_ix = np.floor(center_ys).astype(int)\n top_ix = np.floor(center_xs).astype(int)\n rw = center_ys - left_ix # Relative weight of right-hand cells.\n bw = center_xs - top_ix # Relative weight of bottom cells.\n left_ix = np.mod(left_ix, rows) # Wrap around edges of simulation.\n right_ix = np.mod(left_ix + 1, rows)\n top_ix = np.mod(top_ix, cols)\n bot_ix = np.mod(top_ix + 1, cols)\n\n # A linearly-weighted sum of the 4 surrounding cells.\n flat_f = (1 - rw) * ((1 - bw) * f[left_ix, top_ix] + bw * f[left_ix, bot_ix]) + rw * (\n (1 - bw) * f[right_ix, top_ix] + bw * f[right_ix, bot_ix]\n )\n return np.reshape(flat_f, (rows, cols))\n\n\ndef make_continuous(f, occlusion):\n non_occluded = 1 - occlusion\n num = (\n np.roll(f, 1, axis=0) * np.roll(non_occluded, 1, axis=0)\n + np.roll(f, -1, axis=0) * np.roll(non_occluded, -1, axis=0)\n + np.roll(f, 1, axis=1) * np.roll(non_occluded, 1, axis=1)\n + np.roll(f, -1, axis=1) * np.roll(non_occluded, -1, axis=1)\n )\n den = (\n np.roll(non_occluded, 1, axis=0)\n + np.roll(non_occluded, -1, axis=0)\n + np.roll(non_occluded, 1, axis=1)\n + np.roll(non_occluded, -1, axis=1)\n )\n return f * non_occluded + (1 - non_occluded) * num / (den + 0.001)\n\n\ndef sigmoid(x):\n return 0.5 * (np.tanh(x) + 1.0) # Output ranges from 0 to 1.\n\n\ndef simulate(vx, vy, num_time_steps, occlusion, ax=None, render=False):\n occlusion = sigmoid(occlusion)\n\n # Disallow occlusion outside a certain area.\n mask = np.zeros((rows, cols))\n mask[10:30, 10:30] = 1.0\n occlusion = occlusion * mask\n\n # Initialize smoke bands.\n red_smoke = np.zeros((rows, cols))\n red_smoke[rows // 4 : rows // 2] = 1\n blue_smoke = np.zeros((rows, cols))\n blue_smoke[rows // 2 : 3 * rows // 4] = 1\n\n print(\"Running simulation...\")\n vx, vy = project(vx, vy, occlusion)\n for t in range(num_time_steps):\n plot_matrix(ax, red_smoke, occlusion, blue_smoke, t, render)\n vx_updated = advect(vx, vx, vy)\n vy_updated = advect(vy, vx, vy)\n vx, vy = project(vx_updated, vy_updated, occlusion)\n red_smoke = advect(red_smoke, vx, vy)\n red_smoke = occlude(red_smoke, occlusion)\n blue_smoke = advect(blue_smoke, vx, vy)\n blue_smoke = occlude(blue_smoke, occlusion)\n plot_matrix(ax, red_smoke, occlusion, blue_smoke, num_time_steps, render)\n return vx, vy\n\n\ndef plot_matrix(ax, r, g, b, t, render=False):\n if ax:\n plt.cla()\n ax.imshow(np.concatenate((r[..., np.newaxis], g[..., np.newaxis], b[..., np.newaxis]), axis=2))\n ax.set_xticks([])\n ax.set_yticks([])\n plt.draw()\n if render:\n plt.savefig(f\"step{t:03d}.png\", bbox_inches=\"tight\")\n plt.pause(0.001)\n\n\nif __name__ == \"__main__\":\n simulation_timesteps = 20\n\n print(\"Loading initial and target states...\")\n init_vx = np.ones((rows, cols))\n init_vy = np.zeros((rows, cols))\n\n # Initialize the occlusion to be a block.\n init_occlusion = -np.ones((rows, cols))\n init_occlusion[15:25, 15:25] = 0.0\n init_occlusion = init_occlusion.ravel()\n\n def drag(vx):\n return np.mean(init_vx - vx)\n\n def lift(vy):\n return np.mean(vy - init_vy)\n\n def objective(params):\n cur_occlusion = np.reshape(params, (rows, cols))\n final_vx, final_vy = simulate(init_vx, init_vy, simulation_timesteps, cur_occlusion)\n return -lift(final_vy) / drag(final_vx)\n\n # Specify gradient of objective function using autograd.\n objective_with_grad = value_and_grad(objective)\n\n fig = plt.figure(figsize=(8, 8))\n ax = fig.add_subplot(111, frameon=False)\n\n def callback(weights):\n cur_occlusion = np.reshape(weights, (rows, cols))\n simulate(init_vx, init_vy, simulation_timesteps, cur_occlusion, ax)\n\n print(\"Rendering initial flow...\")\n callback(init_occlusion)\n\n print(\"Optimizing initial conditions...\")\n result = minimize(\n objective_with_grad,\n init_occlusion,\n jac=True,\n method=\"CG\",\n options={\"maxiter\": 50, \"disp\": True},\n callback=callback,\n )\n\n print(\"Rendering optimized flow...\")\n final_occlusion = np.reshape(result.x, (rows, cols))\n simulate(init_vx, init_vy, simulation_timesteps, final_occlusion, ax, render=True)\n\n print(\"Converting frames to an animated GIF...\") # Using imagemagick.\n os.system(f\"convert -delay 5 -loop 0 step*.png -delay 250 step{simulation_timesteps:03d}.png wing.gif\")\n os.system(\"rm step*.png\")\n"
},
{
"path": "examples/gaussian_process.py",
"content": "import matplotlib.pyplot as plt\nfrom scipy.optimize import minimize\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nimport autograd.scipy.stats.multivariate_normal as mvn\nfrom autograd import value_and_grad\nfrom autograd.numpy.linalg import solve\n\n\ndef make_gp_funs(cov_func, num_cov_params):\n \"\"\"Functions that perform Gaussian process regression.\n cov_func has signature (cov_params, x, x')\"\"\"\n\n def unpack_kernel_params(params):\n mean = params[0]\n cov_params = params[2:]\n noise_scale = np.exp(params[1]) + 0.0001\n return mean, cov_params, noise_scale\n\n def predict(params, x, y, xstar):\n \"\"\"Returns the predictive mean and covariance at locations xstar,\n of the latent function value f (without observation noise).\"\"\"\n mean, cov_params, noise_scale = unpack_kernel_params(params)\n cov_f_f = cov_func(cov_params, xstar, xstar)\n cov_y_f = cov_func(cov_params, x, xstar)\n cov_y_y = cov_func(cov_params, x, x) + noise_scale * np.eye(len(y))\n pred_mean = mean + np.dot(solve(cov_y_y, cov_y_f).T, y - mean)\n pred_cov = cov_f_f - np.dot(solve(cov_y_y, cov_y_f).T, cov_y_f)\n return pred_mean, pred_cov\n\n def log_marginal_likelihood(params, x, y):\n mean, cov_params, noise_scale = unpack_kernel_params(params)\n cov_y_y = cov_func(cov_params, x, x) + noise_scale * np.eye(len(y))\n prior_mean = mean * np.ones(len(y))\n return mvn.logpdf(y, prior_mean, cov_y_y)\n\n return num_cov_params + 2, predict, log_marginal_likelihood\n\n\n# Define an example covariance function.\ndef rbf_covariance(kernel_params, x, xp):\n output_scale = np.exp(kernel_params[0])\n lengthscales = np.exp(kernel_params[1:])\n diffs = np.expand_dims(x / lengthscales, 1) - np.expand_dims(xp / lengthscales, 0)\n return output_scale * np.exp(-0.5 * np.sum(diffs**2, axis=2))\n\n\ndef build_toy_dataset(D=1, n_data=20, noise_std=0.1):\n rs = npr.RandomState(0)\n inputs = np.concatenate([np.linspace(0, 3, num=n_data // 2), np.linspace(6, 8, num=n_data // 2)])\n targets = (np.cos(inputs) + rs.randn(n_data) * noise_std) / 2.0\n inputs = (inputs - 4.0) / 2.0\n inputs = inputs.reshape((len(inputs), D))\n return inputs, targets\n\n\nif __name__ == \"__main__\":\n D = 1\n\n # Build model and objective function.\n num_params, predict, log_marginal_likelihood = make_gp_funs(rbf_covariance, num_cov_params=D + 1)\n\n X, y = build_toy_dataset(D=D)\n objective = lambda params: -log_marginal_likelihood(params, X, y)\n\n # Set up figure.\n fig = plt.figure(figsize=(12, 8), facecolor=\"white\")\n ax = fig.add_subplot(111, frameon=False)\n plt.show(block=False)\n\n def callback(params):\n print(f\"Log likelihood {-objective(params)}\")\n plt.cla()\n\n # Show posterior marginals.\n plot_xs = np.reshape(np.linspace(-7, 7, 300), (300, 1))\n pred_mean, pred_cov = predict(params, X, y, plot_xs)\n marg_std = np.sqrt(np.diag(pred_cov))\n ax.plot(plot_xs, pred_mean, \"b\")\n ax.fill(\n np.concatenate([plot_xs, plot_xs[::-1]]),\n np.concatenate([pred_mean - 1.96 * marg_std, (pred_mean + 1.96 * marg_std)[::-1]]),\n alpha=0.15,\n fc=\"Blue\",\n ec=\"None\",\n )\n\n # Show samples from posterior.\n rs = npr.RandomState(0)\n sampled_funcs = rs.multivariate_normal(pred_mean, pred_cov, size=10)\n ax.plot(plot_xs, sampled_funcs.T)\n\n ax.plot(X, y, \"kx\")\n ax.set_ylim([-1.5, 1.5])\n ax.set_xticks([])\n ax.set_yticks([])\n plt.draw()\n plt.pause(1.0 / 60.0)\n\n # Initialize covariance parameters\n rs = npr.RandomState(0)\n init_params = 0.1 * rs.randn(num_params)\n\n print(\"Optimizing covariance parameters...\")\n cov_params = minimize(value_and_grad(objective), init_params, jac=True, method=\"CG\", callback=callback)\n plt.pause(10.0)\n"
},
{
"path": "examples/generative_adversarial_net.py",
"content": "# Implements a Generative Adversarial Network, from\n# arxiv.org/abs/1406.2661\n# but, it always collapses to generating a single image.\n# Let me know if you can get it to work! - David Duvenaud\n\nfrom data import load_mnist, save_images\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.misc import flatten\n\n### Define geneerator, discriminator, and objective ###\n\n\ndef relu(x):\n return np.maximum(0, x)\n\n\ndef sigmoid(x):\n return 0.5 * (np.tanh(x) + 1.0)\n\n\ndef logsigmoid(x):\n return x - np.logaddexp(0, x)\n\n\ndef init_random_params(scale, layer_sizes, rs=npr.RandomState(0)):\n \"\"\"Build a list of (weights, biases) tuples,\n one for each layer in the net.\"\"\"\n return [\n (\n scale * rs.randn(m, n), # weight matrix\n scale * rs.randn(n),\n ) # bias vector\n for m, n in zip(layer_sizes[:-1], layer_sizes[1:])\n ]\n\n\ndef batch_normalize(activations):\n mbmean = np.mean(activations, axis=0, keepdims=True)\n return (activations - mbmean) / (np.std(activations, axis=0, keepdims=True) + 1)\n\n\ndef neural_net_predict(params, inputs):\n \"\"\"Params is a list of (weights, bias) tuples.\n inputs is an (N x D) matrix.\"\"\"\n inpW, inpb = params[0]\n inputs = relu(np.dot(inputs, inpW) + inpb)\n for W, b in params[1:-1]:\n outputs = batch_normalize(np.dot(inputs, W) + b)\n inputs = relu(outputs)\n outW, outb = params[-1]\n outputs = np.dot(inputs, outW) + outb\n return outputs\n\n\ndef generate_from_noise(gen_params, num_samples, noise_dim, rs):\n noise = rs.rand(num_samples, noise_dim)\n samples = neural_net_predict(gen_params, noise)\n return sigmoid(samples)\n\n\ndef gan_objective(gen_params, dsc_params, real_data, num_samples, noise_dim, rs):\n fake_data = generate_from_noise(gen_params, num_samples, noise_dim, rs)\n logprobs_fake = logsigmoid(neural_net_predict(dsc_params, fake_data))\n logprobs_real = logsigmoid(neural_net_predict(dsc_params, real_data))\n return np.mean(logprobs_real) - np.mean(logprobs_fake)\n\n\n### Define minimax version of adam optimizer ###\n\n\ndef adam_minimax(\n grad_both,\n init_params_max,\n init_params_min,\n callback=None,\n num_iters=100,\n step_size_max=0.001,\n step_size_min=0.001,\n b1=0.9,\n b2=0.999,\n eps=10**-8,\n):\n \"\"\"Adam modified to do minimiax optimization, for instance to help with\n training generative adversarial networks.\"\"\"\n\n x_max, unflatten_max = flatten(init_params_max)\n x_min, unflatten_min = flatten(init_params_min)\n\n m_max = np.zeros(len(x_max))\n v_max = np.zeros(len(x_max))\n m_min = np.zeros(len(x_min))\n v_min = np.zeros(len(x_min))\n for i in range(num_iters):\n g_max_uf, g_min_uf = grad_both(unflatten_max(x_max), unflatten_min(x_min), i)\n g_max, _ = flatten(g_max_uf)\n g_min, _ = flatten(g_min_uf)\n\n if callback:\n callback(\n unflatten_max(x_max), unflatten_min(x_min), i, unflatten_max(g_max), unflatten_min(g_min)\n )\n\n m_max = (1 - b1) * g_max + b1 * m_max # First moment estimate.\n v_max = (1 - b2) * (g_max**2) + b2 * v_max # Second moment estimate.\n mhat_max = m_max / (1 - b1 ** (i + 1)) # Bias correction.\n vhat_max = v_max / (1 - b2 ** (i + 1))\n x_max = x_max + step_size_max * mhat_max / (np.sqrt(vhat_max) + eps)\n\n m_min = (1 - b1) * g_min + b1 * m_min # First moment estimate.\n v_min = (1 - b2) * (g_min**2) + b2 * v_min # Second moment estimate.\n mhat_min = m_min / (1 - b1 ** (i + 1)) # Bias correction.\n vhat_min = v_min / (1 - b2 ** (i + 1))\n x_min = x_min - step_size_min * mhat_min / (np.sqrt(vhat_min) + eps)\n return unflatten_max(x_max), unflatten_min(x_min)\n\n\n### Setup and run on MNIST ###\n\nif __name__ == \"__main__\":\n # Model hyper-parameters\n noise_dim = 10\n gen_layer_sizes = [noise_dim, 200, 784]\n dsc_layer_sizes = [784, 200, 1]\n\n # Training parameters\n param_scale = 0.001\n batch_size = 100\n num_epochs = 50\n step_size_max = 0.01\n step_size_min = 0.01\n\n print(\"Loading training data...\")\n N, train_images, _, test_images, _ = load_mnist()\n\n init_gen_params = init_random_params(param_scale, gen_layer_sizes)\n init_dsc_params = init_random_params(param_scale, dsc_layer_sizes)\n\n num_batches = int(np.ceil(len(train_images) / batch_size))\n\n def batch_indices(iter):\n idx = iter % num_batches\n return slice(idx * batch_size, (idx + 1) * batch_size)\n\n # Define training objective\n seed = npr.RandomState(0)\n\n def objective(gen_params, dsc_params, iter):\n idx = batch_indices(iter)\n return gan_objective(gen_params, dsc_params, train_images[idx], batch_size, noise_dim, seed)\n\n # Get gradients of objective using autograd.\n both_objective_grad = grad(objective, argnum=(0, 1))\n\n print(\" Epoch | Objective | Fake probability | Real Probability \")\n\n def print_perf(gen_params, dsc_params, iter, gen_gradient, dsc_gradient):\n if iter % 10 == 0:\n ability = np.mean(objective(gen_params, dsc_params, iter))\n fake_data = generate_from_noise(gen_params, 20, noise_dim, seed)\n real_data = train_images[batch_indices(iter)]\n probs_fake = np.mean(sigmoid(neural_net_predict(dsc_params, fake_data)))\n probs_real = np.mean(sigmoid(neural_net_predict(dsc_params, real_data)))\n print(f\"{iter // num_batches:15}|{ability:20}|{probs_fake:20}|{probs_real:20}\")\n save_images(fake_data, \"gan_samples.png\", vmin=0, vmax=1)\n\n # The optimizers provided can optimize lists, tuples, or dicts of parameters.\n optimized_params = adam_minimax(\n both_objective_grad,\n init_gen_params,\n init_dsc_params,\n step_size_max=step_size_max,\n step_size_min=step_size_min,\n num_iters=num_epochs * num_batches,\n callback=print_perf,\n )\n"
},
{
"path": "examples/gmm.py",
"content": "\"\"\"Implements a Gaussian mixture model, in which parameters are fit using\ngradient descent. This example runs on 2-dimensional data, but the model\nworks on arbitrarily-high dimension.\"\"\"\n\nimport matplotlib.pyplot as plt\nfrom data import make_pinwheel\nfrom scipy.optimize import minimize\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nimport autograd.scipy.stats.multivariate_normal as mvn\nfrom autograd import grad, hessian_vector_product\nfrom autograd.misc.flatten import flatten_func\nfrom autograd.scipy.special import logsumexp\n\n\ndef init_gmm_params(num_components, D, scale, rs=npr.RandomState(0)):\n return {\n \"log proportions\": rs.randn(num_components) * scale,\n \"means\": rs.randn(num_components, D) * scale,\n \"lower triangles\": np.zeros((num_components, D, D)) + np.eye(D),\n }\n\n\ndef log_normalize(x):\n return x - logsumexp(x)\n\n\ndef unpack_gmm_params(params):\n normalized_log_proportions = log_normalize(params[\"log proportions\"])\n return normalized_log_proportions, params[\"means\"], params[\"lower triangles\"]\n\n\ndef gmm_log_likelihood(params, data):\n cluster_lls = []\n for log_proportion, mean, cov_sqrt in zip(*unpack_gmm_params(params)):\n cov = np.dot(cov_sqrt.T, cov_sqrt)\n cluster_lls.append(log_proportion + mvn.logpdf(data, mean, cov))\n return np.sum(logsumexp(np.vstack(cluster_lls), axis=0))\n\n\ndef plot_ellipse(ax, mean, cov_sqrt, alpha, num_points=100):\n angles = np.linspace(0, 2 * np.pi, num_points)\n circle_pts = np.vstack([np.cos(angles), np.sin(angles)]).T * 2.0\n cur_pts = mean + np.dot(circle_pts, cov_sqrt)\n ax.plot(cur_pts[:, 0], cur_pts[:, 1], \"-\", alpha=alpha)\n\n\ndef plot_gaussian_mixture(params, ax):\n for log_proportion, mean, cov_sqrt in zip(*unpack_gmm_params(params)):\n alpha = np.minimum(1.0, np.exp(log_proportion) * 10)\n plot_ellipse(ax, mean, cov_sqrt, alpha)\n\n\nif __name__ == \"__main__\":\n init_params = init_gmm_params(num_components=10, D=2, scale=0.1)\n\n data = make_pinwheel(radial_std=0.3, tangential_std=0.05, num_classes=3, num_per_class=100, rate=0.4)\n\n def objective(params):\n return -gmm_log_likelihood(params, data)\n\n flattened_obj, unflatten, flattened_init_params = flatten_func(objective, init_params)\n\n fig = plt.figure(figsize=(12, 8), facecolor=\"white\")\n ax = fig.add_subplot(111, frameon=False)\n plt.show(block=False)\n\n def callback(flattened_params):\n params = unflatten(flattened_params)\n print(f\"Log likelihood {-objective(params)}\")\n ax.cla()\n ax.plot(data[:, 0], data[:, 1], \"k.\")\n ax.set_xticks([])\n ax.set_yticks([])\n plot_gaussian_mixture(params, ax)\n plt.draw()\n plt.pause(1.0 / 60.0)\n\n minimize(\n flattened_obj,\n flattened_init_params,\n jac=grad(flattened_obj),\n hessp=hessian_vector_product(flattened_obj),\n method=\"Newton-CG\",\n callback=callback,\n )\n"
},
{
"path": "examples/gplvm.py",
"content": "# Implements a Gaussian process latent-variable model.\n# The (high-dimensional) data, Y is explained by some low-dimensional latent\n# data X, warped by a function drawn from a GP prior (f). So Y = f(X), but\n# we don't know X or f.\n#\n# In this example, we optimize X and the hyperparameters of the GP, but\n# we integrate over all possible functions f.\n#\n# Normally the observed data would be high-dimensional.\n#\n# David Duvenaud (duvenaud@gmail.com)\n\n\nimport matplotlib.pyplot as plt\nfrom data import make_pinwheel\nfrom gaussian_process import make_gp_funs, rbf_covariance\nfrom scipy.optimize import minimize\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import value_and_grad\nfrom autograd.scipy.stats import norm\n\nif __name__ == \"__main__\":\n data_dimension = 2 # Normally the data dimension would be much higher.\n latent_dimension = 2\n\n # Build model and objective function.\n params_per_gp, predict, log_marginal_likelihood = make_gp_funs(\n rbf_covariance, num_cov_params=latent_dimension + 1\n )\n total_gp_params = data_dimension * params_per_gp\n\n data = make_pinwheel(radial_std=0.3, tangential_std=0.05, num_classes=3, num_per_class=30, rate=0.4)\n datalen = data.shape[0]\n\n num_latent_params = datalen * latent_dimension\n\n def unpack_params(params):\n gp_params = np.reshape(params[:total_gp_params], (data_dimension, params_per_gp))\n latents = np.reshape(params[total_gp_params:], (datalen, latent_dimension))\n return gp_params, latents\n\n def objective(params):\n gp_params, latents = unpack_params(params)\n gp_likelihood = sum(\n [log_marginal_likelihood(gp_params[i], latents, data[:, i]) for i in range(data_dimension)]\n )\n latent_prior_likelihood = np.sum(norm.logpdf(latents))\n return -gp_likelihood - latent_prior_likelihood\n\n # Set up figure.\n fig = plt.figure(figsize=(12, 8), facecolor=\"white\")\n latent_ax = fig.add_subplot(121, frameon=False)\n data_ax = fig.add_subplot(122, frameon=False)\n plt.show(block=False)\n\n def callback(params):\n print(f\"Log likelihood {-objective(params)}\")\n gp_params, latents = unpack_params(params)\n\n data_ax.cla()\n data_ax.plot(data[:, 0], data[:, 1], \"bx\")\n data_ax.set_xticks([])\n data_ax.set_yticks([])\n data_ax.set_title(\"Observed Data\")\n\n latent_ax.cla()\n latent_ax.plot(latents[:, 0], latents[:, 1], \"kx\")\n latent_ax.set_xticks([])\n latent_ax.set_yticks([])\n latent_ax.set_xlim([-2, 2])\n latent_ax.set_ylim([-2, 2])\n latent_ax.set_title(\"Latent coordinates\")\n\n plt.draw()\n plt.pause(1.0 / 60.0)\n\n # Initialize covariance parameters\n rs = npr.RandomState(1)\n init_params = rs.randn(total_gp_params + num_latent_params) * 0.1\n\n print(\"Optimizing covariance parameters and latent variable locations...\")\n minimize(value_and_grad(objective), init_params, jac=True, method=\"CG\", callback=callback)\n"
},
{
"path": "examples/hmm_em.py",
"content": "import string\nfrom functools import partial\nfrom os.path import dirname, join\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import value_and_grad as vgrad\nfrom autograd.scipy.special import logsumexp\n\n\ndef EM(init_params, data, callback=None):\n def EM_update(params):\n natural_params = list(map(np.log, params))\n loglike, E_stats = vgrad(log_partition_function)(natural_params, data) # E step\n if callback:\n callback(loglike, params)\n return list(map(normalize, E_stats)) # M step\n\n def fixed_point(f, x0):\n x1 = f(x0)\n while different(x0, x1):\n x0, x1 = x1, f(x1)\n return x1\n\n def different(params1, params2):\n allclose = partial(np.allclose, atol=1e-3, rtol=1e-3)\n return not all(map(allclose, params1, params2))\n\n return fixed_point(EM_update, init_params)\n\n\ndef normalize(a):\n def replace_zeros(a):\n return np.where(a > 0.0, a, 1.0)\n\n return a / replace_zeros(a.sum(-1, keepdims=True))\n\n\ndef log_partition_function(natural_params, data):\n if isinstance(data, list):\n return sum(map(partial(log_partition_function, natural_params), data))\n\n log_pi, log_A, log_B = natural_params\n\n log_alpha = log_pi\n for y in data:\n log_alpha = logsumexp(log_alpha[:, None] + log_A, axis=0) + log_B[:, y]\n\n return logsumexp(log_alpha)\n\n\ndef initialize_hmm_parameters(num_states, num_outputs):\n init_pi = normalize(npr.rand(num_states))\n init_A = normalize(npr.rand(num_states, num_states))\n init_B = normalize(npr.rand(num_states, num_outputs))\n return init_pi, init_A, init_B\n\n\ndef build_dataset(filename, max_lines=-1):\n \"\"\"Loads a text file, and turns each line into an encoded sequence.\"\"\"\n encodings = dict(list(map(reversed, enumerate(string.printable))))\n digitize = lambda char: encodings[char] if char in encodings else len(encodings)\n encode_line = lambda line: np.array(list(map(digitize, line)))\n nonblank_line = lambda line: len(line) > 2\n\n with open(filename) as f:\n lines = f.readlines()\n\n encoded_lines = list(map(encode_line, list(filter(nonblank_line, lines))[:max_lines]))\n num_outputs = len(encodings) + 1\n\n return encoded_lines, num_outputs\n\n\nif __name__ == \"__main__\":\n np.random.seed(0)\n np.seterr(divide=\"ignore\")\n\n # callback to print log likelihoods during training\n print_loglike = lambda loglike, params: print(loglike)\n\n # load training data\n lstm_filename = join(dirname(__file__), \"lstm.py\")\n train_inputs, num_outputs = build_dataset(lstm_filename, max_lines=60)\n\n # train with EM\n num_states = 20\n init_params = initialize_hmm_parameters(num_states, num_outputs)\n pi, A, B = EM(init_params, train_inputs, print_loglike)\n"
},
{
"path": "examples/ica.py",
"content": "import matplotlib.cm as cm\nimport matplotlib.pyplot as plt\nfrom scipy.optimize import minimize\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nimport autograd.scipy.stats.t as t\nfrom autograd import value_and_grad\n\n\ndef make_ica_funs(observed_dimension, latent_dimension):\n \"\"\"These functions implement independent component analysis.\n\n The model is:\n latents are drawn i.i.d. for each data point from a product of student-ts.\n weights are the same across all datapoints.\n each data = latents * weghts + noise.\"\"\"\n\n def sample(weights, n_samples, noise_std, rs):\n latents = rs.randn(latent_dimension, n_samples)\n latents = np.array(sorted(latents.T, key=lambda a_entry: a_entry[0])).T\n noise = rs.randn(n_samples, observed_dimension) * noise_std\n observed = predict(weights, latents) + noise\n return latents, observed\n\n def predict(weights, latents):\n return np.dot(weights, latents).T\n\n def logprob(weights, latents, noise_std, observed):\n preds = predict(weights, latents)\n log_lik = np.sum(t.logpdf(preds, 2.4, observed, noise_std))\n return log_lik\n\n num_weights = observed_dimension * latent_dimension\n\n def unpack_weights(weights):\n return np.reshape(weights, (observed_dimension, latent_dimension))\n\n return num_weights, sample, logprob, unpack_weights\n\n\ndef color_scatter(ax, xs, ys):\n colors = cm.rainbow(np.linspace(0, 1, len(ys)))\n for x, y, c in zip(xs, ys, colors):\n ax.scatter(x, y, color=c)\n\n\nif __name__ == \"__main__\":\n observed_dimension = 100\n latent_dimension = 2\n true_noise_var = 1.0\n n_samples = 200\n\n num_weights, sample, logprob, unpack_weights = make_ica_funs(observed_dimension, latent_dimension)\n\n num_latent_params = latent_dimension * n_samples\n total_num_params = num_weights + num_latent_params + 1\n\n def unpack_params(params):\n weights = unpack_weights(params[:num_weights])\n latents = np.reshape(\n params[num_weights : num_weights + num_latent_params], (latent_dimension, n_samples)\n )\n noise_std = np.exp(params[-1])\n return weights, latents, noise_std\n\n rs = npr.RandomState(0)\n true_weights = np.zeros((observed_dimension, latent_dimension))\n for i in range(latent_dimension):\n true_weights[:, i] = np.sin(np.linspace(0, 4 + i * 3.2, observed_dimension))\n\n true_latents, data = sample(true_weights, n_samples, true_noise_var, rs)\n\n # Set up figure.\n fig2 = plt.figure(figsize=(6, 6), facecolor=\"white\")\n ax_data = fig2.add_subplot(111, frameon=False)\n ax_data.matshow(data)\n\n fig1 = plt.figure(figsize=(12, 16), facecolor=\"white\")\n ax_true_latents = fig1.add_subplot(411, frameon=False)\n ax_est_latents = fig1.add_subplot(412, frameon=False)\n ax_true_weights = fig1.add_subplot(413, frameon=False)\n ax_est_weights = fig1.add_subplot(414, frameon=False)\n\n plt.show(block=False)\n ax_true_weights.scatter(true_weights[:, 0], true_weights[:, 1])\n ax_true_weights.set_title(\"True weights\")\n color_scatter(ax_true_latents, true_latents[0, :], true_latents[1, :])\n ax_true_latents.set_title(\"True latents\")\n ax_true_latents.set_xticks([])\n ax_true_weights.set_xticks([])\n ax_true_latents.set_yticks([])\n ax_true_weights.set_yticks([])\n\n def objective(params):\n weight_matrix, latents, noise_std = unpack_params(params)\n return -logprob(weight_matrix, latents, noise_std, data) / n_samples\n\n def callback(params):\n weights, latents, noise_std = unpack_params(params)\n print(f\"Log likelihood {-objective(params)}, noise_std {noise_std}\")\n ax_est_weights.cla()\n ax_est_weights.scatter(weights[:, 0], weights[:, 1])\n ax_est_weights.set_title(\"Estimated weights\")\n ax_est_latents.cla()\n color_scatter(ax_est_latents, latents[0, :], latents[1, :])\n ax_est_latents.set_title(\"Estimated latents\")\n ax_est_weights.set_yticks([])\n ax_est_latents.set_yticks([])\n ax_est_weights.set_xticks([])\n ax_est_latents.set_xticks([])\n plt.draw()\n plt.pause(1.0 / 60.0)\n\n # Initialize and optimize model.\n rs = npr.RandomState(0)\n init_params = rs.randn(total_num_params)\n minimize(value_and_grad(objective), init_params, jac=True, method=\"CG\", callback=callback)\n plt.pause(20)\n"
},
{
"path": "examples/logistic_regression.py",
"content": "import autograd.numpy as np\nfrom autograd import grad\nfrom autograd.test_util import check_grads\n\n\ndef sigmoid(x):\n return 0.5 * (np.tanh(x) + 1)\n\n\ndef logistic_predictions(weights, inputs):\n # Outputs probability of a label being true according to logistic model.\n return sigmoid(np.dot(inputs, weights))\n\n\ndef training_loss(weights):\n # Training loss is the negative log-likelihood of the training labels.\n preds = logistic_predictions(weights, inputs)\n label_probabilities = preds * targets + (1 - preds) * (1 - targets)\n return -np.sum(np.log(label_probabilities))\n\n\n# Build a toy dataset.\ninputs = np.array([[0.52, 1.12, 0.77], [0.88, -1.08, 0.15], [0.52, 0.06, -1.30], [0.74, -2.49, 1.39]])\ntargets = np.array([True, True, False, True])\n\n# Build a function that returns gradients of training loss using autograd.\ntraining_gradient_fun = grad(training_loss)\n\n# Check the gradients numerically, just to be safe.\nweights = np.array([0.0, 0.0, 0.0])\ncheck_grads(training_loss, modes=[\"rev\"])(weights)\n\n# Optimize weights using gradient descent.\nprint(\"Initial loss:\", training_loss(weights))\nfor i in range(100):\n weights -= training_gradient_fun(weights) * 0.01\n\nprint(\"Trained loss:\", training_loss(weights))\n"
},
{
"path": "examples/lstm.py",
"content": "\"\"\"Implements the long-short term memory character model.\nThis version vectorizes over multiple examples, but each string\nhas a fixed length.\"\"\"\n\nfrom os.path import dirname, join\n\nfrom rnn import build_dataset, concat_and_multiply, one_hot_to_string, sigmoid, string_to_one_hot\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.misc.optimizers import adam\nfrom autograd.scipy.special import logsumexp\n\n\ndef init_lstm_params(input_size, state_size, output_size, param_scale=0.01, rs=npr.RandomState(0)):\n def rp(*shape):\n return rs.randn(*shape) * param_scale\n\n return {\n \"init cells\": rp(1, state_size),\n \"init hiddens\": rp(1, state_size),\n \"change\": rp(input_size + state_size + 1, state_size),\n \"forget\": rp(input_size + state_size + 1, state_size),\n \"ingate\": rp(input_size + state_size + 1, state_size),\n \"outgate\": rp(input_size + state_size + 1, state_size),\n \"predict\": rp(state_size + 1, output_size),\n }\n\n\ndef lstm_predict(params, inputs):\n def update_lstm(input, hiddens, cells):\n change = np.tanh(concat_and_multiply(params[\"change\"], input, hiddens))\n forget = sigmoid(concat_and_multiply(params[\"forget\"], input, hiddens))\n ingate = sigmoid(concat_and_multiply(params[\"ingate\"], input, hiddens))\n outgate = sigmoid(concat_and_multiply(params[\"outgate\"], input, hiddens))\n cells = cells * forget + ingate * change\n hiddens = outgate * np.tanh(cells)\n return hiddens, cells\n\n def hiddens_to_output_probs(hiddens):\n output = concat_and_multiply(params[\"predict\"], hiddens)\n return output - logsumexp(output, axis=1, keepdims=True) # Normalize log-probs.\n\n num_sequences = inputs.shape[1]\n hiddens = np.repeat(params[\"init hiddens\"], num_sequences, axis=0)\n cells = np.repeat(params[\"init cells\"], num_sequences, axis=0)\n\n output = [hiddens_to_output_probs(hiddens)]\n for input in inputs: # Iterate over time steps.\n hiddens, cells = update_lstm(input, hiddens, cells)\n output.append(hiddens_to_output_probs(hiddens))\n return output\n\n\ndef lstm_log_likelihood(params, inputs, targets):\n logprobs = lstm_predict(params, inputs)\n loglik = 0.0\n num_time_steps, num_examples, _ = inputs.shape\n for t in range(num_time_steps):\n loglik += np.sum(logprobs[t] * targets[t])\n return loglik / (num_time_steps * num_examples)\n\n\nif __name__ == \"__main__\":\n num_chars = 128\n\n # Learn to predict our own source code.\n text_filename = join(dirname(__file__), \"lstm.py\")\n train_inputs = build_dataset(text_filename, sequence_length=30, alphabet_size=num_chars, max_lines=60)\n\n init_params = init_lstm_params(input_size=128, output_size=128, state_size=40, param_scale=0.01)\n\n def print_training_prediction(weights):\n print(\"Training text Predicted text\")\n logprobs = np.asarray(lstm_predict(weights, train_inputs))\n for t in range(logprobs.shape[1]):\n training_text = one_hot_to_string(train_inputs[:, t, :])\n predicted_text = one_hot_to_string(logprobs[:, t, :])\n print(training_text.replace(\"\\n\", \" \") + \"|\" + predicted_text.replace(\"\\n\", \" \"))\n\n def training_loss(params, iter):\n return -lstm_log_likelihood(params, train_inputs, train_inputs)\n\n def callback(weights, iter, gradient):\n if iter % 10 == 0:\n print(\"Iteration\", iter, \"Train loss:\", training_loss(weights, 0))\n print_training_prediction(weights)\n\n # Build gradient of loss function using autograd.\n training_loss_grad = grad(training_loss)\n\n print(\"Training LSTM...\")\n trained_params = adam(training_loss_grad, init_params, step_size=0.1, num_iters=1000, callback=callback)\n\n print()\n print(\"Generating text from LSTM...\")\n num_letters = 30\n for t in range(20):\n text = \"\"\n for i in range(num_letters):\n seqs = string_to_one_hot(text, num_chars)[:, np.newaxis, :]\n logprobs = lstm_predict(trained_params, seqs)[-1].ravel()\n text += chr(npr.choice(len(logprobs), p=np.exp(logprobs)))\n print(text)\n"
},
{
"path": "examples/mixture_variational_inference.py",
"content": "# Implements black-box variational inference, where the variational\n# distribution is a mixture of Gaussians.\n#\n# This trick was written up by Alex Graves in this note:\n# http://arxiv.org/abs/1607.05690\n\nimport matplotlib.pyplot as plt\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nimport autograd.scipy.stats.norm as norm\nfrom autograd import grad\nfrom autograd.misc.optimizers import adam\nfrom autograd.scipy.special import logsumexp\n\n\ndef diag_gaussian_log_density(x, mu, log_std):\n return np.sum(norm.logpdf(x, mu, np.exp(log_std)), axis=-1)\n\n\ndef unpack_gaussian_params(params):\n # Variational dist is a diagonal Gaussian.\n D = np.shape(params)[0] // 2\n mean, log_std = params[:D], params[D:]\n return mean, log_std\n\n\ndef variational_log_density_gaussian(params, x):\n mean, log_std = unpack_gaussian_params(params)\n return diag_gaussian_log_density(x, mean, log_std)\n\n\ndef sample_diag_gaussian(params, num_samples, rs):\n mean, log_std = unpack_gaussian_params(params)\n D = np.shape(mean)[0]\n return rs.randn(num_samples, D) * np.exp(log_std) + mean\n\n\ndef variational_lower_bound(params, t, logprob, sampler, log_density, num_samples, rs):\n \"\"\"Provides a stochastic estimate of the variational lower bound,\n for any variational family and model density.\"\"\"\n samples = sampler(params, num_samples, rs)\n log_qs = log_density(params, samples)\n log_ps = logprob(samples, t)\n log_ps = np.reshape(log_ps, (num_samples, -1))\n log_qs = np.reshape(log_qs, (num_samples, -1))\n return np.mean(log_ps - log_qs)\n\n\ndef init_gaussian_var_params(D, mean_mean=-1, log_std_mean=-5, scale=0.1, rs=npr.RandomState(0)):\n init_mean = mean_mean * np.ones(D) + rs.randn(D) * scale\n init_log_std = log_std_mean * np.ones(D) + rs.randn(D) * scale\n return np.concatenate([init_mean, init_log_std])\n\n\ndef log_normalize(x):\n return x - logsumexp(x)\n\n\ndef build_mog_bbsvi(logprob, num_samples, k=10, rs=npr.RandomState(0)):\n init_component_var_params = init_gaussian_var_params\n component_log_density = variational_log_density_gaussian\n component_sample = sample_diag_gaussian\n\n def unpack_mixture_params(mixture_params):\n log_weights = log_normalize(mixture_params[:k])\n var_params = np.reshape(mixture_params[k:], (k, -1))\n return log_weights, var_params\n\n def init_var_params(D, rs=npr.RandomState(0), **kwargs):\n log_weights = np.ones(k)\n component_weights = [init_component_var_params(D, rs=rs, **kwargs) for i in range(k)]\n return np.concatenate([log_weights] + component_weights)\n\n def sample(var_mixture_params, num_samples, rs):\n \"\"\"Sample locations aren't a continuous function of parameters\n due to multinomial sampling.\"\"\"\n log_weights, var_params = unpack_mixture_params(var_mixture_params)\n samples = np.concatenate(\n [component_sample(params_k, num_samples, rs)[:, np.newaxis, :] for params_k in var_params],\n axis=1,\n )\n ixs = np.random.choice(k, size=num_samples, p=np.exp(log_weights))\n return np.array([samples[i, ix, :] for i, ix in enumerate(ixs)])\n\n def mixture_log_density(var_mixture_params, x):\n \"\"\"Returns a weighted average over component densities.\"\"\"\n log_weights, var_params = unpack_mixture_params(var_mixture_params)\n component_log_densities = np.vstack(\n [component_log_density(params_k, x) for params_k in var_params]\n ).T\n return logsumexp(component_log_densities + log_weights, axis=1, keepdims=False)\n\n def mixture_elbo(var_mixture_params, t):\n # We need to only sample the continuous component parameters,\n # and integrate over the discrete component choice\n\n def mixture_lower_bound(params):\n \"\"\"Provides a stochastic estimate of the variational lower bound.\"\"\"\n samples = component_sample(params, num_samples, rs)\n log_qs = mixture_log_density(var_mixture_params, samples)\n log_ps = logprob(samples, t)\n log_ps = np.reshape(log_ps, (num_samples, -1))\n log_qs = np.reshape(log_qs, (num_samples, -1))\n return np.mean(log_ps - log_qs)\n\n log_weights, var_params = unpack_mixture_params(var_mixture_params)\n component_elbos = np.stack([mixture_lower_bound(params_k) for params_k in var_params])\n return np.sum(component_elbos * np.exp(log_weights))\n\n return init_var_params, mixture_elbo, mixture_log_density, sample\n\n\nif __name__ == \"__main__\":\n # Specify an inference problem by its unnormalized log-density.\n D = 2\n\n def log_density(x, t):\n mu, log_sigma = x[:, 0], x[:, 1]\n sigma_density = norm.logpdf(log_sigma, 0, 1.35)\n mu_density = norm.logpdf(mu, -0.5, np.exp(log_sigma))\n sigma_density2 = norm.logpdf(log_sigma, 0.1, 1.35)\n mu_density2 = norm.logpdf(mu, 0.5, np.exp(log_sigma))\n return np.logaddexp(sigma_density + mu_density, sigma_density2 + mu_density2)\n\n init_var_params, elbo, variational_log_density, variational_sampler = build_mog_bbsvi(\n log_density, num_samples=40, k=10\n )\n\n def objective(params, t):\n return -elbo(params, t)\n\n # Set up plotting code\n def plot_isocontours(ax, func, xlimits=[-2, 2], ylimits=[-4, 2], numticks=101, cmap=None):\n x = np.linspace(*xlimits, num=numticks)\n y = np.linspace(*ylimits, num=numticks)\n X, Y = np.meshgrid(x, y)\n zs = func(np.concatenate([np.atleast_2d(X.ravel()), np.atleast_2d(Y.ravel())]).T)\n Z = zs.reshape(X.shape)\n plt.contour(X, Y, Z, cmap=cmap)\n ax.set_yticks([])\n ax.set_xticks([])\n\n fig = plt.figure(figsize=(8, 8), facecolor=\"white\")\n ax = fig.add_subplot(111, frameon=False)\n plt.ion()\n plt.show(block=False)\n\n num_plotting_samples = 51\n\n def callback(params, t, g):\n print(f\"Iteration {t} lower bound {-objective(params, t)}\")\n\n plt.cla()\n target_distribution = lambda x: np.exp(log_density(x, t))\n var_distribution = lambda x: np.exp(variational_log_density(params, x))\n plot_isocontours(ax, target_distribution)\n plot_isocontours(ax, var_distribution, cmap=plt.cm.bone)\n ax.set_autoscale_on(False)\n\n rs = npr.RandomState(0)\n samples = variational_sampler(params, num_plotting_samples, rs)\n plt.plot(samples[:, 0], samples[:, 1], \"x\")\n\n plt.draw()\n plt.pause(1.0 / 30.0)\n\n print(\"Optimizing variational parameters...\")\n variational_params = adam(\n grad(objective), init_var_params(D), step_size=0.1, num_iters=2000, callback=callback\n )\n"
},
{
"path": "examples/natural_gradient_black_box_svi.py",
"content": "import matplotlib.pyplot as plt\n\n# same BBSVI function!\nfrom black_box_svi import black_box_variational_inference\n\nimport autograd.numpy as np\nimport autograd.scipy.stats.norm as norm\nfrom autograd.misc.optimizers import adam, sgd\n\nif __name__ == \"__main__\":\n # Specify an inference problem by its unnormalized log-density.\n # it's difficult to see the benefit in low dimensions\n # model parameters are a mean and a log_sigma\n np.random.seed(42)\n obs_dim = 20\n Y = np.random.randn(obs_dim, obs_dim).dot(np.random.randn(obs_dim))\n\n def log_density(x, t):\n mu, log_sigma = x[:, :obs_dim], x[:, obs_dim:]\n sigma_density = np.sum(norm.logpdf(log_sigma, 0, 1.35), axis=1)\n mu_density = np.sum(norm.logpdf(Y, mu, np.exp(log_sigma)), axis=1)\n return sigma_density + mu_density\n\n # Build variational objective.\n D = obs_dim * 2 # dimension of our posterior\n objective, gradient, unpack_params = black_box_variational_inference(log_density, D, num_samples=2000)\n\n # Define the natural gradient\n # The natural gradient of the ELBO is the gradient of the elbo,\n # preconditioned by the inverse Fisher Information Matrix. The Fisher,\n # in the case of a diagonal gaussian, is a diagonal matrix that is a\n # simple function of the variance. Intuitively, statistical distance\n # created by perturbing the mean of an independent Gaussian is\n # determined by how wide the distribution is along that dimension ---\n # the wider the distribution, the less sensitive statistical distances is\n # to perturbations of the mean; the narrower the distribution, the more\n # the statistical distance changes when you perturb the mean (imagine\n # an extremely narrow Gaussian --- basically a spike. The KL between\n # this Gaussian and a Gaussian $\\epsilon$ away in location can be big ---\n # moving the Gaussian could significantly reduce overlap in support\n # which corresponds to a greater statistical distance).\n #\n # When we want to move in directions of steepest ascent, we multiply by\n # the inverse fisher --- that way we make quicker progress when the\n # variance is wide, and we scale down our step size when the variance\n # is small (which leads to more robust/less chaotic ascent).\n def fisher_diag(lam):\n mu, log_sigma = unpack_params(lam)\n return np.concatenate([np.exp(-2.0 * log_sigma), np.ones(len(log_sigma)) * 2])\n\n # simple! basically free!\n natural_gradient = lambda lam, i: (1.0 / fisher_diag(lam)) * gradient(lam, i)\n\n # function for keeping track of callback ELBO values (for plotting below)\n def optimize_and_lls(optfun):\n num_iters = 200\n elbos = []\n\n def callback(params, t, g):\n elbo_val = -objective(params, t)\n elbos.append(elbo_val)\n if t % 50 == 0:\n print(f\"Iteration {t} lower bound {elbo_val}\")\n\n init_mean = -1 * np.ones(D)\n init_log_std = -5 * np.ones(D)\n init_var_params = np.concatenate([init_mean, init_log_std])\n variational_params = optfun(num_iters, init_var_params, callback)\n return np.array(elbos)\n\n # let's optimize this with a few different step sizes\n elbo_lists = []\n step_sizes = [0.1, 0.25, 0.5]\n for step_size in step_sizes:\n # optimize with standard gradient + adam\n optfun = lambda n, init, cb: adam(gradient, init, step_size=step_size, num_iters=n, callback=cb)\n standard_lls = optimize_and_lls(optfun)\n\n # optimize with natural gradient + sgd, no momentum\n optnat = lambda n, init, cb: sgd(\n natural_gradient, init, step_size=step_size, num_iters=n, callback=cb, mass=0.001\n )\n natural_lls = optimize_and_lls(optnat)\n elbo_lists.append((standard_lls, natural_lls))\n\n # visually compare the ELBO\n plt.figure(figsize=(12, 8))\n colors = [\"b\", \"k\", \"g\"]\n for col, ss, (stand_lls, nat_lls) in zip(colors, step_sizes, elbo_lists):\n plt.plot(\n np.arange(len(stand_lls)),\n stand_lls,\n \"--\",\n label=\"standard (adam, step-size = %2.2f)\" % ss,\n alpha=0.5,\n c=col,\n )\n plt.plot(np.arange(len(nat_lls)), nat_lls, \"-\", label=\"natural (sgd, step-size = %2.2f)\" % ss, c=col)\n\n llrange = natural_lls.max() - natural_lls.min()\n plt.ylim((natural_lls.max() - llrange * 0.1, natural_lls.max() + 10))\n plt.xlabel(\"optimization iteration\")\n plt.ylabel(\"ELBO\")\n plt.legend(loc=\"lower right\")\n plt.title(\"%d dimensional posterior\" % D)\n plt.show()\n"
},
{
"path": "examples/negative_binomial_maxlike.py",
"content": "import scipy.optimize\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.scipy.special import gammaln\n\n# The code in this example implements a method for finding a stationary point of\n# the negative binomial likelihood via Newton's method, described here:\n# https://en.wikipedia.org/wiki/Negative_binomial_distribution#Maximum_likelihood_estimation\n\n\ndef newton(f, x0):\n # wrap scipy.optimize.newton with our automatic derivatives\n return scipy.optimize.newton(f, x0, fprime=grad(f), fprime2=grad(grad(f)))\n\n\ndef negbin_loglike(r, p, x):\n # the negative binomial log likelihood we want to maximize\n return gammaln(r + x) - gammaln(r) - gammaln(x + 1) + x * np.log(p) + r * np.log(1 - p)\n\n\ndef negbin_sample(r, p, size):\n # a negative binomial is a gamma-compound-Poisson\n return npr.poisson(npr.gamma(r, p / (1 - p), size=size))\n\n\ndef fit_maxlike(x, r_guess):\n # follows Wikipedia's section on negative binomial max likelihood\n assert np.var(x) > np.mean(x), \"Likelihood-maximizing parameters don't exist!\"\n loglike = lambda r, p: np.sum(negbin_loglike(r, p, x))\n p = lambda r: np.sum(x) / np.sum(r + x)\n rprime = lambda r: grad(loglike)(r, p(r))\n r = newton(rprime, r_guess)\n return r, p(r)\n\n\nif __name__ == \"__main__\":\n # generate data\n npr.seed(0)\n data = negbin_sample(r=5, p=0.5, size=1000)\n\n # fit likelihood-extremizing parameters\n r, p = fit_maxlike(data, r_guess=1)\n\n # report fit\n print(\"Fit parameters:\")\n print(f\"r={r}, p={p}\")\n\n print(\"Check that we are at a local stationary point:\")\n loglike = lambda r, p: np.sum(negbin_loglike(r, p, data))\n grad_both = grad(loglike, argnum=(0, 1))\n print(grad_both(r, p))\n\n import matplotlib.pyplot as plt\n\n xm = data.max()\n plt.figure()\n plt.hist(data, bins=np.arange(xm + 1) - 0.5, density=True, label=\"normed data counts\")\n plt.xlim(0, xm)\n plt.plot(np.arange(xm), np.exp(negbin_loglike(r, p, np.arange(xm))), label=\"maxlike fit\")\n plt.xlabel(\"k\")\n plt.ylabel(\"p(k)\")\n plt.legend(loc=\"best\")\n plt.show()\n"
},
{
"path": "examples/neural_net.py",
"content": "\"\"\"A multi-layer perceptron for classification of MNIST handwritten digits.\"\"\"\n\nfrom data import load_mnist\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.misc.flatten import flatten\nfrom autograd.misc.optimizers import adam\nfrom autograd.scipy.special import logsumexp\n\n\ndef init_random_params(scale, layer_sizes, rs=npr.RandomState(0)):\n \"\"\"Build a list of (weights, biases) tuples,\n one for each layer in the net.\"\"\"\n return [\n (\n scale * rs.randn(m, n), # weight matrix\n scale * rs.randn(n),\n ) # bias vector\n for m, n in zip(layer_sizes[:-1], layer_sizes[1:])\n ]\n\n\ndef neural_net_predict(params, inputs):\n \"\"\"Implements a deep neural network for classification.\n params is a list of (weights, bias) tuples.\n inputs is an (N x D) matrix.\n returns normalized class log-probabilities.\"\"\"\n for W, b in params:\n outputs = np.dot(inputs, W) + b\n inputs = np.tanh(outputs)\n return outputs - logsumexp(outputs, axis=1, keepdims=True)\n\n\ndef l2_norm(params):\n \"\"\"Computes l2 norm of params by flattening them into a vector.\"\"\"\n flattened, _ = flatten(params)\n return np.dot(flattened, flattened)\n\n\ndef log_posterior(params, inputs, targets, L2_reg):\n log_prior = -L2_reg * l2_norm(params)\n log_lik = np.sum(neural_net_predict(params, inputs) * targets)\n return log_prior + log_lik\n\n\ndef accuracy(params, inputs, targets):\n target_class = np.argmax(targets, axis=1)\n predicted_class = np.argmax(neural_net_predict(params, inputs), axis=1)\n return np.mean(predicted_class == target_class)\n\n\nif __name__ == \"__main__\":\n # Model parameters\n layer_sizes = [784, 200, 100, 10]\n L2_reg = 1.0\n\n # Training parameters\n param_scale = 0.1\n batch_size = 256\n num_epochs = 5\n step_size = 0.001\n\n print(\"Loading training data...\")\n N, train_images, train_labels, test_images, test_labels = load_mnist()\n\n init_params = init_random_params(param_scale, layer_sizes)\n\n num_batches = int(np.ceil(len(train_images) / batch_size))\n\n def batch_indices(iter):\n idx = iter % num_batches\n return slice(idx * batch_size, (idx + 1) * batch_size)\n\n # Define training objective\n def objective(params, iter):\n idx = batch_indices(iter)\n return -log_posterior(params, train_images[idx], train_labels[idx], L2_reg)\n\n # Get gradient of objective using autograd.\n objective_grad = grad(objective)\n\n print(\" Epoch | Train accuracy | Test accuracy \")\n\n def print_perf(params, iter, gradient):\n if iter % num_batches == 0:\n train_acc = accuracy(params, train_images, train_labels)\n test_acc = accuracy(params, test_images, test_labels)\n print(f\"{iter // num_batches:15}|{train_acc:20}|{test_acc:20}\")\n\n # The optimizers provided can optimize lists, tuples, or dicts of parameters.\n optimized_params = adam(\n objective_grad,\n init_params,\n step_size=step_size,\n num_iters=num_epochs * num_batches,\n callback=print_perf,\n )\n"
},
{
"path": "examples/neural_net_regression.py",
"content": "import matplotlib.pyplot as plt\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nimport autograd.scipy.stats.norm as norm\nfrom autograd import grad\nfrom autograd.misc import flatten\nfrom autograd.misc.optimizers import adam\n\n\ndef init_random_params(scale, layer_sizes, rs=npr.RandomState(0)):\n \"\"\"Build a list of (weights, biases) tuples, one for each layer.\"\"\"\n return [\n (\n rs.randn(insize, outsize) * scale, # weight matrix\n rs.randn(outsize) * scale,\n ) # bias vector\n for insize, outsize in zip(layer_sizes[:-1], layer_sizes[1:])\n ]\n\n\ndef nn_predict(params, inputs, nonlinearity=np.tanh):\n for W, b in params:\n outputs = np.dot(inputs, W) + b\n inputs = nonlinearity(outputs)\n return outputs\n\n\ndef log_gaussian(params, scale):\n flat_params, _ = flatten(params)\n return np.sum(norm.logpdf(flat_params, 0, scale))\n\n\ndef logprob(weights, inputs, targets, noise_scale=0.1):\n predictions = nn_predict(weights, inputs)\n return np.sum(norm.logpdf(predictions, targets, noise_scale))\n\n\ndef build_toy_dataset(n_data=80, noise_std=0.1):\n rs = npr.RandomState(0)\n inputs = np.concatenate([np.linspace(0, 3, num=n_data // 2), np.linspace(6, 8, num=n_data // 2)])\n targets = np.cos(inputs) + rs.randn(n_data) * noise_std\n inputs = (inputs - 4.0) / 2.0\n inputs = inputs[:, np.newaxis]\n targets = targets[:, np.newaxis] / 2.0\n return inputs, targets\n\n\nif __name__ == \"__main__\":\n init_scale = 0.1\n weight_prior_variance = 10.0\n init_params = init_random_params(init_scale, layer_sizes=[1, 4, 4, 1])\n\n inputs, targets = build_toy_dataset()\n\n def objective(weights, t):\n return -logprob(weights, inputs, targets) - log_gaussian(weights, weight_prior_variance)\n\n print(grad(objective)(init_params, 0))\n\n # Set up figure.\n fig = plt.figure(figsize=(12, 8), facecolor=\"white\")\n ax = fig.add_subplot(111, frameon=False)\n plt.show(block=False)\n\n def callback(params, t, g):\n print(f\"Iteration {t} log likelihood {-objective(params, t)}\")\n\n # Plot data and functions.\n plt.cla()\n ax.plot(inputs.ravel(), targets.ravel(), \"bx\", ms=12)\n plot_inputs = np.reshape(np.linspace(-7, 7, num=300), (300, 1))\n outputs = nn_predict(params, plot_inputs)\n ax.plot(plot_inputs, outputs, \"r\", lw=3)\n ax.set_ylim([-1, 1])\n plt.draw()\n plt.pause(1.0 / 60.0)\n\n print(\"Optimizing network parameters...\")\n optimized_params = adam(grad(objective), init_params, step_size=0.01, num_iters=1000, callback=callback)\n"
},
{
"path": "examples/ode_net.py",
"content": "# A demo of gradients through scipy.integrate.odeint,\n# estimating the dynamics of a system given a trajectory.\n\n\nimport matplotlib.pyplot as plt\nimport numpy as npo\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.builtins import tuple\nfrom autograd.misc.optimizers import adam\nfrom autograd.scipy.integrate import odeint\n\nN = 30 # Dataset size\nD = 2 # Data dimension\nmax_T = 1.5\n\n\n# Two-dimensional damped oscillator\ndef func(y, t0, A):\n return np.dot(y**3, A)\n\n\ndef nn_predict(inputs, t, params):\n for W, b in params:\n outputs = np.dot(inputs, W) + b\n inputs = np.maximum(0, outputs)\n return outputs\n\n\ndef init_nn_params(scale, layer_sizes, rs=npr.RandomState(0)):\n \"\"\"Build a list of (weights, biases) tuples, one for each layer.\"\"\"\n return [\n (\n rs.randn(insize, outsize) * scale, # weight matrix\n rs.randn(outsize) * scale,\n ) # bias vector\n for insize, outsize in zip(layer_sizes[:-1], layer_sizes[1:])\n ]\n\n\n# Define neural ODE model.\ndef ode_pred(params, y0, t):\n return odeint(nn_predict, y0, t, tuple((params,)), rtol=0.01)\n\n\ndef L1_loss(pred, targets):\n return np.mean(np.abs(pred - targets))\n\n\nif __name__ == \"__main__\":\n # Generate data from true dynamics.\n true_y0 = np.array([2.0, 0.0]).T\n t = np.linspace(0.0, max_T, N)\n true_A = np.array([[-0.1, 2.0], [-2.0, -0.1]])\n true_y = odeint(func, true_y0, t, args=(true_A,))\n\n def train_loss(params, iter):\n pred = ode_pred(params, true_y0, t)\n return L1_loss(pred, true_y)\n\n # Set up figure\n fig = plt.figure(figsize=(12, 4), facecolor=\"white\")\n ax_traj = fig.add_subplot(131, frameon=False)\n ax_phase = fig.add_subplot(132, frameon=False)\n ax_vecfield = fig.add_subplot(133, frameon=False)\n plt.show(block=False)\n\n # Plots data and learned dynamics.\n def callback(params, iter, g):\n pred = ode_pred(params, true_y0, t)\n\n print(f\"Iteration {iter:d} train loss {L1_loss(pred, true_y):.6f}\")\n\n ax_traj.cla()\n ax_traj.set_title(\"Trajectories\")\n ax_traj.set_xlabel(\"t\")\n ax_traj.set_ylabel(\"x,y\")\n ax_traj.plot(t, true_y[:, 0], \"-\", t, true_y[:, 1], \"g-\")\n ax_traj.plot(t, pred[:, 0], \"--\", t, pred[:, 1], \"b--\")\n ax_traj.set_xlim(t.min(), t.max())\n ax_traj.set_ylim(-2, 2)\n ax_traj.xaxis.set_ticklabels([])\n ax_traj.yaxis.set_ticklabels([])\n ax_traj.legend()\n\n ax_phase.cla()\n ax_phase.set_title(\"Phase Portrait\")\n ax_phase.set_xlabel(\"x\")\n ax_phase.set_ylabel(\"y\")\n ax_phase.plot(true_y[:, 0], true_y[:, 1], \"g-\")\n ax_phase.plot(pred[:, 0], pred[:, 1], \"b--\")\n ax_phase.set_xlim(-2, 2)\n ax_phase.set_ylim(-2, 2)\n ax_phase.xaxis.set_ticklabels([])\n ax_phase.yaxis.set_ticklabels([])\n\n ax_vecfield.cla()\n ax_vecfield.set_title(\"Learned Vector Field\")\n ax_vecfield.set_xlabel(\"x\")\n ax_vecfield.set_ylabel(\"y\")\n ax_vecfield.xaxis.set_ticklabels([])\n ax_vecfield.yaxis.set_ticklabels([])\n\n # vector field plot\n y, x = npo.mgrid[-2:2:21j, -2:2:21j]\n dydt = nn_predict(np.stack([x, y], -1).reshape(21 * 21, 2), 0, params).reshape(-1, 2)\n mag = np.sqrt(dydt[:, 0] ** 2 + dydt[:, 1] ** 2).reshape(-1, 1)\n dydt = dydt / mag\n dydt = dydt.reshape(21, 21, 2)\n\n ax_vecfield.streamplot(x, y, dydt[:, :, 0], dydt[:, :, 1], color=\"black\")\n ax_vecfield.set_xlim(-2, 2)\n ax_vecfield.set_ylim(-2, 2)\n\n fig.tight_layout()\n plt.draw()\n plt.pause(0.001)\n\n # Train neural net dynamics to match data.\n init_params = init_nn_params(0.1, layer_sizes=[D, 150, D])\n optimized_params = adam(grad(train_loss), init_params, num_iters=1000, callback=callback)\n"
},
{
"path": "examples/print_trace.py",
"content": "\"\"\"Demonstrates how to use the tracer module, independent of autodiff, by\ncreating a trace that prints out functions and their arguments as they're being\nevaluated\"\"\"\n\nimport autograd.numpy as np # autograd has already wrapped numpy for us\nfrom autograd.tracer import Node, trace\n\n\nclass PrintNode(Node):\n def __init__(self, value, fun, args, kwargs, parent_argnums, parents):\n self.varname_generator = parents[0].varname_generator\n self.varname = next(self.varname_generator)\n args_or_vars = list(args)\n for argnum, parent in zip(parent_argnums, parents):\n args_or_vars[argnum] = parent.varname\n print(\"{} = {}({}) = {}\".format(self.varname, fun.__name__, \",\".join(map(str, args_or_vars)), value))\n\n def initialize_root(self, x):\n self.varname_generator = make_varname_generator()\n self.varname = next(self.varname_generator)\n print(f\"{self.varname} = {x}\")\n\n\ndef make_varname_generator():\n for i in range(65, 91):\n yield chr(i)\n raise Exception(\"Ran out of alphabet!\")\n\n\ndef print_trace(f, x):\n start_node = PrintNode.new_root(x)\n trace(start_node, f, x)\n print()\n\n\ndef avg(x, y):\n return (x + y) / 2\n\n\ndef fun(x):\n y = np.sin(x + x)\n return avg(y, y)\n\n\nprint_trace(fun, 1.23)\n\n# Traces can be nested, so we can also trace through grad(fun)\nfrom autograd import grad\n\nprint_trace(grad(fun), 1.0)\n"
},
{
"path": "examples/rkhs.py",
"content": "\"\"\"\nInferring a function from a reproducing kernel Hilbert space (RKHS) by taking\ngradients of eval with respect to the function-valued argument\n\"\"\"\n\nfrom itertools import chain\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.extend import Box, VSpace, defvjp, primitive\nfrom autograd.util import func\n\n\nclass RKHSFun:\n def __init__(self, kernel, alphas={}):\n self.alphas = alphas\n self.kernel = kernel\n self.vs = RKHSFunVSpace(self)\n\n @primitive\n def __call__(self, x):\n return sum([a * self.kernel(x, x_repr) for x_repr, a in self.alphas.items()], 0.0)\n\n def __add__(self, f):\n return self.vs.add(self, f)\n\n def __mul__(self, a):\n return self.vs.scalar_mul(self, a)\n\n\n# TODO: add vjp of __call__ wrt x (and show it in action)\ndefvjp(func(RKHSFun.__call__), lambda ans, f, x: lambda g: RKHSFun(f.kernel, {x: 1}) * g)\n\n\nclass RKHSFunBox(Box, RKHSFun):\n @property\n def kernel(self):\n return self._value.kernel\n\n\nRKHSFunBox.register(RKHSFun)\n\n\nclass RKHSFunVSpace(VSpace):\n def __init__(self, value):\n self.kernel = value.kernel\n\n def zeros(self):\n return RKHSFun(self.kernel)\n\n def randn(self):\n # These arbitrary vectors are not analogous to randn in any meaningful way\n N = npr.randint(1, 3)\n return RKHSFun(self.kernel, dict(zip(npr.randn(N), npr.randn(N))))\n\n def _add(self, f, g):\n assert f.kernel is g.kernel\n return RKHSFun(f.kernel, add_dicts(f.alphas, g.alphas))\n\n def _scalar_mul(self, f, a):\n return RKHSFun(f.kernel, {x: a * a_cur for x, a_cur in f.alphas.items()})\n\n def _inner_prod(self, f, g):\n assert f.kernel is g.kernel\n return sum(\n [a1 * a2 * f.kernel(x1, x2) for x1, a1 in f.alphas.items() for x2, a2 in g.alphas.items()], 0.0\n )\n\n\nRKHSFunVSpace.register(RKHSFun)\n\n\ndef add_dicts(d1, d2):\n d = {}\n for k, v in chain(d1.items(), d2.items()):\n d[k] = d[k] + v if k in d else v\n return d\n\n\nif __name__ == \"__main__\":\n\n def sq_exp_kernel(x1, x2):\n return np.exp(-((x1 - x2) ** 2))\n\n xs = range(5)\n ys = [1, 2, 3, 2, 1]\n\n def logprob(f, xs, ys):\n return -sum((f(x) - y) ** 2 for x, y in zip(xs, ys))\n\n f = RKHSFun(sq_exp_kernel)\n for i in range(100):\n f = f + grad(logprob)(f, xs, ys) * 0.01\n\n for x, y in zip(xs, ys):\n print(f\"{x}\\t{y}\\t{f(x)}\")\n"
},
{
"path": "examples/rnn.py",
"content": "\"\"\"Implements the long-short term memory character model.\nThis version vectorizes over multiple examples, but each string\nhas a fixed length.\"\"\"\n\nfrom os.path import dirname, join\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.misc.optimizers import adam\nfrom autograd.scipy.special import logsumexp\n\n### Helper functions #################\n\n\ndef sigmoid(x):\n return 0.5 * (np.tanh(x) + 1.0) # Output ranges from 0 to 1.\n\n\ndef concat_and_multiply(weights, *args):\n cat_state = np.hstack(args + (np.ones((args[0].shape[0], 1)),))\n return np.dot(cat_state, weights)\n\n\n### Define recurrent neural net #######\n\n\ndef create_rnn_params(input_size, state_size, output_size, param_scale=0.01, rs=npr.RandomState(0)):\n return {\n \"init hiddens\": rs.randn(1, state_size) * param_scale,\n \"change\": rs.randn(input_size + state_size + 1, state_size) * param_scale,\n \"predict\": rs.randn(state_size + 1, output_size) * param_scale,\n }\n\n\ndef rnn_predict(params, inputs):\n def update_rnn(input, hiddens):\n return np.tanh(concat_and_multiply(params[\"change\"], input, hiddens))\n\n def hiddens_to_output_probs(hiddens):\n output = concat_and_multiply(params[\"predict\"], hiddens)\n return output - logsumexp(output, axis=1, keepdims=True) # Normalize log-probs.\n\n num_sequences = inputs.shape[1]\n hiddens = np.repeat(params[\"init hiddens\"], num_sequences, axis=0)\n output = [hiddens_to_output_probs(hiddens)]\n\n for input in inputs: # Iterate over time steps.\n hiddens = update_rnn(input, hiddens)\n output.append(hiddens_to_output_probs(hiddens))\n return output\n\n\ndef rnn_log_likelihood(params, inputs, targets):\n logprobs = rnn_predict(params, inputs)\n loglik = 0.0\n num_time_steps, num_examples, _ = inputs.shape\n for t in range(num_time_steps):\n loglik += np.sum(logprobs[t] * targets[t])\n return loglik / (num_time_steps * num_examples)\n\n\n### Dataset setup ##################\n\n\ndef string_to_one_hot(string, maxchar):\n \"\"\"Converts an ASCII string to a one-of-k encoding.\"\"\"\n ascii = np.array([ord(c) for c in string]).T\n return np.array(ascii[:, None] == np.arange(maxchar)[None, :], dtype=int)\n\n\ndef one_hot_to_string(one_hot_matrix):\n return \"\".join([chr(np.argmax(c)) for c in one_hot_matrix])\n\n\ndef build_dataset(filename, sequence_length, alphabet_size, max_lines=-1):\n \"\"\"Loads a text file, and turns each line into an encoded sequence.\"\"\"\n with open(filename) as f:\n content = f.readlines()\n content = content[:max_lines]\n content = [line for line in content if len(line) > 2] # Remove blank lines\n seqs = np.zeros((sequence_length, len(content), alphabet_size))\n for ix, line in enumerate(content):\n padded_line = (line + \" \" * sequence_length)[:sequence_length]\n seqs[:, ix, :] = string_to_one_hot(padded_line, alphabet_size)\n return seqs\n\n\nif __name__ == \"__main__\":\n num_chars = 128\n\n # Learn to predict our own source code.\n text_filename = join(dirname(__file__), \"rnn.py\")\n train_inputs = build_dataset(text_filename, sequence_length=30, alphabet_size=num_chars, max_lines=60)\n\n init_params = create_rnn_params(input_size=128, output_size=128, state_size=40, param_scale=0.01)\n\n def print_training_prediction(weights):\n print(\"Training text Predicted text\")\n logprobs = np.asarray(rnn_predict(weights, train_inputs))\n for t in range(logprobs.shape[1]):\n training_text = one_hot_to_string(train_inputs[:, t, :])\n predicted_text = one_hot_to_string(logprobs[:, t, :])\n print(training_text.replace(\"\\n\", \" \") + \"|\" + predicted_text.replace(\"\\n\", \" \"))\n\n def training_loss(params, iter):\n return -rnn_log_likelihood(params, train_inputs, train_inputs)\n\n def callback(weights, iter, gradient):\n if iter % 10 == 0:\n print(\"Iteration\", iter, \"Train loss:\", training_loss(weights, 0))\n print_training_prediction(weights)\n\n # Build gradient of loss function using autograd.\n training_loss_grad = grad(training_loss)\n\n print(\"Training RNN...\")\n trained_params = adam(training_loss_grad, init_params, step_size=0.1, num_iters=1000, callback=callback)\n\n print()\n print(\"Generating text from RNN...\")\n num_letters = 30\n for t in range(20):\n text = \"\"\n for i in range(num_letters):\n seqs = string_to_one_hot(text, num_chars)[:, np.newaxis, :]\n logprobs = rnn_predict(trained_params, seqs)[-1].ravel()\n text += chr(npr.choice(len(logprobs), p=np.exp(logprobs)))\n print(text)\n"
},
{
"path": "examples/rosenbrock.py",
"content": "from scipy.optimize import minimize\n\nimport autograd.numpy as np\nfrom autograd import value_and_grad\n\n\ndef rosenbrock(x):\n return 100 * (x[1] - x[0] ** 2) ** 2 + (1 - x[0]) ** 2\n\n\n# Build a function that also returns gradients using autograd.\nrosenbrock_with_grad = value_and_grad(rosenbrock)\n\n# Optimize using conjugate gradients.\nresult = minimize(rosenbrock_with_grad, x0=np.array([0.0, 0.0]), jac=True, method=\"CG\")\nprint(f\"Found minimum at {result.x}\")\n"
},
{
"path": "examples/sinusoid.py",
"content": "import matplotlib.pyplot as plt\n\nimport autograd.numpy as np\nfrom autograd import grad\n\n\ndef fun(x):\n return np.sin(x)\n\n\nd_fun = grad(fun) # First derivative\ndd_fun = grad(d_fun) # Second derivative\n\nx = np.linspace(-10, 10, 100)\nplt.plot(x, list(map(fun, x)), x, list(map(d_fun, x)), x, list(map(dd_fun, x)))\n\nplt.xlim([-10, 10])\nplt.ylim([-1.2, 1.2])\nplt.axis(\"off\")\nplt.savefig(\"sinusoid.png\")\nplt.clf()\n\n\n# Taylor approximation to sin function\ndef fun(x):\n currterm = x\n ans = currterm\n for i in range(1000):\n print(i, end=\" \")\n currterm = -currterm * x**2 / ((2 * i + 3) * (2 * i + 2))\n ans = ans + currterm\n if np.abs(currterm) < 0.2:\n break # (Very generous tolerance!)\n\n return ans\n\n\nd_fun = grad(fun)\ndd_fun = grad(d_fun)\n\nx = np.linspace(-10, 10, 100)\nplt.plot(x, list(map(fun, x)), x, list(map(d_fun, x)), x, list(map(dd_fun, x)))\n\nplt.xlim([-10, 10])\nplt.ylim([-1.2, 1.2])\nplt.axis(\"off\")\nplt.savefig(\"sinusoid_taylor.png\")\nplt.clf()\n"
},
{
"path": "examples/tanh.py",
"content": "import matplotlib.pyplot as plt\n\nimport autograd.numpy as np\nfrom autograd import elementwise_grad as egrad\n\n\"\"\"\nMathematically we can only take gradients of scalar-valued functions, but\nautograd's elementwise_grad function also handles numpy's familiar vectorization\nof scalar functions, which is used in this example.\n\nTo be precise, elementwise_grad(fun)(x) always returns the value of a\nvector-Jacobian product, where the Jacobian of fun is evaluated at x and the\nvector is an all-ones vector with the same size as the output of fun. When\nvectorizing a scalar-valued function over many arguments, the Jacobian of the\noverall vector-to-vector mapping is diagonal, and so this vector-Jacobian\nproduct simply returns the diagonal elements of the Jacobian, which is the\n(elementwise) gradient of the function at each input value over which the\nfunction is vectorized.\n\"\"\"\n\n\ndef tanh(x):\n return (1.0 - np.exp(-2 * x)) / (1.0 + np.exp(-(2 * x)))\n\n\n### Plotting\nplt.figure(figsize=(12, 8))\nx = np.linspace(-7, 7, 700)\nplt.plot(x, tanh(x), label=\"tanh(x)\")\nplt.plot(x, egrad(tanh)(x), label=\"1st derivative\")\nplt.plot(x, egrad(egrad(tanh))(x), label=\"2nd derivative\")\nplt.plot(x, egrad(egrad(egrad(tanh)))(x), label=\"3rd derivative\")\nplt.plot(x, egrad(egrad(egrad(egrad(tanh))))(x), label=\"4th derivative\")\nplt.xlabel(\"x\")\nplt.ylabel(\"y\")\nplt.ylim(-5, 5)\nplt.yticks(np.arange(-5, 6, 1))\nplt.legend()\nplt.grid(True)\nplt.title(\"tanh(x) and its derivatives\")\nplt.savefig(\"tanh.png\")\nplt.show()\n"
},
{
"path": "examples/variational_autoencoder.py",
"content": "# Implements auto-encoding variational Bayes.\n\nfrom data import load_mnist, save_images\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nimport autograd.scipy.stats.norm as norm\nfrom autograd import grad\nfrom autograd.misc.optimizers import adam\nfrom autograd.scipy.special import expit as sigmoid\n\n\ndef diag_gaussian_log_density(x, mu, log_std):\n return np.sum(norm.logpdf(x, mu, np.exp(log_std)), axis=-1)\n\n\ndef unpack_gaussian_params(params):\n # Params of a diagonal Gaussian.\n D = np.shape(params)[-1] // 2\n mean, log_std = params[:, :D], params[:, D:]\n return mean, log_std\n\n\ndef sample_diag_gaussian(mean, log_std, rs):\n return rs.randn(*mean.shape) * np.exp(log_std) + mean\n\n\ndef bernoulli_log_density(targets, unnormalized_logprobs):\n # unnormalized_logprobs are in R\n # Targets must be -1 or 1\n label_probabilities = -np.logaddexp(0, -unnormalized_logprobs * targets)\n return np.sum(label_probabilities, axis=-1) # Sum across pixels.\n\n\ndef relu(x):\n return np.maximum(0, x)\n\n\ndef init_net_params(scale, layer_sizes, rs=npr.RandomState(0)):\n \"\"\"Build a (weights, biases) tuples for all layers.\"\"\"\n return [\n (\n scale * rs.randn(m, n), # weight matrix\n scale * rs.randn(n),\n ) # bias vector\n for m, n in zip(layer_sizes[:-1], layer_sizes[1:])\n ]\n\n\ndef batch_normalize(activations):\n mbmean = np.mean(activations, axis=0, keepdims=True)\n return (activations - mbmean) / (np.std(activations, axis=0, keepdims=True) + 1)\n\n\ndef neural_net_predict(params, inputs):\n \"\"\"Params is a list of (weights, bias) tuples.\n inputs is an (N x D) matrix.\n Applies batch normalization to every layer but the last.\"\"\"\n for W, b in params[:-1]:\n outputs = batch_normalize(np.dot(inputs, W) + b) # linear transformation\n inputs = relu(outputs) # nonlinear transformation\n outW, outb = params[-1]\n outputs = np.dot(inputs, outW) + outb\n return outputs\n\n\ndef nn_predict_gaussian(params, inputs):\n # Returns means and diagonal variances\n return unpack_gaussian_params(neural_net_predict(params, inputs))\n\n\ndef generate_from_prior(gen_params, num_samples, noise_dim, rs):\n latents = rs.randn(num_samples, noise_dim)\n return sigmoid(neural_net_predict(gen_params, latents))\n\n\ndef p_images_given_latents(gen_params, images, latents):\n preds = neural_net_predict(gen_params, latents)\n return bernoulli_log_density(images, preds)\n\n\ndef vae_lower_bound(gen_params, rec_params, data, rs):\n # We use a simple Monte Carlo estimate of the KL\n # divergence from the prior.\n q_means, q_log_stds = nn_predict_gaussian(rec_params, data)\n latents = sample_diag_gaussian(q_means, q_log_stds, rs)\n q_latents = diag_gaussian_log_density(latents, q_means, q_log_stds)\n p_latents = diag_gaussian_log_density(latents, 0, 0)\n likelihood = p_images_given_latents(gen_params, data, latents)\n return np.mean(p_latents + likelihood - q_latents)\n\n\nif __name__ == \"__main__\":\n # Model hyper-parameters\n latent_dim = 10\n data_dim = 784 # How many pixels in each image (28x28).\n gen_layer_sizes = [latent_dim, 300, 200, data_dim]\n rec_layer_sizes = [data_dim, 200, 300, latent_dim * 2]\n\n # Training parameters\n param_scale = 0.01\n batch_size = 200\n num_epochs = 15\n step_size = 0.001\n\n print(\"Loading training data...\")\n N, train_images, _, test_images, _ = load_mnist()\n\n def binarise(images):\n on = images > 0.5\n images = images * 0 - 1\n images[on] = 1.0\n return images\n\n print(\"Binarising training data...\")\n train_images = binarise(train_images)\n test_images = binarise(test_images)\n\n init_gen_params = init_net_params(param_scale, gen_layer_sizes)\n init_rec_params = init_net_params(param_scale, rec_layer_sizes)\n combined_init_params = (init_gen_params, init_rec_params)\n\n num_batches = int(np.ceil(len(train_images) / batch_size))\n\n def batch_indices(iter):\n idx = iter % num_batches\n return slice(idx * batch_size, (idx + 1) * batch_size)\n\n # Define training objective\n seed = npr.RandomState(0)\n\n def objective(combined_params, iter):\n data_idx = batch_indices(iter)\n gen_params, rec_params = combined_params\n return -vae_lower_bound(gen_params, rec_params, train_images[data_idx], seed) / data_dim\n\n # Get gradients of objective using autograd.\n objective_grad = grad(objective)\n\n print(\" Epoch | Objective | Test ELBO \")\n\n def print_perf(combined_params, iter, grad):\n if iter % 10 == 0:\n gen_params, rec_params = combined_params\n bound = np.mean(objective(combined_params, iter))\n message = f\"{iter // num_batches:15}|{bound:20}|\"\n if iter % 100 == 0:\n test_bound = -vae_lower_bound(gen_params, rec_params, test_images, seed) / data_dim\n message += f\"{test_bound:20}\"\n print(message)\n\n fake_data = generate_from_prior(gen_params, 20, latent_dim, seed)\n save_images(fake_data, \"vae_samples.png\", vmin=0, vmax=1)\n\n # The optimizers provided can optimize lists, tuples, or dicts of parameters.\n optimized_params = adam(\n objective_grad,\n combined_init_params,\n step_size=step_size,\n num_iters=num_epochs * num_batches,\n callback=print_perf,\n )\n"
},
{
"path": "license.txt",
"content": "The MIT License (MIT)\n\nCopyright (c) 2025 by the President and Fellows of Harvard University\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
},
{
"path": "noxfile.py",
"content": "import platform\n\nimport nox\n\nNIGHTLY_INDEX_URL = \"https://pypi.anaconda.org/scientific-python-nightly-wheels/simple\"\nUV_NIGHTLY_ENV_VARS = {\n \"UV_INDEX_URL\": NIGHTLY_INDEX_URL,\n \"UV_PRERELEASE\": \"allow\",\n \"UV_INDEX_STRATEGY\": \"first-index\",\n}\n\nnox.needs_version = \">=2024.4.15\"\nnox.options.default_venv_backend = \"uv|virtualenv\"\nnox.options.reuse_existing_virtualenvs = False\nnox.options.error_on_external_run = True\n# nox.options.sessions = [\"lint\", \"validate-package\", \"tests\"]\nnox.options.sessions = [\"tests\"]\n\n\n@nox.session(name=\"validate-package\")\ndef check(session):\n \"\"\"Build source distribution, wheel, and check their metadata\"\"\"\n session.install(\"build\", \"twine\", silent=False)\n session.run(\"python\", \"-m\", \"build\")\n session.run(\"twine\", \"check\", \"--strict\", \"dist/*\")\n\n\n@nox.session(name=\"tests\", tags=[\"tests\"])\ndef run_tests(session):\n \"\"\"Run unit tests and generate a coverage report\"\"\"\n pyproject = nox.project.load_toml(\"pyproject.toml\")\n session.install(*nox.project.dependency_groups(pyproject, \"test\"))\n # SciPy doesn't have wheels on PyPy\n if platform.python_implementation() == \"PyPy\":\n session.install(\"-e.\", silent=False)\n else:\n session.install(\"-e\", \".[scipy]\", silent=False)\n session.run(\"pytest\", \"--cov=autograd\", \"--cov-report=xml\", \"--cov-append\", *session.posargs)\n\n\n@nox.session(name=\"lint\", reuse_venv=True)\ndef ruff(session):\n \"\"\"Lightning-fast linting for Python\"\"\"\n session.install(\"pre-commit\", silent=False)\n session.run(\"pre-commit\", \"run\", \"--all-files\", \"--show-diff-on-failure\")\n\n\n@nox.session(name=\"nightly-tests\", tags=[\"tests\"])\ndef run_nightly_tests(session):\n \"\"\"Run tests against nightly versions of dependencies\"\"\"\n session.install(\"-e.\", silent=False)\n pyproject = nox.project.load_toml(\"pyproject.toml\")\n session.install(*nox.project.dependency_groups(pyproject, \"test\"))\n # SciPy doesn't have wheels on PyPy\n if platform.python_implementation() == \"PyPy\":\n session.install(\n \"numpy\", \"--upgrade\", \"--only-binary\", \":all:\", silent=False, env=UV_NIGHTLY_ENV_VARS\n )\n else:\n session.install(\n \"numpy\", \"scipy\", \"--upgrade\", \"--only-binary\", \":all:\", silent=False, env=UV_NIGHTLY_ENV_VARS\n )\n session.run(\"pytest\", \"--cov=autograd\", \"--cov-report=xml\", \"--cov-append\", *session.posargs)\n"
},
{
"path": "pyproject.toml",
"content": "[build-system]\nrequires = [\"hatchling\"]\nbuild-backend = \"hatchling.build\"\n\n[project]\nname = \"autograd\"\nversion = \"1.8.0\"\nrequires-python = \">=3.10\"\ndescription = \"Efficiently computes derivatives of NumPy code.\"\nreadme = \"README.md\"\nlicense = {file = \"license.txt\"}\nauthors = [\n {name = \"Dougal Maclaurin\", email = \"maclaurin@physics.harvard.edu\"},\n {name = \"David Duvenaud\", email = \"duvenaud@cs.toronto.edu\"},\n {name = \"Matthew Johnson\", email = \"mattjj@csail.mit.edu\"},\n {name = \"Jamie Townsend\", email = \"j.h.n.townsend@uva.nl\"},\n]\nmaintainers = [\n {name = \"Jamie Townsend\", email = \"j.h.n.townsend@uva.nl\"},\n {name = \"Fabian Joswig\", email = \"fabian.joswig@uni-muenster.de\"},\n {name = \"Agriya Khetarpal\", email = \"agriyakhetarpal@outlook.com\"},\n]\nclassifiers = [\n \"Development Status :: 4 - Beta\",\n \"Intended Audience :: Information Technology\",\n \"Intended Audience :: Science/Research\",\n \"License :: OSI Approved :: MIT License\",\n \"Programming Language :: Python :: 3.10\",\n \"Programming Language :: Python :: 3.11\",\n \"Programming Language :: Python :: 3.12\",\n \"Programming Language :: Python :: 3.13\",\n \"Programming Language :: Python :: 3.14\",\n \"Topic :: Scientific/Engineering\",\n]\nkeywords = [\n \"Automatic differentiation\",\n \"backpropagation\",\n \"gradients\",\n \"machine learning\",\n \"optimization\",\n \"neural networks\",\n \"Python\",\n \"NumPy\",\n \"SciPy\",\n]\ndependencies = [\n \"numpy<3\",\n]\n# dynamic = [\"version\"]\n\n[project.urls]\nSource = \"https://github.com/HIPS/autograd\"\n\n[project.optional-dependencies]\nscipy = [\n \"scipy\",\n]\n\n[dependency-groups]\ntest = [\n \"pytest\",\n \"pytest-cov\",\n \"pytest-xdist\",\n]\nexamples = [\"matplotlib\"]\n\n[tool.coverage.run]\nsource = [\"autograd\"]\n\n[tool.coverage.report]\nshow_missing = true\n\n[tool.pytest.ini_options]\nrequired_plugins = [\"pytest-cov\", \"pytest-xdist\"]\n# TODO: generate HTML report, upload to CodeCov\naddopts = \"--color=yes -sra -n auto --cov=autograd --cov-report=xml --cov-report=term\"\n\n[tool.ruff]\nextend-exclude = []\n# TODO: not ignore them\nlint.extend-ignore = [\n \"E731\",\n \"F401\",\n \"F403\",\n \"F841\",\n \"F821\",\n \"E721\",\n \"E722\",\n \"E741\",\n \"E402\",\n \"F811\"\n]\nlint.extend-select = [\"I\", \"W\"]\nline-length = 109\n"
},
{
"path": "tests/_test_complexity.py",
"content": "import time\nimport warnings\n\nimport autograd.numpy as np\nfrom autograd import deriv, grad\nfrom autograd.builtins import list as make_list\n\n\ndef timefunction(f):\n t = time.time()\n f()\n return time.time() - t\n\n\ndef assert_linear_time(f):\n t = timefunction(lambda: f(1))\n t10 = timefunction(lambda: f(10))\n assert t10 > 5 * t, f\"Too fast: f(1) takes {t}, f(10) takes {t10}\"\n assert t10 < 20 * t, f\"Too slow: f(1) takes {t}, f(10) takes {t10}\"\n if not (8 * t < t10 < 12 * t):\n warnings.warn(\"Borderline linearity. May fail on different hardware\")\n\n\ndef test_array_creation():\n def fun(x, N):\n arr = [x for i in range(N)]\n return np.sum(np.array(arr))\n\n assert_linear_time(lambda N: grad(fun)(1.0, 200 * N))\n\n\ndef test_array_indexing():\n def fun(x):\n return sum([x[i] for i in range(len(x))])\n\n assert_linear_time(lambda N: grad(fun)(np.zeros(200 * N)))\n\n\ndef test_list_indexing():\n def fun(x):\n return sum([x[i] for i in range(len(x))])\n\n assert_linear_time(lambda N: grad(fun)([0.0 for i in range(50 * N)]))\n\n\ndef test_list_creation():\n def fun(x, N):\n return make_list(*[x for _ in range(N)])\n\n assert_linear_time(lambda N: deriv(fun)(0.0, 20 * N))\n\n\n# This fails. Need to figure out why\ndef test_array_creation_fwd():\n def fun(x, N):\n arr = [x for i in range(N)]\n return np.sum(np.array(arr))\n\n assert_linear_time(lambda N: deriv(fun)(1.0, 400 * N))\n"
},
{
"path": "tests/check_examples_run.sh",
"content": "#!/bin/bash\n\nPYTHONPATH=\".:$PYTHONPATH\"\ntrap 'kill -INT -$pid && exit 1' INT\n\nworking=()\nfailing=()\n\nexamples=$(find examples -name '*.py' -not -name '__init__.py')\n\necho 'Running all the examples...'\nfor f in $examples; do\n timeout 15s python2 $f > /dev/null 2>&1 & pid=$!\n wait $pid\n status=$?\n if [ $status -eq 0 -o $status -eq 124 ]; then\n echo $f \"seems to work\"\n working+=($f)\n elif [ $status -eq 137 ]; then\n echo $f \"might be working, but had to be killed\"\n working+=($f)\n else\n echo $f \"seems broken, try running manually\"\n failing+=($f)\n fi\ndone\n\nif [ ! ${#working[@]} -eq 0 ]; then\n echo -e '\\033[01;36m'\n echo \"These seemed to WORK:\"\n echo -en '\\033[00m'\n printf '%s\\n' \"${working[@]}\"\n echo\nfi\nif [ ! ${#failing[@]} -eq 0 ]; then\n echo -e '\\033[01;31m'\n echo \"These seemed to FAIL:\"\n echo -en '\\033[00m'\n printf '%s\\n' \"${failing[@]}\"\n echo\nfi\n"
},
{
"path": "tests/conftest.py",
"content": "import numpy as np\nimport pytest\n\n\n@pytest.fixture(autouse=True)\ndef random_seed():\n np.random.seed(42)\n"
},
{
"path": "tests/numpy_utils.py",
"content": "import autograd.numpy.random as npr\nfrom autograd.test_util import combo_check\n\n\ndef stat_check(fun, test_complex=True, **kwargs):\n # Tests functions that compute statistics, like sum, mean, etc\n x = 3.5\n A = npr.randn()\n B = npr.randn(3)\n C = npr.randn(2, 3)\n D = npr.randn(1, 3)\n check = combo_check(fun, (0,), **kwargs)\n check([x, A])\n check([B, C, D], axis=[None, 0], keepdims=[True, False])\n check([C, D], axis=[None, 0, 1], keepdims=[True, False])\n if test_complex:\n c = npr.randn() + 0.1j * npr.randn()\n E = npr.randn(2, 3) + 0.1j * npr.randn(2, 3)\n check([x, c, A])\n check([B, C, D, E], axis=[None, 0], keepdims=[True, False])\n\n\ndef unary_ufunc_check(fun, lims=[-2, 2], test_complex=True, **kwargs):\n scalar = transform(lims, 0.4)\n vector = transform(lims, npr.rand(2))\n mat = transform(lims, npr.rand(3, 2))\n mat2 = transform(lims, npr.rand(1, 2))\n check = combo_check(fun, (0,), **kwargs)\n check([scalar, vector, mat, mat2])\n if test_complex:\n comp = transform(lims, 0.4) + 0.1j * transform(lims, 0.3)\n matc = transform(lims, npr.rand(3, 2)) + 0.1j * npr.rand(3, 2)\n check([comp, matc])\n\n\ndef binary_ufunc_check(fun, lims_A=[-2, 2], lims_B=[-2, 2], test_complex=True, **kwargs):\n T_A = lambda x: transform(lims_A, x)\n T_B = lambda x: transform(lims_B, x)\n scalar = 0.6\n vector = npr.rand(2)\n mat = npr.rand(3, 2)\n mat2 = npr.rand(1, 2)\n check = combo_check(fun, (0, 1), **kwargs)\n check([T_A(scalar), T_A(vector), T_A(mat), T_A(mat2)], [T_B(scalar), T_B(vector), T_B(mat), T_B(mat2)])\n if test_complex:\n comp = 0.6 + 0.3j\n matc = npr.rand(3, 2) + 0.1j * npr.rand(3, 2)\n check(\n [T_A(scalar), T_A(comp), T_A(vector), T_A(matc), T_A(mat2)],\n [T_B(scalar), T_B(comp), T_B(vector), T_B(matc), T_B(mat2)],\n )\n\n\ndef binary_ufunc_check_no_same_args(fun, lims_A=[-2, 2], lims_B=[-2, 2], test_complex=True, **kwargs):\n T_A = lambda x: transform(lims_A, x)\n T_B = lambda x: transform(lims_B, x)\n scalar1 = 0.6\n scalar2 = 0.7\n vector1 = npr.rand(2)\n vector2 = npr.rand(2)\n mat11 = npr.rand(3, 2)\n mat12 = npr.rand(3, 2)\n mat21 = npr.rand(1, 2)\n mat22 = npr.rand(1, 2)\n check = combo_check(fun, (0, 1), **kwargs)\n check(\n [T_A(scalar1), T_A(vector1), T_A(mat11), T_A(mat21)],\n [T_B(scalar2), T_B(vector2), T_B(mat12), T_B(mat22)],\n )\n if test_complex:\n comp1 = 0.6 + 0.3j\n comp2 = 0.1 + 0.2j\n matc1 = npr.rand(3, 2) + 0.1j * npr.rand(3, 2)\n matc2 = npr.rand(3, 2) + 0.1j * npr.rand(3, 2)\n check(\n [T_A(scalar1), T_A(comp1), T_A(vector1), T_A(matc1), T_A(mat21)],\n [T_B(scalar2), T_B(comp2), T_B(vector2), T_B(matc2), T_B(mat22)],\n )\n\n\ndef transform(lims, x):\n return x * (lims[1] - lims[0]) + lims[0]\n"
},
{
"path": "tests/profiling.py",
"content": "from contextlib import contextmanager\nfrom time import time\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\n\n\n@contextmanager\ndef tictoc(text=\"\"):\n print(\"--- Start clock ---\")\n t1 = time()\n yield\n dt = time() - t1\n print(f\"--- Stop clock {text}: {dt} seconds elapsed ---\")\n\n\ndef fan_out_fan_in():\n \"\"\"The 'Pearlmutter test'\"\"\"\n\n def fun(x):\n for i in range(10**4):\n x = (x + x) / 2.0\n return np.sum(x)\n\n with tictoc():\n grad(fun)(1.0)\n\n\ndef convolution():\n # MNIST-scale convolution operation\n import autograd.scipy.signal\n\n convolve = autograd.scipy.signal.convolve\n dat = npr.randn(256, 3, 28, 28)\n kernel = npr.randn(3, 5, 5)\n with tictoc():\n convolve(dat, kernel, axes=([2, 3], [1, 2]), dot_axes=([1], [0]))\n\n\ndef dot_equivalent():\n # MNIST-scale convolution operation\n\n dat = npr.randn(256, 3, 24, 5, 24, 5)\n kernel = npr.randn(3, 5, 5)\n with tictoc():\n np.tensordot(dat, kernel, axes=[(1, 3, 5), (0, 1, 2)])\n\n\n# fan_out_fan_in()\n# convolution()\ndot_equivalent()\n"
},
{
"path": "tests/test_binary_ops.py",
"content": "import itertools as it\nimport warnings\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad, value_and_grad\nfrom autograd.test_util import check_grads\n\nrs = npr.RandomState(0)\n\n\ndef arg_pairs():\n scalar = 2.0\n vector = rs.randn(4)\n mat = rs.randn(3, 4)\n mat2 = rs.randn(1, 4)\n allargs = [scalar, vector, mat, mat2]\n yield from it.product(allargs, allargs)\n\n\ndef test_mul():\n fun = lambda x, y: x * y\n for arg1, arg2 in arg_pairs():\n check_grads(fun)(arg1, arg2)\n\n\ndef test_add():\n fun = lambda x, y: x + y\n for arg1, arg2 in arg_pairs():\n check_grads(fun)(arg1, arg2)\n\n\ndef test_sub():\n fun = lambda x, y: x - y\n for arg1, arg2 in arg_pairs():\n check_grads(fun)(arg1, arg2)\n\n\ndef test_div():\n fun = lambda x, y: x / y\n make_gap_from_zero = lambda x: np.sqrt(x**2 + 0.5)\n for arg1, arg2 in arg_pairs():\n arg1 = make_gap_from_zero(arg1)\n arg2 = make_gap_from_zero(arg2)\n check_grads(fun)(arg1, arg2)\n\n\ndef test_mod():\n fun = lambda x, y: x % y\n make_gap_from_zero = lambda x: np.sqrt(x**2 + 0.5)\n for arg1, arg2 in arg_pairs():\n if arg1 is not arg2: # Gradient undefined at x == y\n arg1 = make_gap_from_zero(arg1)\n arg2 = make_gap_from_zero(arg2)\n check_grads(fun)(arg1, arg2)\n\n\ndef test_pow():\n fun = lambda x, y: x**y\n make_positive = lambda x: np.abs(x) + 1.1 # Numeric derivatives fail near zero\n for arg1, arg2 in arg_pairs():\n arg1 = make_positive(arg1)\n check_grads(fun)(arg1, arg2)\n\n\ndef test_arctan2():\n for arg1, arg2 in arg_pairs():\n check_grads(np.arctan2)(arg1, arg2)\n\n\ndef test_hypot():\n for arg1, arg2 in arg_pairs():\n check_grads(np.hypot, modes=[\"rev\"])(arg1, arg2)\n\n\ndef test_comparison_grads():\n compare_funs = [\n lambda x, y: np.sum(x < x) + 0.0,\n lambda x, y: np.sum(x <= y) + 0.0,\n lambda x, y: np.sum(x > y) + 0.0,\n lambda x, y: np.sum(x >= y) + 0.0,\n lambda x, y: np.sum(x == y) + 0.0,\n lambda x, y: np.sum(x != y) + 0.0,\n ]\n\n with warnings.catch_warnings(record=True) as w:\n for arg1, arg2 in arg_pairs():\n zeros = (arg1 + arg2) * 0 # get correct shape\n for fun in compare_funs:\n assert np.all(grad(fun)(arg1, arg2) == zeros)\n assert np.all(grad(fun, argnum=1)(arg1, arg2) == zeros)\n\n\ndef test_comparison_values():\n compare_funs = [\n lambda x, y: np.sum(x < x) + 0.0,\n lambda x, y: np.sum(x <= y) + 0.0,\n lambda x, y: np.sum(x > y) + 0.0,\n lambda x, y: np.sum(x >= y) + 0.0,\n lambda x, y: np.sum(x == y) + 0.0,\n lambda x, y: np.sum(x != y) + 0.0,\n ]\n\n for arg1, arg2 in arg_pairs():\n for fun in compare_funs:\n fun_val = fun(arg1, arg2)\n fun_val_from_grad, _ = value_and_grad(fun)(arg1, arg2)\n assert fun_val == fun_val_from_grad, (fun_val, fun_val_from_grad)\n"
},
{
"path": "tests/test_builtins.py",
"content": "import autograd.numpy as np\nfrom autograd import grad\nfrom autograd.builtins import isinstance\n\n\ndef test_isinstance():\n def checker(ex, type_, truthval):\n assert isinstance(ex, type_) == truthval\n return 1.0\n\n examples = [\n [list, [[]], [()]],\n [np.ndarray, [np.zeros(1)], [[]]],\n [(tuple, list), [[], ()], [np.zeros(1)]],\n ]\n\n for type_, positive_examples, negative_examples in examples:\n for ex in positive_examples:\n checker(ex, type_, True)\n grad(checker)(ex, type_, True)\n\n for ex in negative_examples:\n checker(ex, type_, False)\n grad(checker)(ex, type_, False)\n"
},
{
"path": "tests/test_complex.py",
"content": "import autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n\ndef test_real_type():\n fun = lambda x: np.sum(np.real(x))\n df = grad(fun)\n assert np.isrealobj(df(2.0))\n assert np.iscomplexobj(df(1.0j))\n\n\ndef test_real_if_close_type():\n fun = lambda x: np.sum(np.real(x))\n df = grad(fun)\n assert np.isrealobj(df(1.0))\n assert np.iscomplexobj(df(1.0j))\n\n\ndef test_angle_real():\n fun = lambda x: np.angle(x)\n d_fun = lambda x: grad(fun)(x)\n check_grads(fun)(npr.rand())\n check_grads(d_fun)(npr.rand())\n\n\ndef test_angle_complex():\n fun = lambda x: np.angle(x)\n d_fun = lambda x: grad(fun)(x)\n check_grads(fun)(npr.rand() + 1j * npr.rand())\n check_grads(d_fun)(npr.rand() + 1j * npr.rand())\n\n\ndef test_abs_real():\n fun = lambda x: np.abs(x)\n d_fun = lambda x: grad(fun)(x)\n check_grads(fun)(1.1)\n check_grads(d_fun)(2.1)\n\n\ndef test_abs_complex():\n fun = lambda x: np.abs(x)\n d_fun = lambda x: grad(fun)(x)\n check_grads(fun)(1.1 + 1.2j)\n check_grads(d_fun)(1.1 + 1.3j)\n"
},
{
"path": "tests/test_core.py",
"content": "\"\"\"This file doesn't import the numpy wrapper, to check if core works\non basic operations even without numpy.\"\"\"\n\nimport warnings\n\nfrom autograd.core import make_vjp\nfrom autograd.wrap_util import unary_to_nary\n\n\n@unary_to_nary\ndef grad(fun, x):\n vjp, _ = make_vjp(fun, x)\n return vjp(1.0)\n\n\n# Non-numpy gradient checking functions.\ndef nd(f, x, eps=1e-4):\n return (f(x + eps / 2) - f(x - eps / 2)) / eps\n\n\ndef check_close(a, b, atol=1e-4, rtol=1e-4):\n assert abs(a - b) < atol + rtol * abs(b), f\"Diffs are: {a - b}\"\n\n\ndef check_binary_func(fun, independent=False):\n with warnings.catch_warnings(record=independent) as w:\n x, y = 0.7, 1.8\n a = grad(fun)(x, y)\n b = nd(lambda x: fun(x, y), x)\n check_close(a, b)\n\n a = grad(fun, 1)(x, y)\n b = nd(lambda y: fun(x, y), y)\n check_close(a, b)\n\n\ndef test_add():\n check_binary_func(lambda x, y: x + y)\n\n\ndef test_sub():\n check_binary_func(lambda x, y: x - y)\n\n\ndef test_div():\n check_binary_func(lambda x, y: x / y)\n\n\ndef test_mul():\n check_binary_func(lambda x, y: x * y)\n\n\ndef test_pow():\n check_binary_func(lambda x, y: x**y)\n\n\ndef test_mod():\n check_binary_func(lambda x, y: x % y)\n\n\ndef test_eq():\n check_binary_func(lambda x, y: x == y, independent=True)\n\n\ndef test_neq():\n check_binary_func(lambda x, y: x != y, independent=True)\n\n\ndef test_leq():\n check_binary_func(lambda x, y: x <= y, independent=True)\n\n\ndef test_geq():\n check_binary_func(lambda x, y: x >= y, independent=True)\n\n\ndef test_lt():\n check_binary_func(lambda x, y: x < y, independent=True)\n\n\ndef test_gt():\n check_binary_func(lambda x, y: x > y, independent=True)\n"
},
{
"path": "tests/test_dict.py",
"content": "import operator as op\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import dict as ag_dict\nfrom autograd import grad\nfrom autograd import isinstance as ag_isinstance\nfrom autograd.test_util import check_grads\n\nnpr.seed(0)\n\n\ndef test_getter():\n def fun(input_dict):\n A = np.sum(input_dict[\"item_1\"])\n B = np.sum(input_dict[\"item_2\"])\n C = np.sum(input_dict[\"item_2\"])\n return A + B + C\n\n d_fun = grad(fun)\n input_dict = {\"item_1\": npr.randn(5, 6), \"item_2\": npr.randn(4, 3), \"item_X\": npr.randn(2, 4)}\n\n result = d_fun(input_dict)\n assert np.allclose(result[\"item_1\"], np.ones((5, 6)))\n assert np.allclose(result[\"item_2\"], 2 * np.ones((4, 3)))\n assert np.allclose(result[\"item_X\"], np.zeros((2, 4)))\n\n\ndef test_grads():\n def fun(input_dict):\n A = np.sum(np.sin(input_dict[\"item_1\"]))\n B = np.sum(np.cos(input_dict[\"item_2\"]))\n return A + B\n\n def d_fun(input_dict):\n g = grad(fun)(input_dict)\n A = np.sum(g[\"item_1\"])\n B = np.sum(np.sin(g[\"item_1\"]))\n C = np.sum(np.sin(g[\"item_2\"]))\n return A + B + C\n\n input_dict = {\"item_1\": npr.randn(5, 6), \"item_2\": npr.randn(4, 3), \"item_X\": npr.randn(2, 4)}\n\n check_grads(fun)(input_dict)\n check_grads(d_fun)(input_dict)\n\n\ndef test_iter():\n def fun(input_dict):\n A = 0.0\n B = 0.0\n for i, k in enumerate(sorted(input_dict)):\n A = A + np.sum(np.sin(input_dict[k])) * (i + 1.0)\n B = B + np.sum(np.cos(input_dict[k]))\n return A + B\n\n def d_fun(input_dict):\n g = grad(fun)(input_dict)\n A = np.sum(g[\"item_1\"])\n B = np.sum(np.sin(g[\"item_1\"]))\n C = np.sum(np.sin(g[\"item_2\"]))\n return A + B + C\n\n input_dict = {\"item_1\": npr.randn(5, 6), \"item_2\": npr.randn(4, 3), \"item_X\": npr.randn(2, 4)}\n\n check_grads(fun)(input_dict)\n check_grads(d_fun)(input_dict)\n\n\ndef test_items_values_keys():\n def fun(input_dict):\n A = 0.0\n B = 0.0\n for i, (k, v) in enumerate(sorted(input_dict.items(), key=op.itemgetter(0))):\n A = A + np.sum(np.sin(v)) * (i + 1.0)\n B = B + np.sum(np.cos(v))\n for v in input_dict.values():\n A = A + np.sum(np.sin(v))\n for k in sorted(input_dict.keys()):\n A = A + np.sum(np.cos(input_dict[k]))\n return A + B\n\n def d_fun(input_dict):\n g = grad(fun)(input_dict)\n A = np.sum(g[\"item_1\"])\n B = np.sum(np.sin(g[\"item_1\"]))\n C = np.sum(np.sin(g[\"item_2\"]))\n return A + B + C\n\n input_dict = {\"item_1\": npr.randn(5, 6), \"item_2\": npr.randn(4, 3), \"item_X\": npr.randn(2, 4)}\n\n check_grads(fun)(input_dict)\n check_grads(d_fun)(input_dict)\n\n\ndef test_get():\n def fun(d, x):\n return d.get(\"item_1\", x) ** 2\n\n check_grads(fun, argnum=(0, 1))({\"item_1\": 3.0}, 2.0)\n check_grads(fun, argnum=(0, 1))({\"item_2\": 4.0}, 2.0)\n check_grads(fun, argnum=(0, 1))({}, 2.0)\n\n\ndef test_make_dict():\n def fun(x):\n return ag_dict([(\"a\", x)], b=x)\n\n check_grads(fun, modes=[\"rev\"])(1.0)\n\n def fun(x):\n return ag_dict({\"a\": x})\n\n check_grads(fun, modes=[\"rev\"])(1.0)\n\n # check some other forms of the constructor\n ag_dict()\n ag_dict(())\n ag_dict({})\n\n\ndef test_isinstance():\n def fun(x):\n assert ag_isinstance(x, dict)\n assert ag_isinstance(x, ag_dict)\n return x[\"x\"]\n\n fun({\"x\": 1.0})\n grad(fun)({\"x\": 1.0})\n"
},
{
"path": "tests/test_direct.py",
"content": "\"\"\"\nSet of tests that are as explicit as possible, in case the test helpers like\nautograd.test_util break and start letting everything pass\n\"\"\"\n\nimport numpy as onp\nimport pytest\n\nimport autograd.numpy as np\nfrom autograd import deriv, grad, holomorphic_grad\n\n\ndef test_grad():\n def fun(x):\n return (x + np.sin(x**2)) * x\n\n assert 3.190948746871 - 1e-6 < grad(fun)(1.3) < 3.190948746871 + 1e-6\n\n\ndef test_deriv():\n def fun(x):\n return (x + np.sin(x**2)) * x\n\n assert 3.190948746871 - 1e-6 < deriv(fun)(1.3) < 3.190948746871 + 1e-6\n\n\ndef test_grad_complex_output():\n def fun(x):\n return x * (1.0 + 0.2j)\n\n with pytest.raises(TypeError):\n grad(fun)(1.0)\n\n\ndef test_holomorphic_grad():\n def fun(x):\n return x * (1.0 + 0.2j)\n\n g = holomorphic_grad(fun)(1.0 + 0.0j)\n assert 0.9999 < onp.real(g) < 1.0001\n assert 0.1999 < onp.imag(g) < 0.2001\n"
},
{
"path": "tests/test_fft.py",
"content": "from functools import partial\n\nimport pytest\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n### fwd mode not yet implemented\ncheck_grads = partial(check_grads, modes=[\"rev\"])\n\n\ndef test_fft():\n def fun(x):\n return np.fft.fft(x)\n\n D = 5\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_fft_ortho():\n def fun(x):\n return np.fft.fft(x, norm=\"ortho\")\n\n D = 5\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_fft_axis():\n def fun(x):\n return np.fft.fft(x, axis=0)\n\n D = 5\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef match_complex(fft_fun, mat):\n # ensure hermitian by doing a fft\n if fft_fun.__name__.startswith(\"ir\"):\n return getattr(np.fft, fft_fun.__name__[1:])(mat)\n else:\n return mat\n\n\ndef check_fft_n(fft_fun, D, n):\n def fun(x):\n return fft_fun(x, D + n)\n\n mat = npr.randn(D, D)\n mat = match_complex(fft_fun, mat)\n check_grads(fun)(mat)\n\n\ndef test_fft_n_smaller():\n check_fft_n(np.fft.fft, 5, -2)\n\n\ndef test_fft_n_bigger():\n check_fft_n(np.fft.fft, 5, 2)\n\n\ndef test_ifft_n_smaller():\n check_fft_n(np.fft.ifft, 5, -2)\n\n\ndef test_ifft_n_bigger():\n check_fft_n(np.fft.ifft, 5, 2)\n\n\ndef test_rfft_n_smaller():\n check_fft_n(np.fft.rfft, 4, -2)\n\n\ndef test_rfft_n_bigger():\n check_fft_n(np.fft.rfft, 4, 2)\n\n\ndef test_irfft_n_smaller():\n check_fft_n(np.fft.irfft, 4, -2)\n\n\ndef test_irfft_n_bigger():\n check_fft_n(np.fft.irfft, 4, 2)\n\n\ndef check_fft_s(fft_fun, D):\n def fun(x):\n return fft_fun(x, s=s, axes=axes)\n\n mat = npr.randn(D, D, D) / 10.0\n mat = match_complex(fft_fun, mat)\n s = [D + 2, D - 2]\n axes = [0, 2]\n check_grads(fun)(mat)\n\n\ndef test_fft2_s():\n check_fft_s(np.fft.fft2, 5)\n\n\ndef test_ifft2_s():\n check_fft_s(np.fft.ifft2, 5)\n\n\ndef test_fftn_s():\n check_fft_s(np.fft.fftn, 5)\n\n\ndef test_ifftn_s():\n check_fft_s(np.fft.ifftn, 5)\n\n\ndef test_rfft2_s():\n check_fft_s(np.fft.rfft2, 4)\n\n\ndef test_irfft2_s():\n check_fft_s(np.fft.irfft2, 4)\n\n\ndef test_rfftn_s():\n check_fft_s(np.fft.rfftn, 4)\n\n\ndef test_irfftn_s():\n check_fft_s(np.fft.irfftn, 4)\n\n\n## TODO: fft gradient not implemented for repeated axes\n# def test_fft_repeated_axis():\n# D = 5\n# for fft_fun in (np.fft.fft2,np.fft.ifft2,np.fft.fftn, np.fft.ifftn):\n# def fun(x): return fft_fun(x, s=s, axes=axes)\n\n# mat = npr.randn(D,D,D) / 10.0\n# s = [D + 2, D - 2]\n# axes = [0,0]\n\n# check_grads(rad)(fun)\n\n\ndef test_ifft():\n def fun(x):\n return np.fft.ifft(x)\n\n D = 5\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_fft2():\n def fun(x):\n return np.fft.fft2(x)\n\n D = 5\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_ifft2():\n def fun(x):\n return np.fft.ifft2(x)\n\n D = 5\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_fftn():\n def fun(x):\n return np.fft.fftn(x)\n\n D = 5\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_ifftn():\n def fun(x):\n return np.fft.ifftn(x)\n\n D = 5\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_rfft():\n def fun(x):\n return np.fft.rfft(x)\n\n D = 4\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_rfft_ortho():\n def fun(x):\n return np.fft.rfft(x, norm=\"ortho\")\n\n D = 4\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_rfft_axes():\n def fun(x):\n return np.fft.rfft(x, axis=0)\n\n D = 4\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_irfft():\n def fun(x):\n return np.fft.irfft(x)\n\n D = 4\n mat = npr.randn(D, D) / 10.0\n # ensure hermitian by doing a fft\n mat = np.fft.rfft(mat)\n check_grads(fun)(mat)\n\n\ndef test_irfft_ortho():\n def fun(x):\n return np.fft.irfft(x, norm=\"ortho\")\n\n D = 4\n mat = npr.randn(D, D) / 10.0\n # ensure hermitian by doing a fft\n mat = np.fft.rfft(mat)\n check_grads(fun)(mat)\n\n\ndef test_rfft2():\n def fun(x):\n return np.fft.rfft2(x)\n\n D = 4\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_irfft2():\n def fun(x):\n return np.fft.irfft2(x)\n\n D = 4\n mat = npr.randn(D, D) / 10.0\n # ensure hermitian by doing a fft\n mat = np.fft.rfft2(mat)\n check_grads(fun)(mat)\n\n\ndef test_rfftn():\n def fun(x):\n return np.fft.rfftn(x)\n\n D = 4\n mat = npr.randn(D, D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_rfftn_odd_not_implemented():\n def fun(x):\n return np.fft.rfftn(x)\n\n D = 5\n mat = npr.randn(D, D, D) / 10.0\n with pytest.raises(NotImplementedError):\n check_grads(fun)(mat)\n\n\ndef test_rfftn_subset():\n def fun(x):\n return np.fft.rfftn(x)[(0, 1, 0), (3, 3, 2)]\n\n D = 4\n mat = npr.randn(D, D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_rfftn_axes():\n def fun(x):\n return np.fft.rfftn(x, axes=(0, 2))\n\n D = 4\n mat = npr.randn(D, D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_irfftn():\n def fun(x):\n return np.fft.irfftn(x)\n\n D = 4\n mat = npr.randn(D, D, D) / 10.0\n # ensure hermitian by doing a fft\n mat = np.fft.rfftn(mat)\n check_grads(fun)(mat)\n\n\ndef test_irfftn_subset():\n def fun(x):\n return np.fft.irfftn(x)[(0, 1, 0), (3, 3, 2)]\n\n D = 4\n mat = npr.randn(D, D, D) / 10.0\n # ensure hermitian by doing a fft\n mat = np.fft.rfftn(mat)\n check_grads(fun)(mat)\n\n\ndef test_fftshift():\n def fun(x):\n return np.fft.fftshift(x)\n\n D = 5\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_fftshift_even():\n def fun(x):\n return np.fft.fftshift(x)\n\n D = 4\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_fftshift_axes():\n def fun(x):\n return np.fft.fftshift(x, axes=1)\n\n D = 5\n mat = npr.randn(D, D) / 10.0\n check_grads(fun)(mat)\n\n\ndef test_ifftshift():\n def fun(x):\n return np.fft.ifftshift(x)\n\n D = 5\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_ifftshift_even():\n def fun(x):\n return np.fft.ifftshift(x)\n\n D = 4\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_ifftshift_axes():\n def fun(x):\n return np.fft.ifftshift(x, axes=1)\n\n D = 5\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n"
},
{
"path": "tests/test_graphs.py",
"content": "import warnings\n\nimport pytest\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n\ndef test_grad_fanout():\n fun = lambda x: np.sin(np.sin(x) + np.sin(x))\n df = grad(fun)\n check_grads(fun)(npr.randn())\n check_grads(df)(npr.rand())\n\n\ndef test_grad_const():\n fun = lambda x: 1.0\n with warnings.catch_warnings(record=True) as w:\n warnings.simplefilter(\"ignore\")\n df = grad(fun)\n assert np.allclose(df(2.0), 0.0)\n\n\ndef test_grad_identity():\n fun = lambda x: x\n df = grad(fun)\n ddf = grad(df)\n assert np.allclose(df(2.0), 1.0)\n assert np.allclose(ddf(2.0), 0.0)\n\n\ndef test_hess_vector_prod():\n npr.seed(1)\n randv = npr.randn(10)\n\n def fun(x):\n return np.sin(np.dot(x, randv))\n\n df = grad(fun)\n\n def vector_product(x, v):\n return np.sin(np.dot(v, df(x)))\n\n ddf = grad(vector_product)\n A = npr.randn(10)\n B = npr.randn(10)\n check_grads(fun)(A)\n check_grads(vector_product)(A, B)\n\n\ndef test_enclosing_scope_ref():\n def fun(x):\n inner_fun = lambda y: x * y\n return x * grad(inner_fun)(2.0)\n\n check_grads(fun)(1.0)\n\n\ndef test_enclosing_scope_ref_2():\n def fun(x):\n inner_fun = lambda y: y * x\n return x * grad(inner_fun)(2.0)\n\n check_grads(fun)(1.0)\n\n\ndef test_mutating_outgrad():\n def fun(a):\n b = a + 1.0\n c = b + 1.5\n d = a + b\n e = d + c\n return e\n\n A = npr.randn(5)\n check_grads(fun)(A)\n\n\ndef test_mutating_outgrad_from_indexing():\n def fun(a):\n b = a + 1.0\n c = b[0] + 1.5\n d = a + b\n e = d + c\n return e\n\n A = npr.randn(5)\n check_grads(fun)(A)\n\n\ndef test_complex_mutating_outgrad_from_indexing():\n def fun(a):\n b = a + 1.0j\n c = b[0] + 1.5\n d = a + b\n e = d + c\n return np.sum(np.sin(np.real(e)))\n\n A = npr.randn(5)\n check_grads(fun)(A)\n d_fun = lambda x: grad(fun)(x)\n check_grads(d_fun)(A)\n\n\ndef test_complex_separate_real_and_imaginary():\n def fun(a):\n r, i = np.real(a), np.imag(a)\n a = np.abs(r) ** 1.4 + np.abs(i) ** 1.3\n return np.sum(np.sin(a))\n\n d_fun = lambda x: grad(fun)(x)\n A = npr.randn(5, 3) + 0.1j * npr.randn(5, 3)\n check_grads(fun)(A)\n check_grads(d_fun)(A)\n\n\ndef test_third_derivative():\n fun = lambda x: np.sin(np.sin(x) + np.sin(x))\n df = grad(fun)\n ddf = grad(fun)\n dddf = grad(fun)\n check_grads(fun)(npr.randn())\n check_grads(df)(npr.rand())\n check_grads(ddf)(npr.rand())\n check_grads(dddf)(npr.rand())\n\n\ndef test_third_derivative_other_args():\n fun = lambda x, y: np.sin(np.sin(x) + np.sin(y))\n df = grad(fun)\n ddf = grad(fun, 1)\n dddf = grad(fun)\n check_grads(fun)(npr.randn(), npr.randn())\n check_grads(df)(npr.randn(), npr.randn())\n check_grads(ddf)(npr.randn(), npr.randn())\n check_grads(dddf)(npr.randn(), npr.randn())\n\n\ndef test_third_derivative_other_args2():\n fun = lambda x, y: np.sin(np.sin(x) + np.sin(y))\n df = grad(fun, 1)\n ddf = grad(fun)\n dddf = grad(fun, 1)\n check_grads(fun)(npr.randn(), npr.randn())\n check_grads(df)(npr.randn(), npr.randn())\n check_grads(ddf)(npr.randn(), npr.randn())\n check_grads(dddf)(npr.randn(), npr.randn())\n\n\ndef test_singleton_array_output():\n fun = lambda x: np.sum(np.sin(x), keepdims=True)\n check_grads(fun)(npr.randn(3, 3))\n check_grads(lambda x: np.sum(grad(fun)(x)))(npr.randn(3, 3))\n\n\ndef test_singleton_array_output_axis0():\n fun = lambda x: np.sum(np.sin(x), axis=0, keepdims=False)\n check_grads(fun)(npr.randn(3, 1))\n check_grads(lambda x: np.sum(grad(fun)(x)))(npr.randn(3, 1))\n\n\ndef test_singleton_array_output_axis1():\n fun = lambda x: np.sum(np.sin(x), axis=1, keepdims=False)\n check_grads(fun)(npr.randn(1, 3))\n check_grads(lambda x: np.sum(grad(fun)(x)))(npr.randn(1, 3))\n\n\ndef test_singleton_array_output_axis0_keepdims():\n fun = lambda x: np.sum(np.sin(x), axis=0, keepdims=True)\n check_grads(fun)(npr.randn(3, 1))\n check_grads(lambda x: np.sum(grad(fun)(x)))(npr.randn(3, 1))\n\n\ndef test_singleton_array_output_axis1_keepdims():\n fun = lambda x: np.sum(np.sin(x), axis=1, keepdims=True)\n check_grads(fun)(npr.randn(1, 3))\n check_grads(lambda x: np.sum(grad(fun)(x)))(npr.randn(1, 3))\n\n\ndef test_assignment_raises_error():\n def fun(A, b):\n A[1] = b\n return A\n\n A = npr.randn(5)\n with pytest.raises(TypeError):\n check_grads(fun)(A, 3.0)\n\n\n# def test_nonscalar_output_1():\n# with pytest.raises(TypeError):\n# grad(lambda x: x * 2)(np.zeros(2))\n\n# def test_nonscalar_output_2():\n# with pytest.raises(TypeError):\n# grad(lambda x: x * 2)(np.zeros(2))\n\n# TODO:\n# Diamond patterns\n# Taking grad again after returning const\n# Empty functions\n# 2nd derivatives with fanout, thinking about the outgrad adder\n"
},
{
"path": "tests/test_jacobian.py",
"content": "import autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad, jacobian\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n\ndef test_jacobian_against_grad():\n fun = lambda x: np.sum(np.sin(x), axis=1, keepdims=True)\n A = npr.randn(1, 3)\n assert np.allclose(grad(fun)(A), jacobian(fun)(A))\n\n\ndef test_jacobian_scalar_to_vector():\n fun = lambda x: np.array([x, x**2, x**3])\n val = npr.randn()\n assert np.allclose(jacobian(fun)(val), np.array([1.0, 2 * val, 3 * val**2]))\n\n\ndef test_jacobian_against_stacked_grads():\n scalar_funs = [\n lambda x: np.sum(x**3),\n lambda x: np.prod(np.sin(x) + np.sin(x)),\n lambda x: grad(lambda y: np.exp(y) * np.tanh(x[0]))(x[1]),\n ]\n\n vector_fun = lambda x: np.array([f(x) for f in scalar_funs])\n\n x = npr.randn(5)\n jac = jacobian(vector_fun)(x)\n grads = [grad(f)(x) for f in scalar_funs]\n\n assert np.allclose(jac, np.vstack(grads))\n\n\ndef test_jacobian_higher_order():\n fun = lambda x: np.sin(np.outer(x, x)) + np.cos(np.dot(x, x))\n\n assert jacobian(fun)(npr.randn(2)).shape == (2, 2, 2)\n assert jacobian(jacobian(fun))(npr.randn(2)).shape == (2, 2, 2, 2)\n # assert jacobian(jacobian(jacobian(fun)))(npr.randn(2)).shape == (2,2,2,2,2)\n\n check_grads(lambda x: np.sum(np.sin(jacobian(fun)(x))))(npr.randn(2))\n check_grads(lambda x: np.sum(np.sin(jacobian(jacobian(fun))(x))))(npr.randn(2))\n"
},
{
"path": "tests/test_linalg.py",
"content": "from functools import partial\n\nimport numpy as onp\nimport pytest\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import tuple\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n# Fwd mode not yet implemented\ncheck_grads = partial(check_grads, modes=[\"rev\"])\n\n\ndef check_symmetric_matrix_grads(fun, **grad_check_kwargs):\n def check(*args):\n def symmetrize(A):\n L = np.tril(A)\n return (L + T(L)) / 2.0\n\n new_fun = lambda *args: fun(symmetrize(args[0]), *args[1:])\n check_grads(new_fun, **grad_check_kwargs)(*args)\n\n return check\n\n\nT = lambda A: np.swapaxes(A, -1, -2)\n\n\ndef rand_psd(D):\n mat = npr.randn(D, D)\n return np.dot(mat, mat.T)\n\n\ndef test_inv():\n def fun(x):\n return np.linalg.inv(x)\n\n D = 8\n mat = npr.randn(D, D)\n mat = np.dot(mat, mat) + 1.0 * np.eye(D)\n check_grads(fun)(mat)\n\n\ndef test_pinv():\n def fun(x):\n return np.linalg.pinv(x)\n\n N = 5\n D = 2\n ## Non-square matrices:\n for M in range(N // 2, N + N // 2 + 1):\n mat = npr.randn(N, M)\n check_grads(fun)(mat)\n # Stacked\n mat = npr.randn(D, N, M)\n check_grads(fun)(mat)\n\n ## Square, low (fixed) rank matrices\n def fun_low_rank(x):\n return np.linalg.pinv(np.linalg._dot(np.linalg.T(x), x))\n\n for M in range(N // 2, N + N // 2 + 1):\n mat = npr.randn(N, M)\n check_grads(fun_low_rank)(mat)\n # Stacked\n mat = npr.randn(D, N, M)\n check_grads(fun_low_rank)(mat)\n\n\ndef test_inv_3d():\n fun = lambda x: np.linalg.inv(x)\n\n D = 4\n mat = npr.randn(D, D, D) + 5 * np.eye(D)\n check_grads(fun)(mat)\n\n mat = npr.randn(D, D, D, D) + 5 * np.eye(D)\n check_grads(fun)(mat)\n\n\ndef test_solve_arg1():\n D = 8\n A = npr.randn(D, D) + 10.0 * np.eye(D)\n B = npr.randn(D, D - 1)\n\n def fun(a):\n return np.linalg.solve(a, B)\n\n check_grads(fun)(A)\n\n\ndef test_solve_arg1_1d():\n D = 8\n A = npr.randn(D, D) + 10.0 * np.eye(D)\n B = npr.randn(D)\n\n def fun(a):\n return np.linalg.solve(a, B)\n\n check_grads(fun)(A)\n\n\ndef test_solve_arg2():\n D = 6\n A = npr.randn(D, D) + 1.0 * np.eye(D)\n B = npr.randn(D, D - 1)\n\n def fun(b):\n return np.linalg.solve(A, b)\n\n check_grads(fun)(B)\n\n\ndef test_solve_arg1_3d():\n D = 4\n A = npr.randn(D + 1, D, D) + 5 * np.eye(D)\n B = npr.randn(D + 1, D)\n if onp.lib.NumpyVersion(onp.__version__) < \"2.0.0\":\n fun = lambda A: np.linalg.solve(A, B)\n else:\n fun = lambda A: np.linalg.solve(A, B[..., None])[..., 0]\n check_grads(fun)(A)\n\n\ndef test_solve_arg1_3d_3d():\n D = 4\n A = npr.randn(D + 1, D, D) + 5 * np.eye(D)\n B = npr.randn(D + 1, D, D + 2)\n fun = lambda A: np.linalg.solve(A, B)\n check_grads(fun)(A)\n\n\ndef test_det():\n def fun(x):\n return np.linalg.det(x)\n\n D = 6\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_det_3d():\n fun = lambda x: np.linalg.det(x)\n D = 3\n mat = npr.randn(D, D, D)\n check_grads(fun)(mat)\n\n\ndef test_slogdet():\n def fun(x):\n sign, logdet = np.linalg.slogdet(x)\n return logdet\n\n D = 6\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n check_grads(fun)(-mat)\n\n\ndef test_slogdet_3d():\n fun = lambda x: np.sum(np.linalg.slogdet(x)[1])\n mat = np.concatenate([(rand_psd(5) + 5 * np.eye(5))[None, ...] for _ in range(3)])\n check_grads(fun)(mat)\n\n\ndef test_vector_2norm():\n def fun(x):\n return np.linalg.norm(x)\n\n D = 6\n vec = npr.randn(D)\n check_grads(fun, modes=[\"fwd\", \"rev\"])(vec)\n\n\ndef test_vector_2norm_complex():\n def fun(x):\n return np.linalg.norm(x)\n\n D = 6\n vec = npr.randn(D) + 1j * npr.randn(D)\n check_grads(fun)(vec)\n\n\ndef test_frobenius_norm():\n def fun(x):\n return np.linalg.norm(x)\n\n D = 6\n mat = npr.randn(D, D - 1)\n check_grads(fun, modes=[\"fwd\", \"rev\"])(mat)\n\n\ndef test_frobenius_norm_complex():\n def fun(x):\n return np.linalg.norm(x)\n\n D = 6\n mat = npr.randn(D, D - 1) + 1j * npr.randn(D, D - 1)\n check_grads(fun)(mat)\n\n\ndef test_frobenius_norm_axis():\n def fun(x):\n return np.linalg.norm(x, axis=(0, 1))\n\n D = 6\n mat = npr.randn(D, D - 1, D - 2)\n check_grads(fun, modes=[\"fwd\", \"rev\"])(mat)\n\n\ndef test_frobenius_norm_axis_complex():\n def fun(x):\n return np.linalg.norm(x, axis=(0, 1))\n\n D = 6\n mat = npr.randn(D, D - 1, D - 2) + 1j * npr.randn(D, D - 1, D - 2)\n check_grads(fun)(mat)\n\n\n@pytest.mark.parametrize(\"ord\", range(2, 5))\n@pytest.mark.parametrize(\"size\", [6])\ndef test_vector_norm_ord(size, ord):\n def fun(x):\n return np.linalg.norm(x, ord=ord)\n\n vec = npr.randn(size)\n check_grads(fun, modes=[\"fwd\", \"rev\"])(vec)\n\n\n@pytest.mark.parametrize(\"ord\", range(2, 5))\n@pytest.mark.parametrize(\"size\", [6])\ndef test_vector_norm_ord_complex(size, ord):\n def fun(x):\n return np.linalg.norm(x, ord=ord)\n\n vec = npr.randn(size) + 1j * npr.randn(size)\n check_grads(fun)(vec)\n\n\n@pytest.mark.parametrize(\"axis\", range(3))\n@pytest.mark.parametrize(\"shape\", [(6, 5, 4)])\ndef test_norm_axis(shape, axis):\n def fun(x):\n return np.linalg.norm(x, axis=axis)\n\n arr = npr.randn(*shape)\n check_grads(fun, modes=[\"fwd\", \"rev\"])(arr)\n\n\n@pytest.mark.parametrize(\"axis\", range(3))\n@pytest.mark.parametrize(\"shape\", [(6, 5, 4)])\ndef test_norm_axis_complex(shape, axis):\n def fun(x):\n return np.linalg.norm(x, axis=axis)\n\n arr = npr.randn(*shape) + 1j * npr.randn(*shape)\n check_grads(fun)(arr)\n\n\ndef test_norm_nuclear():\n def fun(x):\n return np.linalg.norm(x, ord=\"nuc\")\n\n D = 6\n mat = npr.randn(D, D - 1)\n # Order 1 because the jvp of the svd is not implemented\n check_grads(fun, modes=[\"fwd\", \"rev\"], order=1)(mat)\n\n\ndef test_norm_nuclear_complex():\n def fun(x):\n return np.linalg.norm(x, ord=\"nuc\")\n\n D = 6\n mat = npr.randn(D, D - 1) + 1j * npr.randn(D, D - 1)\n check_grads(fun)(mat)\n\n\ndef test_norm_nuclear_axis():\n def fun(x):\n return np.linalg.norm(x, ord=\"nuc\", axis=(0, 1))\n\n D = 6\n mat = npr.randn(D, D - 1, D - 2)\n # Order 1 because the jvp of the svd is not implemented\n check_grads(fun, modes=[\"fwd\", \"rev\"], order=1)(mat)\n\n\ndef test_norm_nuclear_axis_complex():\n def fun(x):\n return np.linalg.norm(x, ord=\"nuc\", axis=(0, 1))\n\n D = 6\n mat = npr.randn(D, D - 1, D - 2) + 1j * npr.randn(D, D - 1, D - 2)\n check_grads(fun)(mat)\n\n\ndef test_eigvalh_lower():\n def fun(x):\n w, v = np.linalg.eigh(x)\n return tuple((w, v))\n\n D = 6\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_eigvalh_upper():\n def fun(x):\n w, v = np.linalg.eigh(x, \"U\")\n return tuple((w, v))\n\n D = 6\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\nbroadcast_dot_transpose = partial(np.einsum, \"...ij,...kj->...ik\")\n\n\ndef test_eigvalh_lower_broadcasting():\n def fun(x):\n w, v = np.linalg.eigh(x)\n return tuple((w, v))\n\n D = 6\n mat = npr.randn(2, 3, D, D) + 10 * np.eye(D)[None, None, ...]\n hmat = broadcast_dot_transpose(mat, mat)\n check_grads(fun)(hmat)\n\n\ndef test_eigvalh_upper_broadcasting():\n def fun(x):\n w, v = np.linalg.eigh(x, \"U\")\n return tuple((w, v))\n\n D = 6\n mat = npr.randn(2, 3, D, D) + 10 * np.eye(D)[None, None, ...]\n hmat = broadcast_dot_transpose(mat, mat)\n check_grads(fun)(hmat)\n\n\n# For complex-valued matrices, the eigenvectors could have arbitrary phases (gauge)\n# which makes it impossible to compare to numerical derivatives. So we take the\n# absolute value to get rid of that phase.\ndef test_eigvalh_lower_complex():\n def fun(x):\n w, v = np.linalg.eigh(x)\n return tuple((w, np.abs(v)))\n\n D = 6\n mat = npr.randn(D, D) + 1j * npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_eigvalh_upper_complex():\n def fun(x):\n w, v = np.linalg.eigh(x, \"U\")\n return tuple((w, np.abs(v)))\n\n D = 6\n mat = npr.randn(D, D) + 1j * npr.randn(D, D)\n check_grads(fun)(mat)\n\n\n# Note eigenvalues and eigenvectors for real matrix can still be complex\ndef test_eig_real():\n def fun(x):\n w, v = np.linalg.eig(x)\n return tuple((np.abs(w), np.abs(v)))\n\n D = 8\n mat = npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_eig_complex():\n def fun(x):\n w, v = np.linalg.eig(x)\n return tuple((w, np.abs(v)))\n\n D = 8\n mat = npr.randn(D, D) + 1.0j * npr.randn(D, D)\n check_grads(fun)(mat)\n\n\ndef test_eig_batched():\n def fun(x):\n w, v = np.linalg.eig(x)\n return tuple((w, np.abs(v)))\n\n D = 8\n b = 5\n mat = npr.randn(b, D, D) + 1.0j * npr.randn(b, D, D)\n check_grads(fun)(mat)\n\n\ndef test_cholesky():\n fun = lambda A: np.linalg.cholesky(A)\n check_symmetric_matrix_grads(fun)(rand_psd(6))\n\n\ndef test_cholesky_broadcast():\n fun = lambda A: np.linalg.cholesky(A)\n A = np.concatenate([rand_psd(6)[None, :, :] for i in range(3)], axis=0)\n check_symmetric_matrix_grads(fun)(A)\n\n\ndef test_cholesky_reparameterization_trick():\n def fun(A):\n rng = np.random.RandomState(0)\n z = np.dot(np.linalg.cholesky(A), rng.randn(A.shape[0]))\n return np.linalg.norm(z)\n\n check_symmetric_matrix_grads(fun)(rand_psd(6))\n\n\ndef test_svd_wide_2d():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((u, s, v))\n\n m = 3\n n = 5\n mat = npr.randn(m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_wide_2d_complex():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((np.abs(u), s, np.abs(v)))\n\n m = 3\n n = 5\n mat = npr.randn(m, n) + 1j * npr.randn(m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_wide_3d():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((u, s, v))\n\n k = 4\n m = 3\n n = 5\n mat = npr.randn(k, m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_wide_3d_complex():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((np.abs(u), s, np.abs(v)))\n\n k = 4\n m = 3\n n = 5\n mat = npr.randn(k, m, n) + 1j * npr.randn(k, m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_square_2d():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((u, s, v))\n\n m = 4\n n = 4\n mat = npr.randn(m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_square_2d_complex():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((np.abs(u), s, np.abs(v)))\n\n m = 4\n n = 4\n mat = npr.randn(m, n) + 1j * npr.randn(m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_square_3d():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((u, s, v))\n\n k = 3\n m = 4\n n = 4\n mat = npr.randn(k, m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_square_3d_complex():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((np.abs(u), s, np.abs(v)))\n\n k = 3\n m = 4\n n = 4\n mat = npr.randn(k, m, n) + 1j * npr.randn(k, m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_tall_2d():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((u, s, v))\n\n m = 5\n n = 3\n mat = npr.randn(m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_tall_2d_complex():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((np.abs(u), s, np.abs(v)))\n\n m = 5\n n = 3\n mat = npr.randn(m, n) + 1j * npr.randn(m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_tall_3d():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((u, s, v))\n\n k = 4\n m = 5\n n = 3\n mat = npr.randn(k, m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_tall_3d_complex():\n def fun(x):\n u, s, v = np.linalg.svd(x, full_matrices=False)\n return tuple((np.abs(u), s, np.abs(v)))\n\n k = 4\n m = 5\n n = 3\n mat = npr.randn(k, m, n) + 1j * npr.randn(k, m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_only_s_2d():\n def fun(x):\n s = np.linalg.svd(x, full_matrices=False, compute_uv=False)\n return s\n\n m = 5\n n = 3\n mat = npr.randn(m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_only_s_2d_complex():\n def fun(x):\n s = np.linalg.svd(x, full_matrices=False, compute_uv=False)\n return s\n\n m = 5\n n = 3\n mat = npr.randn(m, n) + 1j * npr.randn(m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_only_s_3d():\n def fun(x):\n s = np.linalg.svd(x, full_matrices=False, compute_uv=False)\n return s\n\n k = 4\n m = 5\n n = 3\n mat = npr.randn(k, m, n)\n check_grads(fun)(mat)\n\n\ndef test_svd_only_s_3d_complex():\n def fun(x):\n s = np.linalg.svd(x, full_matrices=False, compute_uv=False)\n return s\n\n k = 4\n m = 5\n n = 3\n mat = npr.randn(k, m, n) + 1j * npr.randn(k, m, n)\n check_grads(fun)(mat)\n"
},
{
"path": "tests/test_list.py",
"content": "import autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd import isinstance as ag_isinstance\nfrom autograd import list as ag_list\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n\ndef test_getter():\n def fun(input_list):\n A = np.sum(input_list[0])\n B = np.sum(input_list[1])\n C = np.sum(input_list[1])\n return A + B + C\n\n d_fun = grad(fun)\n input_list = [npr.randn(5, 6), npr.randn(4, 3), npr.randn(2, 4)]\n\n result = d_fun(input_list)\n assert np.allclose(result[0], np.ones((5, 6)))\n assert np.allclose(result[1], 2 * np.ones((4, 3)))\n assert np.allclose(result[2], np.zeros((2, 4)))\n\n\ndef test_grads():\n def fun(input_list):\n A = np.sum(np.sin(input_list[0]))\n B = np.sum(np.cos(input_list[1]))\n return A + B\n\n def d_fun(input_list):\n g = grad(fun)(input_list)\n A = np.sum(g[0])\n B = np.sum(np.sin(g[0]))\n C = np.sum(np.sin(g[1]))\n return A + B + C\n\n input_list = [npr.randn(5, 6), npr.randn(4, 3), npr.randn(2, 4)]\n\n check_grads(fun)(input_list)\n check_grads(d_fun)(input_list)\n\n\ndef test_slices():\n def f(x):\n s = slice(None, -1, None)\n y = x[s]\n return y[0]\n\n grad(f)([1.0, 2.0, 3.0])\n\n def f(x):\n y = x[1:3]\n return y[0]\n\n grad(f)([1.0, 2.0, 3.0])\n\n\ndef test_nested_list():\n A = [[1.0], 2.0, 1.5]\n\n def fun(x):\n return x[1:][0]\n\n check_grads(fun)(A)\n\n\ndef test_make_list():\n def fun(x):\n return ag_list((x, x))\n\n check_grads(fun)(1.0)\n\n\ndef test_isinstance():\n def fun(x):\n assert ag_isinstance(x, list)\n assert ag_isinstance(x, ag_list)\n return x[0]\n\n fun([1.0, 2.0, 3.0])\n grad(fun)([1.0, 2.0, 3.0])\n"
},
{
"path": "tests/test_logic.py",
"content": "import warnings\nfrom contextlib import contextmanager\n\nimport pytest\n\nimport autograd.numpy as np\nfrom autograd import deriv, grad\nfrom autograd.core import primitive_vjps\nfrom autograd.extend import primitive\nfrom autograd.test_util import check_grads\n\n\ndef test_assert():\n # from https://github.com/HIPS/autograd/issues/43\n def fun(x):\n assert np.allclose(x, (x * 3.0) / 3.0)\n return np.sum(x)\n\n check_grads(fun)(np.array([1.0, 2.0, 3.0]))\n\n\ndef test_nograd():\n # we want this to raise non-differentiability error\n fun = lambda x: np.allclose(x, (x * 3.0) / 3.0)\n with pytest.raises(TypeError):\n with warnings.catch_warnings(record=True) as w:\n grad(fun)(np.array([1.0, 2.0, 3.0]))\n\n\ndef test_no_vjp_def():\n fun = primitive(lambda x: 2.0 * x)\n with pytest.raises(NotImplementedError):\n grad(fun)(1.0)\n\n\ndef test_no_jvp_def():\n fun = primitive(lambda x: 2.0 * x)\n with pytest.raises(NotImplementedError):\n deriv(fun)(1.0)\n\n\ndef test_falseyness():\n fun = lambda x: np.real(x**2 if np.iscomplex(x) else np.sum(x))\n check_grads(fun)(2.0)\n check_grads(fun)(2.0 + 1j)\n\n\ndef test_unimplemented_falseyness():\n @contextmanager\n def remove_grad_definitions(fun):\n vjpmaker = primitive_vjps.pop(fun, None)\n yield\n if vjpmaker:\n primitive_vjps[fun] = vjpmaker\n\n with remove_grad_definitions(np.iscomplex):\n fun = lambda x: np.real(x**2 if np.iscomplex(x) else np.sum(x))\n check_grads(fun)(5.0)\n check_grads(fun)(2.0 + 1j)\n"
},
{
"path": "tests/test_misc.py",
"content": "import autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad, make_vjp\nfrom autograd.misc import const_graph, flatten\nfrom autograd.test_util import scalar_close\nfrom autograd.tracer import primitive\n\n\ndef test_const_graph():\n L = []\n\n def foo(x, y):\n L.append(None)\n return grad(lambda x: np.sin(x) + x * 2)(x * y)\n\n foo_wrapped = const_graph(foo)\n\n assert len(L) == 0\n assert scalar_close(foo(0.0, 0.0), foo_wrapped(0.0, 0.0))\n assert len(L) == 2\n assert scalar_close(foo(1.0, 0.5), foo_wrapped(1.0, 0.5))\n assert len(L) == 3\n assert scalar_close(foo(1.0, 0.5), foo_wrapped(1.0, 0.5))\n assert len(L) == 4\n\n\ndef test_const_graph_args():\n L = []\n\n @primitive\n def process(var, varname):\n L.append(varname)\n return var\n\n def foo(x, y, z):\n x = process(x, \"x\")\n y = process(y, \"y\")\n z = process(z, \"z\")\n return x + 2 * y + 3 * z\n\n foo_wrapped = const_graph(foo, 1.0, z=3.0)\n\n assert L == []\n assert scalar_close(foo(1.0, 2.0, 3.0), foo_wrapped(2.0))\n assert L == [\"x\", \"y\", \"z\", \"x\", \"y\", \"z\"]\n L = []\n assert scalar_close(foo(1.0, 2.0, 3.0), foo_wrapped(2.0))\n assert L == [\"x\", \"y\", \"z\", \"y\"]\n L = []\n assert scalar_close(foo(1.0, 2.0, 3.0), foo_wrapped(2.0))\n assert L == [\"x\", \"y\", \"z\", \"y\"]\n\n\ndef test_flatten():\n r = np.random.randn\n x = (1.0, r(2, 3), [r(1, 4), {\"x\": 2.0, \"y\": r(4, 2)}])\n x_flat, unflatten = flatten(x)\n assert x_flat.shape == (20,)\n assert x_flat[0] == 1.0\n assert np.all(x_flat == flatten(unflatten(x_flat))[0])\n\n y = (1.0, 2.0, [3.0, {\"x\": 2.0, \"y\": 4.0}])\n y_flat, unflatten = flatten(y)\n assert y_flat.shape == (5,)\n assert y == unflatten(y_flat)\n\n\ndef test_flatten_empty():\n val = (npr.randn(4), [npr.randn(3, 4), 2.5], (), (2.0, [1.0, npr.randn(2)]))\n vect, unflatten = flatten(val)\n val_recovered = unflatten(vect)\n vect_2, _ = flatten(val_recovered)\n assert np.all(vect == vect_2)\n\n\ndef test_flatten_dict():\n val = {\"k\": npr.random((4, 4)), \"k2\": npr.random((3, 3)), \"k3\": 3.0, \"k4\": [1.0, 4.0, 7.0, 9.0]}\n\n vect, unflatten = flatten(val)\n val_recovered = unflatten(vect)\n vect_2, _ = flatten(val_recovered)\n assert np.all(vect == vect_2)\n\n\ndef unflatten_tracing():\n val = [npr.randn(4), [npr.randn(3, 4), 2.5], (), (2.0, [1.0, npr.randn(2)])]\n vect, unflatten = flatten(val)\n\n def f(vect):\n return unflatten(vect)\n\n flatten2, _ = make_vjp(f)(vect)\n assert np.all(vect == flatten2(val))\n\n\ndef test_flatten_nodes_in_containers():\n # see issue #232\n def f(x, y):\n xy, _ = flatten([x, y])\n return np.sum(xy)\n\n grad(f)(1.0, 2.0)\n\n\ndef test_flatten_complex():\n val = 1 + 1j\n flat, unflatten = flatten(val)\n assert np.all(val == unflatten(flat))\n"
},
{
"path": "tests/test_numpy.py",
"content": "import warnings\n\nfrom numpy_utils import combo_check\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n\ndef test_numpy_version():\n import numpy\n\n assert np.__version__ == numpy.__version__\n\n\ndef test_dot():\n def fun(x, y):\n return np.dot(x, y)\n\n mat1 = npr.randn(10, 11)\n mat2 = npr.randn(10, 11)\n vect1 = npr.randn(10)\n vect2 = npr.randn(11)\n vect3 = npr.randn(11)\n\n check_grads(fun)(mat1, vect2)\n check_grads(fun)(mat1, mat2.T)\n check_grads(fun)(vect1, mat1)\n check_grads(fun)(vect2, vect3)\n\n\ndef test_dot_with_floats():\n def fun(x, y):\n return np.dot(x, y)\n\n mat1 = npr.randn(10, 11)\n vect1 = npr.randn(10)\n float1 = npr.randn()\n\n check_grads(fun)(mat1, float1)\n check_grads(fun)(float1, mat1)\n check_grads(fun)(vect1, float1)\n check_grads(fun)(float1, vect1)\n\n\n# No longer supporting this\n# def test_dot_method():\n# def fun(x, y): return x.dot(y)\n\n# mat1 = npr.randn(10, 11)\n# mat2 = npr.randn(10, 11)\n# vect1 = npr.randn(10)\n# vect2 = npr.randn(11)\n# vect3 = npr.randn(11)\n\n# check_grads(fun)(mat1, vect2)\n# check_grads(fun)(mat1, mat2.T)\n# check_grads(fun)(vect1, mat1)\n# check_grads(fun)(vect2, vect3)\n\n\ndef test_outer():\n def fun(x, y):\n return np.outer(x, y)\n\n vect2 = npr.randn(11)\n vect3 = npr.randn(11)\n\n check_grads(fun)(vect2, vect3)\n check_grads(fun)(vect2.T, vect3)\n check_grads(fun)(vect2.T, vect3.T)\n\n\ndef test_max():\n def fun(x):\n return np.max(x)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_max_axis():\n def fun(x):\n return np.max(x, axis=1)\n\n mat = npr.randn(3, 4, 5)\n check_grads(fun)(mat)\n\n\ndef test_max_axis_keepdims():\n def fun(x):\n return np.max(x, axis=1, keepdims=True)\n\n mat = npr.randn(3, 4, 5)\n check_grads(fun)(mat)\n\n\ndef test_min():\n def fun(x):\n return np.min(x)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_min_axis():\n def fun(x):\n return np.min(x, axis=1)\n\n mat = npr.randn(3, 4, 5)\n check_grads(fun)(mat)\n\n\ndef test_min_axis_keepdims():\n def fun(x):\n return np.min(x, axis=1, keepdims=True)\n\n mat = npr.randn(3, 4, 5)\n check_grads(fun)(mat)\n\n\ndef test_sum_1():\n def fun(x):\n return np.sum(x)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_sum_2():\n def fun(x):\n return np.sum(x, axis=0)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_sum_3():\n def fun(x):\n return np.sum(x, axis=0, keepdims=True)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_sum_with_axis_tuple():\n def fun(x):\n return np.sum(x, axis=(1, 2))\n\n mat = npr.randn(10, 11, 7)\n check_grads(fun)(mat)\n\n\ndef test_flipud():\n def fun(x):\n return np.flipud(x)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_fliplr():\n def fun(x):\n return np.fliplr(x)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_rot90():\n def fun(x):\n return np.rot90(x)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_cumsum_axis0():\n def fun(x):\n return np.cumsum(x, axis=0)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_cumsum_axis1():\n def fun(x):\n return np.cumsum(x, axis=1)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_cumsum_1d():\n def fun(x):\n return np.cumsum(x)\n\n mat = npr.randn(10)\n check_grads(fun)(mat)\n\n\ndef test_cumsum_no_axis():\n def fun(x):\n return np.cumsum(x)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_non_numpy_sum():\n def fun(x, y):\n return sum([x, y])\n\n mat1 = npr.randn(10, 11)\n mat2 = npr.randn(10, 11)\n check_grads(fun)(mat1, mat2)\n\n\ndef test_mean_1():\n def fun(x):\n return np.mean(x)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_mean_2():\n def fun(x):\n return np.mean(x, axis=0)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_mean_3():\n def fun(x):\n return np.mean(x, axis=0, keepdims=True)\n\n mat = npr.randn(10, 11)\n check_grads(fun)(mat)\n\n\ndef test_index_ints():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return x[3, 0, 1]\n\n check_grads(fun)(A)\n\n\ndef test_index_slice():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return x[::-1, 2:4, :]\n\n check_grads(fun)(A)\n\n\ndef test_index_lists():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return x[[0, 1, 2], :, :]\n\n check_grads(fun)(A)\n\n\ndef test_index_mixed():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return x[3, 2:, [1, 3]]\n\n check_grads(fun)(A)\n\n\ndef test_vector_slice():\n A = npr.randn(5)\n\n def fun(x):\n return x[2:4]\n\n check_grads(fun)(A)\n\n\ndef test_index_slice_fanout():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n y = x[::-1, 2:4, :]\n z = x[::-1, 3:5, :]\n return y + z\n\n check_grads(fun)(A)\n\n\ndef test_index_multiple_slices():\n A = npr.randn(7)\n\n def fun(x):\n y = x[2:6]\n z = y[1:3]\n return z\n\n check_grads(fun)(A)\n\n\ndef test_reshape_method():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return x.reshape((5 * 4, 6))\n\n check_grads(fun)(A)\n\n\ndef test_reshape_call():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return np.reshape(x, (5 * 4, 6))\n\n check_grads(fun)(A)\n\n\ndef test_reshape_method_nolist():\n # The reshape can be called in two different ways:\n # like A.reshape((5,4)) or A.reshape(5,4).\n # This test checks that we support the second way.\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return x.reshape(5 * 4, 6)\n\n check_grads(fun)(A)\n\n\ndef test_ravel_method():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return x.ravel()\n\n check_grads(fun)(A)\n\n\ndef test_ravel_call():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return np.ravel(x)\n\n check_grads(fun)(A)\n\n\ndef test_flatten_method():\n A = npr.randn(5, 6, 4)\n\n def fun(x):\n return x.flatten()\n\n check_grads(fun)(A)\n\n\ndef test_simple_append_list():\n A = [1.0, 2.0, 3.0]\n b = 4.0\n check_grads(np.append, argnum=(0, 1))(A, b)\n\n\ndef test_simple_append_arr():\n A = np.array([1.0, 2.0, 3.0])\n b = 4.0\n check_grads(np.append, argnum=(0, 1))(A, b)\n\n\ndef test_simple_append_list_2D():\n A = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]\n B = [[7.0, 8.0, 9.0]]\n check_grads(np.append, argnum=(0, 1))(A, B, axis=0)\n\n\ndef test_simple_concatenate():\n A = npr.randn(5, 6, 4)\n B = npr.randn(4, 6, 4)\n\n def fun(x):\n return np.concatenate((A, x))\n\n check_grads(fun)(B)\n\n\ndef test_concatenate_axis_0():\n A = npr.randn(5, 6, 4)\n B = npr.randn(5, 6, 4)\n\n def fun(x):\n return np.concatenate((B, x, B))\n\n check_grads(fun)(A)\n\n\ndef test_concatenate_axis_1():\n A = npr.randn(5, 6, 4)\n B = npr.randn(5, 6, 4)\n\n def fun(x):\n return np.concatenate((B, x, B), axis=1)\n\n check_grads(fun)(A)\n\n\ndef test_concatenate_axis_1_unnamed():\n \"\"\"Tests whether you can specify the axis without saying \"axis=1\".\"\"\"\n A = npr.randn(5, 6, 4)\n B = npr.randn(5, 6, 4)\n\n def fun(x):\n return np.concatenate((B, x, B), 1)\n\n check_grads(fun)(A)\n\n\ndef test_trace():\n def fun(x):\n return np.trace(x, offset=offset)\n\n mat = npr.randn(10, 11)\n offset = npr.randint(-9, 11)\n check_grads(fun)(mat)\n\n\ndef test_trace2():\n def fun(x):\n return np.trace(x, offset=offset)\n\n mat = npr.randn(11, 10)\n offset = npr.randint(-9, 11)\n check_grads(fun)(mat)\n\n\ndef test_trace_extradims():\n def fun(x):\n return np.trace(x, offset=offset)\n\n mat = npr.randn(5, 6, 4, 3)\n offset = npr.randint(-5, 6)\n check_grads(fun)(mat)\n\n\n# TODO: Allow axis1, axis2 args.\n# def test_trace_extradims2():\n# def fun(x): return np.trace(x, offset=offset, axis1=3,axis2=2)\n# mat = npr.randn(5,6,4,3)\n# offset = npr.randint(-5,6)\n# check_grads(fun)(mat)\n\n\ndef test_diag():\n def fun(x):\n return np.diag(x)\n\n mat = npr.randn(10, 10)\n check_grads(fun)(mat)\n\n\ndef test_transpose():\n def fun(x):\n return x.T\n\n mat = npr.randn(8, 8)\n check_grads(fun)(mat)\n\n\ndef test_roll():\n def fun(x):\n return np.roll(x, 2, axis=1)\n\n mat = npr.randn(4, 5)\n check_grads(fun)(mat)\n\n\ndef test_roll_no_axis():\n def fun(x):\n return np.roll(x, 2, axis=1)\n\n mat = npr.randn(4, 5)\n check_grads(fun)(mat)\n\n\ndef test_triu():\n def fun(x):\n return np.triu(x, k=2)\n\n mat = npr.randn(5, 5)\n check_grads(fun)(mat)\n\n\ndef test_tril():\n def fun(x):\n return np.tril(x, k=2)\n\n mat = npr.randn(5, 5)\n check_grads(fun)(mat)\n\n\ndef test_clip():\n def fun(x):\n return np.clip(x, a_min=0.1, a_max=1.1)\n\n mat = npr.randn(5, 5)\n check_grads(fun)(mat)\n\n\ndef test_prod_1():\n def fun(x):\n return np.prod(x)\n\n mat = npr.randn(2, 3) ** 2 / 10.0 + 0.1 # Gradient unstable when zeros are present.\n check_grads(fun)(mat)\n\n\ndef test_prod_2():\n def fun(x):\n return np.prod(x, axis=0)\n\n mat = npr.randn(2, 3) ** 2 + 0.1\n check_grads(fun)(mat)\n\n\ndef test_prod_3():\n def fun(x):\n return np.prod(x, axis=0, keepdims=True)\n\n mat = npr.randn(2, 3) ** 2 + 0.1\n check_grads(fun)(mat)\n\n\ndef test_prod_4():\n def fun(x):\n return np.prod(x)\n\n mat = npr.randn(7) ** 2 + 0.1\n check_grads(fun)(mat)\n\n\ndef test_1d_array():\n def fun(x):\n return np.array([x, x * 1.0, x + 2.5])\n\n check_grads(fun)(3.0)\n\n\ndef test_2d_array():\n def fun(x):\n return np.array([[x, x * 1.0, x + 2.5], [x**2, x, x / 2.0]])\n\n check_grads(fun)(3.0)\n\n\ndef test_1d_array_fanout():\n def fun(x):\n A = np.array([x, x * 1.0, x + 2.5])\n return A + A\n\n check_grads(fun)(3.0)\n\n\ndef test_2d_array_fanout():\n def fun(x):\n A = np.array([[x, x * 1.0, x + 2.5], [x**2, x, x / 2.0]])\n return A + A\n\n check_grads(fun)(3.0)\n\n\ndef test_array_from_scalar():\n def fun(x):\n return np.array(x)\n\n check_grads(fun)(3.0)\n\n\ndef test_array_from_arrays():\n def fun(x):\n return np.array([x, x])\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_array_from_arrays_2():\n def fun(x):\n return np.array([[2 * x, x + 1], [x, x]])\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_len():\n def fun(x):\n assert len(x) == 3\n return x\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_r_basic():\n with warnings.catch_warnings(record=True) as w:\n\n def fun(x):\n c = npr.randn(3, 2)\n b = np.r_[x]\n return b\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_r_double():\n with warnings.catch_warnings(record=True) as w:\n\n def fun(x):\n c = npr.randn(3, 2)\n b = np.r_[x, x]\n return b\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_no_relation():\n with warnings.catch_warnings(record=True) as w:\n c = npr.randn(3, 2)\n\n def fun(x):\n return c\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_r_no_relation():\n with warnings.catch_warnings(record=True) as w:\n c = npr.randn(3, 2)\n\n def fun(x):\n b = np.r_[c]\n return b\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_r_node_and_const():\n with warnings.catch_warnings(record=True) as w:\n c = npr.randn(3, 2)\n\n def fun(x):\n b = np.r_[x, c]\n return b\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_r_mixed():\n with warnings.catch_warnings(record=True) as w:\n c = npr.randn(3, 2)\n\n def fun(x):\n b = np.r_[x, c, x]\n return b\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_r_slicing():\n with warnings.catch_warnings(record=True) as w:\n c = npr.randn(10)\n\n def fun(x):\n b = np.r_[x, c, 1:10]\n return b\n\n A = npr.randn(10)\n check_grads(fun)(A)\n\n\ndef test_c_():\n with warnings.catch_warnings(record=True) as w:\n c = npr.randn(3, 2)\n\n def fun(x):\n b = np.c_[x, c, x]\n return b\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_c_mixed():\n with warnings.catch_warnings(record=True) as w:\n c = npr.randn(3, 2)\n\n def fun(x):\n b = np.c_[x, c, x]\n return b\n\n A = npr.randn(3, 2)\n check_grads(fun)(A)\n\n\ndef test_var_ddof():\n B = npr.randn(3)\n C = npr.randn(3, 4)\n D = npr.randn(1, 3)\n combo_check(np.var, (0,))([B, C, D], axis=[None], keepdims=[True, False], ddof=[0, 1])\n combo_check(np.var, (0,))([C, D], axis=[None, 1], keepdims=[True, False], ddof=[2])\n\n\ndef test_std_ddof():\n B = npr.randn(3)\n C = npr.randn(3, 4)\n D = npr.randn(1, 3)\n combo_check(np.std, (0,))([B, C, D], axis=[None], keepdims=[True, False], ddof=[0, 1])\n combo_check(np.std, (0,))([C, D], axis=[None, 1], keepdims=[True, False], ddof=[2])\n\n\ndef test_where():\n def fun(x, y):\n b = np.where(C, x, y)\n return b\n\n C = npr.randn(4, 5) > 0\n A = npr.randn(4, 5)\n B = npr.randn(4, 5)\n check_grads(fun)(A, B)\n\n\ndef test_squeeze_func():\n A = npr.randn(5, 1, 4)\n\n def fun(x):\n return np.squeeze(x)\n\n check_grads(fun)(A)\n\n\ndef test_squeeze_method():\n A = npr.randn(5, 1, 4)\n\n def fun(x):\n return x.squeeze()\n\n check_grads(fun)(A)\n\n\ndef test_repeat():\n A = npr.randn(5, 3, 4)\n\n def fun(x):\n return np.repeat(x, 2, axis=1)\n\n check_grads(fun)(A)\n\n\ndef test_repeat_axis1_rep1():\n A = npr.randn(5, 3, 4)\n\n def fun(x):\n return np.repeat(x, 1, axis=1)\n\n check_grads(fun)(A)\n\n\ndef test_repeat_axis0():\n A = npr.randn(5, 3)\n\n def fun(x):\n return np.repeat(x, 2, axis=0)\n\n check_grads(fun)(A)\n\n\ndef test_repeat_1d_axis0():\n A = npr.randn(5)\n\n def fun(x):\n return np.repeat(x, 2, axis=0)\n\n check_grads(fun)(A)\n\n\ndef test_repeat_axis0_rep1():\n A = npr.randn(5, 1)\n\n def fun(x):\n return np.repeat(x, 1, axis=0)\n\n check_grads(fun)(A)\n\n\ndef test_expand_dims():\n A = npr.randn(5, 1, 4)\n\n def fun(x):\n return np.expand_dims(x, 2)\n\n check_grads(fun)(A)\n\n\ndef test_tensordot_kwargs_by_position():\n def fun(x):\n return np.tensordot(x * np.ones((2, 2)), x * np.ones((2, 2)), 2)\n\n grad(fun)(1.0)\n\n\ndef test_multi_index():\n A = npr.randn(3)\n fun = lambda x: np.sum(x[[0, 0]])\n check_grads(fun)(A)\n\n\ndef test_multi_index2():\n A = npr.randn(3)\n fun = lambda x: np.sum(x[[0, 1, 0]])\n check_grads(fun)(A)\n\n\ndef test_index_dot_slices():\n A = npr.randn(4)\n\n def fun(x):\n return np.dot(x[:2], x[2:])\n\n check_grads(fun)(A)\n\n\n# def test_index_exp_slicing():\n# def fun(x):\n# b = np.index_exp[x, x]\n# return b\n# A = npr.randn(10, 1)\n# check_grads(fun)(A)\n\n# def test_s_slicing():\n# def fun(x):\n# b = np.s_[x, x]\n# return b\n# A = npr.randn(10, 1)\n# check_grads(fun)(A)\n\n# TODO:\n# getitem\n\n\ndef test_cast_to_int():\n inds = np.ones(5)[:, None]\n\n def fun(W):\n # glue W and inds together\n glued_together = np.concatenate((W, inds), axis=1)\n\n # separate W and inds back out\n new_W = W[:, :-1]\n new_inds = np.int64(W[:, -1])\n\n assert new_inds.dtype == np.int64\n return new_W[new_inds].sum()\n\n W = np.random.randn(5, 10)\n check_grads(fun)(W)\n\n\ndef test_make_diagonal():\n def fun(D):\n return np.make_diagonal(D, axis1=-1, axis2=-2)\n\n D = np.random.randn(4)\n A = np.make_diagonal(D, axis1=-1, axis2=-2)\n assert np.allclose(np.diag(A), D)\n check_grads(fun)(D)\n\n D = np.random.randn(3, 4)\n A = np.make_diagonal(D, axis1=-1, axis2=-2)\n assert all([np.allclose(np.diag(A[i]), D[i]) for i in range(3)])\n check_grads(fun)(D)\n\n\ndef test_diagonal():\n def fun(D):\n return np.diagonal(D, axis1=-1, axis2=-2)\n\n D = np.random.randn(4, 4)\n A = np.make_diagonal(D, axis1=-1, axis2=-2)\n check_grads(fun)(D)\n\n D = np.random.randn(3, 4, 4)\n A = np.make_diagonal(D, axis1=-1, axis2=-2)\n check_grads(fun)(D)\n\n\ndef test_nan_to_num():\n y = np.array([0.0, np.nan, np.inf, -np.inf])\n fun = lambda x: np.sum(np.sin(np.nan_to_num(x + y)))\n\n x = np.random.randn(4)\n check_grads(fun)(x)\n\n\n# TODO(mattjj): np.frexp returns a pair of ndarrays and the second is an int\n# type, for which there is currently no vspace registered\n# def test_frexp():\n# fun = lambda x: np.frexp(x)[0]\n# A = 1.2 #np.random.rand(4,3) * 0.8 + 2.1\n# check_grads(fun)(A)\n\n\ndef test_max_equal_values():\n def fun(x):\n return np.max(np.array([x, x]))\n\n check_grads(fun)(1.0)\n\n\ndef test_max_equal_values_2d():\n def fun(x):\n return np.max(np.array([[x, x], [x, 0.5]]), axis=1)\n\n check_grads(fun)(1.0)\n check_grads(fun)(-1.0)\n\n\ndef test_min_3_way_equality():\n def fun(x):\n return np.min(np.array([[x, x, x], [x, 0.5, 0.5], [0.5, 0.5, 0.5], [x, x, 0.5]]), axis=0)\n\n check_grads(fun)(1.0)\n check_grads(fun)(-1.0)\n\n\ndef test_maximum_equal_values():\n def fun(x):\n return np.maximum(x, x)\n\n check_grads(fun)(1.0)\n\n\ndef test_maximum_equal_values_2d():\n def fun(x):\n return np.maximum(np.array([x, x, 0.5]), np.array([[x, 0.5, x], [x, x, 0.5]]))\n\n check_grads(fun)(1.0)\n check_grads(fun)(-1.0)\n check_grads(fun)(2.0)\n\n\ndef test_linspace():\n for num in [0, 1, 5]:\n\n def fun(x, y):\n return np.linspace(x, y, num)\n\n check_grads(fun)(1.2, 3.4)\n check_grads(fun)(1.2, -3.4)\n check_grads(fun)(1.2, 1.2)\n\n\ndef test_astype():\n x = np.arange(3, dtype=\"float32\")\n\n def f(x):\n return np.sum(np.sin(x.astype(\"float64\")))\n\n assert grad(f)(x).dtype == np.dtype(\"float32\")\n\n\ndef test_gradient():\n check_grads(np.gradient, 0)(npr.randn(10))\n check_grads(np.gradient, 0)(npr.randn(10, 10))\n check_grads(np.gradient, 0)(npr.randn(10, 10, 10))\n\n for a in [None, 0, 1, -1, (0, 1), (0, -1)]:\n check_grads(np.gradient, 0)(npr.randn(10, 10, 10), axis=a)\n"
},
{
"path": "tests/test_performance.py",
"content": "# TODO:\n# Do a huge calculation with trivial primitive computations\n# and lots of diamonds and get a benchmark per-node time and\n# memory cost.\n"
},
{
"path": "tests/test_scalar_ops.py",
"content": "import autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n\ndef test_abs():\n fun = lambda x: 3.0 * np.abs(x)\n check_grads(fun)(1.1)\n check_grads(fun)(-1.1)\n check_grads(fun, order=1)(0.0)\n\n\ndef test_absolute():\n fun = lambda x: 3.0 * np.absolute(x)\n check_grads(fun)(1.1)\n check_grads(fun)(-1.1)\n check_grads(fun, order=1)(0.0)\n\n\ndef test_sin():\n fun = lambda x: 3.0 * np.sin(x)\n check_grads(fun)(npr.randn())\n\n\ndef test_sign():\n fun = lambda x: 3.0 * np.sign(x)\n check_grads(fun)(1.1)\n check_grads(fun)(-1.1)\n\n\ndef test_exp():\n fun = lambda x: 3.0 * np.exp(x)\n check_grads(fun)(npr.randn())\n\n\ndef test_log():\n fun = lambda x: 3.0 * np.log(x)\n check_grads(fun)(abs(npr.randn()))\n\n\ndef test_log2():\n fun = lambda x: 3.0 * np.log2(x)\n check_grads(fun)(abs(npr.randn()))\n\n\ndef test_log10():\n fun = lambda x: 3.0 * np.log10(x)\n check_grads(fun)(abs(npr.randn()))\n\n\ndef test_log1p():\n fun = lambda x: 3.0 * np.log1p(x)\n check_grads(fun)(abs(npr.randn()))\n\n\ndef test_expm1():\n fun = lambda x: 3.0 * np.expm1(x)\n check_grads(fun)(abs(npr.randn()))\n\n\ndef test_exp2():\n fun = lambda x: 3.0 * np.exp2(x)\n check_grads(fun)(abs(npr.randn()))\n\n\ndef test_neg():\n fun = lambda x: 3.0 * -x\n check_grads(fun)(npr.randn())\n\n\ndef test_cos():\n fun = lambda x: 3.0 * np.cos(x)\n check_grads(fun)(npr.randn())\n\n\ndef test_tan():\n fun = lambda x: 3.0 * np.tan(x)\n check_grads(fun)(npr.randn())\n\n\ndef test_cosh():\n fun = lambda x: 3.0 * np.cosh(x)\n check_grads(fun)(npr.randn())\n\n\ndef test_sinh():\n fun = lambda x: 3.0 * np.sinh(x)\n check_grads(fun)(npr.randn())\n\n\ndef test_tanh():\n fun = lambda x: 3.0 * np.tanh(x)\n check_grads(fun)(npr.randn())\n\n\ndef test_arccos():\n fun = lambda x: 3.0 * np.arccos(x)\n check_grads(fun)(0.1)\n\n\ndef test_arcsin():\n fun = lambda x: 3.0 * np.arcsin(x)\n check_grads(fun)(0.1)\n\n\ndef test_arctan():\n fun = lambda x: 3.0 * np.arctan(x)\n check_grads(fun)(0.2)\n\n\ndef test_arccosh():\n fun = lambda x: 3.0 * np.arccosh(x)\n check_grads(fun)(npr.randn() ** 2 + 1.2)\n\n\ndef test_arcsinh():\n fun = lambda x: 3.0 * np.arcsinh(x)\n check_grads(fun)(npr.randn())\n\n\ndef test_arctanh():\n fun = lambda x: 3.0 * np.arctanh(x)\n check_grads(fun)(0.2)\n\n\ndef test_sqrt():\n fun = lambda x: 3.0 * np.sqrt(x)\n check_grads(fun)(10.0 * npr.rand())\n\n\ndef test_power_arg0():\n # the +1.'s here are to avoid regimes where numerical diffs fail\n make_fun = lambda y: lambda x: np.power(x, y)\n fun = make_fun(npr.randn() ** 2 + 1.0)\n check_grads(fun)(npr.rand() ** 2 + 1.0)\n\n # test y == 0. as a special case, c.f. #116\n fun = make_fun(0.0)\n assert grad(fun)(0.0) == 0.0\n assert grad(grad(fun))(0.0) == 0.0\n\n\ndef test_power_arg1():\n x = npr.randn() ** 2\n fun = lambda y: np.power(x, y)\n check_grads(fun)(npr.rand() ** 2)\n\n\ndef test_power_arg1_zero():\n fun = lambda y: np.power(0.0, y)\n check_grads(fun)(npr.rand() ** 2)\n\n\ndef test_mod_arg0():\n fun = lambda x, y: np.mod(x, y)\n check_grads(fun)(npr.rand(), npr.rand())\n\n\ndef test_mod_arg1():\n fun = lambda x, y: np.mod(x, y)\n check_grads(fun)(npr.rand(), npr.rand())\n\n\ndef test_divide_arg0():\n fun = lambda x, y: np.divide(x, y)\n check_grads(fun)(npr.rand(), npr.rand())\n\n\ndef test_divide_arg1():\n fun = lambda x, y: np.divide(x, y)\n check_grads(fun)(npr.rand(), npr.rand())\n\n\ndef test_multiply_arg0():\n fun = lambda x, y: np.multiply(x, y)\n check_grads(fun)(npr.rand(), npr.rand())\n\n\ndef test_multiply_arg1():\n fun = lambda x, y: np.multiply(x, y)\n check_grads(fun)(npr.rand(), npr.rand())\n\n\ndef test_true_divide_arg0():\n fun = lambda x, y: np.true_divide(x, y)\n check_grads(fun)(npr.rand(), npr.rand())\n\n\ndef test_true_divide_arg1():\n fun = lambda x, y: np.true_divide(x, y)\n check_grads(fun)(npr.rand(), npr.rand())\n\n\ndef test_reciprocal():\n fun = lambda x: np.reciprocal(x)\n check_grads(fun)(npr.rand())\n\n\ndef test_negative():\n fun = lambda x: np.negative(x)\n check_grads(fun)(npr.rand())\n\n\ndef test_rad2deg():\n fun = lambda x: 3.0 * np.rad2deg(x)\n check_grads(fun)(10.0 * npr.rand())\n\n\ndef test_deg2rad():\n fun = lambda x: 3.0 * np.deg2rad(x)\n check_grads(fun)(10.0 * npr.rand())\n\n\ndef test_radians():\n fun = lambda x: 3.0 * np.radians(x)\n check_grads(fun)(10.0 * npr.rand())\n\n\ndef test_degrees():\n fun = lambda x: 3.0 * np.degrees(x)\n check_grads(fun)(10.0 * npr.rand())\n\n\ndef test_sinc():\n fun = lambda x: 3.0 * np.sinc(x)\n check_grads(fun)(10.0 * npr.rand())\n"
},
{
"path": "tests/test_scipy.py",
"content": "from functools import partial\n\nimport numpy as npo\n\ntry:\n import scipy\nexcept:\n from warnings import warn\n\n warn(\"Skipping scipy tests.\")\nelse:\n from numpy_utils import unary_ufunc_check\n from scipy.signal import convolve as sp_convolve\n\n import autograd.numpy as np\n import autograd.numpy.random as npr\n import autograd.scipy.integrate as integrate\n import autograd.scipy.linalg as spla\n import autograd.scipy.signal\n import autograd.scipy.special as special\n import autograd.scipy.stats as stats\n import autograd.scipy.stats.multivariate_normal as mvn\n from autograd import grad\n from autograd.test_util import check_grads, combo_check\n\n npr.seed(1)\n R = npr.randn\n U = npr.uniform\n\n # Fwd mode not yet implemented for scipy functions\n combo_check = partial(combo_check, modes=[\"rev\"])\n unary_ufunc_check = partial(unary_ufunc_check, modes=[\"rev\"])\n check_grads = partial(check_grads, modes=[\"rev\"])\n\n def symmetrize_matrix_arg(fun, argnum):\n def T(X):\n return np.swapaxes(X, -1, -2) if np.ndim(X) > 1 else X\n\n def symmetrize(X):\n return 0.5 * (X + T(X))\n\n def symmetrized_fun(*args, **kwargs):\n args = list(args)\n args[argnum] = symmetrize(args[argnum])\n return fun(*args, **kwargs)\n\n return symmetrized_fun\n\n ### Stats ###\n def test_chi2_pdf():\n combo_check(stats.chi2.pdf, [0])([R(4) ** 2 + 1.1], [1, 2, 3])\n\n def test_chi2_cdf():\n combo_check(stats.chi2.cdf, [0])([R(4) ** 2 + 1.1], [1, 2, 3])\n\n def test_chi2_logpdf():\n combo_check(stats.chi2.logpdf, [0])([R(4) ** 2 + 1.1], [1, 2, 3])\n\n def test_beta_cdf():\n combo_check(stats.beta.cdf, [0])([U(0.0, 1.0, 4)], [R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1])\n\n def test_beta_pdf():\n combo_check(stats.beta.pdf, [0, 1, 2])([U(0.0, 1.0, 4)], [R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1])\n\n def test_beta_logpdf():\n combo_check(stats.beta.logpdf, [0, 1, 2])([U(0.0, 1.0, 4)], [R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1])\n\n def test_gamma_cdf():\n combo_check(stats.gamma.cdf, [0])([R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1])\n\n def test_gamma_pdf():\n combo_check(stats.gamma.pdf, [0, 1])([R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1])\n\n def test_gamma_logpdf():\n combo_check(stats.gamma.logpdf, [0, 1])([R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1])\n\n def test_norm_pdf():\n combo_check(stats.norm.pdf, [0, 1, 2])([R(4)], [R(4)], [R(4) ** 2 + 1.1])\n\n def test_norm_cdf():\n combo_check(stats.norm.cdf, [0, 1, 2])([R(4)], [R(4)], [R(4) ** 2 + 1.1])\n\n def test_norm_sf():\n combo_check(stats.norm.sf, [0, 1, 2])([R(4)], [R(4)], [R(4) ** 2 + 1.1])\n\n def test_norm_logpdf():\n combo_check(stats.norm.logpdf, [0, 1, 2])([R(4)], [R(4)], [R(4) ** 2 + 1.1])\n\n def test_norm_logcdf():\n combo_check(stats.norm.logcdf, [0, 1, 2])([R(4)], [R(4)], [R(4) ** 2 + 1.1])\n\n def test_norm_logsf():\n combo_check(stats.norm.logsf, [0, 1, 2])([R(4)], [R(4)], [R(4) ** 2 + 1.1])\n\n def test_norm_pdf_broadcast():\n combo_check(stats.norm.pdf, [0, 1, 2])([R(4, 3)], [R(1, 3)], [R(4, 1) ** 2 + 1.1])\n\n def test_norm_cdf_broadcast():\n combo_check(stats.norm.cdf, [0, 1, 2])([R(4, 3)], [R(1, 3)], [R(4, 1) ** 2 + 1.1])\n\n def test_norm_sf_broadcast():\n combo_check(stats.norm.cdf, [0, 1, 2])([R(4, 3)], [R(1, 3)], [R(4, 1) ** 2 + 1.1])\n\n def test_norm_logpdf_broadcast():\n combo_check(stats.norm.logpdf, [0, 1, 2])([R(4, 3)], [R(1, 3)], [R(4, 1) ** 2 + 1.1])\n\n def test_norm_logcdf_broadcast():\n combo_check(stats.norm.logcdf, [0, 1, 2])([R(4, 3)], [R(1, 3)], [R(4, 1) ** 2 + 1.1])\n\n def test_norm_logsf_broadcast():\n combo_check(stats.norm.logcdf, [0, 1, 2])([R(4, 3)], [R(1, 3)], [R(4, 1) ** 2 + 1.1])\n\n def test_poisson_cdf():\n combo_check(stats.poisson.cdf, [1])([np.round(R(4) ** 2)], [R(4) ** 2 + 1.1])\n\n def test_poisson_logpmf():\n combo_check(stats.poisson.logpmf, [1])([np.round(R(4) ** 2)], [R(4) ** 2 + 1.1])\n\n def test_poisson_pmf():\n combo_check(stats.poisson.pmf, [1])([np.round(R(4) ** 2)], [R(4) ** 2 + 1.1])\n\n def test_poisson_cdf_broadcast():\n combo_check(stats.poisson.cdf, [1])([np.round(R(4, 3) ** 2)], [R(4, 1) ** 2 + 1.1])\n\n def test_poisson_logpmf_broadcast():\n combo_check(stats.poisson.logpmf, [1])([np.round(R(4, 3) ** 2)], [R(4, 1) ** 2 + 1.1])\n\n def test_poisson_pmf_broadcast():\n combo_check(stats.poisson.pmf, [1])([np.round(R(4, 3) ** 2)], [R(4, 1) ** 2 + 1.1])\n\n def test_t_pdf():\n combo_check(stats.t.pdf, [0, 1, 2, 3])([R(4)], [R(4) ** 2 + 2.1], [R(4)], [R(4) ** 2 + 2.1])\n\n def test_t_cdf():\n combo_check(stats.t.cdf, [0, 2])([R(4)], [R(4) ** 2 + 2.1], [R(4)], [R(4) ** 2 + 2.1])\n\n def test_t_logpdf():\n combo_check(stats.t.logpdf, [0, 1, 2, 3])([R(4)], [R(4) ** 2 + 2.1], [R(4)], [R(4) ** 2 + 2.1])\n\n def test_t_logcdf():\n combo_check(stats.t.logcdf, [0, 2])([R(4)], [R(4) ** 2 + 2.1], [R(4)], [R(4) ** 2 + 2.1])\n\n def test_t_pdf_broadcast():\n combo_check(stats.t.pdf, [0, 1, 2, 3])(\n [R(4, 3)], [R(1, 3) ** 2 + 2.1], [R(4, 3)], [R(4, 1) ** 2 + 2.1]\n )\n\n def test_t_cdf_broadcast():\n combo_check(stats.t.cdf, [0, 2])([R(4, 3)], [R(1, 3) ** 2 + 2.1], [R(4, 3)], [R(4, 1) ** 2 + 2.1])\n\n def test_t_logpdf_broadcast():\n combo_check(stats.t.logpdf, [0, 1, 2, 3])(\n [R(4, 3)], [R(1, 3) ** 2 + 2.1], [R(4, 3)], [R(4, 1) ** 2 + 2.1]\n )\n\n def test_t_logcdf_broadcast():\n combo_check(stats.t.logcdf, [0, 2])([R(4, 3)], [R(1, 3) ** 2 + 2.1], [R(4, 3)], [R(4, 1) ** 2 + 2.1])\n\n def make_psd(mat):\n return np.dot(mat.T, mat) + np.eye(mat.shape[0])\n\n def test_mvn_pdf():\n combo_check(symmetrize_matrix_arg(mvn.pdf, 2), [0, 1, 2])(\n [R(4)], [R(4)], [make_psd(R(4, 4))], allow_singular=[False]\n )\n\n def test_mvn_logpdf():\n combo_check(symmetrize_matrix_arg(mvn.logpdf, 2), [0, 1, 2])(\n [R(4)], [R(4)], [make_psd(R(4, 4))], allow_singular=[False]\n )\n\n def test_mvn_entropy():\n combo_check(symmetrize_matrix_arg(mvn.entropy, 1), [0, 1])([10 * R(4)], [make_psd(R(4, 4))])\n\n def test_mvn_sing_cov():\n cov = np.zeros((4, 4))\n cov[0, 0] = cov[1, 1] = 1\n\n # Only allow variations in x along the first two dimensions, because\n # variance is zero in the last two.\n def pdf(x, mean, cov):\n x = np.concatenate([x[:2], mean[2:]])\n return symmetrize_matrix_arg(partial(mvn.pdf, allow_singular=True), 2)(x, mean, cov)\n\n combo_check(pdf, [0, 1])(\n [np.concatenate((R(2), np.zeros(2)))], [np.concatenate((R(2), np.zeros(2)))], [cov]\n )\n\n def logpdf(x, mean, cov):\n x = np.concatenate([x[:2], mean[2:]])\n return symmetrize_matrix_arg(partial(mvn.logpdf, allow_singular=True), 2)(x, mean, cov)\n\n combo_check(logpdf, [0, 1])(\n [np.concatenate((R(2), np.zeros(2)))], [np.concatenate((R(2), np.zeros(2)))], [cov]\n )\n\n def test_mvn_pdf_broadcast():\n combo_check(symmetrize_matrix_arg(mvn.pdf, 2), [0, 1, 2])([R(5, 4)], [R(4)], [make_psd(R(4, 4))])\n\n def test_mvn_logpdf_broadcast():\n combo_check(symmetrize_matrix_arg(mvn.logpdf, 2), [0, 1, 2])([R(5, 4)], [R(4)], [make_psd(R(4, 4))])\n\n alpha = npr.random(4) ** 2 + 1.2\n x = stats.dirichlet.rvs(alpha, size=1)[0, :]\n\n # Need to normalize input so that x's sum to one even when we perturb them to compute numeric gradient.\n def normalize(x):\n return x / sum(x)\n\n def normalized_dirichlet_pdf(x, alpha):\n return stats.dirichlet.pdf(normalize(x), alpha)\n\n def normalized_dirichlet_logpdf(x, alpha):\n return stats.dirichlet.logpdf(normalize(x), alpha)\n\n def test_dirichlet_pdf_x():\n combo_check(normalized_dirichlet_pdf, [0])([x], [alpha])\n\n def test_dirichlet_pdf_alpha():\n combo_check(stats.dirichlet.pdf, [1])([x], [alpha])\n\n def test_dirichlet_logpdf_x():\n combo_check(normalized_dirichlet_logpdf, [0])([x], [alpha])\n\n def test_dirichlet_logpdf_alpha():\n combo_check(stats.dirichlet.logpdf, [1])([x], [alpha])\n\n ### Misc ###\n def test_logsumexp1():\n combo_check(special.logsumexp, [0], modes=[\"fwd\", \"rev\"])(\n [np.array([1.1]), R(4), R(3, 4)], axis=[None, 0], keepdims=[True, False]\n )\n\n def test_logsumexp2():\n combo_check(special.logsumexp, [0], modes=[\"fwd\", \"rev\"])(\n [R(3, 4), R(4, 5, 6), R(1, 5)], axis=[None, 0, 1], keepdims=[True, False]\n )\n\n def test_logsumexp3():\n combo_check(special.logsumexp, [0], modes=[\"fwd\", \"rev\"])(\n [R(4)], b=[np.exp(R(4))], axis=[None, 0], keepdims=[True, False]\n )\n\n def test_logsumexp4():\n combo_check(special.logsumexp, [0], modes=[\"fwd\", \"rev\"])(\n [\n R(3, 4),\n ],\n b=[np.exp(R(3, 4))],\n axis=[None, 0, 1],\n keepdims=[True, False],\n )\n\n def test_logsumexp5():\n combo_check(special.logsumexp, [0], modes=[\"fwd\", \"rev\"])(\n [R(2, 3, 4)], b=[np.exp(R(2, 3, 4))], axis=[None, 0, 1], keepdims=[True, False]\n )\n\n def test_logsumexp6():\n x = npr.randn(1, 5)\n\n def f(a):\n return special.logsumexp(a, axis=1, keepdims=True)\n\n check_grads(f, modes=[\"fwd\", \"rev\"])(x)\n check_grads(lambda a: grad(f)(a), modes=[\"fwd\", \"rev\"])(x)\n\n ### Signal ###\n def test_convolve_generalization():\n ag_convolve = autograd.scipy.signal.convolve\n A_35 = R(3, 5)\n A_34 = R(3, 4)\n A_342 = R(3, 4, 2)\n A_2543 = R(2, 5, 4, 3)\n A_24232 = R(2, 4, 2, 3, 2)\n\n for mode in [\"valid\", \"full\"]:\n assert npo.allclose(\n ag_convolve(A_35, A_34, axes=([1], [0]), mode=mode)[1, 2],\n sp_convolve(A_35[1, :], A_34[:, 2], mode),\n )\n assert npo.allclose(\n ag_convolve(A_35, A_34, axes=([], []), dot_axes=([0], [0]), mode=mode),\n npo.tensordot(A_35, A_34, axes=([0], [0])),\n )\n assert npo.allclose(\n ag_convolve(A_35, A_342, axes=([1], [2]), dot_axes=([0], [0]), mode=mode)[2],\n sum([sp_convolve(A_35[i, :], A_342[i, 2, :], mode) for i in range(3)]),\n )\n assert npo.allclose(\n ag_convolve(A_2543, A_24232, axes=([1, 2], [2, 4]), dot_axes=([0, 3], [0, 3]), mode=mode)[2],\n sum(\n [\n sum(\n [sp_convolve(A_2543[i, :, :, j], A_24232[i, 2, :, j, :], mode) for i in range(2)]\n )\n for j in range(3)\n ]\n ),\n )\n\n def test_convolve():\n combo_check(autograd.scipy.signal.convolve, [0, 1])(\n [R(4), R(5), R(6)], [R(2), R(3), R(4)], mode=[\"full\", \"valid\"]\n )\n\n def test_convolve_2d():\n combo_check(autograd.scipy.signal.convolve, [0, 1])(\n [R(4, 3), R(5, 4), R(6, 7)], [R(2, 2), R(3, 2), R(4, 2), R(4, 1)], mode=[\"full\", \"valid\"]\n )\n\n def test_convolve_ignore():\n combo_check(autograd.scipy.signal.convolve, [0, 1])(\n [R(4, 3)],\n [R(3, 2)],\n axes=[([0], [0]), ([1], [1]), ([0], [1]), ([1], [0]), ([0, 1], [0, 1]), ([1, 0], [1, 0])],\n mode=[\"full\", \"valid\"],\n )\n\n def test_convolve_ignore_dot():\n combo_check(autograd.scipy.signal.convolve, [0, 1])(\n [R(3, 3, 2)],\n [R(3, 2, 3)],\n axes=[([1], [1])],\n dot_axes=[([0], [2]), ([0], [0])],\n mode=[\"full\", \"valid\"],\n )\n\n ### Special ###\n def test_beta():\n combo_check(special.beta, [0, 1])([R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1])\n\n def test_betainc():\n combo_check(special.betainc, [2])([R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1], [U(0.0, 1.0, 4)])\n\n def test_betaln():\n combo_check(special.betaln, [0, 1])([R(4) ** 2 + 1.1], [R(4) ** 2 + 1.1])\n\n def test_gammainc():\n combo_check(special.gammainc, [1])([1], R(4) ** 2 + 1.3)\n\n def test_gammaincc():\n combo_check(special.gammaincc, [1])([1], R(4) ** 2 + 1.3)\n\n def test_polygamma():\n combo_check(special.polygamma, [1])([0], R(4) ** 2 + 1.3)\n\n def test_jn():\n combo_check(special.jn, [1])([2], R(4) ** 2 + 1.3)\n\n def test_yn():\n combo_check(special.yn, [1])([2], R(4) ** 2 + 1.3)\n\n def test_psi():\n unary_ufunc_check(special.psi, lims=[0.3, 2.0], test_complex=False)\n\n def test_digamma():\n unary_ufunc_check(special.digamma, lims=[0.3, 2.0], test_complex=False)\n\n def test_gamma():\n unary_ufunc_check(special.gamma, lims=[0.3, 2.0], test_complex=False)\n\n def test_gammaln():\n unary_ufunc_check(special.gammaln, lims=[0.3, 2.0], test_complex=False)\n\n def test_gammasgn():\n unary_ufunc_check(special.gammasgn, lims=[0.3, 2.0], test_complex=False)\n\n def test_rgamma():\n unary_ufunc_check(special.rgamma, lims=[0.3, 2.0], test_complex=False)\n\n def test_multigammaln():\n combo_check(special.multigammaln, [0])([U(4.0, 5.0), U(4.0, 5.0, (2, 3))], [1, 2, 3])\n\n def test_j0():\n unary_ufunc_check(special.j0, lims=[0.2, 20.0], test_complex=False)\n\n def test_j1():\n unary_ufunc_check(special.j1, lims=[0.2, 20.0], test_complex=False)\n\n def test_y0():\n unary_ufunc_check(special.y0, lims=[0.2, 20.0], test_complex=False)\n\n def test_y1():\n unary_ufunc_check(special.y1, lims=[0.2, 20.0], test_complex=False)\n\n def test_i0():\n unary_ufunc_check(special.i0, lims=[0.2, 20.0], test_complex=False)\n\n def test_i1():\n unary_ufunc_check(special.i1, lims=[0.2, 20.0], test_complex=False)\n\n def test_iv():\n combo_check(special.iv, [1])(U(1.0, 50.0, 4), R(4) ** 2 + 1.3)\n\n def test_ive():\n combo_check(special.ive, [1])(U(1.0, 50.0, 4), R(4) ** 2 + 1.3)\n\n def test_erf():\n unary_ufunc_check(special.erf, lims=[-3.0, 3.0], test_complex=True)\n\n def test_erfc():\n unary_ufunc_check(special.erfc, lims=[-3.0, 3.0], test_complex=True)\n\n def test_erfinv():\n unary_ufunc_check(special.erfinv, lims=[-0.95, 0.95], test_complex=False)\n\n def test_erfcinv():\n unary_ufunc_check(special.erfcinv, lims=[0.05, 1.95], test_complex=False)\n\n def test_logit():\n unary_ufunc_check(special.logit, lims=[0.10, 0.90], test_complex=False)\n\n def test_expit():\n unary_ufunc_check(special.expit, lims=[-4.05, 4.95], test_complex=False)\n\n ### ODE integrator ###\n def func(y, t, arg1, arg2):\n return -np.sqrt(t) - y + arg1 - np.mean((y + arg2) ** 2)\n\n def test_odeint():\n combo_check(integrate.odeint, [1, 2, 3])([func], [R(3)], [np.linspace(0.1, 0.2, 4)], [(R(3), R(3))])\n\n ## Linalg\n def test_sqrtm():\n combo_check(spla.sqrtm, modes=[\"fwd\"], order=2)([R(3, 3)])\n\n def test_sqrtm():\n combo_check(symmetrize_matrix_arg(spla.sqrtm, 0), modes=[\"fwd\", \"rev\"], order=2)([R(3, 3)])\n\n def test_solve_sylvester():\n combo_check(spla.solve_sylvester, [0, 1, 2], modes=[\"rev\", \"fwd\"], order=2)(\n [R(3, 3)], [R(3, 3)], [R(3, 3)]\n )\n\n def test_solve_banded():\n combo_check(spla.solve_banded, [1, 2], modes=[\"rev\"], order=1)([(1, 1)], [R(3, 5)], [R(5)])\n"
},
{
"path": "tests/test_systematic.py",
"content": "import operator as op\n\nimport numpy as onp\nfrom numpy_utils import binary_ufunc_check, binary_ufunc_check_no_same_args, stat_check, unary_ufunc_check\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd.test_util import combo_check\n\nnpr.seed(0)\n\n\n# Array statistics functions\ndef test_max():\n stat_check(np.max)\n\n\n# def test_all(): stat_check(np.all)\n# def test_any(): stat_check(np.any)\ndef test_max():\n stat_check(np.max)\n\n\ndef test_mean():\n stat_check(np.mean)\n\n\ndef test_min():\n stat_check(np.min)\n\n\ndef test_sum():\n stat_check(np.sum)\n\n\ndef test_prod():\n stat_check(np.prod)\n\n\ndef test_var():\n stat_check(np.var)\n\n\ndef test_std():\n stat_check(np.std)\n\n\n# Unary ufunc tests\n\n\ndef test_sin():\n unary_ufunc_check(np.sin)\n\n\ndef test_abs():\n unary_ufunc_check(np.abs, lims=[0.1, 4.0])\n\n\ndef test_absolute():\n unary_ufunc_check(np.absolute, lims=[0.1, 4.0])\n\n\ndef test_arccosh():\n unary_ufunc_check(np.arccosh, lims=[1.1, 4.0])\n\n\ndef test_arcsinh():\n unary_ufunc_check(np.arcsinh, lims=[-0.9, 0.9])\n\n\ndef test_arctanh():\n unary_ufunc_check(np.arctanh, lims=[-0.9, 0.9])\n\n\ndef test_ceil():\n unary_ufunc_check(np.ceil, lims=[-1.5, 1.5], test_complex=False)\n\n\ndef test_cos():\n unary_ufunc_check(np.cos)\n\n\ndef test_cosh():\n unary_ufunc_check(np.cosh)\n\n\ndef test_deg2rad():\n unary_ufunc_check(np.deg2rad, test_complex=False)\n\n\ndef test_degrees():\n unary_ufunc_check(lambda x: np.degrees(x) / 50.0, test_complex=False)\n\n\ndef test_exp():\n unary_ufunc_check(np.exp)\n\n\ndef test_exp2():\n unary_ufunc_check(np.exp2)\n\n\ndef test_expm1():\n unary_ufunc_check(np.expm1)\n\n\ndef test_fabs():\n unary_ufunc_check(np.fabs, test_complex=False)\n\n\ndef test_floor():\n unary_ufunc_check(np.floor, lims=[-1.5, 1.5], test_complex=False)\n\n\ndef test_log():\n unary_ufunc_check(np.log, lims=[0.2, 2.0])\n\n\ndef test_log10():\n unary_ufunc_check(np.log10, lims=[0.2, 2.0])\n\n\ndef test_log1p():\n unary_ufunc_check(np.log1p, lims=[0.2, 2.0])\n\n\ndef test_log2():\n unary_ufunc_check(np.log2, lims=[0.2, 2.0])\n\n\ndef test_rad2deg():\n unary_ufunc_check(lambda x: np.rad2deg(x) / 50.0, test_complex=False)\n\n\ndef test_radians():\n unary_ufunc_check(np.radians, test_complex=False)\n\n\ndef test_sign():\n unary_ufunc_check(np.sign, test_complex=False)\n\n\ndef test_sin():\n unary_ufunc_check(np.sin)\n\n\ndef test_sinh():\n unary_ufunc_check(np.sinh)\n\n\ndef test_sqrt():\n unary_ufunc_check(np.sqrt, lims=[1.0, 3.0])\n\n\ndef test_square():\n unary_ufunc_check(np.square, test_complex=False)\n\n\ndef test_tan():\n unary_ufunc_check(np.tan, lims=[-1.1, 1.1])\n\n\ndef test_tanh():\n unary_ufunc_check(np.tanh)\n\n\ndef test_real():\n unary_ufunc_check(np.real)\n\n\ndef test_real_ic():\n unary_ufunc_check(np.real_if_close)\n\n\ndef test_imag():\n unary_ufunc_check(np.imag)\n\n\ndef test_conj():\n unary_ufunc_check(np.conj)\n\n\ndef test_conjugate():\n unary_ufunc_check(np.conjugate)\n\n\ndef test_angle():\n unary_ufunc_check(np.angle)\n\n\n# Binary ufunc tests\n\n\ndef test_add():\n binary_ufunc_check(np.add)\n\n\ndef test_logaddexp():\n binary_ufunc_check(np.logaddexp, test_complex=False)\n\n\ndef test_logaddexp2():\n binary_ufunc_check(np.logaddexp2, test_complex=False)\n\n\ndef test_remainder():\n binary_ufunc_check_no_same_args(np.remainder, lims_A=[-0.9, 0.9], lims_B=[0.7, 1.9], test_complex=False)\n\n\ndef test_true_divide():\n binary_ufunc_check(np.true_divide, lims_B=[0.8, 1.2], test_complex=False)\n\n\ndef test_mod():\n binary_ufunc_check_no_same_args(np.mod, lims_B=[0.8, 2.1], test_complex=False)\n\n\ndef test_true_divide_neg():\n binary_ufunc_check(np.true_divide, lims_B=[-0.3, -2.0], test_complex=False)\n\n\ndef test_mod_neg():\n binary_ufunc_check_no_same_args(np.mod, lims_B=[-0.3, -2.0], test_complex=False)\n\n\ndef test_op_mul():\n binary_ufunc_check(op.mul)\n\n\ndef test_op_add():\n binary_ufunc_check(op.add)\n\n\ndef test_op_sub():\n binary_ufunc_check(op.sub)\n\n\ndef test_op_mod():\n binary_ufunc_check_no_same_args(op.mod, lims_B=[0.3, 2.0], test_complex=False)\n\n\ndef test_op_mod_neg():\n binary_ufunc_check_no_same_args(op.mod, lims_B=[-0.3, -2.0], test_complex=False)\n\n\n# Misc tests\n\nR = npr.randn\nC = lambda *shape: npr.randn(*shape) + 1j * npr.randn(*shape)\n\n\ndef test_transpose():\n combo_check(np.transpose, [0])(\n [R(2, 3, 4)], axes=[None, [0, 1, 2], [0, 2, 1], [2, 0, 1], [2, 1, 0], [1, 0, 2], [1, 2, 0]]\n )\n\n\ndef test_moveaxis():\n combo_check(np.moveaxis, [0])([R(2, 3, 4)], source=[0, 1, 2], destination=[0, 1, 2])\n\n\ndef test_repeat():\n combo_check(np.repeat, [0])([R(2, 3, 4), R(3, 1)], repeats=[0, 1, 2], axis=[None, 0, 1])\n\n\ndef test_diff():\n combo_check(np.diff, [0])([R(5, 5), R(5, 5, 5)], n=[1, 2], axis=[0, 1])\n combo_check(np.diff, [0])([R(1), R(1, 1)], axis=[0])\n combo_check(np.diff, [0])([R(1, 1), R(3, 1)], axis=[1])\n\n\ndef test_gradient():\n combo_check(np.gradient, [0])([R(5, 5), R(5, 5, 5)], axis=[None, 0, 1, -1])\n combo_check(np.gradient, [0])([R(5, 5, 5)], axis=[(0, 1), (0, -1)])\n\n\ndef test_tile():\n combo_check(np.tile, [0])([R(2, 1, 3, 1)], reps=[(1, 4, 1, 2)])\n combo_check(np.tile, [0])([R(1, 2)], reps=[(1, 2), (2, 3), (3, 2, 1)])\n combo_check(np.tile, [0])([R(1)], reps=[(2,), 2])\n\n\ndef test_kron():\n combo_check(np.kron, [0, 1])(\n [R(5, 5), R(4, 4), R(5), R(5, 1), R(1, 5), R(), C(5, 5)],\n [R(3, 3), R(2, 2), R(3), R(1, 3), R(3, 1), R(), C(3, 3)],\n )\n\n\ndef test_inner():\n combo_check(np.inner, [0, 1])([1.5, R(3), R(2, 3)], [0.3, R(3), R(4, 3)])\n\n\ndef test_dot():\n combo_check(np.dot, [0, 1], order=3)(\n [1.5, R(3), R(2, 3), R(2, 2, 3), C(3), C(2, 3)], [0.3, R(3), R(3, 4), R(2, 3, 4), C(3)]\n )\n\n\ndef test_outer():\n combo_check(np.outer, [0, 1], order=3)([R(3), C(3)], [R(3), C(3)])\n\n\ndef test_matmul():\n combo_check(np.matmul, [0, 1])(\n [R(3), R(2, 3), R(2, 2, 3), C(3), C(2, 3)], [R(3), R(3, 4), R(2, 3, 4), C(3), C(3, 4)]\n )\n\n\ndef test_matmul_broadcast():\n combo_check(np.matmul, [0, 1])([R(1, 2, 2)], [R(3, 2, 1)])\n\n\ndef test_tensordot_1():\n combo_check(np.tensordot, [0, 1], order=3)(\n [R(1, 3), R(2, 3, 2), C(1, 3)], [R(3), R(3, 1), R(3, 4, 2), C(3)], axes=[[(1,), (0,)]]\n )\n\n\ndef test_tensordot_2():\n combo_check(np.tensordot, [0, 1], order=3)(\n [R(3), R(3, 1), R(3, 4, 2)], [R(1, 3), R(2, 3, 2)], axes=[[(0,), (1,)]]\n )\n\n\ndef test_tensordot_3():\n combo_check(np.tensordot, [0, 1], order=3)(\n [R(2, 3), R(2, 3, 4)], [R(1, 2, 3), R(2, 2, 3, 4)], axes=[[(0, 1), (1, 2)], [(1, 0), (2, 1)]]\n )\n\n\ndef test_tensordot_4():\n combo_check(np.tensordot, [0, 1], order=3)([R(2, 2), R(4, 2, 2)], [R(2, 2), R(2, 2, 4)], axes=[1, 2])\n\n\ndef test_tensordot_5():\n combo_check(np.tensordot, [0, 1], order=3)([R(4)], [R()], axes=[0])\n\n\ndef test_tensordot_6():\n combo_check(np.tensordot, [0, 1], order=3)([R(2, 6)], [R(6, 3)], axes=[[[-1], [0]]])\n\n\ndef test_tensordot_7():\n combo_check(np.tensordot, [0, 1], order=3)([R(2, 6)], [R(6, 3)], axes=[[-1, 0]])\n\n\ndef test_tensordot_8():\n combo_check(np.tensordot, [0, 1], order=3)([R(2)], [R(2, 2)], axes=[[0, 1]])\n\n\n# Need custom tests because gradient is undefined when arguments are identical.\ndef test_maximum():\n combo_check(np.maximum, [0, 1])([R(1), R(1, 4), R(3, 4)], [R(1), R(1, 4), R(3, 4)])\n\n\ndef test_fmax():\n combo_check(np.fmax, [0, 1])([R(1), R(1, 4), R(3, 4)], [R(1), R(1, 4), R(3, 4)])\n\n\ndef test_minimum():\n combo_check(np.minimum, [0, 1])([R(1), R(1, 4), R(3, 4)], [R(1), R(1, 4), R(3, 4)])\n\n\ndef test_fmin():\n combo_check(np.fmin, [0, 1])([R(1), R(1, 4), R(3, 4)], [R(1), R(1, 4), R(3, 4)])\n\n\ndef test_sort():\n combo_check(np.sort, [0])([R(1), R(7)])\n\n\nif onp.lib.NumpyVersion(onp.__version__) < \"2.0.0\":\n\n def test_msort():\n combo_check(np.msort, [0])([R(1), R(7)])\n\n\ndef test_partition():\n combo_check(np.partition, [0])([R(7), R(14)], kth=[0, 3, 6])\n\n\ndef test_atleast_1d():\n combo_check(np.atleast_1d, [0])([1.2, R(1), R(7), R(1, 4), R(2, 4), R(2, 4, 5)])\n\n\ndef test_atleast_2d():\n combo_check(np.atleast_2d, [0])([1.2, R(1), R(7), R(1, 4), R(2, 4), R(2, 4, 5)])\n\n\ndef test_atleast_3d():\n combo_check(np.atleast_3d, [0])([1.2, R(1), R(7), R(1, 4), R(2, 4), R(2, 4, 5), R(2, 4, 3, 5)])\n\n\ndef test_einsum_transpose():\n combo_check(np.einsum, [1])([\"ij->ji\"], [R(1, 1), R(4, 4), R(3, 4)])\n\n\ndef test_einsum_matmult():\n combo_check(np.einsum, [1, 2])([\"ij,jk->ik\"], [R(2, 3), C(2, 3)], [R(3, 4), C(3, 4)])\n\n\ndef test_einsum_matmult_broadcast():\n combo_check(np.einsum, [1, 2])([\"...ij,...jk->...ik\"], [R(2, 3), R(2, 2, 3)], [R(3, 4), R(2, 3, 4)])\n\n\ndef test_einsum_matmult_broadcast_leadzero():\n combo_check(np.einsum, [1, 2])([\"...ij,...jk->...ik\"], [R(0, 2, 3)], [R(0, 3, 4)])\n\n\ndef test_einsum_covsum():\n combo_check(np.einsum, [1, 2])([\"ijk,lji->lki\"], [R(3, 4, 4)], [R(4, 4, 3)])\n\n\ndef test_einsum_ellipses():\n combo_check(np.einsum, [1, 2])(\n [\"...jk,...lj->...lk\", \"...,...->...\"], [R(4, 4), R(3, 4, 4)], [R(4, 4), R(3, 4, 4)]\n )\n\n\ndef test_einsum_ellipses_tail():\n combo_check(np.einsum, [1, 2])([\"jk...,lj...->lk...\"], [R(3, 2), R(3, 2, 4)], [R(2, 3), R(2, 3, 4)])\n\n\ndef test_einsum_ellipses_center():\n combo_check(np.einsum, [1, 2])([\"j...k,lj...->lk...\"], [R(2, 2), R(2, 2, 2)], [R(2, 2), R(2, 2, 2)])\n\n\ndef test_einsum_three_args():\n combo_check(np.einsum, [1, 2])([\"ijk,lji,lli->lki\"], [R(3, 4, 4)], [R(4, 4, 3)], [R(4, 4, 3)])\n\n\ndef test_einsum2_transpose():\n combo_check(np.einsum, [0])([R(1, 1), R(4, 4), R(3, 4)], [(0, 1)], [(1, 0)])\n\n\ndef test_einsum2_matmult():\n combo_check(np.einsum, [0, 2])([R(2, 3)], [(0, 1)], [R(3, 4)], [(1, 2)], [(0, 2)])\n\n\ndef test_einsum2_matmult_broadcast():\n combo_check(np.einsum, [0, 2])(\n [R(2, 3), R(2, 2, 3)],\n [(Ellipsis, 0, 1)],\n [R(3, 4), R(2, 3, 4)],\n [(Ellipsis, 1, 2)],\n [(Ellipsis, 0, 2)],\n )\n\n\ndef test_einsum2_covsum():\n combo_check(np.einsum, [0, 2])([R(3, 4, 4)], [(0, 1, 2)], [R(4, 4, 3)], [(3, 1, 0)], [(3, 2, 0)])\n\n\ndef test_einsum2_three_args():\n combo_check(np.einsum, [0, 2])(\n [R(3, 4, 4)], [(0, 1, 2)], [R(4, 4, 3)], [(3, 1, 0)], [R(4, 4, 3)], [(3, 3, 0)], [(3, 2, 0)]\n )\n\n\ndef test_einsum_naked_sum():\n combo_check(np.einsum, [1, 2])([\"k,nk->\"], [R(5)], [R(10, 5)])\n\n\ndef test_einsum_naked_sum2():\n combo_check(np.einsum, [1])([\"abcd->bd\"], [R(3, 2, 3, 2)])\n\n\ndef test_einsum_naked_sum_ellipsis():\n combo_check(np.einsum, [1, 2])([\"...k,...nk->...\"], [R(3, 5)], [R(3, 10, 5)])\n\n\ndef test_einsum_no_output_indices():\n combo_check(np.einsum, [1, 2])([\"ij,k\"], [R(3, 4)], [R(3)])\n\n\ndef test_trace():\n combo_check(np.trace, [0])([R(5, 5), R(4, 5), R(5, 4), R(3, 4, 5)], offset=[-1, 0, 1])\n\n\ndef test_diag():\n combo_check(np.diag, [0])([R(5, 5)], k=[-1, 0, 1])\n\n\ndef test_diag_flat():\n combo_check(np.diag, [0])([R(5)], k=[-1, 0, 1])\n\n\ndef test_tril():\n combo_check(np.tril, [0])([R(5, 5)], k=[-1, 0, 1])\n\n\ndef test_triu():\n combo_check(np.triu, [0])([R(5, 5)], k=[-1, 0, 1])\n\n\ndef test_tril_3d():\n combo_check(np.tril, [0])([R(5, 5, 4)], k=[-1, 0, 1])\n\n\ndef test_triu_3d():\n combo_check(np.triu, [0])([R(5, 5, 4)], k=[-1, 0, 1])\n\n\ndef test_swapaxes():\n combo_check(np.swapaxes, [0])([R(3, 4, 5)], axis1=[0, 1, 2], axis2=[0, 1, 2])\n\n\ndef test_rollaxis():\n combo_check(np.rollaxis, [0])([R(2, 3, 4)], axis=[0, 1, 2], start=[0, 1, 2])\n\n\ndef test_cross():\n combo_check(np.cross, [0, 1])(\n [R(3, 3)], [R(3, 3)], axisa=[-1, 0, 1], axisb=[-1, 0, 1], axisc=[-1, 0, 1], axis=[None, -1, 0, 1]\n )\n\n\ndef test_vsplit_2d():\n combo_check(np.vsplit, [0])([R(4, 8)], [4, [1, 2]])\n\n\ndef test_vsplit_3d():\n combo_check(np.vsplit, [0])([R(4, 4, 4)], [2, [1, 2]])\n\n\ndef test_hsplit_2d():\n combo_check(np.hsplit, [0])([R(4, 8)], [4, [1, 2]])\n\n\ndef test_hsplit_3d():\n combo_check(np.hsplit, [0])([R(4, 4, 4)], [2, [1, 2]])\n\n\ndef test_dsplit_3d():\n combo_check(np.dsplit, [0])([R(4, 4, 4)], [2, [1, 2]])\n\n\ndef test_split_1d():\n combo_check(np.split, [0])([R(1), R(7)], [1], axis=[0])\n\n\ndef test_split_2d():\n combo_check(np.split, [0])([R(4, 8)], [4, [1, 2]], axis=[0, 1])\n\n\ndef test_split_3d():\n combo_check(np.split, [0])([R(4, 4, 4)], [2, [1, 2]], axis=[0, 1, 2])\n\n\ndef test_array_split_1d():\n combo_check(np.array_split, [0])([R(1), R(7)], [1, 3], axis=[0])\n\n\ndef test_array_split_2d():\n combo_check(np.array_split, [0])([R(7, 7)], [4, [3, 5]], axis=[0, 1])\n\n\ndef test_array_split_3d():\n combo_check(np.array_split, [0])([R(7, 7, 7)], [4, [3, 5]], axis=[0, 1, 2])\n\n\ndef test_concatenate_1ist():\n combo_check(np.concatenate, [0])([(R(1), R(3))], axis=[0])\n\n\ndef test_concatenate_tuple():\n combo_check(np.concatenate, [0])([[R(1), R(3)]], axis=[0])\n\n\ndef test_concatenate_2d():\n combo_check(np.concatenate, [0])([(R(2, 2), R(2, 2))], axis=[0, 1])\n\n\ndef test_concatenate_3d():\n combo_check(np.concatenate, [0])([(R(2, 2, 2), R(2, 2, 2))], axis=[0, 1, 2])\n\n\ndef test_vstack_1d():\n combo_check(np.vstack, [0])([R(2), (R(2), R(2))])\n\n\ndef test_vstack_2d():\n combo_check(np.vstack, [0])([R(2, 3), (R(2, 4), R(1, 4))])\n\n\ndef test_vstack_3d():\n combo_check(np.vstack, [0])([R(2, 3, 4), (R(2, 3, 4), R(5, 3, 4))])\n\n\ndef test_hstack_1d():\n combo_check(np.hstack, [0])([R(2), (R(2), R(2))])\n\n\ndef test_hstack_2d():\n combo_check(np.hstack, [0])([R(3, 2), (R(3, 4), R(3, 5))])\n\n\ndef test_hstack_3d():\n combo_check(np.hstack, [0])([R(2, 3, 4), (R(2, 1, 4), R(2, 5, 4))])\n\n\ndef test_stack_1d():\n combo_check(np.stack, [0])([(R(2),), (R(2), R(2))], axis=[0, 1])\n\n\ndef test_row_stack_1d():\n combo_check(np.row_stack, [0])([R(2), (R(2), R(2))])\n\n\ndef test_row_stack_2d():\n combo_check(np.row_stack, [0])([R(2, 3), (R(2, 4), R(1, 4))])\n\n\ndef test_column_stack_1d():\n combo_check(np.column_stack, [0])([R(2), (R(2), R(2))])\n\n\ndef test_column_stack_2d():\n combo_check(np.column_stack, [0])([R(2, 2), (R(2, 2), R(2, 2))])\n\n\ndef test_select():\n combo_check(np.select, [1])(\n [[R(3, 4, 5) > 0, R(3, 4, 5) > 0, R(3, 4, 5) > 0]],\n [[R(3, 4, 5), R(3, 4, 5), R(3, 4, 5)]],\n default=[0.0, 1.1],\n )\n\n\ndef test_pad():\n combo_check(np.pad, [0])(\n [R(2, 2)], [0, 3, (3,), (3, 2), ((3, 2),), ((1, 2), (3, 4)), ((0, 0), (0, 0))], [\"constant\"]\n )\n"
},
{
"path": "tests/test_tests.py",
"content": "from pytest import raises\n\nfrom autograd.extend import defvjp\nfrom autograd.test_util import check_grads\nfrom autograd.tracer import primitive\n\n\ndef test_check_vjp_1st_order_fail():\n @primitive\n def foo(x):\n return x * 2.0\n\n defvjp(foo, lambda ans, x: lambda g: g * 2.001)\n\n with raises(AssertionError, match=\"\\\\(VJP\\\\) check of foo failed\"):\n check_grads(foo, modes=[\"rev\"])(1.0)\n\n\ndef test_check_vjp_2nd_order_fail():\n @primitive\n def foo(x):\n return x * 2.0\n\n defvjp(foo, lambda ans, x: lambda g: bar(g) * 2)\n\n @primitive\n def bar(x):\n return x\n\n defvjp(bar, lambda ans, x: lambda g: g * 1.001)\n\n with raises(AssertionError, match=\"\\\\(VJP\\\\) check of vjp_foo failed\"):\n check_grads(foo, modes=[\"rev\"])(1.0)\n"
},
{
"path": "tests/test_truediv.py",
"content": "# This file is to check that future division works.\n\nfrom test_binary_ops import arg_pairs\n\nimport autograd.numpy as np\nfrom autograd.test_util import check_grads\n\n\ndef test_div():\n fun = lambda x, y: x / y\n make_gap_from_zero = lambda x: np.sqrt(x**2 + 0.5)\n for arg1, arg2 in arg_pairs():\n arg1 = make_gap_from_zero(arg1)\n arg2 = make_gap_from_zero(arg2)\n check_grads(fun)(arg1, arg2)\n"
},
{
"path": "tests/test_tuple.py",
"content": "import autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import grad\nfrom autograd import isinstance as ag_isinstance\nfrom autograd import tuple as ag_tuple\nfrom autograd.test_util import check_grads\n\nnpr.seed(1)\n\n\ndef test_getter():\n def fun(input_tuple):\n A = np.sum(input_tuple[0])\n B = np.sum(input_tuple[1])\n C = np.sum(input_tuple[1])\n return A + B + C\n\n d_fun = grad(fun)\n input_tuple = (npr.randn(5, 6), npr.randn(4, 3), npr.randn(2, 4))\n\n result = d_fun(input_tuple)\n assert np.allclose(result[0], np.ones((5, 6)))\n assert np.allclose(result[1], 2 * np.ones((4, 3)))\n assert np.allclose(result[2], np.zeros((2, 4)))\n\n\ndef test_grads():\n def fun(input_tuple):\n A = np.sum(np.sin(input_tuple[0]))\n B = np.sum(np.cos(input_tuple[1]))\n return A + B\n\n def d_fun(input_tuple):\n g = grad(fun)(input_tuple)\n A = np.sum(g[0])\n B = np.sum(np.sin(g[0]))\n C = np.sum(np.sin(g[1]))\n return A + B + C\n\n input_tuple = (npr.randn(5, 6), npr.randn(4, 3), npr.randn(2, 4))\n\n check_grads(fun)(input_tuple)\n check_grads(d_fun)(input_tuple)\n\n\ndef test_nested_higher_order():\n def outer_fun(x):\n def inner_fun(y):\n return y[0] * y[1]\n\n return np.sum(np.sin(np.array(grad(inner_fun)(ag_tuple((x, x))))))\n\n check_grads(outer_fun)(5.0)\n check_grads(grad(outer_fun))(10.0)\n check_grads(grad(grad(outer_fun)))(10.0)\n\n\ndef test_isinstance():\n def fun(x):\n assert ag_isinstance(x, tuple)\n assert ag_isinstance(x, ag_tuple)\n return x[0]\n\n fun((1.0, 2.0, 3.0))\n grad(fun)((1.0, 2.0, 3.0))\n"
},
{
"path": "tests/test_vspaces.py",
"content": "import itertools as it\nfrom functools import reduce\n\nimport numpy as np\n\nfrom autograd.core import vspace\nfrom autograd.test_util import check_grads, scalar_close\n\n\ndef check_vspace(value):\n vs = vspace(value)\n # --- required attributes ---\n size = vs.size\n add = vs.add\n scalar_mul = vs.scalar_mul\n inner_prod = vs.inner_prod\n randn = vs.randn\n zeros = vs.zeros\n ones = vs.ones\n standard_basis = vs.standard_basis\n\n # --- util ---\n def randns(N=2):\n return [randn() for i in range(N)]\n\n def rand_scalar():\n return float(np.random.randn())\n\n def rand_scalars(N=2):\n return [rand_scalar() for i in range(N)]\n\n def vector_close(x, y):\n z = randn()\n return scalar_close(inner_prod(z, x), inner_prod(z, y))\n\n # --- vector space axioms ---\n def associativity_of_add(x, y, z):\n return vector_close(add(x, add(y, z)), add(add(x, y), z))\n\n def commutativity_of_add(x, y):\n return vector_close(add(x, y), add(y, x))\n\n def identity_element_of_add(x):\n return vector_close(add(zeros(), x), x)\n\n def inverse_elements_of_add(x):\n return vector_close(zeros(), add(x, scalar_mul(x, -1.0)))\n\n def compatibility_of_scalar_mul_with_field_mul(x, a, b):\n return vector_close(scalar_mul(x, a * b), scalar_mul(scalar_mul(x, a), b))\n\n def identity_element_of_scalar_mul(x):\n return vector_close(scalar_mul(x, 1.0), x)\n\n def distributivity_of_scalar_mul_wrt_vector_add(x, y, a):\n return vector_close(scalar_mul(add(x, y), a), add(scalar_mul(x, a), scalar_mul(y, a)))\n\n def distributivity_of_scalar_mul_wrt_scalar_add(x, a, b):\n return vector_close(scalar_mul(x, a + b), add(scalar_mul(x, a), scalar_mul(x, b)))\n\n # --- closure ---\n def add_preserves_vspace(x, y):\n return vs == vspace(add(x, y))\n\n def scalar_mul_preserves_vspace(x, a):\n return vs == vspace(scalar_mul(x, a))\n\n # --- inner product axioms ---\n def symmetry(x, y):\n return scalar_close(inner_prod(x, y), inner_prod(y, x))\n\n def linearity(x, y, a):\n return scalar_close(inner_prod(scalar_mul(x, a), y), a * inner_prod(x, y))\n\n def positive_definitive(x):\n return 0 < inner_prod(x, x)\n\n def inner_zeros():\n return scalar_close(0, inner_prod(zeros(), zeros()))\n\n # --- basis vectors and special vectors---\n def basis_orthonormality():\n return all(\n [\n scalar_close(inner_prod(x, y), 1.0 * (ix == iy))\n for (ix, x), (iy, y) in it.product(enumerate(standard_basis()), enumerate(standard_basis()))\n ]\n )\n\n def ones_sum_of_basis_vects():\n return vector_close(reduce(add, standard_basis()), ones())\n\n def basis_correct_size():\n return len(list(standard_basis())) == size\n\n def basis_correct_vspace():\n return (vs == vspace(x) for x in standard_basis())\n\n def zeros_correct_vspace():\n return vs == vspace(zeros())\n\n def ones_correct_vspace():\n return vs == vspace(ones())\n\n def randn_correct_vspace():\n return vs == vspace(randn())\n\n assert associativity_of_add(*randns(3))\n assert commutativity_of_add(*randns())\n assert identity_element_of_add(randn())\n assert inverse_elements_of_add(randn())\n assert compatibility_of_scalar_mul_with_field_mul(randn(), *rand_scalars())\n assert identity_element_of_scalar_mul(randn())\n assert distributivity_of_scalar_mul_wrt_vector_add(randn(), randn(), rand_scalar())\n assert distributivity_of_scalar_mul_wrt_scalar_add(randn(), *rand_scalars())\n assert add_preserves_vspace(*randns())\n assert scalar_mul_preserves_vspace(randn(), rand_scalar())\n assert symmetry(*randns())\n assert linearity(randn(), randn(), rand_scalar())\n assert positive_definitive(randn())\n assert inner_zeros()\n assert basis_orthonormality()\n assert ones_sum_of_basis_vects()\n assert basis_correct_size()\n assert basis_correct_vspace()\n assert zeros_correct_vspace()\n assert ones_correct_vspace()\n assert randn_correct_vspace()\n\n # --- grads of basic operations ---\n check_grads(add)(*randns())\n check_grads(scalar_mul)(randn(), rand_scalar())\n check_grads(inner_prod)(*randns())\n\n\ndef test_array_vspace():\n check_vspace(np.zeros((3, 2)))\n\n\ndef test_array_vspace_0_dim():\n check_vspace(0.0)\n\n\ndef test_array_vspace_complex():\n check_vspace(1.0j * np.zeros((2, 1)))\n\n\ndef test_list_vspace():\n check_vspace([1.0, np.zeros((2, 1))])\n\n\ndef test_tuple_vspace():\n check_vspace((1.0, np.zeros((2, 1))))\n\n\ndef test_dict_vspace():\n check_vspace({\"a\": 1.0, \"b\": np.zeros((2, 1))})\n\n\ndef test_mixed_vspace():\n check_vspace({\"x\": [0.0, np.zeros((3, 1))], \"y\": ({\"a\": 0.0}, [0.0])})\n"
},
{
"path": "tests/test_wrappers.py",
"content": "import warnings\nfrom functools import partial\n\nimport pytest\n\nimport autograd.numpy as np\nimport autograd.numpy.random as npr\nfrom autograd import (\n checkpoint,\n elementwise_grad,\n grad,\n grad_and_aux,\n hessian,\n hessian_tensor_product,\n jacobian,\n make_ggnvp,\n make_hvp,\n make_jvp,\n tensor_jacobian_product,\n value_and_grad,\n)\nfrom autograd.test_util import check_equivalent, check_grads # , nd\nfrom autograd.tracer import isbox\n\nnpr.seed(1)\n\n\ndef test_return_both():\n fun = lambda x: 3.0 * x**3.2\n d_fun = grad(fun)\n f_and_d_fun = value_and_grad(fun)\n\n test_x = 1.7\n f, d = f_and_d_fun(test_x)\n assert f == fun(test_x)\n assert d == d_fun(test_x)\n\n\ndef test_value_and_grad():\n fun = lambda x: np.sum(np.sin(x) ** 2)\n dfun = grad(fun)\n dfun_both = value_and_grad(fun)\n x = npr.randn(5)\n assert not isbox(dfun_both(x)[0])\n check_equivalent(fun(x), dfun_both(x)[0])\n check_equivalent(dfun(x), dfun_both(x)[1])\n\n def fun2(x):\n return dfun_both(x)[0]\n\n check_grads(fun2)(x)\n\n\ndef test_hessian():\n # Check Hessian of a quadratic function.\n D = 5\n H = npr.randn(D, D)\n\n def fun(x):\n return np.dot(np.dot(x, H), x)\n\n hess = hessian(fun)\n x = npr.randn(D)\n check_equivalent(hess(x), H + H.T)\n\n\ndef test_multigrad():\n def complicated_fun(a, b, c, d, e, f=1.1, g=9.0):\n return a + np.sin(b) + np.cosh(c) + np.cos(d) + np.tan(e) + f + g\n\n def complicated_fun_3_1(d_b):\n d, b = d_b\n return complicated_fun(A, b, C, d, E, f=F, g=G)\n\n A = 0.5\n B = -0.3\n C = 0.2\n D = -1.1\n E = 0.7\n F = 0.6\n G = -0.1\n\n wrapped = grad(complicated_fun, argnum=[3, 1])(A, B, C, D, E, f=F, g=G)\n explicit = grad(complicated_fun_3_1)((D, B))\n check_equivalent(wrapped, explicit)\n\n\ndef test_value_and_multigrad():\n def complicated_fun(a, b, c, d, e, f=1.1, g=9.0):\n return a + np.sin(b) + np.cosh(c) + np.cos(d) + np.tan(e) + f + g\n\n A = 0.5\n B = -0.3\n C = 0.2\n D = -1.1\n E = 0.7\n F = 0.6\n G = -0.1\n\n dfun = grad(complicated_fun, argnum=[3, 1])\n dfun_both = value_and_grad(complicated_fun, argnum=[3, 1])\n\n check_equivalent(complicated_fun(A, B, C, D, E, f=F, g=G), dfun_both(A, B, C, D, E, f=F, g=G)[0])\n\n check_equivalent(dfun(A, B, C, D, E, f=F, g=G), dfun_both(A, B, C, D, E, f=F, g=G)[1])\n\n\ndef test_multigrad_onearg():\n fun = lambda x, y: np.sum(x + np.sin(y))\n packed_fun = lambda xy: np.sum(xy[0] + np.sin(xy[1]))\n A, B = npr.randn(3), npr.randn(3)\n check_equivalent(grad(fun, argnum=[0])(A, B), (grad(packed_fun)((A, B))[0],))\n\n\ndef test_elementwise_grad():\n def simple_fun(a):\n return a + np.sin(a) + np.cosh(a)\n\n A = npr.randn(10)\n\n wrapped = elementwise_grad(simple_fun)(A)\n explicit = np.array([grad(simple_fun)(A[i]) for i in range(len(A))])\n check_equivalent(wrapped, explicit)\n\n\ndef test_elementwise_grad_multiple_args():\n def simple_fun(a, b):\n return a + np.sin(a) + np.cosh(b)\n\n A = 0.9\n B = npr.randn(10)\n argnum = 1\n\n wrapped = elementwise_grad(simple_fun, argnum)(A, B)\n explicit = np.array([grad(simple_fun, argnum)(A, B[i]) for i in range(len(B))])\n check_equivalent(wrapped, explicit)\n\n\ndef test_hessian_tensor_product():\n fun = lambda a: np.sum(np.sin(a))\n a = npr.randn(5)\n v = npr.randn(5)\n H = hessian(fun)(a)\n check_equivalent(np.dot(H, v), hessian_tensor_product(fun)(a, v))\n\n\ndef test_hvp():\n fun = lambda a: np.sum(np.sin(a))\n a = npr.randn(5)\n v = npr.randn(5)\n H = hessian(fun)(a)\n hvp = make_hvp(fun)(a)[0]\n check_equivalent(np.dot(H, v), hvp(v))\n\n\ndef test_hessian_matrix_product():\n fun = lambda a: np.sum(np.sin(a))\n a = npr.randn(5, 4)\n V = npr.randn(5, 4)\n H = hessian(fun)(a)\n check_equivalent(np.tensordot(H, V), hessian_tensor_product(fun)(a, V))\n\n\ndef test_hessian_tensor_product_3d():\n fun = lambda a: np.sum(np.sin(a))\n a = npr.randn(5, 4, 3)\n V = npr.randn(5, 4, 3)\n H = hessian(fun)(a)\n check_equivalent(np.tensordot(H, V, axes=np.ndim(V)), hessian_tensor_product(fun)(a, V))\n\n\ndef test_tensor_jacobian_product():\n # This function will have an asymmetric jacobian matrix.\n fun = lambda a: np.roll(np.sin(a), 1)\n a = npr.randn(5)\n V = npr.randn(5)\n J = jacobian(fun)(a)\n check_equivalent(np.dot(V.T, J), tensor_jacobian_product(fun)(a, V))\n\n\ndef test_matrix_jacobian_product():\n fun = lambda a: np.roll(np.sin(a), 1)\n a = npr.randn(5, 4)\n V = npr.randn(5, 4)\n J = jacobian(fun)(a)\n check_equivalent(np.tensordot(V, J), tensor_jacobian_product(fun)(a, V))\n\n\ndef test_tensor_jacobian_product():\n fun = lambda a: np.roll(np.sin(a), 1)\n a = npr.randn(5, 4, 3)\n V = npr.randn(5, 4)\n J = jacobian(fun)(a)\n check_equivalent(np.tensordot(V, J, axes=np.ndim(V)), tensor_jacobian_product(fun)(a, V))\n\n\ndef test_deprecated_defgrad_wrapper():\n from autograd.core import primitive\n\n @primitive\n def new_mul(x, y):\n return x * y\n\n with warnings.catch_warnings(record=True) as w:\n new_mul.defgrad(lambda ans, x, y: lambda g: y * g)\n new_mul.defgrad(lambda ans, x, y: lambda g: x * g, argnum=1)\n\n def fun(x, y):\n return new_mul(x, y)\n\n mat1 = npr.randn(2, 2)\n mat2 = npr.randn(2, 2)\n check_grads(fun, modes=[\"rev\"])(mat1, mat2)\n\n\ndef test_deprecated_defvjp_wrapper():\n from autograd.core import primitive\n\n @primitive\n def new_mul(x, y):\n return x * y\n\n with warnings.catch_warnings(record=True) as w:\n new_mul.defvjp(lambda g, ans, vs, gvs, x, y: y * g)\n new_mul.defvjp(lambda g, ans, vs, gvs, x, y: x * g, argnum=1)\n\n def fun(x, y):\n return new_mul(x, y)\n\n mat1 = npr.randn(2, 2)\n mat2 = npr.randn(2, 2)\n check_grads(fun, modes=[\"rev\"])(mat1, mat2)\n\n\ndef test_deprecated_defvjp_is_zero_wrapper():\n from autograd.core import primitive\n\n @primitive\n def new_mul(x, y):\n return 0 * x * y\n\n with warnings.catch_warnings(record=True) as w:\n new_mul.defvjp_is_zero([0, 1])\n\n def fun(x, y):\n return new_mul(x, y)\n\n mat1 = npr.randn(2, 2)\n mat2 = npr.randn(2, 2)\n with warnings.catch_warnings(record=True) as w:\n check_grads(fun, modes=[\"rev\"])(mat1, mat2)\n\n\ndef test_deprecated_quick_grad_check_wrapper():\n from autograd.util import quick_grad_check\n\n with warnings.catch_warnings(record=True) as w:\n quick_grad_check(lambda x, y: x**2 + y, 1.0, (2.0,))\n\n\ndef test_partial():\n def f(x, y):\n return x\n\n grad(partial(f, y=1))\n\n\n@pytest.mark.skip(reason=\"fails with NumPy nightlies\")\ndef test_dtypes():\n def f(x):\n return np.real(np.sum(x**2))\n\n # Array y with dtype np.float32\n y = np.random.randn(10, 10).astype(np.float32)\n assert grad(f)(y).dtype.type is np.float32\n\n y = np.random.randn(10, 10).astype(np.float16)\n assert grad(f)(y).dtype.type is np.float16\n\n y = np.random.randn(10, 10).astype(np.longdouble)\n grad(f)(y)\n\n y = np.random.randn(10, 10).astype(np.clongdouble)\n grad(f)(y)\n\n\ndef test_checkpoint_correctness():\n bar = lambda x, y: 2 * x + y + 5\n checkpointed_bar = checkpoint(bar)\n foo = lambda x: bar(x, x / 3.0) + bar(x, x**2)\n foo2 = lambda x: checkpointed_bar(x, x / 3.0) + checkpointed_bar(x, x**2)\n assert np.allclose(foo(3.0), foo2(3.0))\n assert np.allclose(grad(foo)(3.0), grad(foo2)(3.0))\n\n baz = lambda *args: sum(args)\n checkpointed_baz = checkpoint(baz)\n foobaz = lambda x: baz(x, x / 3.0)\n foobaz2 = lambda x: checkpointed_baz(x, x / 3.0)\n assert np.allclose(foobaz(3.0), foobaz2(3.0))\n assert np.allclose(grad(foobaz)(3.0), grad(foobaz2)(3.0))\n\n\ndef checkpoint_memory():\n \"\"\"This test is meant to be run manually, since it depends on\n memory_profiler and its behavior may vary.\"\"\"\n try:\n from memory_profiler import memory_usage\n except ImportError:\n return\n\n def f(a):\n for _ in range(10):\n a = np.sin(a**2 + 1)\n return a\n\n checkpointed_f = checkpoint(f)\n\n def testfun(f, x):\n for _ in range(5):\n x = f(x)\n return np.sum(x)\n\n gradfun = grad(testfun, 1)\n\n A = npr.RandomState(0).randn(100000)\n max_usage = max(memory_usage((gradfun, (f, A))))\n max_checkpointed_usage = max(memory_usage((gradfun, (checkpointed_f, A))))\n\n assert max_checkpointed_usage < max_usage / 2.0\n\n\ndef test_make_jvp():\n A = npr.randn(3, 5)\n x = npr.randn(5)\n v = npr.randn(5)\n fun = lambda x: np.tanh(np.dot(A, x))\n\n jvp_explicit = lambda x: lambda v: np.dot(jacobian(fun)(x), v)\n jvp = make_jvp(fun)\n\n check_equivalent(jvp_explicit(x)(v), jvp(x)(v)[1])\n\n\ndef _make_explicit_ggnvp(f, g=lambda x: 1.0 / 2 * np.dot(x, x)):\n def ggnvp_maker(x):\n J = jacobian(f)(x)\n H = hessian(g)(f(x))\n\n def ggnvp(v):\n return np.dot(J.T, np.dot(H, np.dot(J, v)))\n\n return ggnvp\n\n return ggnvp_maker\n\n\ndef test_make_ggnvp():\n A = npr.randn(5, 4)\n x = npr.randn(4)\n v = npr.randn(4)\n\n fun = lambda x: np.dot(A, x)\n check_equivalent(make_ggnvp(fun)(x)(v), _make_explicit_ggnvp(fun)(x)(v))\n\n fun2 = lambda x: np.tanh(np.dot(A, x))\n check_equivalent(make_ggnvp(fun2)(x)(v), _make_explicit_ggnvp(fun2)(x)(v))\n\n\ndef test_make_ggnvp_nondefault_g():\n A = npr.randn(5, 4)\n x = npr.randn(4)\n v = npr.randn(4)\n\n g = lambda y: np.sum(2.0 * y**2 + y**4)\n\n fun = lambda x: np.dot(A, x)\n check_equivalent(make_ggnvp(fun, g)(x)(v), _make_explicit_ggnvp(fun, g)(x)(v))\n\n fun2 = lambda x: np.tanh(np.dot(A, x))\n check_equivalent(make_ggnvp(fun2, g)(x)(v), _make_explicit_ggnvp(fun2, g)(x)(v))\n\n\ndef test_grad_and_aux():\n A = npr.randn(5, 4)\n x = npr.randn(4)\n\n f = lambda x: (np.sum(np.dot(A, x)), x**2)\n g = lambda x: np.sum(np.dot(A, x))\n\n assert len(grad_and_aux(f)(x)) == 2\n\n check_equivalent(grad_and_aux(f)(x)[0], grad(g)(x))\n check_equivalent(grad_and_aux(f)(x)[1], x**2)\n\n\n## No longer support this behavior\n# def test_make_ggnvp_broadcasting():\n# A = npr.randn(4, 5)\n# x = npr.randn(10, 4)\n# v = npr.randn(10, 4)\n\n# fun = lambda x: np.tanh(np.dot(x, A))\n# res1 = np.stack([_make_explicit_ggnvp(fun)(xi)(vi) for xi, vi in zip(x, v)])\n# res2 = make_ggnvp(fun)(x)(v)\n# check_equivalent(res1, res2)\n\n\ndef test_wrapped_name_and_docs():\n def foo(x):\n pass\n\n assert grad.__name__ == \"grad\"\n # Python 3.13: Compiler now strip indents from docstrings.\n # https://docs.python.org/3.13/whatsnew/3.13.html#other-language-changes\n assert grad.__doc__.startswith(tuple(f\"\\n{indent}Returns a function which\" for indent in (\" \", \"\")))\n assert grad(foo, 1).__name__ == \"grad_of_foo_wrt_argnum_1\"\n assert grad(foo, 1).__doc__.startswith(\" grad of function foo with\")\n"
}
]