JIT compile inference func in LLM pretrain tutorial

[windmaple]

Rendered preview: https://jax-ai-stack--197.org.readthedocs.build/en/197/JAX_for_LLM_pretraining.html

[jakevdp]

Can you describe the changes you made? The GitHub UI can't show the diff because it's too large. Is that expected? Given the title of the PR, I would expect a change of only a few lines.

[windmaple]

Apparently loading tensorboard completely blows up the notebook in size (110KB -> 83MB); I've avoided doing that so that the notebook can render on GitHub now.

Previously for inference the input wasn't padded so JIT wasn't used. Now this is done properly and speed is much faster. Also streaming tokens out instead of batching them.

I had to re-run the notebook, which trigged a lot of metadata and output changes in the notebook.

Btw, have you considered using ReviewNB so that nb diff is much easier to review? We used it for another project last year and it worked like a charm.

[jakevdp]

The lint is failing because the CI expects jupytext version 1.15.2. Locally you can use pre-commit run --all-files to ensure that you're using the exact same package versions and commands as the CI. More info at https://jax-ai-stack--197.org.readthedocs.build/en/197/contributing.html#editing-an-existing-notebook