JAX: 고성능 배열 컴퓨팅#
JAX는 고성능 수치 계산을 위해 Autograd와 XLA를 결합한 것입니다.
신경망을 훈련시키고자 한다면, Flax를 사용하고 그 문서부터 시작하세요. 관련 도구로는 Optax와 Orbax가 있습니다. JAX 위에 빌드된 엔드-투-엔드 트랜스포머 라이브러리를 보려면, MaxText를 확인하세요.
익숙한 API
JAX는 연구자와 엔지니어들이 쉽게 적응할 수 있도록 NumPy 스타일의 익숙한 API를 제공합니다.
변환
JAX는 컴파일, 배치 처리, 자동 미분 및 병렬화를 위한 조립 가능한 함수 변환 기능이 포함되어 있습니다.
어디서든 실행
동일한 코드가 CPU, GPU, 및 TPU를 포함한 여러 백엔드에서 실행됩니다.