Flax — open-source библиотека глубокого обучения, построенная на JAX. Она предоставляет инструменты для определения, обучения, оценки и экспериментирования с нейронными сетями, особенно в исследовательских и высокопроизводительных рабочих процессах машинного обучения.
A high-performance open-source library for deep learning built on top of JAX. Flax provides functions for training neural networks, as well as methods for evaluating their performance.