Yash Katariya 7de522c5a3 Enter into auto mode for .at[...].get(...) a bit earlier so that all ops inside _gather are in auto mode. ...

Fix select's batching rule where `explicit_mesh_axis` that we capture in `axis_data` was not propagated properly to the `broadcast` happening in `bdim_at_front`.

PiperOrigin-RevId: 748867490

2025-04-17 17:42:13 -07:00

..

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

array_api_metadata.py

[array_api] update array_api_version to 2024.12

2025-04-02 13:46:07 -07:00

array_creation.py

DOC: add documentation note about default dtypes

2025-03-28 15:20:58 -07:00

array_methods.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

einsum.py

jnp.einsum: add support for __jax_array__

2025-04-01 12:26:26 -07:00

error.py

Add OOB checks to jax.numpy array indexing

2025-04-01 19:17:57 -07:00

fft.py

refactor: move scalar type defs out of lax_numpy.py

2025-02-06 14:48:10 -08:00

index_tricks.py

Set __module__ attribute for objects in jax.numpy

2024-11-15 06:03:54 -08:00

indexing.py

Enter into auto mode for .at[...].get(...) a bit earlier so that all ops inside _gather are in auto mode.

2025-04-17 17:42:13 -07:00

lax_numpy.py

jnp.repeat: don't cast repeats to array, as they must be static.

2025-04-08 10:32:03 -07:00

linalg.py

jnp.linalg: add symmetrize_input argument & docs

2025-04-07 14:46:38 -07:00

polynomial.py

Migrated to mypy 1.14.1 with --allow_redefinition

2025-02-13 15:38:28 +00:00

reductions.py

Added lax.axis_size and switched all existing usage of psum(1, ...) to it

2025-04-17 02:22:25 -07:00

scalar_types.py

cleanup now that we depend on ml_dtypes>=0.5

2025-03-28 07:44:38 -07:00

setops.py

Update jnp.unique to support upstream interface changes.

2025-03-04 05:24:52 -05:00

sorting.py

refactor: move sorting ops out of lax_numpy

2025-02-07 08:18:04 -08:00

tensor_contractions.py

refactor: move lax_numpy tensor contractions into their own file

2025-02-10 18:56:18 -08:00

ufunc_api.py

refactor: move lax_numpy indexing routines to their own submodule

2025-02-13 12:03:07 -08:00

ufuncs.py

jnp.power: better docs for invalid input

2025-04-14 10:42:29 -07:00

util.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

vectorize.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

window_functions.py

refactor: move lax_numpy window functions into their own file

2025-02-07 11:21:38 -08:00