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

..

control_flow

Use _make_lengths_same for explicit mode too.

2025-04-17 10:48:46 -07:00

__init__.py

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

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

ann.py

Make sure vma on ShapedArray exists by default to make development easier. The field is populated inside shard_map guarded on the varying_axes_in_types config though.

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

convolution.py

[NFC] Rename standard_insert_pbroadcast to standard_insert_pvary

2025-04-15 11:02:45 -07:00

eigh.py

refactor: import numpy objects directly in jax.numpy

2025-02-14 12:47:58 -08:00

fft.py

Make sure vma on ShapedArray exists by default to make development easier. The field is populated inside shard_map guarded on the varying_axes_in_types config though.

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

lax.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

linalg.py

[NFC] Rename standard_insert_pbroadcast to standard_insert_pvary

2025-04-15 11:02:45 -07:00

other.py

Simplify and consolidate dot algorithm control in lax.

2024-10-07 13:21:34 -07:00

parallel.py

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

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

qdwh.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

slicing.py

[NFC] Rename standard_insert_pbroadcast to standard_insert_pvary

2025-04-15 11:02:45 -07:00

special.py

[NFC] Rename standard_insert_pbroadcast to standard_insert_pvary

2025-04-15 11:02:45 -07:00

stack.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

svd.py

Enable pivoted QR on CPU devices.

2025-01-09 20:44:45 +00:00

utils.py

Make sure vma on ShapedArray exists by default to make development easier. The field is populated inside shard_map guarded on the varying_axes_in_types config though.

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

windowed_reductions.py

Use _make_lengths_same for explicit mode too.

2025-04-17 10:48:46 -07:00