Improve JAX for pretraining notebook. - #293

Merged

emilyfertig merged 11 commits into

jax-ml:revamp-2025from

f33ac:revamp-2025

Mar 10, 2026

Commits on Jan 6, 2026

Improve JAX for LLM pretraining notebook.
f33ac
committed
78c083d

Commits on Jan 29, 2026

Update the notebook. Fix size.
f33ac
committed
1741da0
Remove widgets
f33ac
committed
c42601a

Commits on Jan 30, 2026

Sync notebook with Markdown format.
f33ac
committed
f32942b
Validate outputs with pre-commit.
f33ac
committed
82223fd

Commits on Feb 3, 2026

readd nbval tags; update markdown
f33-kt
committed
40a764d

Commits on Feb 4, 2026

Validate outputs with pre-commit.
f33ac
committed
35727cc
Clean metadata.
f33ac
committed
91306ee
Remove Google Colab metadata.
f33ac
committed
eed8c50

Commits on Feb 6, 2026

Fix minor issues.
f33ac
committed
0d20f18

Commits on Feb 10, 2026

Add proper nbval tags.
f33ac
committed
cee9e4e