Matthew Johnson 9a8523603c Add experimental rematerialization decorator ...

We want to allow users to control how reverse-mode autodiff saves values
from the forward pass. In particular, we want it to be easy to signal
that a function shouldn't have any of its intermediate residuals stored
for the backward pass, and instead those values should be recomputed
from the function's saved inputs. (This feature is especially handy for
accelerators on which memory access is much more expensive than FLOPs
are.) In JAX terms, since we implement reverse-mode as a composition of
forward-mode, partial evaluation, and transposition, we want users to
control how partial evaluation behaves.

See https://github.com/google/jax/pull/1749 for more.

Co-authored-by: Dougal Maclaurin <dougalm@google.com>

2019-11-27 19:52:24 -08:00

..

__init__.py

Add a handwritten ThreeFry2x32 CUDA kernel. (#1756)

2019-11-24 13:06:23 -05:00

lax_control_flow.py

Add experimental rematerialization decorator

2019-11-27 19:52:24 -08:00

lax_fft.py

Change test tolerance logic not to choose tolerance values based on f… (#1701)

2019-11-16 13:51:42 -05:00

lax_parallel.py

Cast perm to tuple

2019-11-16 14:39:13 -08:00

lax.py

Add experimental rematerialization decorator

2019-11-27 19:52:24 -08:00