Jake VanderPlas 8b46e53a4f jax.lax: improve docs for several APIs

2025-03-18 08:55:38 -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

jax.lax: improve docs for several APIs

2025-03-18 08:55:38 -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

Update gammainc and gammaincc against scipy 1.16: return nan whenever one of operands is nan.

2025-03-12 17:48:45 +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

Raise a better error with more info when we see duplicate axis in a PartitionSpec resulting from a sharding rule.

2025-03-13 15:24:10 -07:00

windowed_reductions.py

[sharding_in_types] Add sharding rules for the following primitives:

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