* Fix minor typo in cell
One of the arguments to `hvp` wasn't being used, which made the example slightly confusing.
* Fix both definitions of hvp in the autodiff cookbook.
Co-authored-by: Peter Hawkins <phawkins@google.com>
* Implement jax.ops.index_mul.
* Add index_mul to documentation.
* Fix RHS JVP rule for scatter_mul, fix test bug that meant it was not tested.
* Fix typo in docstring.
* Implement np.roots.
* Expose jit-compatible variant of np.roots.
General np.roots implementation has a value dependent output shape.
If the input coefficients are guaranteed to have no leading zeros,
output shape is independent of values. Skip checking for leading
zeros by setting a keyword argument.
* Fix typo.
* Make roots jit-argument keyword only.
Co-Authored-By: Stephan Hoyer <shoyer@google.com>
* Format docstring to enable parsing.
Co-Authored-By: Stephan Hoyer <shoyer@google.com>
* Add np.roots function to documentation.
* Add more tests for np.roots function.
- Include length 0 polynomial coefficients
- Test strip_zeros=False argument
- Test jit compiled version (only on cpu due to eigvals)
- Confirm that adding leading zeros while skipping check
for them results in nan's (expected behavior)
* Fix bug in np.roots test.
The polynomial with coefficents [0] never fails because the number of
roots is 0.
* Avoid bug in eigvals and adjust test accuracy.
The parameters of the test that was changed are non-essential
since they test for how the code behaves given invalid inputs.
The accuracy in comparing to the numpy result is changed because
the algorithm in those cases is slightly changed with respect to
the original numpy algorithm (to allow jit).
Co-authored-by: Stephan Hoyer <shoyer@google.com>
Introduced two new constructors for PartialVal: unknown and known.
These should make it easier to read the code where we construct
PartialVal:
* instead of PartialVal((aval, core.unit) we use PartialVal.unknown(aval)
* instead of PartialVal((None, pval)) we use PartialVal.known(pval)
Also disabled some new tests in random_tests.py on Mac. They segfault,
apparently due to the same issue #432.
* Minor update to docs; trigger readthedocs
* Updated Common Gotchas notebook
Handle errors explicitly, otherwise it is too hard to test the notebook by 'Run all'
* Added a section about pure functions to Common Gotchas
