Jake VanderPlas d8bdd9a19d lax.sort: regularize handling of -0.0 and -NaN

2022-01-13 13:03:41 -08:00

..

__init__.py

register the lax module for stack trace filtering

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

control_flow.py

Resolves issue 8744

2022-01-12 21:10:45 +00:00

convolution.py

[MLIR] Change signature of lowering rules.

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

fft.py

Port remaining translation rules inside JAX to new style.

2021-10-19 09:48:37 -07:00

lax.py

lax.sort: regularize handling of -0.0 and -NaN

2022-01-13 13:03:41 -08:00

linalg.py

Fix type signature of fallback translation rule for lax.linalg.eig.

2021-11-29 15:26:49 -08:00

other.py

Split convolution functions out of jax._src.lax.lax and into a separate module (jax._src.lax.convolution).

2021-11-23 12:35:50 -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

[MLIR] Change signature of lowering rules.

2021-12-15 19:06:58 -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