jax authors d9f82f7b9b [JAX] Move experimental.ann.approx_*_k into lax. ...

Updated docs, tests and the example code snippets.

PiperOrigin-RevId: 431781401

2022-03-01 14:46:33 -08:00

..

__init__.py

register the lax module for stack trace filtering

2021-02-22 21:11:41 -08:00

ann.py

[JAX] Move experimental.ann.approx_*_k into lax.

2022-03-01 14:46:33 -08:00

control_flow.py

Handle jaxpr constants correctly in MLIR lowering of conditional branches.

2022-03-01 08:56:08 -08:00

convolution.py

[MLIR] Change signature of lowering rules.

2021-12-15 19:06:58 -08:00

fft.py

Support string fft_type values in lax.fft.

2022-02-03 08:52:38 -05:00

lax.py

Add separate mechanism for threading name stacks to the lowering

2022-02-23 09:59:09 -08:00

linalg.py

linalg_test: disable implicit rank promotion

2022-01-26 09:29:06 -08:00

other.py

Update type annotations and use the new convolution.py file

2022-01-07 14:54:41 -08:00

parallel.py

Ensure that pdot parameters are hashable.

2021-12-16 07:22:59 -08:00

polar.py

Document jax.lax.Precision

2021-11-08 14:15:31 -08:00

qdwh.py

Enable QDWH TPU tests.

2021-11-30 15:47:50 -08:00

slicing.py

Index update operators: add scatter_apply()

2022-02-18 09:44:40 -08:00

svd.py

[JAX] Added jit-able singular value decomposition.

2022-02-03 11:16:55 -08:00

utils.py

add staging logic for polymorphic shapes in jaxprs

2022-01-05 14:11:12 -08:00

windowed_reductions.py

[MLIR] Add a @mlir.cache_lowering decorator that lowers a primitive out-of-line as a reusable function.

2021-12-16 08:34:52 -08:00