Yash Katariya bcfe95e98e 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. ...

There will be more improvements and semantics clarification coming in the future as we integrate it more into JAX.

Co-authored-by: Dougal Maclaurin <dougalm@google.com>
PiperOrigin-RevId: 668991384

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

..

control_flow

refactor lax.loops to avoid importing from jax.numpy

2024-08-28 14:41:59 -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

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

Return early from eigh for small matrices. This was accidentally removed cl/577222219.

2024-04-22 15:10:38 -07:00

fft.py

Bumped the minimum jaxlib to 0.4.23

2024-04-13 08:18:33 +01:00

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

linalg.py

Activate Eigenvalue Decompositions to XLA's FFI

2024-08-28 03:53:49 -07:00

other.py

refactor lax.loops to avoid importing from jax.numpy

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

parallel.py

Merge pull request #22812 from superbobry:maint

2024-08-02 04:43:33 -07:00

qdwh.py

Get rid of the is_hermitian argument for lax.qdwh. If it was known that H was also positive semi-definite, the polar decomposition would be I*H. But for indefinite H, the QDWH algorithm does not differ from the general case for Hermitian inputs.

2024-06-13 15:33:49 -07:00

slicing.py

Don't wrap singleton ir.Types during HLO lowering.

2024-07-08 12:24:45 -04:00

special.py

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

2023-11-17 11:47:14 -08:00

stack.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

svd.py

Fix Typos

2024-06-17 13:55:46 +05:30

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