🧬

JAX Quest

최근 수정: 2026-07-29

첫 jit부터 미분 가능한 물리 시뮬레이션까지

가속기에서 실행되는 수치 코드를 함수 변환으로 다루는 법을 배운다. 14개 트랙, 73개 레슨, 약 28시간에 걸쳐 "jax.numpy는 그냥 NumPy일까?"라는 질문에서 출발해 미분 가능한 시뮬레이터와 사용자 정의 XLA 커널까지 나아간다.

14 tracks · 73 lessons · ~28h · difficulty: intermediate-to-advanced

Level 0호기심

0 XP0/73 lessons0/17 achievements

0/100 XP to next level100 XP to go0% complete

JAX는 또 하나의 딥러닝 프레임워크라기보다는 강력한 수치 계산 시스템으로, 머신러닝에도 널리 쓰인다. 이 퀘스트는 첫 jax.grad 호출에서 출발해 SPMD 병렬 처리, custom_vjp, 그래디언트 체크포인팅, 메타러닝, 미분 가능한 물리 시뮬레이션 같은 연구 수준의 패턴으로 나아간다. 네 가지 핵심 변환인 jit, grad, vmap, pmap을 자유롭게 합성하는 법을 익히고, Flax·Equinox·Optax·Diffrax·NumPyro·Brax가 그 철학을 어떻게 이어받는지도 살펴본다. 완주할 즈음에는 현대 수치 계산 연구의 코드를 읽고 직접 설계할 토대가 생긴다.

Tracks

