JAX — библиотека численных вычислений и машинного обучения, поддерживающая высокопроизводительные операции с массивами, автоматическое дифференцирование, компиляцию и выполнение на ускорителях. Она часто используется для исследований и разработки моделей, требующих гибких, компонуемых преобразований численных программ.
An array computing library, bringing together XLA (Accelerated Linear Algebra) and automatic differentiation for high-performance numerical computing. JAX provides a simple and powerful API for writing accelerated numerical code with composable transformations. JAX provides features such as: - (automatic differentiation) - (just-in-time compilation) - (automatic vectorization or batching) - (parallelization) JAX is a language for expressing and composing transformations of numerical code, analogous---but much larger in scope---to Python's NumPy library. (In fact, the .numpy library under JAX is a functionally equivalent, but entirely rewritten version of the Python NumPy library.) JAX is particularly well-suited for speeding up many machine learning tasks by transforming the models and data into a form suitable for parallelism across GPU and TPU accelerator chips. Flax, Optax, Pax, and many other libraries are built on the JAX infrastructure.