jax authors 0b3f0e11fb Reverts ebb75db8a523150c48376d15391f84380a2bb110 ...

PiperOrigin-RevId: 688477769

2024-10-22 03:29:32 -07:00

..

control_flow

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -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

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

fft.py

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

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

lax.py

Reverts ebb75db8a523150c48376d15391f84380a2bb110

2024-10-22 03:29:32 -07:00

linalg.py

Activate the FFI implementation of SVD on GPU.

2024-10-17 17:57:06 -07:00

other.py

Simplify and consolidate dot algorithm control in lax.

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

parallel.py

Reverts 6cf09f8c24c67ff650b95d174501fff3cb59db0d

2024-10-04 13:56:27 -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

Improve docs for lax.gather & lax.scatter

2024-10-15 16:42:44 -07:00

special.py

Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.

2023-11-17 11:47:14 -08: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

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

utils.py

[sharding_in_types] Add dot_general sharding rule. We only handle the simple cases and rely on xla to insert the collectives.

2024-10-11 16:05:13 -07:00

windowed_reductions.py

Clean up and fix primal type to tangent type mapping

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