[
  {
    "question": "What is the fundamental design difference between PyTorch and JAX?",
    "options": ["PyTorch is faster", "PyTorch mutates tensors eagerly while JAX compiles pure functions without side effects", "PyTorch doesn't support automatic differentiation", "JAX doesn't support GPUs"],
    "correct": 1,
    "explanation": "PyTorch uses eager execution with mutable state (e.g., tensor.grad is modified in place). JAX embraces functional programming: functions are pure, state is explicit, and jit compilation optimizes entire computation graphs.",
    "stage": "pre"
  },
  {
    "question": "What does jax.jit do?",
    "options": ["It adds dropout regularization", "It compiles a Python function into optimized XLA code that runs much faster than interpreted Python", "It computes gradients", "It initializes model weights"],
    "correct": 1,
    "explanation": "jax.jit traces a function and compiles it to XLA (Accelerated Linear Algebra) machine code. The first call is slow (compilation), but subsequent calls run the optimized compiled version.",
    "stage": "pre"
  },
  {
    "question": "What does jax.vmap do?",
    "options": ["Computes second-order gradients", "Manages GPU memory", "Validates model architecture", "Vectorizes a function to run over a batch dimension without writing explicit loops"],
    "correct": 3,
    "explanation": "jax.vmap automatically vectorizes a function written for a single example to process an entire batch. You write code for one sample and vmap handles batching, often more efficiently than manual batch loops.",
    "stage": "post"
  },
  {
    "question": "How does JAX handle model state (weights) differently than PyTorch?",
    "options": ["JAX requires explicit state passing -- weights are function arguments, not mutable object attributes", "JAX and PyTorch handle state identically", "JAX doesn't support trainable weights", "JAX stores weights on CPU only"],
    "correct": 0,
    "explanation": "In PyTorch, weights live inside nn.Module objects and are mutated in place. In JAX, weights are passed explicitly as function arguments and returned as new values. No mutation, no hidden state.",
    "stage": "post"
  },
  {
    "question": "When would you choose JAX over PyTorch?",
    "options": ["When you need the largest ecosystem of pretrained models", "For deployment on mobile devices", "For quick prototyping of small models", "When training at massive scale on TPU pods, where compilation and functional transforms provide significant speedups"],
    "correct": 3,
    "explanation": "JAX excels at scale: jit compilation eliminates Python overhead, pmap handles multi-device parallelism naturally, and XLA optimization across the full computation graph gives significant speedups on TPU clusters.",
    "stage": "post"
  }
]