Séance 20 du groupe de travail PyR présentant le calcul tensoriel avec Jax dans Python:
Le notebook est visible dans le repo suivant:
https://forgemia.inra.fr/gdtpyr/gdt_pyr/-/tree/main/GDT_PYR_20_Jax?ref_type=heads
https://jax.readthedocs.io/en/latest/notebooks/quickstart.html
Jax est un outil de calcul tensoriel permettant d’utiliser le compilateur XLA sur Python.
Ses principaux avantages sont :
- Acceleration du calcul tensoriel.
- Gestion automatique de l’utilisation du CPU, GPU (et TPU).
- Opérations complexes automatisées telles que l’autodérivation (utile pour faire de l’optimisation).
- Construit exactement comme Numpy pour faciliter sa prise en main.
https://jax.readthedocs.io/en/latest/notebooks/quickstart.html