01🧬JAX는 왜 만들어졌고 무엇이 다른가
0/6 lessons
JAX가 필요한 이유와 PyTorch·TensorFlow와의 차이
JAX의 뿌리를 따라가며 이것이 왜 또 하나의 딥러닝 프레임워크가 아닌지 살펴본다. 합성 가능한 함수 변환이라는 핵심 명제가 수치 코드를 작성하는 방식을 어떻게 바꾸는지도 배운다.
Lesson list (6)
퀴즈 · 3 문제→
02🔢NumPy처럼 쓰는 JAX, 기초 다지기
0/6 lessons
가속기에서 실행되고 미분할 수 있는 NumPy
가속기에서 실행되는 NumPy라는 말의 뜻을 알아본다. jax.numpy가 NumPy와 닮은 부분, 의도적으로 다른 부분, 초보자가 자주 놓치는 기본 dtype까지 짚는다.
Lesson list (6)
퀴즈 · 2 문제→
03✨순수 함수, 모든 변환을 여는 규칙
0/5 lessons
JAX의 변환을 가능하게 하는 하나의 원칙
JAX가 순수 함수를 요구하는 이유와 추적 시점에서 순수성이 실제로 무엇을 뜻하는지 배운다. 부수 효과가 있는 NumPy 코드를 JAX에 맞는 형태로 리팩터링하는 방법도 익힌다.
Lesson list (5)
퀴즈 · 2 문제→
04⚡jit, XLA 컴파일
0/5 lessons
추적하고 컴파일한 뒤 캐시하는 과정
jit이 함수를 추적하고 컴파일한 뒤 결과를 캐시하는 과정을 배운다. 정적 인자와 추상 값의 의미를 익히고, 추적의 원리를 모를 때 마주치는 대표적인 오류도 해석한다.
Lesson list (5)
퀴즈 · 2 문제→
05📐grad, 함수형 자동 미분
0/5 lessons
함수를 받아 그래디언트 함수를 만드는 변환
JAX의 함수형 자동 미분을 배운다. grad, value_and_grad, has_aux에서 출발해 고차 그래디언트와 Jacobian으로 나아가고, grad를 jit·vmap과 합성하는 표준 패턴을 익힌다.
Lesson list (5)
퀴즈 · 2 문제→
06🔄vmap, 자동 벡터화
0/4 lessons
배치 반복문 없이 함수를 벡터화하기
배치 반복문을 직접 쓰지 않고 함수를 벡터화한다. in_axes와 out_axes, 예제별 그래디언트, 중첩 vmap을 익혀 배치 차원을 손으로 처리하는 코드를 없앤다.
Lesson list (4)
퀴즈 · 2 문제→
07🌐pmap, 여러 장치에 걸친 병렬화
0/5 lessons
여러 장치에서 실행하는 SPMD 병렬 처리
여러 장치에서 SPMD 방식으로 계산하는 법을 배운다. 집합 통신과 분산 학습 패턴을 익히고, jax.sharding과 Mesh가 순수한 pmap 사용을 대체하는 현대적인 선택지가 된 이유도 살펴본다.
Lesson list (5)
퀴즈 · 2 문제→
08🎲난수, 함수형으로 다루기
0/6 lessons
명시적 PRNG 키와 재현 가능한 난수
JAX가 난수를 명시적으로 다루는 이유를 배운다. PRNG 키를 안전하게 나누고 배분하는 법과, 같은 seed를 계속 재사용하는 실수를 막는 패턴을 익힌다.
Lesson list (6)
퀴즈 · 2 문제→
09🌳Pytree, JAX의 보편 자료구조
0/5 lessons
중첩된 자료구조를 한 방식으로 다루는 추상화
Pytree는 매개변수 묶음, 상태 딕셔너리, 모델처럼 중첩된 구조를 JAX가 한 방식으로 다루게 해주는 추상화다. 매핑과 평탄화, 사용자 정의 노드 등록을 익힌다.
Lesson list (5)
퀴즈 · 2 문제→
10🧠신경망 만들기, Flax와 Equinox
0/6 lessons
Flax NNX와 Equinox로 구성하는 JAX 네이티브 신경망
JAX 코어에 nn.Module이 없는 이유와 Flax NNX·Equinox가 그 자리를 어떻게 채우는지 알아본다. 같은 CNN을 두 라이브러리로 구현해 비교하고 Transformer 블록도 구성한다.
Lesson list (6)
퀴즈 · 2 문제→
11🔄학습 루프, Optax와 JAX의 방식
0/5 lessons
Optax, 학습률 스케줄, 명시적인 학습 루프
Optax의 그래디언트 변환을 레고 블록처럼 합성한다. 학습률 스케줄과 표준 학습 단계를 구성하고, scan과 Orbax 체크포인트로 오래 실행해도 견고한 학습 루프를 만든다.
Lesson list (5)
퀴즈 · 2 문제→
12⚡고급 변환과 패턴
0/5 lessons
custom_vjp, remat, scan, 메타러닝
사용자 정의 그래디언트 규칙, 메모리를 아끼는 그래디언트 체크포인팅, jit 안에서 구조적인 제어 흐름을 만드는 scan, jax.debug, 메타러닝의 그래디언트-위-그래디언트 패턴을 배운다.
Lesson list (5)
퀴즈 · 2 문제→
13🔬머신러닝 너머의 JAX, 과학 계산
0/5 lessons
Diffrax, NumPyro, 미분 가능한 물리 시뮬레이션
머신러닝 밖에서 JAX가 어떻게 쓰이는지 살펴본다. Diffrax로 미분방정식을 풀고 NumPyro로 베이지안 추론을 수행하며, 물리 시뮬레이터를 끝까지 미분하는 예제를 만든다.
Lesson list (5)
퀴즈 · 2 문제→
14🚀생태계, 배포, 다음 단계
0/5 lessons
프로덕션 배포와 프로파일링, 그리고 다음 단계
모델을 프로덕션 형식으로 내보내고 TPU와 GPU에서 실행하는 법을 살펴본다. 프로파일링 도구와 Pallas 사용자 정의 커널을 익히고, 2026년 현재 JAX 생태계의 위치와 다음 단계를 정리한다.
Lesson list (5)
퀴즈 · 2 문제→

이 페이지에서 버그를 발견하셨거나 피드백이 있으세요?문제 신고

🔔 답글 알림 (로그인 필요)

로그인 — 댓글을 남기려면 로그인해 주세요.

아직 댓글이 없어요. 첫 댓글을 남겨보세요.

Tracks

01🧬JAX는 왜 만들어졌고 무엇이 다른가

02🔢NumPy처럼 쓰는 JAX, 기초 다지기

03✨순수 함수, 모든 변환을 여는 규칙

04⚡jit, XLA 컴파일

05📐grad, 함수형 자동 미분

06🔄vmap, 자동 벡터화

07🌐pmap, 여러 장치에 걸친 병렬화

08🎲난수, 함수형으로 다루기

09🌳Pytree, JAX의 보편 자료구조

10🧠신경망 만들기, Flax와 Equinox

11🔄학습 루프, Optax와 JAX의 방식

12⚡고급 변환과 패턴

13🔬머신러닝 너머의 JAX, 과학 계산

14🚀생태계, 배포, 다음 단계