9 ms·
Jax – Composable transformations of Python and NumPy programs
- xapata 7y agoAnyone from the Numba team care to comment?
- albertzeyer 7y agoWhile there is an overlap in how they work technically, I think the overlap is not too big. So if you ask about a Numba vs JAX comparison here, I'm not sure if such a comparison makes too much sense. JAX core is an extensible system for transforming numerical Python functions. This core is used to implement automatic differentiation, translation to TF XLA, etc. Numba does not have such a generic function transformation framework - it just supports a single transformation, that is from numerical Python function to machine code.
- cl3misch 7y agoA big difference is jax having a tracing compiler while numba does not. In general, jax is more suited towards vectorizable code (so using numpy functions broadcasted over axes) while numba is better for accelerating manual for loops. (I'm not from either jax or numba, but a keen jax user for non-ML research.)
- hedgehog 7y agoOut of curiosity what are you using JAX for?
- emu 7y agoWe (JAX) see JAX and Numba as mostly complementary; they have different strengths and they are focused on different things. We haven't tried combining them yet, but we think it would be fun to explore (https://github.com/google/jax/issues/1870 https://github.com/google/jax/issues/1870). For example, you could use Numba to hand write a numerical kernel that then participates in a machine learning model that uses JAX automatic differentiation.
- chrisaycock 7y agoHere is an earlier submission of JAX from December 2018: https://news.ycombinator.com/item?id=18636054 https://news.ycombinator.com/item?id=18636054 JAX is pretty neat because it is effectively a derivatives compiler: it can automatically differentiate a function and JIT compile the result. This makes training in machine learning both fast and easy because gradient descent no longer has to be written by hand.
- superpermutat0r 7y agoFunction getting differentiated regularly is a loss function defined as (DesiredOutput(x) - HugeNumberOfParametersAppliedTo(x))^2. Are you saying that the symbolic expression gets transformed and is then used to represent the gradient? I thought that PyTorch, Tensorflow and similar already do that.
- tel 7y agoMany frameworks already compute derivatives, but they don't use a symbolic representation. Instead they use a method called "automatic differentiation" which does something along the lines of (a) extracts a trace of the algorithm by executing the code with dummy arguments, then (b) uses the chain rule to compute component derivatives at each node in the execution tree and combine them into the final answer. These methods are much faster than perturbation-based derivatives and much more applicable than symbolic methods (which cannot be automatically extracted from a program).
- p1esk 7y agoNot sure what you mean by “automatically extracted from a program”, all DL frameworks manually write backward pass for each op.
- tel 7y agoI mean the tracing operation that produces a structure appropriate for AD computation. I agree with you that there's work needed to specify the node derivatives. Although, honestly, I misspoke. The difference between AD and symbolic differentiation is more subtle. Really AD is profiting because it uses AST representations to keep a graph of intermediate values while symbolic methods can blow up exponentially (or require clever, difficult to generalize tricks to reconstruct that graph).
- xiaodai 7y agoI 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 agoJAX 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 agoWhy would this be 'the next big thing' if Google seems to be committing to Swift as its DL language of choice?
- UncleOxidant 7y agoIt 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 agoMore 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.
- throwawayiionqz 7y agoI am wondering what is the state of things for fixing gradients are kinks, or if there is no hope for general automatic differentiation libraries. import jax.numpy as np from jax import grad, jit, vmap from jax import random def relu(x): return np.where(x>0, x, np.zeros_like(x)) def identity(x): return relu(x) - relu(-x) derivative_identity = grad(identity) derivative_identity(0.0) returns 1.0 or 0.0? It currently returns 0.0. (Edit: typo in sign)
- jl2718 7y ago0.0>0 is false.
- ajtulloch 7y agoSee https://en.m.wikipedia.org/wiki/Subderivative https://en.m.wikipedia.org/wiki/Subderivative
- throwawayiionqz 7y agoThe derivative of x->x (identity) should be 1 regardless of your definition of derivative. Subgradients are only applicable when summing convex functions. Here relu(x) - relu(-x) is a sum of a convex function and a concave function.
- tel 7y agoFor many kinks, it doesn't in practice seem to matter that much. Most applications are computing large stochastic derivatives which smooth out kinks through averaging. Critical kinks are those that affect the geometry of the gradient in systematic ways. For instance, a model with a mixture of discrete and continuous parameters. These are serious blockers and require more complex methods to solve such as Rao-Blackwellization (marginalizing out the discrete parameters). Generally this appears as model bias or substantially increased, often fatal variance in loss curves.
- bionhoward 7y agoJax is compelling because it’s a lot simpler than other solutions. It’s hard to quantify the value of simplicity over new features, but I’d argue we usually undervalue simplicity; it helps move way faster and makes debugging & maintenance exponentially easier. Helps focus on the 80/20 of helping customers instead of internal bullshit...
- tel 7y agoJAX is similar to many other autodifferentiation frameworks being developed for machine learning. It's a spiritual successor to the popular autograd framework (same authors). The advantages JAX brings are * Numpy-adjacent interface * Auto-vectorization * Great tie-in with XLA * Higher-order derivatives (critical for some applications) * Simplified, functional interface The disadvantage is that it's younger and not trying to be a fully-fledged competitive neural network framework and thus is behind and less resourced compared with other libraries implementing auto-differentiation. I think it's really exciting to see more people learning about JAX and beginning to use it for serious projects.
- jlebar 7y ago> The disadvantage is that it's [...] not trying to be a fully-fledged competitive neural network framework Perhaps this is an unpopular opinion, but to me this is an advantage. JAX is a library upon which you can build a neural-network framework -- or a framework for something else. It seems clear that we're still figuring out the Right Way to write code that defines a neural network. The fact that JAX lets you write different, competing libraries/APIs/DSLs -- however you want to think of it -- lets us innovate more freely.
- tel 7y agoHonestly, I don't disagree, but I feel like it's harder to get resources if you're not aiming for neural networks somewhat directly.
- polskibus 7y agoDoes the resulting chain of transformation pull all data into main memory at once, or can JAX handle larger than memory datasets?
- tbenst 7y agoHas anyone benchmarked Jax? Curious how it compares to PyTorch for nontrivial networks, say ResNet.
- komuher 7y agoIt is compiled to XLA so should be a lot faster then pure PyTorch but probably will be slower then TVM (https://tvm.apache.org/ https://tvm.apache.org/) i can prepare some benchmarks in next few days if u are interested :)
- tbenst 7y agoThat would be wonderful if you’re able to! Also doubles as a good intro to Jax :). Please feel free to tweet at me (@tbenst) or email [same username at stanford dot edu] if you do get around to it.
- komuher 7y agoSure (dont have twitter yet) but will post it here on hacker news in next week probably :)
- p1esk 7y agoI'm very interested in seeing Resnet results. Especially FP16 precision running on tensor cores (V100 cards). Please use synthetic input vectors so that the input pipeline is not a bottleneck (as it is often the case, and it varies per framework).
- m0zg 7y ago> XLA so should be a lot faster I've yet to see anything get "a lot faster" because of XLA. It's a ton of complicated code, but then you end up spending the vast majority of time in NVIDIA's cuDNN anyway, so any benefits you might have hoped for will be marginal at best.
- p1esk 7y agohttps://blog.exxactcorp.com/nvidia-quadro-rtx-8000-deep-learning-performance-benchmarks-for-tensorflow-2019/ https://blog.exxactcorp.com/nvidia-quadro-rtx-8000-deep-lear... Almost double speedup for FP16 Resnet-50.
- p1esk 7y agoHow is it different from ChainerX?
- emu 7y agoHello from the JAX team! We'd like to take this opportunity to give a shout out to some of the awesome projects folks are building on top of JAX, e.g., * Flax, a neural network library for JAX (https://github.com/google/flax https://github.com/google/flax) * Haiku, a neural network library for JAX inspired by Sonnet (https://github.com/deepmind/dm-haiku https://github.com/deepmind/dm-haiku) * RLax, a library for building reinforcement learning agents (https://github.com/deepmind/rlax https://github.com/deepmind/rlax) * NumPyro, a probabilistic programming library on top of JAX (https://github.com/pyro-ppl/numpyro https://github.com/pyro-ppl/numpyro) * JAX-MD, a differentiable molecular dynamics package built on top of JAX (https://github.com/google/jax-md https://github.com/google/jax-md)
- jimfleming 7y agoI've noticed that many JAX libraries (including those from Google) seem to adopt an object-oriented style more similar to Torch/Keras rather than JAX's functional style demonstrated in modules like jax.experimental.stax. This is disappointing since stax is quite clean and these libraries seem to use a lot of hacks to make OO work with JAX. Is there an effort to implement and maintain more full-featured functional libraries in the jax/stax style?
- alevskaya 7y agoI've been involved w. jax/stax/trax/flax - I think the real issue w. the stax-like functional form is that it gets unwieldy very quickly when dealing w. more complicated models that are natively general-graphs as opposed to simple sequential pipelines that can be trivially mapped to a combinator expression tree. Of course there are many solutions here, but ultimately if you're building an NN library you need to build something that ML researchers actually want to use daily, and that tends to look closer to hackable pytorch-like DSLs rather than higher-order functional code - which often looks elegant but tends to hurt readability and rework speed.
- p1esk 7y agoWhat about Trax (https://github.com/google/trax https://github.com/google/trax) and how does it compare with Flax or Haiku? Interesting that googlers who are supposed to use Tensorflow are now actively developing a new autograd engine and at least three new DL frameworks on top of it. What do you think about this segmentation?
- choeger 7y agoI assume that this only works for total derivatives (e.g., with respect to time) not general partial derivatives (i.e., the n-th total derivative with respect to another variable), right? Otherwise how would one avoid generating massive amounts of code, and temporary data along with it, if I want to
- nmca 7y agoYou can run a k steps of adam-sgd and differentiate through the learning rate and scalarising parameters of a composite loss function in order to meta-learn them. Support for general n-th total derivatives is rather good :)