Machine Learning Jax

This are a series of notebooks in the use of Jax/Flax for Machine Learning Models for the Class at Cinvestav GDL. We took the decision of using Jax given

1. The jit abilities to improve code even at CPU level
2. The GPU support
3. The easiness to stop the catching of GPU memory

Slides

These are the basic slides for the theory to be reviewed at the class.

Structure of the notebooks

Under Construction