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>
* WIP: linear solvers
* Draft of lax.linear_solve
* Refactor pytree munging inside lax.root.
The primitive's implementation and JVP rules are now 100% pytree free.
* Fixup linear_solve
* Linearize multiple times in _root_jvp to avoid zeros
* fix deftraced
* add a symmetric argument
* Fixup float64; add a test for symmetric/non-symmetric
* test zeros in linear_solve_jvp
* Revisions per review
* Adjust signature of linear_solve
* restore botched test
* variable names
* WIP: root solve jaxpr
* WIP more tests
* rewrite root
* Root works with jaxprs
* root -> custom_root
* WIP undefined tangent
* Delayed undefined JVP errors
* use raise_on_undefined_tangents inside define_implicit_gradient
* more tests on jvps with undefined tangents
* Remove define_implicit_gradient
* Support closures in custom_root
* revert api-test
* another test
* jit tests
* spelling
This change creates a new fft primitive in lax, and uses it to implement numpy's np.fft.fftn function.
Not-yet-implemented functionality:
- vmap
- 's' argument of fftn
- other numpy np.fft functions
Resolves#505.
