|
|
|
|
|
|
by mattjjatgoogle
992 days ago
|
|
|
Actually that never changed. The README has always had an example of differentiating through native Python control flow: 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! |
|
|