Некоторые функции JAX:
1. Автоматическое дифференцирование — позволяет получить производную любой функции, написанной на Python
2. JIT-компиляция — функция jax.jit() берёт Python-код, анализирует его и компилирует в специализированный машинный код для GPU или TPU
3. Автоматическая векторизация — автоматически превращает функцию для одного элемента в функцию для целого массива, при этом оптимизирует вычисления на уровне ядер процессора
4. Параллелизм на нескольких устройствах — позволяет параллельное выполнение функций на нескольких устройствах, в основном для параллелизма данных
Языки: Python (86.5%), C++ (10.5%), Starlark (2.1%), ...
⭐️ Star 35k
https://github.com/jax-ml/jax
⚡️ @becaps // max.ru/becaps