Repolode
Предложить
jax-ml

jax-ml/jax-ai-stack

Устанавливаемый через pip набор библиотек JAX с проверенными совместимыми версиями Flax, Optax, Orbax, Chex и Grain для ИИ и численных вычислений.

Открыть на GitHub
Звёзды
315
Форки
51
Открытые issues
38
Последний коммит
25 дней назад

Что это?

JAX AI Stack — единая точка входа для установки набора совместимых библиотек на базе JAX. Команда `pip install jax-ai-stack` устанавливает JAX и такие компоненты, как Flax для создания нейронных сетей, Optax для обработки градиентов и оптимизации, Orbax для сохранения чекпойнтов, Chex для написания надёжного кода на JAX, Grain для загрузки данных и ml_dtypes для расширения набора типов данных NumPy. Версии зафиксированы в сочетаниях, проверенных интеграционными тестами репозитория.

Сам JAX — библиотека для вычислений с массивами, включающая такие преобразования программ, как jit, vmap и grad. В нём почти нет средств, специфичных для нейронных сетей, а окружающая экосистема развивается в виде отдельных пакетов. Этот набор объединяет библиотеки, которые разработчики Google используют для генеративных ИИ-моделей, поэтому пользователям не нужно вручную подбирать совместимые версии.

Дополнительные опции позволяют установить TensorFlow Datasets, а также аппаратно-специфичные сборки JAX для GPU с CUDA или TPU — всё одной командой.

Ключевые возможности

  • 01Устанавливает JAX вместе с Flax, Optax, Orbax, Chex, Grain и ml_dtypes в версиях, проверенных на совместную работу.
  • 02Предоставляет преобразования JAX, такие как jit, vmap и grad.
  • 03Позволяет установить tensorflow-datasets как дополнительную опцию через pip.
  • 04Поддерживает аппаратно-специфичные установки JAX для GPU с CUDA и TPU.

Сценарии использования

  • Настройка среды JAX для обучения нейронных сетей с помощью Flax и Optax.
  • Сохранение и восстановление чекпойнтов моделей с помощью Orbax в проектах на JAX.
  • Запуск задач JAX на GPU NVIDIA или Google TPU с проверенным набором зависимостей.

Быстрый старт

  1. 1

    Установите стек JAX AI из PyPI

    pip install jax-ai-stack
  2. 2

    При необходимости добавьте extra tfds для поддержки TensorFlow Datasets

    pip install jax-ai-stack[tfds]
  3. 3

    При необходимости установите версию с поддержкой GPU, добавив extra JAX CUDA в ту же команду

    pip install jax-ai-stack "jax[cuda]"

Что важно знать

  • Это набор других библиотек, поэтому он не предоставляет самостоятельных функций помимо тех, которые есть в зафиксированных пакетах.
  • Согласно README, документация и учебные материалы всё ещё находятся в разработке.

В подборках

Обучение и запуск ИИ на GPU

Темы

#jax#machine-learning#deep-learning#python#numerical-computing#gpu#tpu

Частые вопросы

Что такое jax-ai-stack?
Устанавливаемый через pip набор библиотек JAX с проверенными совместимыми версиями Flax, Optax, Orbax, Chex и Grain для ИИ и численных вычислений. JAX AI Stack — единая точка входа для установки набора совместимых библиотек на базе JAX. Команда `pip install jax-ai-stack` устанавливает JAX и такие компоненты, как Flax для создания нейронных сетей, Optax для обработки градиентов и оптимизации, Orbax для сохранения чекпойнтов, Chex для написания надёжного кода на JAX, Grain для загрузки данных и ml_dtypes для расширения набора типов данных NumPy. Версии зафиксированы в сочетаниях, проверенных интеграционными тестами репозитория.
jax-ai-stack можно использовать бесплатно? Какая у него лицензия?
jax-ai-stack распространяется с открытым исходным кодом под лицензией Apache-2.0. Перед использованием в коммерческих продуктах проверьте условия лицензии.
jax-ai-stack всё ещё поддерживается?
Активно поддерживается. Последний коммит в jax-ai-stack — 13 окт. 2026 г..
Для чего можно использовать jax-ai-stack?
Настройка среды JAX для обучения нейронных сетей с помощью Flax и Optax. Сохранение и восстановление чекпойнтов моделей с помощью Orbax в проектах на JAX. Запуск задач JAX на GPU NVIDIA или Google TPU с проверенным набором зависимостей.
Как установить jax-ai-stack?
Установите стек JAX AI из PyPI: pip install jax-ai-stack При необходимости добавьте extra tfds для поддержки TensorFlow Datasets: pip install jax-ai-stack[tfds] При необходимости установите версию с поддержкой GPU, добавив extra JAX CUDA в ту же команду: pip install jax-ai-stack "jax[cuda]"
Где скачать jax-ai-stack?
Исходный код и инструкции по установке — на GitHub: https://github.com/jax-ml/jax-ai-stack

Описание создано ИИ на основе README проекта. Статистика GitHub обновляется ежедневно.

Похожие проекты