Peter Hawkins f94c42b271 Fix rendering of cube root in docs.

2022-02-04 11:30:10 -05:00

..

__init__.py

register the lax module for stack trace filtering

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

control_flow.py

NotImplementedError for new_checkpoint-of-xmap

2022-01-19 20:40:38 -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

Fix rendering of cube root in docs.

2022-02-04 11:30:10 -05: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

Disallow scatter_mul gradients if unique_indices=False.

2022-01-24 14:55:02 -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