JAX는 Google에서 출시한 차별화 가능한 프로그래밍 프레임워크입니다. NumPy API와 자동 차등 XLA 컴파일 및 하드웨어 가속 기능을 제공하여 최첨단 ML 연구를 위한 중요한 인프라가 됩니다.

JAX

JAX의 핵심 매개변수 및 통계

JAX는 주류 딥 러닝 프레임워크 중에서 독특한 경로를 택했습니다. 즉, 스스로를 "신경망 라이브러리"라고 부르지 않고 "미분 가능한 수치 계산 프레임워크"라고 부릅니다. DeepMind의 많은 핵심 연구(AlphaFold, Gemini 부분 인프라 AlphaGo 개선)가 JAX를 기반으로 하는 것은 바로 이러한 기본 디자인입니다. PyTorch 및 TensorFlow와 달리 JAX는 높은 수준의 신경망 API를 제공하지 않습니다. 대신, 개발자가 순수 함수형 스타일로 계산을 표현한 다음 XLA 컴파일러를 통해 효율적인 GPU/TPU 커널로 컴파일할 수 있는 구성 가능한 함수 변환 세트를 제공합니다.

프로젝트	잭스	파이토치	텐서플로우
공식 포지셔닝	고성능 미분 가능 프로그래밍 프레임워크	딥러닝 연구 프레임워크	엔드투엔드 ML 플랫폼
프로그래밍 패러다임	기능성(순수함수+변환기)	명령형(기본적으로 열망)	선언적 + 명령형 하이브리드
자동 차별화	grad(역방향 모드)/jacfwd(정방향 모드)	autograd(역방향 모드)	GradientTape(역방향 모드)
컴파일 메커니즘	XLA(jit 데코레이터)	토치다이나모/인덕터	XLA(tf.function)
병렬 전략	pmap/pjit/shard_map	DDP/FSDP	미러링전략/FSDP
하드웨어 지원	엔비디아 GPU, AMD GPU, 구글 TPU	NVIDIA GPU, AMD GPU, Apple MPS	엔비디아 GPU, AMD GPU, TPU
신경망 라이브러리	아마/하이쿠(제3자)	내장형 torch.nn	내장 tf.keras
오픈 소스 라이센스	아파치 2.0	BSD	아파치 2.0
GitHub 스타	33,000+	87,000+	188,000+
첫 번째 릴리스	2018-12	2016-09	2015-11
주요 사용자	최첨단 ML 연구(DeepMind 등)	학계 + 산업	엔터프라이즈급 프로덕션 배포

핵심 차이점: JAX의 기능적 설계는 PyTorch/TensorFlow와의 근본적인 차이점입니다. JAX에는 "모델 객체" 및 "훈련 주기"라는 개념이 없지만 순수 함수와 변환 함수(jit, grad, vmap, pmap)의 조합을 사용하여 계산을 표현합니다. 이 디자인은 대규모 병렬 훈련 및 사용자 정의 과학 연구 컴퓨팅 시나리오에서 JAX 고유의 이점을 제공하지만 학습 곡선도 더욱 가파르게 됩니다.

JAX의 사용자 및 시장 인지도

연구 기관 채택: JAX는 최고의 ML 연구 기관 사이에서 보급률이 매우 높습니다. DeepMind는 2020년부터 JAX를 핵심 연구 프레임워크로 사용했습니다. AlphaFold 2/3, Gemini 시리즈 모델 Chinchilla 및 Gopher와 같은 마일스톤 성과는 모두 JAX 또는 상위 계층 라이브러리를 기반으로 구현되었습니다. Google Brain(현재 Google DeepMind) 내의 대규모 실험 인프라에서도 JAX를 기본 컴퓨팅 엔진으로 사용합니다.

오픈 소스 커뮤니티: GitHub의 JAX 코어 저장소는 33,000개 이상의 별표를 받았으며 포크 수가 3,100개를 초과했습니다. 신경망 라이브러리(Flax, Haiku), 최적화 도구(Optax), 강화 학습(RLax, Acme), 그래프 신경망(Jraph), 베이지안 추론(NumPyro, TensorFlow Probability for JAX) 및 기타 방향을 다루는 JAX를 중심으로 구축된 200개 이상의 생태학적 프로젝트가 있습니다.

엔터프라이즈 애플리케이션: Google 외에도 NVIDIA(CUDA 및 cuDNN을 통해 JAX 성능을 심층적으로 최적화), Hugging Face(Transformers에서 JAX/Flax 백엔드 지원), Cohere, Anthropic 및 기타 회사에서도 일부 교육 또는 추론 작업에 JAX를 사용하고 있습니다. Hugging Face에는 이미 모델 라이브러리에 JAX/Flax를 지원하는 수천 개의 사전 훈련된 모델이 있습니다.

업계 벤치마킹: NeurIPS, ICML, ICLR과 같은 상위 컨퍼런스 논문에서 JAX의 사용 비율은 2020년 5% 미만에서 2025년 약 35~40%로 증가할 것이며 연구 방법론의 중요한 인프라가 되었습니다. 대학 과정에서 교수 도구로 사용되는 JAX의 비율도 해마다 증가하고 있습니다.

JAX의 비용 이점: 라이선스 비용이 없는 고성능 컴퓨팅 인프라

JAX의 비용 구조는 프레임워크 자체와 실행 중인 하드웨어라는 두 가지 차원에서 독립적으로 평가되어야 합니다.

C 측/개인 개발자:

