NNX is a Neural Networks library for JAX that provides a simple yet powerful module system that adheres to standard Python semantics. Its aim is to combine the robustness of Flax with a simplified, Pythonic API akin to that of PyTorch.
I think we have to make a distiction here:
- On one hand, having access to these large scale language models that can do few-shot learning is incredibly useful for the industry as in can be easily deployed to solve thosands of simple tasks.
- On the other hand, this approach will not solve harder problems (as Yann points out) and "just" creating bigger models using the same techniques is probably not the path forward in those domains.
One thing not mentioned in the original "Why Swift for Tensorflow" document and was a mayor source of conflict when the differentiable programming feature was formally proposed by the S4TF as a standard Swift feature: Swift has no mechanisms for metaprogramming. The reason is that Automatic Differentiation can be implemented 100% using metaprogramming, instead the S4TF team had to create internally certain features for this, that is probably one of the reasons why it took so long to get the most basic stuff working.
In retrospective you can really say Swift was a bad choice for the project because the time to market was much slower than it could be vs e.g choosing Julia. The other thing they didn't take into account was the actual market, that is, the Data Science ecosystem in Swift is non-existente, you have an excellent Deep Learning library standing alone without a numpy, a pandas, a scipy, a opencv, a pillow, ect, which makes doing real application with it nearly impossible.
That said, Swift as a language is amazing, doing parallel computation is so easy, not having a garbage collector makes it super efficient. Its the kind of thing we need, but the language right now is not in the right state.
I think the ML community really needs a better language than Python but not because of the ML part, that works really good, its because of the Data Engineering part (which is 80-90% of most projects) where python really struggles for being slow and not having true parallelism (multiprocessing is suboptimal).
That said I love Python as a language, but if it doesn't fix its issues, on the (very) long run its inevitable the data science community will move to a better solution. Python 4 should focus 100% of JIT compilation.