With JAX you _can_ just use python control flow with gradients, so long as you don't try to jit-compile the control flow. This ends up being just like PyTorch /

With JAX you _can_ just use python control flow with gradients, so long as you don't try to jit-compile the control flow. This ends up being just like PyTorch / eager execution. However, you have the option of compiling the largest possible blocks of code that are convenient. When writing code in JAX, I often start by write code eagerly and then try to jit increasingly large pieces to get more performance. You might also find the experimental control flow interesting since it tries to be more pythonic while also being possible to jit-compile. (https://github.com/google/jax/blob/master/jax/experimental/loops.py https://github.com/google/jax/blob/master/jax/experimental/l...)