Adam Paszke f86bf12b5a Add support for axis names in jnp.{sum,min,max} ...

Similarly to `jnp.einsum`, whenever we encounter an extension to the
positional NumPy API (in the case of reductions, the extension is
whenever a non-integer axis is specified), we reroute the call to a
parallel primitive instead of the standard lax reductions.

Note that this makes the parallel primitives implement a strict subset
of functionality of the lax reductions so in the future (when we decide
that we want axes to be truly first class) we can always swap out the
implementation for the parallel version. But, it makes sense to keep
them separate for the ease of prototyping in the near future.

2021-02-01 11:41:05 +00:00

..

optimize

all changes plus test verifcation on TPU squashed

2020-10-26 23:58:09 +01:00

sparse

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

stats

add support for scipy.stats.poisson.cdf

2021-01-24 16:15:31 +00:00

__init__.py

Move most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

linalg.py

Fix some doc build warnings

2021-01-25 14:08:57 -08:00

ndimage.py

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

signal.py

Fix #4775 + additional fixes

2020-11-09 10:40:14 +05:30

special.py

Add support for axis names in jnp.{sum,min,max}

2021-02-01 11:41:05 +00:00