Document context_model as coordinates-only with label context

README.md

@@ -11,8 +11,8 @@ tags:
 # STIP tutorial checkpoints
 
 Small checkpoints used by the [`stip`](https://github.com/instadeepai/stip) tutorial
-notebooks to demonstrate sampling methods and capapbilities of the repository. 
-They are toy models (a two-layer MLP, ~50k parameters,
+notebooks, so that a tutorial can demonstrate sampling without spending ten
+minutes training first. They are toy models (a two-layer MLP, ~50k parameters,
 trained for 1500 steps on a 4-component 2D Gaussian mixture) and have no
 use outside the notebooks.
 
@@ -23,13 +23,13 @@ Checkpoints are [Orbax](https://orbax.readthedocs.io) directories written by
 ## `conditioning_and_guidance/`
 
 Used by `tutorials/notebooks/4.conditioning_and_guidance.ipynb`. Both models are
-`VelocityGenerativeModel`s over two modalities — `coordinates` (continuous, 2D)
-and `index` (discrete, 4 categories) — with a `FlowMatchingOneSidedInterpolant`.
+`VelocityGenerativeModel`s with a `FlowMatchingOneSidedInterpolant`, but over
+different modalities:
 
-| Path | Model | Role in the notebook |
-|---|---|---|
-| `conditioning_and_guidance/joint_model` | Unconditional cross-modal MLP | Intrinsic guidance (Section 4): conditioning a model using a model trained unconditionally |
-| `conditioning_and_guidance/context_model` | The same MLP plus a label context path, trained with 20% context dropout | Context conditioning and classifier-free guidance (Sections 5-7) |
+| Path | Model | Modalities | Role in the notebook |
+|---|---|---|---|
+| `conditioning_and_guidance/joint_model` | Unconditional cross-modal MLP | `coordinates` (continuous, 2D) and `index` (discrete, 4 categories) | Intrinsic guidance (Section 4): conditioning a model that was never trained to be conditional |
+| `conditioning_and_guidance/context_model` | The same MLP plus a label context path, trained with 20% context dropout | `coordinates` only; the corner label is passed as `context_data` instead of as a modality | Context conditioning and classifier-free guidance (Sections 5-7) |
 
 ### Loading
 
@@ -54,3 +54,13 @@ gen_model = nnx.merge(graphdef, checkpointer.restore_ema(params))
 
 `restore_ema` reads only `ema_params` and `extra`, and applies the same bias
 correction the training loop uses for evaluation.
+
+## Reproducing
+
+```bash
+uv run python tutorials/scripts/train_conditioning_checkpoints.py
+```
+
+The script mirrors the notebook's model definitions and PRNG chain, so it
+reproduces these exact weights. A checkpoint pins the parameter structure: if a
+notebook's network changes, re-run the script and re-upload.