simple_decoding_loop

A simple decoding loop for the transformer model.

This can be used to sample from a transformer in decoding mode, and can also be used as a starting point for more sophisticated sampling algorithms.

Functions

temperature_sample_pyloop(model, prompt, rng)

Runs temperature sampling in a Python for loop.