2026의 JAX vs PyTorch, 그리고 다음 단계

# JAX shines when you need:
# 1. Composable transformations (grad of grad, vmap of grad)
# 2. TPU support (JAX is first-class on TPUs)
# 3. Functional programming style
# 4. Scientific computing + ML hybrid workloads
# 5. Custom differentiation rules
# 6. Research that pushes framework boundaries

# PyTorch is better when you need:
# 1. Largest community and ecosystem
# 2. Dynamic computation graphs with easy debugging
# 3. Rapid prototyping with less boilerplate
# 4. Production deployment tooling (TorchServe, etc.)
# 5. Most SOTA model implementations available first

# Key JAX milestones:
# JAX 0.5.0 (Feb 2025) — Partitionable PRNG by default
# JAX 0.6.0 (Apr 2025) — CUDA 12.8+ required, API cleanup
# JAX 0.7.0 (Jul 2025) — Migrated from GSPMD to Shardy,
#                         direct linearization, jax.P alias,
#                         minimum Python 3.11
# JAX 0.8.x (Dec 2025) — Continued stabilization
# JAX 0.9.x (Jan 2026) — Latest release as of early 2026

# Resources for continuing your JAX journey:

# Official documentation
# - jax.readthedocs.io (core JAX)
# - flax.readthedocs.io (Flax NNX)
# - optax.readthedocs.io (Optax)
# - docs.kidger.site/equinox (Equinox)
# - docs.kidger.site/diffrax (Diffrax)
# - num.pyro.ai (NumPyro)

# Learning projects to try:
# 1. Train a small Transformer on a text dataset
# 2. Implement a variational autoencoder (VAE)
# 3. Solve a differential equation with Diffrax
# 4. Build a Bayesian neural network with NumPyro
# 5. Write a differentiable physics simulation
# 6. Fine-tune a model with LoRA using Flax NNX

# Community
# - github.com/jax-ml/jax (source + issues)
# - JAX Discussions on GitHub
# - r/MachineLearning on Reddit
# - JAX Discord / Slack communities