Peter Hawkins 74346f464b [JAX] Change jnp.take_along_axis to return invalid (e.g. NaN) values for out-of-bounds indices. ...

Previously, out-of-bounds indices were clipped into range, but that behavior is error prone. We would rather fail in a more visible way when out-of-bounds indices are used. Future changes will migrate other JAX indexing operations to have the same semantics.

PiperOrigin-RevId: 443390170

2022-04-21 08:52:14 -07: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

[JAX] Change jnp.take_along_axis to return invalid (e.g. NaN) values for out-of-bounds indices.

2022-04-21 08:52:14 -07:00

linalg.py

remove numpy.linalg._promote_arg_dtypes

2022-04-21 00:23:56 +08: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