Peter Hawkins a52f07a21b Add an optional mode= argument to jnp.take_along_axis. ...

This allows users of jnp.take_along_axis to override the out-of-bounds indexing behavior.
Default to "clip", which for the forward computation is identical to the current behavior. In a future change, we will change this to "fill".

2022-04-19 16:07:00 -04:00

..

__init__.py

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

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

fft.py

Fix fft dtype for norm='ortho'

2022-03-10 10:39:52 -05:00

index_tricks.py

simplify _IndexGrid

2022-04-15 04:17:27 +08:00

lax_numpy.py

Add an optional mode= argument to jnp.take_along_axis.

2022-04-19 16:07:00 -04:00

linalg.py

Replace reshape with expand_dims if possible

2022-03-31 01:34:26 +01:00

ndarray.py

lax_numpy.py: factor ndarray into its own module

2022-03-02 12:18:46 -08:00

polynomial.py

Add polydiv to jax.numpy

2022-04-13 18:31:27 +00:00

reductions.py

maint: remove numpy reductions dependence on ufuncs

2022-03-23 15:56:59 -07:00

setops.py

lax_numpy: refactor set operations into separate private submodule

2022-03-21 09:38:11 -07:00

ufuncs.py

lax_numpy.py: factor ufuncs into their own private submodule

2022-03-14 19:14:33 -07:00

util.py

docs: use intersphinx links for wrapped functions

2022-03-29 14:43:59 -07:00

vectorize.py

djax: let make_jaxpr build dyn shape jaxprs

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