🧬

JAX Quest

최근 수정: 2026-05-03

첫 jit 부터 미분 가능한 physics 까지

가속기 위의 합성 가능한 함수 변환. 14 track, 73 lesson, ~28 시간. 'jax.numpy 가 그냥 NumPy?' 에서 미분 가능 simulator 와 custom XLA kernel 작성까지.

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 는 또 하나의 deep learning framework 가 아니라 — ML 도 잘하는 numerical computing 시스템. 이 quest 는 첫 jax.grad 호출부터 research-grade 패턴까지: SPMD 병렬, custom_vjp, gradient checkpointing, meta-learning, differentiable physics. 4 대 변환 (jit, grad, vmap, pmap) 이 자유롭게 합성, 그리고 나머지 생태계 (Flax, Equinox, Optax, Diffrax, NumPyro, Brax) 가 같은 DNA 를 재사용. 끝내고 나면 — 현대 numerical research 에 fluent.

Tracks

01🧬JAX 가 존재하는 이유와 무엇이 다른가
0/6 lessons
JAX 가 왜 존재하는지, PyTorch 와 TensorFlow 와 무엇이 다른지
JAX가 어디서 왔는지, 왜 또 하나의 deep learning framework가 아닌지, 그리고 functional-transform 명제가 numerics 작성 방식에 의미하는 것.
Lesson list (6)
퀴즈 · 3 문제→
02🔢JAX as NumPy — 기반
0/6 lessons
가속기 위의 미분 가능한 NumPy 로서의 JAX
가속기 위 NumPy 의미. jax.numpy 가 NumPy 와 같은 곳, 의도적으로 다른 곳, 모든 신참을 무는 dtype default.
Lesson list (6)
퀴즈 · 2 문제→
03✨Pure Function — 모든 걸 푸는 규칙
0/5 lessons
pure 함수 — 모든 변환을 unlock 하는 규칙
JAX 가 pure 함수를 강요하는 이유, 'pure' 가 trace-time 에서 실제로 의미하는 것, impure NumPy 코드를 JAX-clean 모양으로 refactor 하는 법.
Lesson list (5)
퀴즈 · 2 문제→
04⚡jit — XLA Just-In-Time Compilation
0/5 lessons
XLA 의 just-in-time compilation
jit 이 어떻게 trace, compile, cache 하는지. static argument, abstract value, tracing 을 익히기 전에 부딪힐 정석 에러.
Lesson list (5)
퀴즈 · 2 문제→
05📐grad — 자동 미분, functional 식
0/5 lessons
functional 식 자동 미분
JAX 의 functional autodiff: grad, value_and_grad, has_aux, 고차 gradient, Jacobian, 그리고 정석 compose-grad-with-jit-and-vmap 패턴.
Lesson list (5)
퀴즈 · 2 문제→
06🔄vmap — 자동 vectorization
0/4 lessons
batch loop 없는 자동 vectorization
batch loop 작성 없이 vectorize. in_axes, out_axes, per-example gradient, nested vmap — batch dimension 손코딩이 끝나는 변환.
Lesson list (4)
퀴즈 · 2 문제→
07🌐pmap — 다중 device 병렬화
0/5 lessons
다중 device SPMD 병렬화
여러 device 의 SPMD 병렬화. collective, distributed training 패턴, 그리고 jax.sharding + Mesh 가 pure pmap 의 modern 대체인 이유.
Lesson list (5)
퀴즈 · 2 문제→
08🎲Random Number — functional 방식
0/6 lessons
functional PRNG key 와 reproducible randomness
JAX 가 randomness 를 명시적으로 만든 이유, PRNG key 가 어떻게 깨끗히 split / partition 되는지, 같은 seed 로 영원히 학습하는 사고를 막는 패턴.
Lesson list (6)
퀴즈 · 2 문제→
09🌳Pytree — JAX 의 보편 자료구조
0/5 lessons
JAX 의 보편 nested-structure 추상화
Pytree 는 JAX 의 nested-structure 추상화 — 모든 parameter pack, state dict, model. map, flatten, custom node 등록.
Lesson list (5)
퀴즈 · 2 문제→
10🧠Neural Network 만들기 — Flax, Equinox 등
0/6 lessons
Flax NNX, Equinox, JAX-native neural network
JAX core 에 nn.Module 이 없는 이유, 그 빈자리를 채우는 것 (Flax NNX, Equinox), 같은 CNN 이 두 라이브러리에서 어떻게 보이는지. plus Transformer block.
Lesson list (6)
퀴즈 · 2 문제→
11🔄Training Loop — Optax 와 JAX 식
0/5 lessons
Optax, schedule, 명시적 training loop
Optax 가 Lego 처럼 gradient transformation 합성. schedule, 정석 training step, robust run 을 위해 scan + Orbax 로 checkpoint 하는 법.
Lesson list (5)
퀴즈 · 2 문제→
12⚡Advanced Transformation 과 패턴
0/5 lessons
custom_vjp, remat, scan, meta-learning 패턴
Custom gradient rule, memory 를 위한 gradient checkpointing, jit 아래 structured control flow 의 scan, jax.debug, meta-learning grad-of-grad 패턴.
Lesson list (5)
퀴즈 · 2 문제→
13🔬ML 너머의 JAX — 과학 계산
0/5 lessons
ML 너머의 JAX — Diffrax, NumPyro, differentiable physics
ML 외부의 JAX: Diffrax 로 differential equation, NumPyro 로 Bayesian inference, physics simulator, fully differentiable simulation 예제.
Lesson list (5)
퀴즈 · 2 문제→
14🚀생태계, 배포, 다음 단계
0/5 lessons
production deploy, profiling, 앞으로의 길
production model export, TPU/GPU runtime, profiling 도구, Pallas custom kernel, 그리고 다음 단계의 솔직한 2026 현재.
Lesson list (5)
퀴즈 · 2 문제→

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

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

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

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

Tracks

01🧬JAX 가 존재하는 이유와 무엇이 다른가

02🔢JAX as NumPy — 기반

03✨Pure Function — 모든 걸 푸는 규칙

04⚡jit — XLA Just-In-Time Compilation

05📐grad — 자동 미분, functional 식

06🔄vmap — 자동 vectorization

07🌐pmap — 다중 device 병렬화

08🎲Random Number — functional 방식

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

10🧠Neural Network 만들기 — Flax, Equinox 등

11🔄Training Loop — Optax 와 JAX 식

12⚡Advanced Transformation 과 패턴

13🔬ML 너머의 JAX — 과학 계산

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