Peter Hawkins 9c86fdec02 Make optimization_barrier a public lax API.

2024-09-06 00:18:57 +00:00

..

control_flow

Make optimization_barrier a public lax API.

2024-09-06 00:18:57 +00: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

Removed various ununsed functions

2024-08-01 11:18:19 +01:00

convolution.py

Clean up version switches from dense array migration

2024-05-28 10:58:51 -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

refactor jax.lax to not depend on jax.numpy

2024-09-01 07:49:49 -07:00

lax.py

Make optimization_barrier a public lax API.

2024-09-06 00:18:57 +00:00

linalg.py

Merge pull request #23353 from jakevdp:lax-deps

2024-09-03 05:59:26 -07:00

other.py

refactor jax.lax to not depend on jax.numpy

2024-09-01 07:49:49 -07:00

parallel.py

refactor jax.lax to not depend on jax.numpy

2024-09-01 07:49:49 -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

refactor jax.lax to not depend on jax.numpy

2024-09-01 07:49:49 -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

Initial integration of sharding in types in JAX. Currently we just support nary ops in forward only sharding propagation. Currently this functionality is experimental and hidden behind jax_sharding_in_types config flag.

2024-08-29 10:50:04 -07:00

windowed_reductions.py

refactor lax.loops to avoid importing from jax.numpy

2024-08-28 14:41:59 -07:00