CHECKPOINT token-chunked response-only mean logp from a PRE-COMPUTED
post-finalNorm hidden h [1,T,hidden] (the segmented-backward head variant of
chunkedLogpMeanB1). Each token-chunk's logitsFromHidden is wrapped in a
Checkpoint, so the per-chunk [chunk,V] logits are RECOMPUTED in the
backward and freed sequentially — the recompute boundary is what actually
bounds the head term under MLX's lazy eval when nested inside the segmented
mlx_vjp (unlike the analytic CustomVjp, whose per-chunk graph isn't freed
incrementally inside an outer vjp). Checkpoints + slices go in sink.
CHECKPOINT token-chunked response-only mean logp from a PRE-COMPUTED post-finalNorm hidden
h [1,T,hidden](the segmented-backward head variant ofchunkedLogpMeanB1). Each token-chunk'slogitsFromHiddenis wrapped in aCheckpoint, so the per-chunk[chunk,V]logits are RECOMPUTED in the backward and freed sequentially — the recompute boundary is what actually bounds the head term under MLX's lazy eval when nested inside the segmentedmlx_vjp(unlike the analytic CustomVjp, whose per-chunk graph isn't freed incrementally inside an outer vjp). Checkpoints + slices go insink.