[BUG] The implementation bug of the ’predict‘ function in class ChronosPipeline

[tychenn]

Describe the bug

Below is a segment of the prediction code. My description of the bug and my questions are written in the comments of the code below.

        predictions = []
        remaining = prediction_length

        while remaining > 0:
            token_ids, attention_mask, scale,mean,std = self.tokenizer.context_input_transform(
                context_tensor
            )
            samples = self.model(
                token_ids.to(self.model.device),
                attention_mask.to(self.model.device),
                min(remaining, self.model.config.prediction_length),
                num_samples,
                temperature,
                top_k,
                top_p,
            )
# After we execute the first round of the model's prediction code in the loop, 
# the shape of the |samples| variable is (batchsize, num_samples, pred_len), e.g. (32, 20 ,25).
            prediction = self.tokenizer.output_transform(
                samples.to(token_ids.device), scale
            )

            predictions.append(prediction)
            remaining -= prediction.shape[-1]

            if remaining <= 0:
                break
# So the entire loop will only execute one iteration, and the code will definitely exit the loop at this point. 
# Since the second iteration will not be executed, how will the effect of the following line of code be reflected?
            context_tensor = torch.cat(
                [context_tensor, prediction.median(dim=1).values], dim=-1
            )
        return torch.cat(predictions, dim=-1)
# So, is the autoregression during the prediction process in Chronos completed by calling the T5 model, 
# rather than by this segment of code?

Expected behavior

The code can run without any issues, but I believe the current execution of the program is not what the author originally intended.

Expected reply If my understanding is incorrect, please correct me. If there is indeed an issue with the code, I hope the author can fix this bug. Thank you very much! 😊

[abdulfatir]

Hi @tychenn! The autoregressive generation works as follows:

• If your prediction length is <= 64 steps, the while loop runs only once.
• In case your prediction length is > 64 steps, the loop generates first 64 predictions and then concatenates the median prediction as "ground truth" to the context tensor and feeds it into the model again to generate the next 64 steps until the prediction length is reached. We do this instead of directly unrolling the model beyond 64 steps because the model was trained with a prediction length of 64 and direct unrolling may lead to poor performance.

[tychenn]

Hi @abdulfatir ! Thank you for your patient responses. I have carefully reviewed the overall process of the code again, and now I fully understand its operating mechanism. Thank you very much!