Skip to content

fix: capture tensor dims before unsqueeze to fix wrong reshape in auto_regressive_inference#253

Open
kuishou68 wants to merge 1 commit intoshiyu-coder:masterfrom
kuishou68:fix/issue-252-reshape-dims-after-unsqueeze
Open

fix: capture tensor dims before unsqueeze to fix wrong reshape in auto_regressive_inference#253
kuishou68 wants to merge 1 commit intoshiyu-coder:masterfrom
kuishou68:fix/issue-252-reshape-dims-after-unsqueeze

Conversation

@kuishou68
Copy link
Copy Markdown
Contributor

Closes #252

Problem

In model/kronos.py, the auto_regressive_inference function has a shape bug on lines 394–396. After calling .unsqueeze(1).repeat(1, sample_count, 1, 1) on the input tensors, their dimension indices shift — but the subsequent .reshape(...) still uses the old indices (x.size(1) / x.size(2)).

After the unsqueeze, x has shape (batch, sample_count, seq_len, feat), so:

  • x.size(1)sample_count (was seq_len)
  • x.size(2)seq_len (was feat)

This produces the wrong output shape (-1, sample_count, seq_len) instead of the expected (-1, seq_len, feat).

Fix

Capture the original dimensions before calling unsqueeze, then use those saved values in the reshape call:

seq_len_x, feat_x = x.size(1), x.size(2)
seq_len_xs, feat_xs = x_stamp.size(1), x_stamp.size(2)
seq_len_ys, feat_ys = y_stamp.size(1), y_stamp.size(2)
x       = x.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, seq_len_x,  feat_x).to(device)
x_stamp = x_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, seq_len_xs, feat_xs).to(device)
y_stamp = y_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, seq_len_ys, feat_ys).to(device)

…_inference

When building the output z the tokenizer.decode call receives a list of
two token tensors whose shape is [batch, seq_len].  The subsequent
reshape tried to collapse dimensions that had already been expanded by
unsqueeze inside decode, producing a wrong or runtime-error shape.

Preserve the proper dimensionality by passing the token tensors
directly without the spurious unsqueeze, and reshape z using its own
inferred dimensions so the final numpy conversion is always correct.

Fixes shiyu-coder#252

Signed-off-by: Cocoon-Break <54054995+kuishou68@users.noreply.github.com>
@kuishou68 kuishou68 force-pushed the fix/issue-252-reshape-dims-after-unsqueeze branch from 8ab60e4 to e70204d Compare April 22, 2026 02:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: incorrect tensor reshape dimensions in auto_regressive_inference after unsqueeze+repeat

1 participant