jax — JAX — это библиотека Python для вычислений с массивами и трансформации программ, ориентированная на ускорители и предназначенная для высокопроизводительных численных вычислений и масштабного машинного обучения.

JAX — это библиотека Python для вычислений с массивами и трансформации программ, ориентированная на ускорители и предназначенная для высокопроизводительных численных вычислений и масштабного машинного обучения.

Об этом проекте

JAX предоставляет систему композируемых трансформаций функций для программ на Python и NumPy. Она разработана для масштабирования численных вычислений на аппаратных ускорителях, включая GPU и TPU, используя XLA (Accelerated Linear Algebra) для компиляции. Ключевые возможности включают: - Автоматическое дифференцирование: с помощью `jax.grad` поддерживается дифференцирование в обратном и прямом режимах через циклы Python, ветвления, рекурсию и замыкания, что позволяет вычислять производные высших порядков. - JIT-компиляция (Just-In-Time): с помощью `jax.jit` чистые функции компилируются через XLA для повышения производительности выполнения на ускорителях. - Авто-векторизация: с помощью `jax.vmap` функции отображаются вдоль осей массивов путем переноса циклов на уровень примитивных операций, что избавляет от необходимости ручного управления размерностями батчей. - Масштабирование и параллелизация: поддерживаются различные режимы масштабирования, включая автоматическую параллелизацию на базе компилятора, явный шардинг и ручное программирование для каждого устройства с использованием явных коллективов. JAX совместим с несколькими платформами, включая Linux, macOS и Windows, с поддержкой CPU, NVIDIA GPU, Google TPU, AMD GPU, а также экспериментальной поддержкой GPU от Apple и Intel.