JAX is a sophisticated and well implemented high level tracer, combined with a backend compiler that makes use of XLA. High level tracing has some fundamental draw backs in key features we want (efficient scalar AD for example), so the approach being advocated here is doing AD as a compiler transform on the original source language. Mike has some details in an earlier paper: https://arxiv.org/abs/1810.07951. As a shameless plug, you can plug (heh) Julia into XLA as well, e.g. for targeting TPUs (https://arxiv.org/abs/1810.09868). In fact, Julia and JAX are probably better at that than TensorFlow, because the TF execution model is a bit of a mismatch for what XLA expects.
Does this support user defined data structures? This is a thing holding me back in Stan and other frameworks with AD (they only want scalars or arrays)
"Without compromising on performance, Zygote supports the full flexibility and dynamism of the Julia language, including control flow, recursion, closures, structs, dictionaries, and more."
Cool and amazing. This thing should be used to enrich base Julia, not simply be an extension.
Julia is designed in such a way that third party 'extensions' (packages) are first class citizens. Very little actually needs to be in Base or the standard library. When one types `using Zygote` you're not using a programming language with a siloed AD exension. You've radically extenended julia into a differentiable programming language. The beauty of Zygote.jl being a package is that we don't force one AD approach o…
Yes, this is exactly correct. We try very hard to extend Julia with "mechanisms" not "features", such that if you don't like the exact way we're doing it, you can just do your own (with blackjack), while still re-using as much of the rest of the system as possible. The same holds true for AD. We think (and we hope) that Zygote (and related packages) will be an extremely good, general purpose AD system that's a good default, but if you're in a specialized domain, you may be able to do a lot better with a specialized AD transform (e.g. your domain might want to do part of the computation in the frequency domain automatically, or have a domain-specific optimizations that are not generally valid, etc.). Our user base is to varied to be able to prescribe a one-size-fits-all approach. The only thing we can do is provide the necessary mechanisms and high quality default implementations (as well as good explanations as to how things work) and let it loose on the world. Everything being in the same language helps enormously here, since even users coming from the scientific side are well versed enough in the language to at least be able to read and understand the semantics of code in the compiler (even if they think it employs dark magic).
This is huge and I have high hopes for zygote,jl. Observation/request: For higher-order derivatives (Hessian, Laplacian, etc), AD libraries typically provide API shortcuts. I have found it difficult to control or predict the memory footprint and the time complexity of these API shortcuts. Laplacian is case in point: it is sometimes computed by first computing the Hessian by forward-over-reverse and then taking the tr…
Unfortunately, this isn't super easy since the time complexity heavily depends on what optimizations the compiler will apply and how the higher order AD is exactly implemented (there's many ways to do so that all give you the same answer and you probably want the system to pick the best for you). What we could do fairly easily however is to have a compiler introspection tool that runs through the computation, but doe…
Does this support user defined data structures? This is a thing holding me back in Stan and other frameworks with AD (they only want scalars or arrays)
Julia doesn't even really have a distinction between user defined and base data structures. Every julia datastructure that's not needed for bootstrap is implemented in julia and all are on an equal footing with a datastructure I define at the repl. Even numbers.
This is why julia is so powerful for so many different use-cases, when I make my own special custom array, number, string or whatever type, it's a first class citizen and with enough optimization will be just as fast, extensible and generic as whatever was provided by base.
Would also be curious to understand how this compares. They seem to account for recursion and if-branches as well.
They're not really all that close. For instance that Autograd library requires that you use its own version of numpy instead of the regular version because it can't differentiate C. Because most non-trivial python code is actually written in C, there's almost no performant programs you can just differentiate out of the box unless autograd itself has a fork of that package. My understanding is that the python Autograd…
Would also be curious to understand how this compares. They seem to account for recursion and if-branches as well.
They're not really all that close. For instance that Autograd library requires that you use its own version of numpy instead of the regular version because it can't differentiate C. Because most non-trivial python code is actually written in C, there's almost no performant programs you can just differentiate out of the box unless autograd itself has a fork of that package. My understanding is that the python Autograd…
JAX (github.com/google/jax), which is being developed by many of the authors of Autograd, is a probably a better comparison. At the cost of requiring you to rewrite control flow in a functional way, it eliminates Python's overhead by compiling NumPy code into XLA.
Does this support user defined data structures? This is a thing holding me back in Stan and other frameworks with AD (they only want scalars or arrays)
Yes https://github.com/FluxML/Zygote.jl "Without compromising on performance, Zygote supports the full flexibility and dynamism of the Julia language, including control flow, recursion, closures, structs, dictionaries, and more."
Keno: first of all, let me give a big public thank you to you and your colleagues. (For those here who don't know, Keno is listed as one of the authors of the paper, and works closely with Mike Innes, lead author and also lead developer of Zygote. Mike is also an active member of HN.) Second, let me bring up what I think is a significant issue. My perception is that most deep learning researchers and practitioners --…
Zygote is an orthogonal piece of technology on this front and relies on a good optimizing compiler behind it to target actual hardware. Its focus is primarily on expressability. We've been talking about automatic kernel generation for a while (and when I saw kernel generation what I mean is basically search for access patterns), but note that it's not quite as bad a problem in julia, because you can use higher order…
In terms of trying to break free of dependence on hand optimized kernels: a few people, myself included, have been working on some theoretical approaches to generating cache-efficient rearrangements for neutral net like problems. We've worked it out for convolution like problems [1] and have some upcoming results generalizing these techniques to other problems. Please feel free to email if you'd like to talk.