• 프레임워크 비용: JAX는 완전한 오픈 소스, Apache 2.0 프로토콜, 라이센스 비용이 없으며 상업적 용도로 무조건 사용할 수 있습니다.
• 하드웨어 비용: 개인은 자신의 GPU(NVIDIA GeForce 시리즈 AMD Radeon 시리즈)에서 JAX를 무료로 실행할 수 있습니다. GPU가 필요하지 않은 소규모 실험의 경우 순수 CPU 실행도 무료입니다. TPU 액세스 요금은 Google Cloud TPU를 통해 시간 단위로 청구되지만 Google은 제한된 무료 TPU 할당량(예: TRC 프로젝트)을 제공합니다.

개발자/API 호출 레이어:

• JAX 자체는 클라우드 API 서비스를 제공하지 않습니다. 개발자는 프레임워크 자체에 대해 비용을 지불할 필요가 없습니다.
• 교육 인프라 비용은 선택한 클라우드 컴퓨팅 플랫폼에 따라 다릅니다. Google Cloud를 예로 들어보겠습니다.
  • GPU 인스턴스(예: A100 80G): 시간당 약 $3.50-$5.00
  • TPU v5p Pod(멀티 칩 슬라이싱): 구성에 따라 시간당 약 $30~$100+
• AWS 및 Azure도 JAX GPU 교육을 지원하며 해당 GPU 인스턴스 가격에 따라 요금이 청구됩니다.

기업/개인 배포:

• 제로 프레임워크 비용: 기업 라이센스 비용, 사용자 제한, API 호출 제한이 없습니다.
• 숨겨진 비용: -인재 확보: JAX 함수형 프로그래밍에 익숙한 ML 엔지니어는 PyTorch 개발자보다 급여 프리미엄이 높아 채용이 더 어렵습니다.
  • 마이그레이션 비용: PyTorch/TensorFlow에서 JAX로 마이그레이션하려면 훈련 파이프라인과 데이터 처리 프로세스를 다시 작성해야 하며 2~6개월의 초기 변환 기간이 있을 수 있습니다.
  • 운영 및 유지 관리 비용: 대규모 JAX 교육에는 Google Cloud TPU 또는 자체 구축 GPU 클러스터 배포가 필요하며 운영 및 유지 관리 복잡성은 규모에 비례합니다.
• 숨겨진 이점: JAX의 XLA 컴파일 및 메모리 관리 최적화는 대규모 교육에서 컴퓨팅 리소스 소비를 15%-30% 줄일 수 있으며(동등한 PyTorch 구현과 비교하여) 장기적으로 마이그레이션 비용을 상쇄할 수 있습니다.

비용 차원	잭스	파이토치	텐서플로우
프레임워크 라이센스 비용	$0	$0	$0
엔터프라이즈 라이센스 모델	없음(아파치 2.0)	없음(BSD)	없음(아파치 2.0)
최소 작동 임계값	CPU는 충분하다(무료)	CPU는 충분하다(무료)	CPU는 충분하다(무료)
일반적인 GPU 훈련 비용	클라우드 GPU 인스턴스별 청구	클라우드 GPU 인스턴스별 청구	클라우드 GPU 인스턴스별 청구
TPU 사용 비용	Google Cloud 필요($30+/h)	TPU를 직접 지원하지 않습니다	Google Cloud 필요(동일 가격)
인재 확보의 어려움	높음(개발자 수가 적음)	낮음(대규모 커뮤니티)	중간
마이그레이션 비용	높음(패러다임 전환)	—	중간(Keras가 이미 존재함)
대규모 교육 자원 효율성	우수(XLA 컴파일 및 최적화)	양호(Dynamo는 지속적으로 개선됨)	좋음(XLA 컴파일 및 최적화)

JAX의 주요 기능

• 자동 미분(grad): 모든 Python 함수에서 파생되며 역방향 모드(가장 일반적으로 사용됨) 및 순방향 모드(jacfwd)를 지원합니다. 과학적 컴퓨팅 및 최적화 문제의 핵심 기능인 고차 도함수(예: 헤세 행렬)를 계산하기 위해 중첩될 수 있습니다. value_and_grad는 함수 값과 기울기를 동시에 반환할 수 있어 반복 계산을 줄일 수 있습니다.
• 적시 컴파일(jit): XLA를 통해 Python 함수를 효율적인 GPU/TPU 커널로 컴파일합니다. 첫 번째 호출은 컴파일을 트리거하고(함수 복잡성에 따라 최대 5~60초) 후속 호출은 컴파일된 고성능 코드를 직접 실행합니다. 컴파일된 함수는 종종 손으로 작성한 CUDA에 가까운 속도로 실행되어 매트릭스 집약적인 작업에서 순수 Python에 비해 50~100배의 속도 향상을 달성합니다.
• 자동 벡터화(vmap): 일괄 처리 논리를 함수에 자동으로 매핑하므로 일괄 루프를 수동으로 작성할 필요가 없습니다. 예를 들어 단일 샘플 추론 함수에 vmap을 적용하면 자동으로 배치 추론 기능이 확보됩니다. 내부적으로 vmap은 배치 차원을 기존 벡터화 작업에 병합하며 성능은 수동 for 루프보다 훨씬 뛰어납니다.
• 교차 장치 병렬 처리(pmap / pjit / shard_map): pmap은 자동으로 계산을 여러 장치에 복사하고 데이터 병렬 처리를 수행합니다. 'pjit'(Partitioned JIT)는 샤딩 사양을 통해 계산 그래프를 장치 배열로 자동으로 분할합니다. shard_map(JAX 0.4.16+)은 사용자 정의 샤딩 전략에 적합한 명시적인 SPMD 프로그래밍 모델을 제공합니다. 세 가지 시나리오는 단순한 데이터 병렬 처리부터 복잡한 모델 병렬 처리까지 모든 시나리오를 다룹니다.
• Pallas 커널 언어: JAX 0.4.20+에 도입된 사용자 정의 GPU 커널 DSL을 사용하면 하위 수준 GPU 커널을 Python(CUDA와 유사하지만 구문이 더 간단함)으로 작성하고 XLA를 통해 컴파일하고 실행할 수 있습니다. Flash Attention의 사용자 정의 구현과 같이 극단적인 성능 요구 사항이 있는 사용자 정의 연산자에 적합합니다.
• 난수 생성(jax.random): 기능적 난수 시스템 - 각 무작위 함수는 암시적 전역 상태를 피하면서 PRNG 키 값을 명시적으로 수신하고 반환합니다. 이 디자인은 재현성을 보장하고

