4 ms·

I really like the UX of JAX, https://github.com/google/jax https://github.com/google/jax The API is numpy plus like four functions. That's the beginning and e

by jlebar 7y ago

I really like the UX of JAX, https://github.com/google/jax https://github.com/google/jax The API is numpy plus like four functions. That's the beginning and end. It does JIT compilation under the hood, so can run quite fast.

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...)
nihilest 7y ago
That's not a fair comparison because JAX does not really provide high-level libraries (and for ones it does, it introduces many more than those four functions). For example, there must be at least 4 different neural network libraries in JAX. (e.g., https://github.com/google/jax/blob/master/jax/experimental/stax.py https://github.com/google/jax/blob/master/jax/experimental/s..., https://github.com/google/trax https://github.com/google/trax).
- jlebar 7y ago
  Sure, but OP's point is that TensorFlow's built-in libraries suck and you're forced to use them. They are confusing, change all the time, etc etc. Since JAX's API is lower-level, you can choose a neural network API that doesn't have the problems of TF's API. Or worst case, you make it yourself; you can see that the higher-level APIs provided by JAX are all very simple.

by jlebar 7y ago

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...)

nihilest 7y ago

That's not a fair comparison because JAX does not really provide high-level libraries (and for ones it does, it introduces many more than those four functions). For example, there must be at least 4 different neural network libraries in JAX. (e.g., https://github.com/google/jax/blob/master/jax/experimental/stax.py https://github.com/google/jax/blob/master/jax/experimental/s..., https://github.com/google/trax https://github.com/google/trax).

jlebar 7y ago
Sure, but OP's point is that TensorFlow's built-in libraries suck and you're forced to use them. They are confusing, change all the time, etc etc. Since JAX's API is lower-level, you can choose a neural network API that doesn't have the problems of TF's API. Or worst case, you make it yourself; you can see that the higher-level APIs provided by JAX are all very simple.