Port np.roots (#2250)
* 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>
