JAX가 등장하기 전에는 Autograd라는 작은 Python 라이브러리가 있었어. 평범한 Python과 NumPy 코드를 그대로 작성하면 그래디언트를 자동으로 구해 주는 도구였지. 수식을 코드로 옮기기만 하면 미분이 따라온다는 발상은 우아했어. 다만 CPU에서만 실행됐고 컴파일할 수 없어서 GPU나 TPU 같은 현대적인 가속기를 활용하지 못했어.
2018년, Google Brain의 Matt Johnson, Roy Frostig, Dougal Maclaurin, Chris Leary는 Autograd의 후계자를 만들기 시작했어. 목표는 Autograd의 단순한 인터페이스, 즉 "평범한 Python 함수를 그대로 미분한다"는 원칙을 유지하면서 Google의 XLA 컴파일러를 결합하는 것이었지. 같은 코드를 GPU와 TPU에서도 빠르게 실행하려는 시도였고, 그 결과가 JAX야. 이름은 원래 "Just After eXecution"의 줄임말이었지만 지금은 그저 JAX라는 이름으로 쓰여.
2018년 말 오픈소스로 공개된 JAX는 TensorFlow의 계산 그래프 구축 방식이나 PyTorch의 객체지향 계층과 다른 중심축을 제안했어. 바로 합성 가능한 함수 변환이야. NumPy를 쓰듯 평범한 함수를 작성한 다음, 그 함수를 한 번 감싸 컴파일하고 미분하고 벡터화하고 병렬화해. 각 변환을 서로 조합할 수도 있고.
💡 왜 이게 중요한가
JAX는 "또 하나의 딥러닝 프레임워크"로 만들어진 게 아니야. 머신러닝에도 강력한 수치 계산 시스템에 가까워. 그래서 계층을 쌓는 방법보다 수학적 함수를 어떻게 표현하고 변환할지에 초점을 둬. 이 관점이 PyTorch나 TensorFlow와 다른 사용감을 만들어.
오늘날 JAX는 Google DeepMind의 주요 연구 프레임워크 가운데 하나야. AlphaFold의 단백질 구조 예측, Gemini 백엔드 연구, 여러 논문에서 JAX가 사용됐어. Stanford, MIT, Berkeley, MILA 같은 연구 기관에서도 과학 계산, 물리 시뮬레이션, 베이지안 추론, 강화학습에 JAX를 활용해.
첫 JAX 프로그램은 이렇게 단순한 모습이야.
import jax
import jax.numpy as jnp
# 평범한 함수 — class 도 decorator 도 필요 없음
def f(x):
return jnp.sum(x ** 2)
x = jnp.array([1.0, 2.0, 3.0])
# 그냥 호출
print(f(x)) # 14.0
# 이제 gradient — 함수 한 번 감싸기만 하면 끝
grad_f = jax.grad(f)
print(grad_f(x)) # [2. 4. 6.]
# Compile 해서 빠르게
fast_f = jax.jit(f)
print(fast_f(x)) # 14.0 (다음 호출부터 훨씬 빠름)
여기에 JAX의 약속이 담겨 있어. 먼저 함수를 작성하고, JAX가 제공하는 도구로 그 함수를 변환해. 프레임워크의 클래스를 상속하거나 새로운 텐서 사용법부터 배울 필요는 없어. 함수를 받아 함수를 돌려주는 방식, 그것이 출발점이야.