JAX can be used for any high-performance scientific computing with arrays, but because it is a good starting point to build deep learning libraries, please find this course in the AI section.
An introduction to JAX
High-performance array computing and differentiation with 