5 ms·

I continue to wonder how does JAX compare to Zygote.jl see https://twitter.com/oxinabox_frames/status/1239495711332478976?s=19 https://twitter.com/oxinabox_fram

I continue to wonder how does JAX compare to Zygote.jl see https://twitter.com/oxinabox_frames/status/1239495711332478976?s=19 https://twitter.com/oxinabox_frames/status/12394957113324789... for a good start

komuher 7y ago
JAX is focusing more on performance and probably will be next big think for google machine learning teams like https://github.com/google/flax https://github.com/google/flax (most of the people hate tensorflow) Dunno about zygote future
- fishmaster 7y ago
  Why would this be 'the next big thing' if Google seems to be committing to Swift as its DL language of choice?
UncleOxidant 7y ago
It would be great to see a comparison of Zygote to Jax (both features and performance). I think the Julia language with the ability to manipulate it's AST via macros makes this kind of source-to-source transformation a lot easier than it would be in Python. I'm not sure how they achieved this in Python, but I'm going to guess it was a lot more work. Also, Zygote is much more ambitious than Jax. Zygote aims to support all of the Julia language whereas Jax is limited to a subset of Python. I wonder if the Zygote folks are biting off too much here? Though currently Zygote doesn't support mutation.
- ddragon 7y ago
  More than AST manipulation (which is probably not a key element in Zygote), the most important aspect is Julia's multi-stage JIT compilation (which is actually an aggressive mixture of JIT and AoT). A Julia program has access to it's own compiler at runtime, which you can call to compile any program down to Julia's Intermediate Representation, which happens to be a Julia structure that you can freely manipulate just like the AST and Static Single Assignment making it already similar to an execution graph. Then the Zygote library traverses the structure systematically writing the backward pass as if it was part of the code in the first place (plus applying any kind of optimization), and finally compiles everything together down to machine code through the LLVM. Since Python is not compiled like Julia, the approach usually involves using custom types that store a separate intermediate representation from python (a custom one or the MLIR for example) and overloading functions (and perhaps other more sophisticated ways of metaprogramming) to map all the required operations while interpreting the Python code before compiling. That's why Zygote says it works on Zygote unaware libraries, as it doesn't store the graph within any of the types or directly overload any of the methods (it does need to know how to reverse them though), and also why it can directly operate on control structures since they are part of the full IR even when you can't overload them. And while Zygote aims to fully support the Julia language, they'll certainly work on getting the most important operations working well enough before focusing on the more complicated stuff like mutation.
