3 ms·

JAX looks really interesting but it seems that it’s not as straightforward to just JIT things. Generally I think you need to use their flow control mechanisms r

by iwill 7y ago

JAX looks really interesting but it seems that it’s not as straightforward to just JIT things. Generally I think you need to use their flow control mechanisms rather than standard Python ones. There are of course technical reasons for doing so, but this makes me wonder if languages like Julia with packages like Zygote.jl would be nicer for some situations in the long run.

schsam 7y ago
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...)