병렬 컴퓨팅에서는 당연히 스레드로부터 안전합니다.

• 선형 대수학 및 NumPy 호환 API(jax.numpy / jax.lax / jax.scipy): jax.numpy는 NumPy와 거의 동일한 인터페이스를 제공하며 GPU/TPU에서 투명하게 가속될 수 있습니다. jax.lax는 낮은 수준의 선형 대수 기본 요소를 제공하고 jax.scipy는 일반적인 과학 계산 기능을 다룹니다.

JAX의 모델 및 버전 진화

JAX는 2018년 12월 Google에서 오픈소스로 공개되었으며 실험적 프레임워크에서 프로덕션급 인프라로 완전히 발전했습니다.

메인라인 출시

버전	날짜	주요 변경 사항
0.1.0	~2019-02	grad, jit, vmap, pmap 코어 변환기를 제공하는 최초의 공개 릴리스
0.2.0	~2020-06	NumPy API를 안정화하고 jax.numpy 완전한 인터페이스를 도입합니다. DeepMind가 완전히 채택하기 시작
0.3.0	~2022-03	다중 머신 및 다중 TPU 교육을 지원하기 위해 pjit 샤드 컴파일이 추가되었습니다. 상당한 성능 개선
0.4.0	~2023-01	API 안정성 이정표 shard_map 명시적 SPMD 도입; AMD GPU 지원 실험 버전
0.4.16	~2024-06	shard_map 안정; 팔라스 커널 언어 베타
0.4.20	~2024-10	팔라스는 공식적으로 출시되었습니다. 디버그 인프라 개선(jax.debug)
0.4.30	~2025-06	AMD GPU ROCm 지원 향상; 컴파일 캐시 최적화; 새로운 MLIR 백엔드 미리보기
0.4.35	~2025-12	AMD GPU 생산 수준 지원; 다중 노드 통신 최적화; 오류 메시지 가독성 향상
0.5.0	~2026-05	XLA 컴파일 성능은 계속해서 향상되고 있습니다. 팔라스 커널 확장; API 정리

버전 하이라이트 해석

0.2.x 시리즈(2020-2021): JAX가 "NumPy + 자동 미분 + XLA"의 삼위일체 포지셔닝을 확립하는 중요한 기간입니다. 이 기간 동안 DeepMind는 핵심 연구 스택을 TensorFlow에서 JAX로 마이그레이션하여 대규모 ML 연구에서 JAX의 타당성을 검증했습니다.

0.3.x 시리즈(2022-2023): pjit의 도입으로 JAX는 "원클릭 파티션 컴파일"을 지원하는 몇 안 되는 프레임워크 중 하나가 되었습니다. 개발자는 각 장치(PartitionSpec)에서 텐서의 배포 의도만 설명하면 되며 pjit는 자동으로 장치 간 실행 계획을 생성합니다. 같은 기간 EasyLM, T5X, PaLM 등 대규모 교육 라이브러리가 JAX를 기반으로 구축되었습니다.

0.4.x 시리즈(2023-2025): JAX 생태계는 성숙도를 가속화합니다. Pallas 커널 언어는 맞춤형 GPU 연산자의 격차를 메웁니다. shard_map은 SPMD 프로그래밍 모델을 암시적에서 명시적으로 변경하여 대규모 교육을 위한 사용자 지정 샤딩 임계값을 낮춥니다. AMD GPU는 실험에서 생산으로의 전환을 지원합니다.

0.5.0(2026-05): 0.5 라인의 첫 번째 버전으로 0.4.x의 안정성 전략을 이어가며 XLA 컴파일 오버헤드 및 Pallas 커널 개발 경험을 최적화하는 데 중점을 둡니다. 아직 공식적인 정확한 날짜는 없습니다.

JAX의 기술적 장점

기능적 디자인: 결정성 + 구성성

JAX의 "순수 기능" 디자인은 PyTorch/TensorFlow와의 근본적인 차이점입니다. 각 JAX 함수는 내부 상태를 보유하지 않으며 모든 입력 및 출력은 매개변수를 통해 명시적으로 전달됩니다. 즉, 동일한 매개변수 및 입력 세트는 항상 동일한 결과를 생성하며(결정성), 기능은 부작용 없이 자유롭게 결합될 수 있습니다(구성성). 이 디자인은 병렬 컴퓨팅에서 특히 중요합니다. 공유 상태의 경쟁 조건을 걱정할 필요 없이 pmap/pjit는 기능을 임의의 장치에 안전하게 배포할 수 있습니다.

