I can tell you about my experience using JAX- a functional Deep Learning framework. My work is full-time PyTorch. I use tflite for hobby work in Edge AI. I re

I can tell you about my experience using JAX- a functional Deep Learning framework. My work is full-time PyTorch. I use tflite for hobby work in Edge AI. I really like using PyTorch and the experience is far superior than using TF. But JAX is on another league altogther. Purely functional approach to Deep Learning makes so much sense, and I personally find this approach really better and gives me more peace of mind. When working with large datasets, and many kinds of transformations, and passing those through functions, thousands of times- one can't help appreciating the whole functional framework. It is so much simple and effective. And I don't need to tell you about the advantages with parallel processing. And it is not with Haskell or an ML language I had my functional enlightenment- but it was with JAX. In my limited exposure, much of Haskell examples and use cases I found are related to some form of text processing- an area in which I am least interested. (Please be kind if I am saying something inaccurate.) So, JAX is not only a good approach to Deep Learning, but finally made me appreciate the functional paradigm. When getting started, I also wrote a Notebook [0] where I wrote a beginners' guide for JAX. Head over there if you are interested. It was written after my first exposure to JAX. [0]: https://www.kaggle.com/code/truthr/jax-0 https://www.kaggle.com/code/truthr/jax-0