3 ms·

XLA is the limiting factor in a lot of these cases, though maybe saying limiting factor is wrong because it's more of a "trade-off factor". XLA wants to know th

XLA is the limiting factor in a lot of these cases, though maybe saying limiting factor is wrong because it's more of a "trade-off factor". XLA wants to know the static size of a lot of arguments so it can build a mathematical description of the compute graph and fuse linear algebra commands freely. What the Julia libraries like Zygote do is say "there is no good mathematical description of this, so I will generate source code instead" (and some programs like Tapenade are similar). For example, while loops are translated into for loops where a stack of the Boolean choices are stored so they can be ran in reverse during the backpass. The Julia libraries can sometimes have more trouble automatically fusing linear algebra commands though, since then they need to say "my IR lets non-static things occur, so therefore I need to prove it's static before doing transformation X". It's much easier to know you can do such transformations if anything written in the IR obeys the rules required for the transform! So it's a trade-off. In the search for allowing differentiation of any program in the language, the Julia AD tools have gone for extreme flexibility (and can rely on the fact that Julia has a compiler that can JIT compile any generated valid Julia code) and I find it really interesting to try and elucidate what you actually gain from that.

awaythrowact 5y ago
If the next machine learning killer-app model requires autodiff'ed dynamic control flow, do you think Google/Facebook will build that capability into XLA/TorchScript? Seems like if NLP SOTA requires dynamic control flow, Google will build it? Maybe they let you declare some subgraph as "dynamic" to avoid static optimizations? But maybe the static graph assumption is so deeply embedded into the XLA architecture, they'd be better off just adopting Julia? (I honestly don't know the answer, asking your opinion!)
- ChrisRackauckas 5y ago
  "Maybe they let you declare some subgraph as 'dynamic' to avoid static optimizations?" What you just described is Tensorflow Eager and why it has some performance issues (but more flexibility!). XLA makes some pretty strong assumptions and I don't think that should change. Tensorflow's ability to automatically generate good automatically parallelized production code stems from the restrictions it has imposed. So I wouldn't even try for a "one true AD to rule them all" since making things more flexible will reduce the amount of compiler optimizations that can be automatically performed. To get the more flexible form, you really would want to do it in a way that uses a full programming language's IR as its target. I think trying to use a fully dynamic programming language IR directly (Python, R, etc.) directly would be pretty insane because it would be hard to enforce rules and get performance. So some language that has a front end over an optimizing compiler (LLVM) would probably make the most sense. Zygote and Diffractor uses Julia's IR, but there are other ways to do this as well. Enzyme (https://github.com/wsmoses/Enzyme.jl https://github.com/wsmoses/Enzyme.jl) uses the LLVM IR directly for doing source-to-source translations. Using some dialect of LLVM (provided by MLIR) might be an interesting place to write a more ML-focused flexible AD system. Swift for Tensorflow used the Swift IR. This mindset starts to show why those tools were chosen.
  - awaythrowact 5y ago
    Makes sense. I don't use TF Eager, but I do use Jax, and Jax lets you arbitrarily compose JITed and non-JITed code, which made me think that might be a viable pattern. I guess I wondered if there might be something like "nonstatic_jit(foo)" that would do "julia style" compiling on function foo, in addition to "jit(foo)" that compiles foo to optimized XLA ops. Probably impractical. Thanks.
    - sillysaurusx 5y ago
      This is entirely possible. You have direct access to XLA, after all: https://jax.readthedocs.io/en/latest/notebooks/XLA_in_Python.html https://jax.readthedocs.io/en/latest/notebooks/XLA_in_Python... And XLA has dynamic shape semantics (currently unused by jax) via SetDimensionSize: https://www.tensorflow.org/xla/operation_semantics#setdimensionsize https://www.tensorflow.org/xla/operation_semantics#setdimens...
6gvONxR4sf7o 5y ago
> and can rely on the fact that Julia has a compiler that can JIT compile any generated valid Julia code This seems to be the key bit. It’s a great data point around the meme of “with a sufficiently advanced compiler…” In this case we have sufficiently advanced compilers to make very different JIT trade offs. XLA is differently powerful compared to Julia. Very cool, thanks for the insight.

awaythrowact 5y ago

If the next machine learning killer-app model requires autodiff'ed dynamic control flow, do you think Google/Facebook will build that capability into XLA/TorchScript? Seems like if NLP SOTA requires dynamic control flow, Google will build it? Maybe they let you declare some subgraph as "dynamic" to avoid static optimizations? But maybe the static graph assumption is so deeply embedded into the XLA architecture, they'd be better off just adopting Julia? (I honestly don't know the answer, asking your opinion!)

ChrisRackauckas 5y ago
"Maybe they let you declare some subgraph as 'dynamic' to avoid static optimizations?" What you just described is Tensorflow Eager and why it has some performance issues (but more flexibility!). XLA makes some pretty strong assumptions and I don't think that should change. Tensorflow's ability to automatically generate good automatically parallelized production code stems from the restrictions it has imposed. So I wouldn't even try for a "one true AD to rule them all" since making things more flexible will reduce the amount of compiler optimizations that can be automatically performed. To get the more flexible form, you really would want to do it in a way that uses a full programming language's IR as its target. I think trying to use a fully dynamic programming language IR directly (Python, R, etc.) directly would be pretty insane because it would be hard to enforce rules and get performance. So some language that has a front end over an optimizing compiler (LLVM) would probably make the most sense. Zygote and Diffractor uses Julia's IR, but there are other ways to do this as well. Enzyme (https://github.com/wsmoses/Enzyme.jl https://github.com/wsmoses/Enzyme.jl) uses the LLVM IR directly for doing source-to-source translations. Using some dialect of LLVM (provided by MLIR) might be an interesting place to write a more ML-focused flexible AD system. Swift for Tensorflow used the Swift IR. This mindset starts to show why those tools were chosen.
- awaythrowact 5y ago
  Makes sense. I don't use TF Eager, but I do use Jax, and Jax lets you arbitrarily compose JITed and non-JITed code, which made me think that might be a viable pattern. I guess I wondered if there might be something like "nonstatic_jit(foo)" that would do "julia style" compiling on function foo, in addition to "jit(foo)" that compiles foo to optimized XLA ops. Probably impractical. Thanks.
  - sillysaurusx 5y ago
    This is entirely possible. You have direct access to XLA, after all: https://jax.readthedocs.io/en/latest/notebooks/XLA_in_Python.html https://jax.readthedocs.io/en/latest/notebooks/XLA_in_Python... And XLA has dynamic shape semantics (currently unused by jax) via SetDimensionSize: https://www.tensorflow.org/xla/operation_semantics#setdimensionsize https://www.tensorflow.org/xla/operation_semantics#setdimens...

6gvONxR4sf7o 5y ago

> and can rely on the fact that Julia has a compiler that can JIT compile any generated valid Julia code This seems to be the key bit. It’s a great data point around the meme of “with a sufficiently advanced compiler…” In this case we have sufficiently advanced compilers to make very different JIT trade offs. XLA is differently powerful compared to Julia. Very cool, thanks for the insight.