ReformerLM output length limit

In ReformerLM_output_gen function, output_gen = trax.supervised.decoding.autoregressive_sample_stream(…) seems not having a limit of the words being generated. Will it be automatically stopped by reaching the maximum window size of sequences in each batch? If so, which boundary in the following setting is chosen?[128, 256, 512, 1024],
batch_sizes=[16, 8, 4, 2, 1]),

Hello @YIHUI!

as said in the docs of autoregressive_sample_stream:
“Inputs and outputs always come in batches, even if size 1. If inputs is present, it must have shape (batch_size, inputs_sequence_length), and each output in the stream has shape (batch_size, 1).”

So, in this case you need to create input tokens using the the tokenize function and this will be your inputs for autoregressive_sample_stream with shape (1, n), batch_size=1. Is this make sense for you?

Best regards,
Wesley P.