* first commit with placeholders for tests
* added tests for the following:
1 - inverse
2 - dtypes
3 - size
4 - axis
added error tests for the following:
1 - multiple axes provided instead of single axis
2 - axis out of bounds
* removed placeholders
added functions to .rst file
* Implement jax.numpy.nonzero.
* Implement the one-argument form of np.where.
* Fix output type and error message.
* Add lax_description strings to where and nonzero.
* Support IntEnum values as arguments to JAX functions.
When abstractifying a Python value, search the method-resolution order (MRO) of the type rather than only looking at the value's own type. IntEnum instances are subclasses of int, so this allows us to correctly handle them as integers, much as NumPy itself does.
Change the JAX type promotion table to prefer inexact types during type promotion.
NumPy's type promotion rules tend to promote aggressively to float64, which isn't a very accelerator-friendly behavior when not all accelerators (e.g., TPUs) support 64-bit floating point types. Even on accelerators that support 64-bit floating point types (e.g., GPUs), promotion to a 64-bit type comes with a significant performance cost.
This change makes JAX type promotion between inexact and exact types closer to PyTorch's promotion semantics, which are a better fit for modern accelerators:
e.g.,
```
import numpy as onp
from jax import numpy as np
In [1]: onp.promote_types(onp.float32, onp.int32)
Out[1]: dtype('float64')
In [2]: onp.promote_types(onp.float16, onp.int64)
Out[2]: dtype('float64')
In [3]: np.promote_types(onp.float32, onp.int32)
Out[3]: dtype('float32')
In [4]: np.promote_types(onp.float16, onp.int64)
Out[4]: dtype('float16')
```
This change is in preparation for enabling x64 mode by default on all platforms.
We want to allow users to control how reverse-mode autodiff saves values
from the forward pass. In particular, we want it to be easy to signal
that a function shouldn't have any of its intermediate residuals stored
for the backward pass, and instead those values should be recomputed
from the function's saved inputs. (This feature is especially handy for
accelerators on which memory access is much more expensive than FLOPs
are.) In JAX terms, since we implement reverse-mode as a composition of
forward-mode, partial evaluation, and transposition, we want users to
control how partial evaluation behaves.
See https://github.com/google/jax/pull/1749 for more.
Co-authored-by: Dougal Maclaurin <dougalm@google.com>
* [doc] Note that building jaxlib from source isn't always necessary
Building jaxlib from source is time-consuming and the source of most
pain for building JAX. It's also not necessary (in my experience) for
pure-Python changes.
This commit adds notes to the 'building from source' documentation to
make this explicit.
* Move ``jaxlib`` skip instructions to top
Also add a new section for the np.fft module. These functions were previously
not appearing in the docs, because fftn is not exposed as np.fftn but only as
np.fft.fftn.
* WIP: linear solvers
* Draft of lax.linear_solve
* Refactor pytree munging inside lax.root.
The primitive's implementation and JVP rules are now 100% pytree free.
* Fixup linear_solve
* Linearize multiple times in _root_jvp to avoid zeros
* fix deftraced
* add a symmetric argument
* Fixup float64; add a test for symmetric/non-symmetric
* test zeros in linear_solve_jvp
* Revisions per review
* Adjust signature of linear_solve
* restore botched test
* variable names
* WIP: root solve jaxpr
* WIP more tests
* rewrite root
* Root works with jaxprs
* root -> custom_root
* WIP undefined tangent
* Delayed undefined JVP errors
* use raise_on_undefined_tangents inside define_implicit_gradient
* more tests on jvps with undefined tangents
* Remove define_implicit_gradient
* Support closures in custom_root
* revert api-test
* another test
* jit tests
* spelling
* Moved all notebooks to docs/notebooks.
Now all notebooks are in the same place, thus all are subject
to auto-doc generation at readthedocs.io and to automated testing
with travis.
Some notebooks are too slow, exclude them at docs/conf.py:exclude_patterns.
Cleanup a bit the section headings in notebooks so that they show
up well in readtehdocs.io.
* Increase the cell timeout for executing notebooks
* Exclude also the neural network notebook from auto-generation (timing out)
* Disable the score_matching notebook from auto-doc (travis does not have sklearn)
Had to extend the docs/requirements.txt file to install
matplotlb (needed by the Gotchas notebook) and ".",
needed by everything. This results in a reduction
of the sphinx warnings from 3300 to 1200!
