Fix shape error in TransformerBlock

[stephen-huan]

Clear shape error that isn't caught because embed_dim = ff_dim = 32 happened to be chosen. Fix the shape error and set ff_dim = 64 (arbitrary) to make sure it doesn't happen again.

[stephen-huan]

How should the notebooks be re-ran? I tried

jupytext --execute docs/source/JAX_transformer_text_classification.ipynb

and

jupyter nbconvert --execute \
    --ClearMetadataPreprocessor.enabled=True \
    --ClearMetadataPreprocessor.clear_notebook_metadata=False \
    --inplace \
    docs/source/JAX_transformer_text_classification.ipynb \

but I get a lot of tqdm in the stderr outputs from the "Training the model" section (tqdm logs every "tick" rather showing the last update as currently on the website). I ended up opening it in jupyter lab and re-running all the cells.

Also, nnx.display(model) now shows an interactive html widget instead of printing to stdout. I opened the website built in CI and it seems to work.

Some tutorials have an explicit id in the markdown source which induces a "id" field in the "metadata" in the respective cell of the corresponding notebook.

jax-ai-stack/docs/source/JAX_for_LLM_pretraining.md

Line 30 in 7f1c736

+++ {"id": "NIOXoY1xgiww"}

jax-ai-stack/docs/source/JAX_for_LLM_pretraining.ipynb

Lines 29 to 31 in 7f1c736

	"metadata": {
	"id": "NIOXoY1xgiww"
	},

However, some cells of this notebook have "id" as a top-level cell field and no "id"''s in the markdown source.

jax-ai-stack/docs/source/JAX_transformer_text_classification.ipynb

Lines 228 to 230 in 7f1c736

	"cell_type": "markdown",
	"id": "deec4652",
	"metadata": {},

This doesn't seem to be valid jupyter and so jupyter removes them when I re-ran the cells.

[NbConvertApp] Converting notebook docs/source/JAX_transformer_text_classification.ipynb to notebook
[NbConvertApp] ERROR | Notebook JSON is invalid: Additional properties are not allowed ('id' was unexpected)

Failed validating 'additionalProperties' in markdown_cell:

On instance['cells'][10]:
{'cell_type': 'markdown',
 'id': 'deec4652',
 'metadata': {},
 'source': '**Note:** We are using a single-device setup, so [device '
           'paralle...'}
[NbConvertApp] Writing 787288 bytes to docs/source/JAX_transformer_text_classification.ipynb

jupyter --version

Selected Jupyter core packages...
IPython          : 9.6.0
ipykernel        : 7.1.0
ipywidgets       : not installed
jupyter_client   : 8.6.3
jupyter_core     : 5.9.1
jupyter_server   : 2.17.0
jupyterlab       : 4.4.10
nbclient         : 0.10.2
nbconvert        : 7.16.6
nbformat         : 5.10.4
notebook         : not installed
qtconsole        : not installed
traitlets        : 5.14.3