Flatten a ``(B, n)``-style shape tensor to per-sample length.

Create a batch of ChatML-formatted sequences.

Create synthetic z_0, z_noisy, timesteps for one sample.

Tokenize, encode through VAE. Returns None if the conversation is too short or has no response.

Preprocess one JSONL file in place. Extract the common prompt prefix, split generate, and reorder keys.

Process one JSON record. Split the ``generate`` field at the first newline whose preceding segment contains an actual word or sentence.

Render a conversation dict into token IDs + roles. Returns (ids: list[int], roles: list[str]) where roles[i] is 'P' or 'R'.

Reorder dictionary keys so id, prompt, generate, and ground_truth come first.

Sample a block size from the configured distribution.

NA-form forward pass — one prior-transport step on block ``b``. Each call evaluates the block-causal vector field ``v_psi(z_t^(b), t;