JAX 편집하기 (부분)

==주요 특징==
===Autograd (자동 미분)===
JAX는 Python 함수를 입력받아 해당 함수의 기울기를 자동 계산한다.
*reverse-mode differentiation
*forward-mode differentiation
*많은 연산에 대한 강력한 합성(differentiable programming)
예시:<syntaxhighlight lang="python">
import jax.numpy as jnp
import jax

f = lambda x: x**2 + 3*x
jax.grad(f)(2.0)  # 7.0
</syntaxhighlight>
===JIT 컴파일===
XLA 기반의 JIT 컴파일 기능(`jax.jit`)은 Python 코드를 고성능 기계 코드로 변환하여 실행 속도를 크게 높인다.
*연산 fusion
*device placement 자동 처리
*GPU/TPU에서 높은 효율
<syntaxhighlight lang="python">
@jax.jit
def compute(x):
    return jnp.sin(x) * jnp.cos(x)
</syntaxhighlight>
===벡터화 (vmap)===
루프를 Python이 아닌 XLA에서 병렬화하여 대규모 데이터 처리 속도를 높인다.
*Python loop → XLA 병렬 연산 자동 변환
===pmap (병렬 처리)===
다중 GPU/TPU 장치에서 동일한 함수를 병렬 실행.
*Data parallelism 구현
*TPU Pod 환경에서 핵심 기능
===함수형 패러다임===
JAX는 순수 함수(pure function)를 강조하며, 상태(state) 없는 계산을 지향한다.  모델 파라미터도 불변성을 유지한 채 함수 인자로 전달하는 방식으로 관리한다.