메커니즘 → 효과: 순수 함수와 변환기의 결합된 아키텍처를 통해 grad, jit, vmap 및 pmap을 임의로 중첩하고 합성할 수 있습니다(예: jit(grad(vmap(fn)))). 각 변환 계층은 한 차원의 계산 의미에만 초점을 맞추고 다른 차원을 방해하지 않습니다. 이것이 표현력 측면에서 JAX의 핵심 장점입니다. PyTorch의 torch.vmap 및 torch.compile은 후속 "캐치업" 기능이며 구성 가능성과 안정성은 JAX의 기본 디자인만큼 좋지 않습니다.

XLA 컴파일: 한 번 컴파일하면 모든 장치에서 실행됩니다.

XLA(Accelerated Linear Algebra)는 Python 함수 수준 계산 그래프를 대상 하드웨어에 최적화된 실행 코드로 컴파일하는 JAX의 기본 컴파일러입니다. PyTorch의 즉시 실행 모드(각 작업이 독립적으로 예약됨)와 비교하여 XLA 컴파일은 다음 메커니즘을 통해 성능 향상을 달성합니다.

• 작업 융합: 연속적인 소규모 작업(예: 'add → relu → matmul → Softmax')을 단일 GPU 커널로 융합하여 메모리 왕복 및 커널 실행 오버헤드를 줄입니다. Transformer 훈련에서 융합은 일반적으로 커널 호출 수를 30%-50% 줄입니다.
• 비디오 메모리 최적화: XLA는 컴파일 단계에서 텐서의 수명 주기를 분석하고 버퍼 재사용 및 삭제 전략을 자동으로 삽입합니다. 수동 관리에 비해 최대 비디오 메모리 사용량을 10%-20% 줄일 수 있습니다.
• 장치 독립적: 동일한 JAX 코드는 수정 없이 CPU, NVIDIA GPU, AMD GPU, Google TPU에서 실행될 수 있으며 XLA는 컴파일 타임에 대상 하드웨어에 자동으로 적응합니다.

대규모 훈련: 단일 카드에서 1만 카드까지 원활한 확장

JAX의 병렬 추상화(pmap → pjit → shard_map)는 단일 머신에서 대규모 TPU Pod까지 점진적인 확장 경로를 형성합니다.

• pmap(데이터 병렬성): 모델을 N개의 장치에 복사하고, 각 장치는 서로 다른 마이크로 배치를 처리하고, all-reduce를 통해 그라디언트를 동기화합니다. 구성 비용이 가장 낮은 단일 시스템 다중 카드 시나리오에 적합합니다.
• pjit(모델 병렬성 + 데이터 병렬성): PartitionSpec을 통해 텐서의 장치 분포를 설명함으로써 컴파일러는 자동으로 장치 간 계산 그래프와 통신 계획을 생성합니다. 모델 매개변수가 단일 장치의 메모리를 초과하는 중규모 및 대규모 교육에 적합합니다.
• shard_map(명시적 SPMD): 0.4.16+에 도입되어 개발자가 샤딩된 데이터에서 실행되는 함수를 직접 작성할 수 있으며 컴파일러는 자동으로 샤드 간 통신을 처리합니다. 맞춤형 샤딩 전략(예: 순차 병렬 처리, 전문가 병렬 처리)에 적합합니다.

효과: DeepMind는 JAX + pjit를 사용하여 6,144개의 TPU v4 칩에서 5,000억 개의 매개변수로 GShard-MoE 모델을 교육하여 선형에 가까운 확장 효율성을 달성했습니다. 이러한 대규모 병렬 처리 기능은 현재 주류 프레임워크의 JAX + TPU 조합을 통해서만 달성할 수 있습니다.

적응 경계(적용 가능 및 적용 가능하지 않은 시나리오)

JAX가 가장 적합한 시나리오:

• 대규모 분산 훈련(100칼로리 ~ 10,000칼로리 수준), 특히 TPU 클러스터에 대한 훈련
• 고차 도함수 또는 맞춤형 기울기 계산이 필요한 과학적 계산(물리적 시뮬레이션, 분자 역학, 기후 모델링)
• 연구 중심의 실험 코드(모델 구조, 맞춤형 손실 함수, 실험 연산자의 빈번한 수정 필요)
• 복잡한 모델 병렬 전략(MoE, 시퀀스 병렬성, 텐서 샤딩 등)을 사용한 대규모 모델 교육

JAX가 좋지 않은 시나리오:

• 신속한 프로토타이핑 및 교육 시작하기(PyTorch보다 훨씬 가파른 학습 곡선)
• 동적 제어 흐름 집약적 모델(예: tree-RNN, 재귀 그래프 네트워크), jax.lax.while_loop/cond가 지원을 제공하지만 표현 및 디버깅은 PyTorch 동적 그래프보다 훨씬 덜 편리합니다.
• Python이 아닌 외부 시스템과 자주 상호 작용해야 하는 프로덕션 추론 파이프라인
• 캐주얼/비연구 ML 프로젝트(커뮤니티 모델 라이브러리 및 도구의 풍부함은 PyTorch보다 훨씬 적음)
• 이미 성숙한 PyTorch 코드 기반과 팀 경험을 보유하고 있으며 마이그레이션 비용이 이점보다 높습니다.

성능 및 처리량

XLA 컴파일을 통해 얻은 JAX의 성능은 다음 차원에서 직접 작성한 최적화 코드와 경쟁할 수 있습니다.

