3 ms·
> With its updated version of Autograd, JAX can automatically differentiate native Python and NumPy code. Nice. When did they make this change? Here is the ol
by sampo 3y ago
> With its updated version of Autograd, JAX can automatically differentiate native Python and NumPy code.
Nice. When did they make this change?
Here is the old way in the docs, where you needed to define functions for the if-true branch and the if-false branch, and feed them to a conditional function, to get the normal if-then-else conditional.
https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#cond https://jax.readthedocs.io/en/latest/notebooks/Common_Gotcha...
- cl3misch 3y agoThis is still the case afaik. For vanilla "if", the condition must be known at compile time. For runtime, you have to use "cond", "where", or "select" (which may be analogous).
- mattjjatgoogle 3y agoActually, that's never been a constraint for JAX autodiff. JAX grew out of the original Autograd (https://github.com/hips/autograd https://github.com/hips/autograd), so differentiating through Python control flow always worked. It's jax.jit and jax.vmap which place constraints on control flow, requiring structured control flow combinators like those.
- mattjjatgoogle 3y agoActually that never changed. The README has always had an example of differentiating through native Python control flow: https://github.com/google/jax/commit/948a8db0adf233f333f3e5f64d324f308c277773#diff-b335630551682c19a781afebcf4d07bf978fb1f8ac04c6bf87428ed5106870f5R309-R322 https://github.com/google/jax/commit/948a8db0adf233f333f3e5f... The constraints on control flow expressions come from jax.jit (because Python control flow can't be staged out) and jax.vmap (because we can't take multiple branches of Python control flow, which we might need to do for different batch elements). But autodiff of Python-native control flow works fine!