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

..

image

Merge pull request #9689 from fehiepsi:image

2022-03-30 06:05:30 -07:00

lax

remove lax._device_put_raw

2022-04-20 14:49:53 -07:00

lib

[MHLO] Remove most XLA translation rules.

2022-04-18 08:28:35 -07:00

nn

Use lax.broadcasted_iota in jax.nn.one_hot.

2022-04-20 09:09:19 -04:00

numpy

[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

ops

lax_numpy: factor out indexing tricks

2022-03-17 11:05:45 -07:00

scipy

Fix typo of #10381

2022-04-21 02:17:09 +08:00

third_party

Add some functions for spectral analysis.

2022-02-17 15:59:24 +09:00

__init__.py

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

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

abstract_arrays.py

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

ad_checkpoint.py

[MHLO] Remove most XLA translation rules.

2022-04-18 08:28:35 -07:00

ad_util.py

Move jax.ad_util to jax._src.ad_util.

2021-06-07 14:51:34 -07:00

api_util.py

Deprecate jax.tree_util.tree_multimap

2022-04-01 14:51:54 -07:00

api.py

[MHLO] Change jax.xla_computation() to use MHLO lowering internally.

2022-04-13 06:28:38 -07:00

checkify.py

Checkify: support checks on data-independent values.

2022-04-19 15:06:00 +01:00

cloud_tpu_init.py

Remove flags set for all v4 TPUs. Topology flags will now be set in libTPU.

2022-04-18 16:39:34 -07:00

config.py

Switch gpu, cpu and jax2tf to use the new OptimizationBarrier op.

2022-04-13 22:51:36 -07:00

custom_api_util.py

enable custom transformation "stacking"

2022-01-11 17:55:08 -08:00

custom_batching.py

[MHLO] Remove most XLA translation rules.

2022-04-18 08:28:35 -07:00

custom_derivatives.py

fix issue #10366

2022-04-19 13:18:00 -07:00

custom_transpose.py

[MHLO] Remove most XLA translation rules.

2022-04-18 08:28:35 -07:00

device_array.py

Re-land changes from https://github.com/google/jax/pull/10069

2022-04-04 12:18:43 -07:00

dispatch.py

Fix name stack bugs

2022-04-19 11:14:41 -07:00

distributed.py

Add jax.distributed and jax.dlpack to the docs.

2022-02-17 16:10:07 -05:00

dlpack.py

Drop dead code now that the minimum jaxlib version is 0.3.2.

2022-04-13 13:34:00 -04:00

dtypes.py

Add comment explaining implementation in promote_types

2022-04-20 08:44:49 -07:00

errors.py

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

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

flatten_util.py

Documentation improvements.

2022-03-08 09:37:33 -05:00

iree.py

add some simple iree tests

2022-04-14 10:55:00 -07:00

lax_reference.py

Improve implementation of cbrt() in JAX.

2021-07-22 14:01:28 -07:00

pretty_printer.py

add flag for jaxpr colorful syntax highlighting

2022-02-12 14:15:28 -08:00

prng.py

Reapply: fff370d78d107ed81431becf9dfe97eba77863fb by Lukas Geiger <lukas.geiger94@gmail.com>:

2022-04-19 07:12:27 -07:00

profiler.py

Add profiler annotations

2021-12-06 15:13:35 -08:00

public_test_util.py

Separate jax.test_util implementations into public and private sources.

2022-04-04 14:43:39 -07:00

random.py

[MHLO] Remove most XLA translation rules.

2022-04-18 08:28:35 -07:00

source_info_util.py

Increase minimum jaxlib version to 0.3.0.

2022-03-04 10:33:03 -05:00

stages.py

Expose the input and output sharding on the compiled object.

2022-04-13 10:18:25 -07:00

test_util.py

[MHLO] Remove most XLA translation rules.

2022-04-18 08:28:35 -07:00

traceback_util.py

Increase minimum jaxlib version to 0.3.0.

2022-03-04 10:33:03 -05:00

tree_util.py

Merge pull request #9569 from GJBoth:tree_flatten_docs

2022-04-14 16:09:47 -07:00

util.py

Comment on implementation of unzip2 & unzip3

2022-04-14 13:41:05 -07:00