- oxinabox 7y ago
  AST manipulation does not do it for AD in general (though XGrad.jl is doing it as a symbolic diff approach https://github.com/dfdx/XGrad.jl/ https://github.com/dfdx/XGrad.jl/), you want to work with the code reduced down to Single assignment form, and work with code blocks and the control flow graph. for a number of reasons: 1. all control flow looks the same now (no 2 different kinds of loops + GOTOs) 2. one expression per line, no need to untangle larger experessions 3. Host of techniques from compiler world can come out to be used. This is how Zgyote works, it used code that has been lowered into this form
eigenspace 7y ago
Yeah I agree with Lyndon, Zygote is a much more ambitious project than Jax, so I would not be surprised if we see them make quicker progress than Zygote in this space.
- UncleOxidant 7y ago
  I wonder if the Zygote folks would be better off focusing on a subset of Julia for now instead of supporting all of Julia. And then focusing on performance for that subset. Maybe this is more difficult because of the way Zygote is implemented (via Julia's macro system) which just lends itself to whole-language support?
  - eigenspace 7y ago
    It's hard to say, I really do think that what the Zygote folks are trying to do is quite achievable, but I also worry that Zygote's prominence may have been bad for the Julia autodiff community. A year or two ago, it at least seemed to be a much more vibrant situation where there was more effort going into lots of different approaches, but then Zygote really picked up steam and it seemed like it was going to solve all our problems and cure cancer so development on the other packages slowed down, but now the 'finish line' for Zygote doesn't seem to be as close as we initially thought.
    - UncleOxidant 7y ago
      On the other hand, there are a lot less Julia developers than Python developers so focusing on Zygote instead of splitting the effort over several AD packages makes some sense, doesn't it? There were already a few AD packages that were fairly mature: https://www.juliadiff.org/ https://www.juliadiff.org/ Also, I think performance is very important because Zygote is used in Flux (Julia's main deep learning framework at this point.) and if it's keeping Flux's performance from matching, say, PyTorch's then that's going to limit adoption of Flux and Julia. One of Julia's main claims to fame is performance so people coming to kick the tires are going to be disappointed if it's actually slower for this domain.

komuher 7y ago

JAX is focusing more on performance and probably will be next big think for google machine learning teams like https://github.com/google/flax https://github.com/google/flax (most of the people hate tensorflow) Dunno about zygote future

fishmaster 7y ago
Why would this be 'the next big thing' if Google seems to be committing to Swift as its DL language of choice?

UncleOxidant 7y ago

It would be great to see a comparison of Zygote to Jax (both features and performance). I think the Julia language with the ability to manipulate it's AST via macros makes this kind of source-to-source transformation a lot easier than it would be in Python. I'm not sure how they achieved this in Python, but I'm going to guess it was a lot more work. Also, Zygote is much more ambitious than Jax. Zygote aims to support all of the Julia language whereas Jax is limited to a subset of Python. I wonder if the Zygote folks are biting off too much here? Though currently Zygote doesn't support mutation.

ddragon 7y ago
More than AST manipulation (which is probably not a key element in Zygote), the most important aspect is Julia's multi-stage JIT compilation (which is actually an aggressive mixture of JIT and AoT). A Julia program has access to it's own compiler at runtime, which you can call to compile any program down to Julia's Intermediate Representation, which happens to be a Julia structure that you can freely manipulate just like the AST and Static Single Assignment making it already similar to an execution graph. Then the Zygote library traverses the structure systematically writing the backward pass as if it was part of the code in the first place (plus applying any kind of optimization), and finally compiles everything together down to machine code through the LLVM. Since Python is not compiled like Julia, the approach usually involves using custom types that store a separate intermediate representation from python (a custom one or the MLIR for example) and overloading functions (and perhaps other more sophisticated ways of metaprogramming) to map all the required operations while interpreting the Python code before compiling. That's why Zygote says it works on Zygote unaware libraries, as it doesn't store the graph within any of the types or directly overload any of the methods (it does need to know how to reverse them though), and also why it can directly operate on control structures since they are part of the full IR even when you can't overload them. And while Zygote aims to fully support the Julia language, they'll certainly work on getting the most important operations working well enough before focusing on the more complicated stuff like mutation.
oxinabox 7y ago
AST manipulation does not do it for AD in general (though XGrad.jl is doing it as a symbolic diff approach https://github.com/dfdx/XGrad.jl/ https://github.com/dfdx/XGrad.jl/), you want to work with the code reduced down to Single assignment form, and work with code blocks and the control flow graph. for a number of reasons: 1. all control flow looks the same now (no 2 different kinds of loops + GOTOs) 2. one expression per line, no need to untangle larger experessions 3. Host of techniques from compiler world can come out to be used. This is how Zgyote works, it used code that has been lowered into this form

eigenspace 7y ago

Yeah I agree with Lyndon, Zygote is a much more ambitious project than Jax, so I would not be surprised if we see them make quicker progress than Zygote in this space.

UncleOxidant 7y ago
I wonder if the Zygote folks would be better off focusing on a subset of Julia for now instead of supporting all of Julia. And then focusing on performance for that subset. Maybe this is more difficult because of the way Zygote is implemented (via Julia's macro system) which just lends itself to whole-language support?
- eigenspace 7y ago
  It's hard to say, I really do think that what the Zygote folks are trying to do is quite achievable, but I also worry that Zygote's prominence may have been bad for the Julia autodiff community. A year or two ago, it at least seemed to be a much more vibrant situation where there was more effort going into lots of different approaches, but then Zygote really picked up steam and it seemed like it was going to solve all our problems and cure cancer so development on the other packages slowed down, but now the 'finish line' for Zygote doesn't seem to be as close as we initially thought.
  - UncleOxidant 7y ago
    On the other hand, there are a lot less Julia developers than Python developers so focusing on Zygote instead of splitting the effort over several AD packages makes some sense, doesn't it? There were already a few AD packages that were fairly mature: https://www.juliadiff.org/ https://www.juliadiff.org/ Also, I think performance is very important because Zygote is used in Flux (Julia's main deep learning framework at this point.) and if it's keeping Flux's performance from matching, say, PyTorch's then that's going to limit adoption of Flux and Julia. One of Julia's main claims to fame is performance so people coming to kick the tires are going to be disappointed if it's actually slower for this domain.