Update jax_intro lecture: restructure, fix bugs, improve exposition

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: jstac