michaelmarien 3e9f8248f2 Expand implementation of lax_numpy._quantile to allow the input of a tuple as axis argument ...

* support and test edge case where axis argument is empty tuple ()
* replace swapaxis + reshape methodology by one call to lax.reshape for computational efficiency's sake
* add check on repeated axis and throw ValueError
* introduced and changed corresponding numpy code to swap and reshape axis to be quantiled
* introduced code to accomodate the reintroduction of those axes if keepdims=True
* added testcases

2022-02-08 21:03:02 +01:00

..

__init__.py

Move jax.numpy internals into jax._src.numpy.

2020-10-16 20:35:19 -04:00

fft.py

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08:00

lax_numpy.py

Expand implementation of lax_numpy._quantile to allow the input of a tuple as axis argument

2022-02-08 21:03:02 +01:00

linalg.py

linalg_test: disable implicit rank promotion

2022-01-26 09:29:06 -08:00

polynomial.py

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08:00

util.py

maint: update pre-commit package versions & fix new mypy errors

2022-01-31 13:39:11 -08:00

vectorize.py

djax: let make_jaxpr build dyn shape jaxprs

2022-02-01 00:10:21 -08:00