• TTFT(Time to First Token): JAX의 jit 컴파일은 전체 계산 그래프 분석 및 하드웨어 코드 생성을 완료해야 하기 때문에 처음으로 오랜 시간(보통 5~60초)이 걸립니다. 매개변수 변경 후 재컴파일 감지를 포함한 후속 호출의 오버헤드가 크게 줄어듭니다. 이에 비해 PyTorch Eager 모드는 컴파일 지연이 전혀 없으며 TorchDynamo의 예열 시간은 약 10~30초입니다.
• 처리량(교육 처리량): 표준 Transformer 교육 작업에서 JAX + TPU 조합의 처리량은 일반적으로 동일한 GPU 구성을 사용하는 PyTorch보다 20%-50% 더 높습니다. GPU의 맥락에서 JAX와 PyTorch 간의 성능 격차는 줄어들고 JAX는 잘 통합된 특정 연산자에서 여전히 선두를 달리고 있습니다. 구체적인 값은 모델 아키텍처, 배치 크기, 하드웨어 유형에 따라 다르며 공식적인 통합 벤치마크는 없습니다.
• TPM/RPM 빈도 제어: 로컬 프레임워크인 JAX에는 API 호출 빈도 제어가 없습니다. Google Cloud TPU를 사용하는 경우 클라우드 리소스 할당량 제한(시간별 TPU 칩 시간 할당량) 및 비API 수준 TPM/RPM 제한이 적용됩니다.

JAX 사용 방법

설치

JAX는 다양한 하드웨어 백엔드에 대한 pip 설치 패키지를 제공합니다.

``배쉬

CPU 버전(범용, GPU가 필요하지 않음)

pip 설치 jax jaxlib

NVIDIA GPU 버전(CUDA 12)

pip 설치 jax[cuda12]

AMD GPU 버전(ROCm)

pip 설치 jax[rocm]

TPU 버전(Google Cloud TPU 환경에서 실행해야 함)

pip 설치 jax[tpu]

설치 후 상황을 확인하십시오: `python -c "import jax; print(jax.devices())"`, 현재 사용 가능한 하드웨어 장치 목록을 출력해야 합니다.

### 핵심 API 코드 예

**자동 미분 예**:

``파이썬
수입 잭스
jax.numpy를 jnp로 가져오기

데프 f(x):
    return jnp.sin(x) * jnp.exp(-x**2)

# 1차 미분
df = jax.grad(f)
print(df(1.0)) # df/dx at x=1.0

#2차 파생(그라디언트 중첩)
d2f = jax.grad(jax.grad(f))
print(d2f(1.0)) # d²f/dx² at x=1.0

# 함수값과 그래디언트를 모두 반환
val_grad = jax.value_and_grad(f)
print(val_grad(1.0)) # (f(1.0), df(1.0))

적시 컴파일 예시:

``파이썬 수입 잭스 jax.numpy를 jnp로 가져오기

행렬 곱셈 함수를 컴파일합니다.

@jax.jit def matmul_fast(A, B): jnp.dot(A, B)를 반환합니다.

첫 번째 호출은 XLA 컴파일을 트리거합니다(약간 더 오래 걸립니다).

A = jnp.ones((4096, 4096)) B = jnp.ones((4096, 4096)) C = matmul_fast(A, B) # 컴파일 + 실행

후속 호출은 컴파일된 코드를 직접 실행합니다.

C = matmul_fast(A, B) # 실행만 가능, 컴파일 오버헤드 없음

정적 매개변수 예: 계산 그래프로 추적할 필요가 없는 매개변수를 지정합니다.

@jax.jit(static_argnums=(2,)) def conv_with_padding(x, w, padding_mode): jnp.convolve(x, w, 모드=padding_mode)를 반환합니다.

**자동 벡터화 예**:

``파이썬
수입 잭스
jax.numpy를 jnp로 가져오기

#단일 샘플 추론 기능
def 예측_단일(매개변수, x):
    jnp.dot(params, x)를 반환합니다.

# 자동 배치 추론
배치_예측 = jax.vmap(predict_single, in_axes=(없음, 0))
# in_axes=(None, 0)은 매개변수가 분할(공유)되지 않고 x가 0번째 차원을 따라 분할됨을 의미합니다.

params = jnp.ones((256, 64))
배치_x = jnp.ones((32, 64)) # 32개 샘플
결과 = 배치_예측(매개변수, 배치_x) # 모양: (32, 256)

교차 장치 병렬 처리 예:

``파이썬 수입 잭스 jax.numpy를 jnp로 가져오기

데이터 병렬성: pmap은 기능을 모든 장치에 복사합니다.

def train_step(매개변수, 배치): 손실 = 계산_손실(매개변수, 배치) grads = jax.grad(compute_loss)(params, 배치) 반사 손실, jax.pmean(grads, axis_name='devices')

num_devices 장치는 각 배치의 일부를 처리합니다.

params = jnp.ones((1024, 512)) 배치 = jnp.ones((64, 512)) # 각 장치에 자동으로 분할됩니다. loss, grads = jax.pmap(train_step, axis_name='devices')(params, 배치)

**주요 매개변수 설명**:
- `jax.jit(fun, static_argnums=(), donate_argnums=())`: `static_argnums`는 계산 그래프로 추적되지 않는 매개변수 인덱스를 지정합니다(모양/구성 매개변수에 적용). `donate_argnums`는 비디오 메모리를 절약하기 위해 입력 버퍼를 덮어쓸 수 있음을 선언합니다.
- `jax.grad(fun, argnums=0, has_aux=False)`: `argnums`는 어떤 매개변수가 차별화되는지 지정합니다. `has_aux=True`인 경우 함수는 `(기본 출력, 보조 데이터)`를 반환하고 grad는 기본 출력만 구별합니다.
- `jax.vmap(fun, in_axes=0, out_axes=0)`: `in_axes`/`out_axes`는 배치 차원에 해당하는 입력/출력 텐서의 차원을 지정합니다.
- `jax.pmap(fun, axis_name, devices=None)`: `axis_name`은 `pmean`/`all_gather`와 같은 집단 통신 작업에 사용되는 명명된 식별자입니다. `devices`는 참여 장치의 하위 집합을 지정할 수 있습니다.
- `jax.lax.with_sharding_constraint(x, sharding)`: pjit에서 텐서 샤딩 전략을 명시적으로 지정합니다.

