On first looking into JAX

Much have I travell'd in the realms of gold, And many goodly states and kingdoms seen; Round many western islands have I been Which bards in fealty to Apollo hold. Oft of one wide expanse had I been told That deep-brow'd Homer ruled as his demesne; Yet did I never breathe its pure serene Till I heard Chapman speak out loud and bold: Then felt I like some watcher of the skies When a new planet swims into his ken; Or like stout Cortez when with eagle eyes He star'd at the Pacific -- and all his men Look'd at each other with a wild surmise -- Silent, upon a peak in Darien. John Keats, On First Looking into Chapman's Homer I've been working with PyTorch quite a lot for the last couple of years, and feel like I've come to a reasonably solid understanding of how it all fits together. Working through Sebastian Raschka's book "Build a Large Language Model (from Scratch)", training my own LLMs locally and in the cloud, rebuilding Andrej Karpathy's 2015-vintage RNNs -- over time, it all adds…