4 ms·
JAX is similar to many other autodifferentiation frameworks being developed for machine learning. It's a spiritual successor to the popular autograd framework (
by tel 7y ago
JAX 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.