Matthew Johnson 66a6eb299e add autodiff rules for jax.lax.ragged_all_to_all collective ...

also update the ragged_all_to_all docstring. pseudocode in the style of the shard_map tutorial would be better and cleaner, but it needs the context of the tutorial to explain; i'll add ra2a to the shmap tutorial in the future.

PiperOrigin-RevId: 735957604

2025-03-11 18:22:02 -07:00

..

control_flow

jax.disable_jit: ensure while_loop behaves similarly to non-disable_jit version

2025-03-11 09:53:34 -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

Clean up and fix primal type to tangent type mapping

2024-09-18 13:43:54 -07:00

convolution.py

Clone RandomAlgorithm into lax.py, instead of using the version from XLA.

2024-10-11 11:03:15 -07:00

eigh.py

refactor: import numpy objects directly in jax.numpy

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

fft.py

Clone RandomAlgorithm into lax.py, instead of using the version from XLA.

2024-10-11 11:03:15 -07:00

lax.py

add autodiff rules for jax.lax.ragged_all_to_all collective

2025-03-11 18:22:02 -07:00

linalg.py

Move (most) jaxlib linalg custom call registration into JAX.

2025-03-10 08:17:44 -07:00

other.py

Simplify and consolidate dot algorithm control in lax.

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

parallel.py

add autodiff rules for jax.lax.ragged_all_to_all collective

2025-03-11 18:22:02 -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

Add a divisibility check so that we make sure that sharding evenly divides the shape (until this restriction is lifted) to make sure we don't create bad shardings.

2025-03-07 07:01:34 -08:00

special.py

Fix scipy.special.gammainc/gammaincc evaluation at boundary points

2025-03-11 21:18:47 +02: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

Start adding primitive registration helper functions to lax.linalg.

2025-02-21 04:05:34 -08:00

windowed_reductions.py

[sharding_in_types] Add sharding rules for the following primitives:

2025-02-22 10:45:58 -08:00