Add tutorial introducing JAX

[dgegen]

Motivated by the list of desired tutorials provided in Notes from jaxoplanet docathon, I wrote an introductory tutorial to JAX, trying to illustrate the most important features of JAX while keeping computationally intensive performance comparisons to a minimum, resulting in an execution time of ~5 seconds on my laptop.

Feedback is greatly appreciated. In particular with regards to

• Length: Do you think the tutorial is too short or too long?
• Content: Are there important features missing that you'd like to see covered?
• Placement: Is the location within the list of tutorials satisfactory, or should it be moved?

[soichiro-hattori]

@dgegen: I think this is an awesome tutorial! I'm wondering if you think there's somewhere where you could include a snippet about how JAX handles out-of-bounds indexing? Or if not even just linking https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#out-of-bounds-indexing?

To answer your questions:

• Length: I think it's slightly on the longer side.
• Content: Just the out-of-bounds indexing I mentioned above!
• Placement: I think it's good!

[dgegen]

Thank you for your suggestion and kind words, @soichiro-hattori! I will revisit the tutorial soon to see where information about out-of-bounds indexing would best fit into the flow of information. Cheers!

[dgegen]

@soichiro-hattori: I added the missing type hints as well as a small section on out-of-bounds indexing at the end of An accelerated NumPy.

Regarding the length of the tutorial: I tried to present the most important functions of JAX in a short and structured way. But of course the length of the tutorial depends on the choice of the target group and the teaching objectives. So by classifying the tutorial on the slightly longer side of the spectrum, do you mean that in comparison with the other tutorials or in reference to the word brief in its title? Or do you suggest covering less features? Or the same features, but more concisely? Or was it just an intuitive classification, perhaps rooted in a subconscious recognition of the tension between our desire for short tutorials and the fact that learning the most important JAX functions takes some time? :))

[soichiro-hattori]

Hi @dgegen, I'm sorry for the late reply there's been a lot happening the past few weeks!

Thanks for adding the changes!

Or was it just an intuitive classification, perhaps rooted in a subconscious recognition of the tension between our desire for short tutorials and the fact that learning the most important JAX functions takes some time? :))

Yes haha It was more of an intuitive classification! I don't see an issue with the length and I think it's great!

I'm happy for this to be merged! @lgrcia, @catrionamurray?

[dfm]

Thanks @dgegen for your contribution, and thanks @soichiro-hattori for the careful review!