### 개발 도구 및 디버깅

- **jax.debug**: 0.4.20+에서는 컴파일된 중간 값을 볼 수 있는 중단점 및 인쇄 도구를 제공합니다.
- **jax.make_jaxpr**: 계산 그래프 구조를 분석하기 위해 함수를 JAX 내부 표현(Jaxpr)으로 변환합니다.
- **jax.profiler**: 커널 시간 소비 및 비디오 메모리 할당을 볼 수 있는 TensorBoard와 통합된 성능 분석 도구입니다.
- **Orbax**: Google의 공식 JAX 체크포인트 라이브러리로, 비동기 저장 및 SPMD 샤딩 체크포인트를 지원합니다.

## JAX 제품 가격

JAX 자체는 완전한 오픈 소스이며 무료이며, 총 비용은 프레임워크 사용 비용과 하드웨어 운영 비용의 두 부분으로 구성됩니다.

**프레임워크 사용 비용**:

| 프로젝트 | 가격 | 설명 |
|---|---|---|
| JAX 프레임워크 | $0 | Apache 2.0 오픈 소스 프로토콜, 무제한 상업적 사용 |
| 아마 / 하이쿠 / Optax | $0 | 상위 수준 라이브러리도 오픈 소스이며 무료입니다 |
| 기업 라이센스 | $0 | 추가 기업 계약이나 라이센스 비용이 필요하지 않습니다 |
| 기술지원 | 커뮤니티 무료/Google Cloud 유료 기술 지원 | 공식적인 무료 지원 계획; Google Cloud 고객은 TPU 관련 지원을 받을 수 있습니다 |

**하드웨어 운영 비용**:

| 하드웨어 유형 | 획득 방법 | 참고가격 |
|---|---|---|
| CPU | 자체 서버 또는 클라우드 CPU 인스턴스 | 기존 컴퓨팅 리소스에 포함됨 |
| NVIDIA GPU(개인용) | 자체 GPU | 일회성 하드웨어 투자($300-$3,000) |
| NVIDIA GPU(클라우드) | Google Cloud/AWS/Azure GPU 인스턴스 | $0.50-$5.00/시간(T4/A100/H100에 따라 다름) |
| AMD GPU(클라우드) | Google Cloud A3 인스턴스/자체 구축 | NVIDIA 클라우드 GPU |
| 구글 클라우드 TPU v5e | Google Cloud 주문형/선점형 | ~$1.50-$4.00/시간(단일 칩) |
| 구글 클라우드 TPU v5p | Google Cloud 주문형/선점형 | ~$12.00-$30.00+/시간(단일 칩) |
| TPU Pod(멀티칩 슬라이싱) | Google Cloud 선점 | 비즈니스 견적 필요, 일반적으로 시간당 $100 이상 |

**무료 할당량**: Google은 학계 연구자들에게 제한된 무료 TPU 액세스 할당량을 제공하는 TPU Research Cloud(TRC) 프로젝트를 제공합니다. 신규 Google Cloud 사용자는 TPU/GPU 인스턴스 테스트를 위해 $300의 평가판 크레딧을 받을 수 있습니다.

**유료 제안**:
- 개인 연구: 자체 GPU 또는 TRC 무료 TPU 할당량을 사용하는 것이 거의 비용이 들지 않는 가장 좋은 방법입니다.
- 중소 규모 팀: NVIDIA GPU 클라우드 인스턴스(A100 80G, ~$4/시간)를 사용하고 월예산 $1,000-$5,000.
- 대규모 교육 팀: TPU 클러스터와 GPU 클러스터의 비용 ​​성능을 평가해야 합니다. TPU Pod는 대규모 병렬 시나리오(칩 256개 이상)에서 더 효율적이지만 초기 구성 비용이 더 높고 Google Cloud에 바인딩되어 있습니다. 결정을 내리기 전 2~4주 동안 소규모로 파일럿 비교를 진행하는 것이 좋습니다.

## JAX 애플리케이션 시나리오

