Yash Katariya 1cb8d31c66 Convert in_shardings to physical shardings in cpp dispatch path because the same happens with prng arrays. ...

Also comment out key reuse check in cpp dispatch since it's True for jax tests which prevent prng keys from taking Cpp dispatch.

PiperOrigin-RevId: 613289252

2024-03-06 11:42:40 -08:00

..

control_flow

[xmap-removal] remove reduce_axes from grad / vjp / backward_pass

2024-02-25 15:50:54 -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

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

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

convolution.py

Emit dense arrays for StableHLO ops migrating to dense arrays

2024-01-24 16:41:37 -08:00

eigh.py

[JAX] Add an option subset_by_index that allows computing a contiguous subset of eigenvalues from eigh.

2023-10-27 09:29:33 -07:00

fft.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

lax.py

Convert in_shardings to physical shardings in cpp dispatch path because the same happens with prng arrays.

2024-03-06 11:42:40 -08:00

linalg.py

[JAX] Add an option subset_by_index that allows computing a contiguous subset of singular components from svd.

2024-02-15 16:33:09 -08:00

other.py

Define lax.Precision directly in Python, rather than inheriting from a C++ type in jaxlib.

2024-03-04 14:01:31 -08:00

parallel.py

Pipe tiled through all_to_all primitive

2024-03-04 16:33:56 -08:00

qdwh.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

slicing.py

Split name_stack out of mlir.ModuleContext.

2024-02-20 07:17:23 -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

Use lower-case PEP 585 names for types.

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

svd.py

[JAX] Add an option subset_by_index that allows computing a contiguous subset of singular components from svd.

2024-02-15 16:33:09 -08:00

utils.py

Avoid a couple more uses of operator.attrgetter(), since it's unnecessarily slow compared to writing a function directly.

2024-02-23 09:20:56 -08:00

windowed_reductions.py

Split name_stack out of mlir.ModuleContext.

2024-02-20 07:17:23 -08:00