#3 JAX Course - Future of ML research in JAX/Flax

:man_dancing::fire:We’re super excited to host Jonathan Heek from the Google Brain team working on JAX.
Drop your questions for them in this thread now!

How to join?
:link:Register here to join live

:movie_camera:Recording: JAX Course - 3. Future of ML research in JAX/Flax - YouTube

🧑‍🏫What will you learn?

  • Jonathan will be talking about the future of ML research in JAX/FLAX
  • Q&A with our hosts and guest speaker
2 Likes

Why is the nn api called Linen? :smiley:

i have a broad question – when should I use jax as opposed to pytorch or tensorflow? what does jax do that pytorch/tensorflow do not do?

4 Likes

Will ask after the presentatin :smiley:

Will it be correct to interpret, that JAX is only (or mostly) suitable for research but is not suitable for industry applications unlike TF and pytorch ?

3 Likes

Can you comment on the design choices for the flax.training.* API. :coffee:

2 Likes

Blog we’re discussing: Supercharged high-resolution ocean simulation with JAX | dionhaefner.github.io

will there be hands-on tutorials with jax during these courses?

2 Likes

Thanks for checking-Yes, the next few ones will have some more “homework” + code walkthroughs

Here are the link from prev sessions @kaczmarj:

1 Like

How can we have more insight in how XLA is choosing the most optimal path? There must be a way to comment on what XLA is doing is some cases though not as definitively as one would want. Is there a repertoire of such possible operations?

2 Likes

Can you share some resources on “pipelining for democratisation of AI”, the way you are explaining it? Sorry if it is already on the JAX official guide, its just that I have not come across it yet.

2 Likes

https://flax.readthedocs.io/en/latest/notebooks/jax_for_the_impatient.html

1 Like