Yash Katariya 473e2bf527 Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip. ...

Also allow users to enter into `Auto`/`User` mode inside jit along all or some axes.

Add checks to make sure that avals inside a context match the surrounding context. This check happens inside `abstract_eval` rules but maybe we need a more central place for it which we can create later on.

PiperOrigin-RevId: 707128096

2024-12-17 09:17:21 -08:00

..

control_flow

Fix incorrect capitalization in scan error message

2024-12-16 11:37:31 -08: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

Use a broadcasted gather in the sort JVP, rather than forming explicit iotas.

2024-12-13 09:23:34 -05:00

linalg.py

Activate Triangular Solve to XLA's FFI

2024-12-11 02:22:37 -08:00

other.py

Simplify and consolidate dot algorithm control in lax.

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

parallel.py

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08: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

Cleanup: replace lax._abstractify with core.get_aval

2024-12-12 14:08:17 -08: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

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08:00

windowed_reductions.py

Cleanup: replace lax._abstractify with core.get_aval

2024-12-12 14:08:17 -08:00