Token Generation using ConditionalAutoregressive2D

[aszala]

Hi, I am trying to use Jukebox to generate my own tokens. The paper mentions that you pass all previously generated tokens to the model as input for generating the next token.

However, I was reading through the code to understand how it works, and I am confused about how it passes the previously generated tokens to the model.

Here is the code snippet from the ConditionalAutoregressive2D class in the sample method I am referring to: https://github.com/openai/jukebox/blob/08efbbc1d4ed1a3cef96e08a931944c8b4d63bb3/jukebox/prior/autoregressive.py#L222-L237

Even the self.get_emb(sample_t, n_samples, x, x_cond, y_cond) line, doesn't retain the previous tokens, it just adds the updated positional embedding.

https://github.com/openai/jukebox/blob/08efbbc1d4ed1a3cef96e08a931944c8b4d63bb3/jukebox/prior/autoregressive.py#L177-L197

Could someone explain how Autoregressively generates and where the previous tokens are used as input?

Thanks!