- **최첨단 ML 연구 및 논문 재발**: NeurIPS/ICML/ICLR 2024~2025년 논문의 약 35%에는 Transformer 변형부터 확산 모델, 강화 학습 알고리즘에 이르기까지 JAX 구현이 포함됩니다. **구현 팁**: JAX 논문을 재현할 때 Flax 또는 Haiku를 기반으로 한 오픈 소스 구현을 찾는 데 우선순위를 두십시오. 상위 수준 라이브러리에 의존하지 않는 순수 JAX 코드는 일반적으로 프로덕션 환경으로 직접 마이그레이션하기가 어렵습니다.
- **대규모 모델 학습 인프라**: JAX를 기반으로 구축된 학습 라이브러리(T5X, EasyLM, PaLM 파이프라인)는 Google의 내부 1000억 개 이상의 매개변수 모델 대부분의 학습을 지원합니다. **구현 팁**: 수백억 개의 매개변수 교육을 시작하기 전에 팀에는 pjit/shard_map 샤딩 의미 체계에 익숙한 최소 1~2명의 엔지니어가 있어야 합니다. 그렇지 않으면 디버깅 주기가 2~4주까지 길어질 수 있습니다.
- **과학 컴퓨팅 및 물리적 시뮬레이션**: JAX의 차별화 가능한 특성은 분자 역학(JAX-MD), 천체 물리학 모델링(JAX-Cosmo) 및 기후 시뮬레이션(JAX-Climate)과 같은 분야에서 고유한 이점을 제공합니다. 기존 과학 컴퓨팅 도구(예: MATLAB 및 Fortran)와 비교하여 JAX는 자동 차별화 및 GPU/TPU 가속 기능을 제공하여 과학 모델 개발을 위한 임계값을 낮춥니다. **구현 팁**: 과학 컴퓨팅 시나리오에서는 JAX(`jax.config.update("jax_enable_x64", True)`)의 64비트 모드를 먼저 사용해야 합니다. 기본 32비트 모드에서는 누적 정밀도 오류가 발생할 수 있습니다.
- **강화 학습 훈련 플랫폼**: DeepMind의 오픈 소스 RL 라이브러리(Acme, RLax, Mava)는 모두 JAX를 기반으로 구축되었으며 vmap 및 pmap을 사용하여 컨텍스트 병렬성과 훈련 병렬성을 달성합니다. **구현 팁**: RL 훈련에는 많은 수의 상황별 상호 작용이 포함되는 경우가 많습니다. JAX의 순수 함수 모델은 RL의 "상태-행동-보상" 주기에 자연스럽게 들어맞습니다. 그러나 vmap이 컨텍스트 병렬일 때 각 컨텍스트의 종료 조건이 다르기 때문에 계산 낭비가 발생한다는 점에 주의해야 합니다.
- **GPU/TPU 커널 개발 및 프로토타입 검증**: Pallas 커널 언어는 GPU 커널 개발을 위해 CUDA보다 높은 추상화 수준을 제공하며 맞춤형 연산자(예: Flash Attention 변형)를 빠르게 검증하는 데 적합합니다. **구현 팁**: Pallas는 현재 NVIDIA GPU 및 TPU만 지원하며, AMD GPU 지원은 아직 제공되지 않습니다.

안정적인; 프로덕션 수준 커널 개발은 미세 조정을 위해 여전히 CUDA로 돌아가야 합니다.

## JAX 적용 그룹

- **최첨단 ML 연구자(핵심 사용자)**: JAX의 주요 대상 그룹입니다. DeepMind, Google Brain, 최고의 AI 연구소 또는 최고의 대학에서 ML 연구를 수행하는 경우 JAX가 "모국어"입니다. JAX 함수형 프로그래밍과 pjit/shard_map 샤딩 전략에 대한 깊은 숙달은 대규모 실험을 진행하는 데 필수적인 기술입니다. **선행조건**: 자동 미분의 원리, 분산 학습의 기본 개념을 이해하고 하나 이상의 딥러닝 프레임워크를 사용해 본 경험이 있어야 합니다.
- **과학 컴퓨팅 및 미분 방정식 연구원**: 물리, 화학, 생물학, 기후 등 분야에서 수치 시뮬레이션 및 미분 방정식 풀이가 필요한 연구원. JAX의 grad/vmap/pmap 조합을 사용하면 수학 공식에서 실행 가능한 시뮬레이션까지의 주기를 크게 단축할 수 있습니다. **전제 조건**: NumPy/SciPy 생태계에 익숙하며 JAX의 수치 계산 부분을 시작하기 위해 딥 러닝 경험이 필요하지 않습니다.
- **대형 모델 교육 엔지니어**: 10B-1T 파라메트릭 규모 모델 교육을 담당하는 엔지니어링 팀입니다. JAX + TPU는 입증된 몇 안 되는 Wanka 수준 교육 솔루션 중 하나입니다. **전제 조건**: SPMD 프로그래밍 모델, 통신 토폴로지(all-reduce/all-gather/reduce-scatter), Google Cloud TPU 운영 및 유지 관리 지식에 대한 심층적인 이해가 필요합니다.
- **기계 학습 엔지니어(신중한 평가 필요)**: 일상 업무가 미세 조정, 배포 및 비즈니스 통합을 위해 사전 훈련된 모델을 사용하는 것이라면 JAX는 최선의 선택이 아닙니다. PyTorch의 커뮤니티 생태계, 배포 도구(TorchServe, ONNX, TensorRT) 및 완전성은 JAX를 훨씬 능가합니다. **부적합한 조건**: 장기적인 연구 요구가 없는 시나리오에서는 팀이 PyTorch를 메인 스택으로 사용하고 프로젝트 제공 주기가 3개월 이내인 경우 JAX를 도입하지 않는 것이 좋습니다.
- **학생 및 초보자(우선 순위로 권장되지 않음)**: JAX의 높은 추상화 및 기능적 디자인은 ML 초보자에게 친숙하지 않습니다. 먼저 PyTorch를 통해 딥러닝의 기본 개념(텐서, 자동 미분, 학습 루프)을 확립한 후, 고성능 컴퓨팅이나 특정 기술을 재현할 때 사용하는 것이 좋습니다.

조사하면서 JAX를 배우십시오. **부적절한 조건**: 딥 러닝을 처음 접한 지 6개월 미만인 학습자의 경우 JAX의 학습 곡선으로 인해 과도한 인지 부하가 ​​발생할 수 있습니다.

## 요약 및 전망

JAX는 "미분 가능 프로그래밍"이라는 기술적 방향에서 정의적인 지위를 가지고 있습니다. JAX의 기능적 설계와 기본 하드웨어의 높은 추상화 수준은 최고 임계값을 가진 최첨단 ML 연구에서 대체할 수 없습니다.

