What is JAX?

[@GoogleResearch]

Subscribe to the Google Research Channel → goo.gle/GoogleResearch

[@gufrankanat3383]

Such a great introduction. Thanks for sharing

[@LaurenceMoroney]

Welcome!

[@khidhrhalab3543]

If you were new to the field what would you learn first?jax or tensorflow ?

[@adityamwagh]

PyTorch 😂 JK JK. Learn JAX

[@arshsharma8627]

Tensorflow

[@stokedfool]

Definitely interested in what JAX has to offer. Appreciate the explanation. Cheers

[@isaacatia-abugbilla2476]

Will there be a specialization course for jax on learning platforms like you did with tensorflow?

[@LaurenceMoroney]

No plans at the moment, sorry! I'm going to refresh the TensorFlow specializations, so may have some room for JAX in that.

[@isaacatia-abugbilla2476]

@@LaurenceMoroney Perfect! Thank you

[@pedropgusmao]

Why is this separate from Tensorflow?

[@thedislikebutton163]

This is for researchers primarily. Different usecase

[@LaurenceMoroney]

It's not. We work on both. But, to be clear, TensorFlow has high-level APIs for defining and training neural nets (amongst other scenarios) -- JAX is primarily optimized for numeric computing -- so it's not limited to deep learning etc.

[@readbyname]

​​@@LaurenceMoroney what about flax? Is flax going to replace tensorflow or tensorflow.keras for ML and DL?

[@gokulakrishnanm]

I heard that you are going to refresh TF specialisation when it will be done.

[@LaurenceMoroney]

We're just getting started, so no end-date is confirmed, but in general, I'm aiming for early summer.

[@gokulakrishnanm]

@@LaurenceMoroney ok I can wait for it 😇

[@user-wr4yl7tx3w]

More videos on Jax. 😊

[@LaurenceMoroney]

This is the first in a series. Let me know what kind of stuff you'd like to see.

[@kartikpodugu]

how to convert a model trained using JAX to ONNX format ?

[@marverickbin]

Im into pytorch. Why should I consider going jax?

[@LaurenceMoroney]

I'll give the answer in two ways, and the same as I give my staff. First: Soft skills --- it's never good to be a one-trick pony. Learn as many of the technologies that are as adjacent to yours as possible. Even if you never adopt them, it will make you understand your own better. I encourage Googlers to kick the tires of Pytorch for those very reasons. Otherwise all one would ever learn is from marketing and hype. Two: Hard skills -- the goal of JAX is to be a high performance numeric computing framework. It makes it a foundational technology in deep learning, but not *only* in deep learning. It's not a straight up apples to apples comparison. It lets you do things at a lower level that higher level frameworks like torch, tensorflow and keras abstract away. There's nothing wrong with abstraction, of course, but, if you want to be better at X, it's always good to have a way of getting deeper.

[@ai.aspirations]

great!

[@ranimsaidi9564]

How can JAX help in HPC applications?

[@LaurenceMoroney]

If you need math done really really fast with JIT compilation that’s optimized for accelerators - that’s how JAX can help HPC 😀

[@TheEditorify]

Just use Julia. JAX is just a copy.

[@st0rm129]

why is Chuck from Better Call Saul teaching me about a ML structure library?

[@FoodReviewerByNusrat]

good content

[@LaurenceMoroney]

Thanks!

[@FoodReviewerByNusrat]

@@LaurenceMoroney I think you have been working on youtube for a long time but not getting the expected success.. If you want I can help you and do something good..

[@LaurenceMoroney]

@@FoodReviewerByNusrat I’m happy with the level of success I’ve had 😀

[@MarcoServetto]

I wonder, is there any support for ML in Java instead? now that we are moving toward SW as a way to combine more and more ML systems seen as black boxes, It would be convenient to have support for other kinds of languages. In particular, python tends to favor a "system" programming style where the machine you are running on and how stuff is installed on it is relevant, and we often end up with a family of connected processes. While Java tend to push toward a large long living process where dynamic class loading with multiple class loaders can take care with more efficiency and portability of those same situations.

[@falcon20243]

The main code for most ML libraries is in C++. There are bindings available for most of the popular languages including Java. Both pytorch and tensorflow support JVM languages. This video is about JAX which is mainly for research purposes. People mostly use JAVA when they want to deploy models in Android.

[@LaurenceMoroney]

There are bindings for Java in TensorFlow, but they're not nearly as popular as JavaScript or Python.

[@Fortyq]

how is this better than PyTorch?

[@LaurenceMoroney]

Apples and Oranges

[@Fortyq]

@@LaurenceMoroney IMO, JAX is just another frontend for XLA. I would be happy to listen to your opinion.

[@LaurenceMoroney]

That’s what the X in JAX stands for! 😊