Jake VanderPlas 34f116c0e0 vmap: preserve weak_type in batching tracer

2022-03-30 11:06:56 -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

vmap: preserve weak_type in batching tracer

2022-03-30 11:06:56 -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

[mhlo] Clean up ops that can use InferTensorTypeWithReify

2022-03-30 10:44:16 -07: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

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

[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