jax-ml/jax-ai-stack

어떤 프로젝트인가요?

핵심 기능

• 01JAX와 함께 Flax, Optax, Orbax, Chex, Grain, ml_dtypes를 서로 호환되는 것으로 테스트된 버전으로 설치합니다.
• 02jit, vmap, grad 같은 JAX 변환 기능을 제공합니다.
• 03pip에서 선택적 tensorflow-datasets extra를 제공합니다.
• 04CUDA GPU 및 TPU용 하드웨어별 JAX 설치를 지원합니다.

활용 사례

• Flax와 Optax로 신경망을 학습할 JAX 환경을 설정합니다.
• JAX 프로젝트에서 Orbax로 모델 체크포인트를 저장하고 복원합니다.
• 검증된 의존성 조합으로 NVIDIA GPU 또는 Google TPU에서 JAX 작업을 실행합니다.

빠른 시작

1. 1

   PyPI에서 JAX AI 스택을 설치합니다

   pip install jax-ai-stack
2. 2

   필요한 경우 tfds extra를 추가해 TensorFlow Datasets 지원을 활성화합니다

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

   필요한 경우 같은 명령에 JAX CUDA extra를 추가해 GPU 지원과 함께 설치합니다

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

사용 전 알아둘 점

• !다른 라이브러리를 묶은 번들이므로, 고정된 패키지가 제공하는 기능 외에 독립적인 기능은 없습니다.
• !README에 따르면 문서와 튜토리얼은 아직 작업 중입니다.

포함된 컬렉션

토픽

자주 묻는 질문

jax-ai-stack은(는) 무엇인가요?

AI 및 수치 계산에 사용할 수 있도록 테스트된 JAX 라이브러리(Flax, Optax, Orbax, Chex, Grain) 버전을 고정해 제공하는 pip 설치형 번들입니다. JAX AI Stack은 함께 사용할 수 있는 JAX 기반 라이브러리 묶음을 설치하는 단일 진입점입니다. `pip install jax-ai-stack`을 실행하면 JAX와 함께 신경망 구축용 Flax, 그래디언트 처리 및 최적화용 Optax, 체크포인트 저장용 Orbax, 신뢰성 있는 JAX 코드 작성을 돕는 Chex, 데이터 로딩용 Grain, NumPy dtype 확장 기능인 ml_dtypes 등의 구성 프로젝트가 설치됩니다. 각 버전은 저장소의 통합 테스트로 검증된 조합에 맞춰 고정되어 있습니다.

jax-ai-stack은(는) 무료로 사용할 수 있나요? 라이선스는 무엇인가요?

jax-ai-stack은(는) Apache-2.0 라이선스의 오픈소스입니다. 상업용 제품에 사용하기 전에 라이선스 조건을 확인하세요.

jax-ai-stack은(는) 아직 유지보수되고 있나요?

활발히 유지보수 중. jax-ai-stack의 마지막 커밋은 2026년 10월 13일입니다.

jax-ai-stack은(는) 어디에 사용할 수 있나요?

Flax와 Optax로 신경망을 학습할 JAX 환경을 설정합니다. JAX 프로젝트에서 Orbax로 모델 체크포인트를 저장하고 복원합니다. 검증된 의존성 조합으로 NVIDIA GPU 또는 Google TPU에서 JAX 작업을 실행합니다.

jax-ai-stack는 어떻게 설치하나요?

PyPI에서 JAX AI 스택을 설치합니다: pip install jax-ai-stack 필요한 경우 tfds extra를 추가해 TensorFlow Datasets 지원을 활성화합니다: pip install jax-ai-stack[tfds] 필요한 경우 같은 명령에 JAX CUDA extra를 추가해 GPU 지원과 함께 설치합니다: pip install jax-ai-stack "jax[cuda]"

jax-ai-stack은(는) 어디서 받을 수 있나요?

소스 코드와 설치 방법은 GitHub에 있습니다: https://github.com/jax-ml/jax-ai-stack