PyImageSearch Tutorial on JAX Part 2 | Power tools of JAX

[ariG23498]

The first part of the tutorial was shared in this discussion.

The second part of the tutorial is up.

Link: https://pyimagesearch.com/2023/02/27/learning-jax-in-2023-part-2-jaxs-power-tools-grad-jit-vmap-and-pmap/

We (Ritwik and I) cover the following:

• 👨‍🏫Important Functional Transformations
  • Pure Functions
  • jaxpr
  • jax.grad
  • jax.jit
  • jax.vmap
  • jax.pmap
• 🎱 How Does JAX Handle Randomness?

We look forward to constructive feedback from the community 🤗

[jakevdp]

Nice to see this tutorial progressing! A few comments:

1. As in the first section, pip install jax will not work in a fresh environment. The only reason it works in Colab is because jaxlib is already installed. You might wish to recommend something like pip install jax[cpu] (this would be preferred over installing jaxlib separately) or refer to the JAX installation guide.
2. For the stateful vs. stateless discussion, the current version makes it sounds like class attributes are the problem. But I think this is a bit misleading: class attributes are fine, the issue is with in-place modification of class attributes. So maybe in place of the Stateless example, you could do something like this: Stateless class, I wonder if it might be more illustrative to have a pure version of a class that returns a copy; e.g. something like this:

   class PureState(NamedTuple):
     state: Any
     def update_state(self, new_state):
       return PureState(new_state)
   
   p1 = PureState(1)
   p2 = p1.update_state(2)
   print(p1)  # un-modified
   print(p2)  # new object

   That's more akin to how state works in many JAX functions, and inheriting from NamedTuple ensures that the state is read-only.
3. In the jit timing example, you time JIT compilation of a single matmul: I wouldn't expect JIT to give much speed improvement in such a case, and I suspect the 30% improvement is more of a statistical fluctuation than a true indication of a speedup. Users re-running your code may find that the JIT version is slower, which may cause confusion. I would suggest JIT compiling a function with a couple steps, which is where JIT will actually start to shine.
4. For the example where JIT fails, you may mention that you can reexpress it as return jnp.where(x > 0, x + 1, x) or similar to avoid giving the impression that JIT can't handle conditional computations at all.

Overall it looks great! Thanks for sharing