Peter Hawkins 3fc1fdb148 Add a JVP rule for the general case of lax.reduce.

2021-03-30 17:31:47 -04:00

..

__init__.py

register the lax module for stack trace filtering

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

control_flow.py

Thread through transform name.

2021-03-30 15:39:11 +01:00

fft.py

Fix transpose rule for jnp.fft.irfft

2021-03-25 18:49:38 -07:00

lax.py

Add a JVP rule for the general case of lax.reduce.

2021-03-30 17:31:47 -04:00

linalg.py

Increase threshold for switching to unbatched triangular solve on GPU

2021-03-22 21:29:39 +00:00

other.py

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

parallel.py

remove disable_omnistaging mechanism

2021-03-29 15:26:57 -07:00