Jake VanderPlas df1ceaeeb1 Deprecate jax.tree_util.tree_multimap

2022-04-01 14:51:54 -07:00

..

__init__.py

register the lax module for stack trace filtering

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

ann.py

[JAX] Update ann to use XLA based fallback ApproxTopK.

2022-03-15 07:50:48 -07:00

control_flow.py

Deprecate jax.tree_util.tree_multimap

2022-04-01 14:51:54 -07: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

Share compare_mhlo function between lax.py and mlir.py.

2022-03-30 17:02:19 -04:00

linalg.py

Fix for hipsparse in ROCm.

2022-03-25 17:41:42 +00:00

other.py

Use real Precision type for lax.PrecisionType

2022-03-04 04:21:25 -08:00

parallel.py

Replace reshape with expand_dims if possible

2022-03-31 01:34:26 +01:00

polar.py

Document jax.lax.Precision

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

qdwh.py

[linalg] Fix type promotion in QDWH.

2022-03-03 20:59:28 -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

fixup from 5415306: remove extraneous lines

2022-03-11 15:19:10 -08:00