PyImageSearch Tutorial on JAX Part 1

[ritwikraha]

The first part of a tutorial series authored by Aritra Roy Gosthipaty and Ritwik Raha on JAX was published at PyImageSearch. This covers:

• Origin of JAX
  • autograd
  • XLA
• 👀 What Is JAX (revisited)?
• ⬇️ Import JAX
• 📚 Understanding the Components: API Layering of JAX
  • 💯 Numerical Computation in JAX
  • jax.numpy
  • jax.lax

We look forward to constructive feedback from the community 🤗

[jakevdp]

This looks great! A couple comments on a read-through of the first section:

• On installation: pip install jax won't work by itself; you also need to install a backend-specific jaxlib. See https://github.com/google/jax#installation for more info
• It seems strange to delve into the details of autograd, because for the most part JAX supersedes that package. If it's a JAX tutorial, why not focus on how to do autodiff with JAX?
• The tutorial says it's written with JAX 0.4.1, but then mentions DeviceArray which was removed in v0.4.1. (see https://jax.readthedocs.io/en/latest/jax_array_migration.html#jax-array-migration). You might check that the JAX version used in your code snippets matches the JAX version used in your text. Everything the tutorial says about DeviceArray is also true about the new Array.
• This note: "The clause does not make changes in place. The mutation creates another DeviceArray with the necessary changes." is only correct outside JIT; in most cases within JIT updates will happen in-place.

I hope that's helpful – excited to see other sections as

[ariG23498]

Hey @jakevdp,

Thanks for the comments. We really appreciate your time and detailed pointers.

On installation: pip install jax won't work by itself; you also need to install a backend-specific jaxlib. See https://github.com/google/jax#installation for more info

We were majorly working on Google Colab which comes pretty much pre-installed with the jax library. You guidance on this matter is appreciated and we will be making this change to our blog post as well.

It seems strange to delve into the details of autograd, because for the most part JAX supersedes that package. If it's a JAX tutorial, why not focus on how to do autodiff with JAX?

The main idea behind this tutorial was to make one familiar with the roots of JAX. We wanted people to get the essence of functional programming, autodiff and much more from its parent autograd. We cover JAX and autodiff with JAX in a later tutorial (Part 2).

The tutorial says it's written with JAX 0.4.1, but then mentions DeviceArray which was removed in v0.4.1. (see https://jax.readthedocs.io/en/latest/jax_array_migration.html#jax-array-migration). You might check that the JAX version used in your code snippets matches the JAX version used in your text. Everything the tutorial says about DeviceArray is also true about the new Array.

Thanks for pointing this out. I think this is a problem with Google Colab. We did pip install --upgrade jax on a colab which gave us the 0.4.1 version, but did not give us the Array class. We will be using 0.3.25 in the blog post so that we guide our readers better. We also would like to add this mishap as our learning from this discussion and point the readers to the link that you have posted.

This note: "The clause does not make changes in place. The mutation creates another DeviceArray with the necessary changes." is only correct outside JIT; in most cases within JIT updates will happen in-place.

Great point! We will add this to the blog as well.

[ariG23498]

We have updated our blog post! Thanks for your help @jakevdp.