Improve JAX for pretraining notebook. by f33ac · Pull Request #293 · jax-ml/jax-ai-stack

✅ All contributors are covered under a CLA with Google

See https://cla.developers.google.com/ for more info about Google's Contributor License Agreement (CLA).

ℹ️ Googlers: Go here to view more details and manage scans for this pull request.

@f33ac

The following contributors were found for this pull request:

✅ cee9e4e Author: @f33ac <acz*****ki@f33.ai>
✅ 40a764d Author: @f33-kt <kt***s@fourteen33.com>

(Only the first commit for a unique contributor is listed.)