**핵심 역량**:
- **패러다임 리더십**: 기능적 + 변환기의 설계는 이론적으로 명령형 프레임워크보다 복잡한 계산을 표현하고 결합하는 데 더 적합합니다. 이러한 이점은 분산 및 다중 장치 시나리오에서 특히 두드러집니다.
- **하드웨어 추상화 깊이**: JAX + XLA의 조합은 CPU에서 TPU Pod까지 통합 프로그래밍 모델을 제공합니다. 이 모델은 한 번 작성하면 다양한 하드웨어 백엔드에서 실행할 수 있으며 이는 현재 주류 프레임워크 중에서 고유합니다.
- **대규모 훈련 검증**: DeepMind와 Google 내에서 수천~수만 개의 칩 규모로 수년간의 생산 검증을 거친 후, 대규모 병렬 훈련에 대한 JAX의 기술적 성숙도를 실제 전투에서 테스트했습니다.

**현재 제한사항**:
- **가파른 학습 곡선**: 기능적 패러다임, 변환기 구성, 샤딩 의미론과 같은 개념에는 전문적인 사고 전환이 필요합니다. 개발자가 PyTorch에서 마이그레이션하는 데 일반적으로 1~3개월이 걸립니다.
- **불충분한 생태적 풍부함**: 커뮤니티 모델 라이브러리, 타사 도구, 배포 솔루션 및 튜토리얼 리소스의 풍부함은 PyTorch보다 훨씬 적습니다. 2026년 중반 현재 PyPI의 JAX 관련 패키지 수는 PyTorch 생태계의 약 1/10입니다.
- **디버깅 어려움**: 컴파일된 함수 오류 메시지는 직관적이지 않으며 `jit` 내부의 Python 디버거(pdb)는 지원이 제한되어 있습니다. `jax.debug` 및 `jax.make_jaxpr`이 상황을 개선하고 있지만 전반적인 디버깅 경험은 여전히 ​​PyTorch Eager 모드보다 뒤떨어져 있습니다.
- **Google의 전략적 위험**: JAX의 핵심 개발은 Google이 주도하며 외부 기여자의 영향력은 제한적입니다. Google 내 TensorFlow/JAX 이중 프레임워크 간에는 유사한 상황이 있으며 기술 로드맵의 장기적인 방향에 대한 불확실성이 있습니다.

**추가 관찰 포인트**:
1. **Google 내부 통합**: Google DeepMind가 향후 2~3년 내에 TensorFlow와 JAX의 기술 경로를 통합할 것인지, 아니면 JAX를 유일한 연구 프레임워크로 명시할 것인지 여부.
2. **생태학적 성장률**: JAX 생태계는 모델 라이브러리(Hugging Face JAX/Flax 모델 비율) 및 도구 체인(디버거 프로파일러, 배포 계획) 측면에서 PyTorch와의 격차를 줄일 수 있습니까?
3. **AMD GPU 및 Apple Silicon 지원**: NVIDIA 이외의 하드웨어에 대한 JAX 지원의 성숙도는 채택 확대에 직접적인 영향을 미칩니다.
4. **커뮤니티 거버넌스 구조**: Google이 단일 회사에 대한 의존도를 줄이기 위해 보다 개방적인 커뮤니티 거버넌스 모델(예: JAX Foundation)을 구축할지 여부입니다.

**조달 및 채택 위험 평가**:
- **최첨단 연구 팀**(최고의 컨퍼런스 논문을 출판하고 새로운 아키텍처를 탐구하는 것을 목표로 함)의 경우: JAX는 반드시 숙달해야 하는 핵심 기술입니다. 1~2명의 엔지니어를 투자해 먼저 학습하고 3~6개월 내에 내부 JAX 기능을 구축하는 것이 좋습니다.
- **대형 모델 학습 팀**(대상 학습 10B+ 매개변수 모델): JAX + TPU 솔루션은 확장 효율성(특히 512+ 칩 크기) 측면에서 여전히 PyTorch + GPU 솔루션보다 앞서 있지만 Google Cloud TPU의 가용성과 비용을 평가해야 합니다. 먼저 Google TRC 무료 TPU 할당량을 신청하고 4~8주 동안 기술 검증을 수행하는 것이 좋습니다.
- **중소규모 ML 팀**(목표 7B 이하의 모델 미세 조정/추론): JAX는 권장되지 않습니다. PyTorch는 더 나은 도구 체인, 커뮤니티 지원 및 인재 풀을 갖추고 있으며 JAX 채택에 따른 숨겨진 비용(채용, 교육, 마이그레이션)이 성능 향상보다 클 수 있습니다. 향후 JAX 생태적 성숙도가 크게 향상되면 2027~2028년에 재평가될 수 있습니다.

관련 도구: <a href="https://www.aistarmap.com/ko-KR/aitool/hugging-face" target="_self" class="tool-link"><img src="https://res.aistarmap.com/uploads/images/tools/zh-CN/hugging-face/logo_1785413327.png" alt="포옹하는 얼굴" class="tool-logo" style="width: 20px; height: 20px; margin-right: 5px; vertical-align: middle;">포옹하는 얼굴</a>, replicate

버전 정보

• JAX 0.5.0 ：아직 공식적인 정확한 날짜는 없습니다. XLA 컴파일 성능 및 Pallas 커널이 지속적으로 개선됩니다. （~2026-05）
• JAX 0.4.35 ：아직 공식적인 정확한 날짜는 없습니다. AMD GPU에 대한 향상된 지원 및 성능 최적화. （~2025-12）