Four defects in partially_AR_model.py surfaced by CodeRabbit on the line-ending PR

While #6668 normalised the line endings of espnet2/asr/partially_AR_model.py, CodeRabbit reviewed the (unchanged) code and reported four findings. #6668 is line endings only, so they are recorded here for a separate fix. The code has been this way since the file was added; git diff --ignore-space-at-eol on #6668 is empty for this file.

The one marked critical is a wrong loop variable (i instead of m in _get_mask_idx), which would make every mask in a batch resolve to the position of the first one; that needs to be verified against a real partially-AR decode before it is fixed, since test_partially_AR.py passes as things stand.

espnet2/asr/partially_AR_model.py line 150

_ Functional Correctness_ | _ Major_ | ⚡ Quick win

max_mask_parallel is overwritten on the first call and keeps a tensor value.

mask_num is a 0-dim tensor from line 143. Line 150 stores mask_num + 1 into the instance attribute, so the sentinel -1 is lost after the first utterance. Every later utterance reuses the first utterance's mask count as the batch cap. An utterance with more masks is then split into batches without the user requesting it, which makes decoding results depend on input order.

Use a local variable and an int.

espnet2/asr/partially_AR_model.py line 156

_ Stability & Availability_ | _ Major_ | ⚡ Quick win

The batch loop runs one empty iteration when mask_num divides evenly.

range((mask_num // self.max_mask_parallel) + 1) always adds a trailing iteration. If mask_num is a multiple of max_mask_parallel, that iteration gets bs_iter == mask_num and max_iter == 0. No mask is registered, and line 175 calls the beam search with an empty masks list. PartiallyARBeamSearch.init_hyp then fails on assert len(self.masks) > 0, "add_mask must be called before init_hyp". Example: max_mask_parallel=2 with 4 masks.

Use ceiling division.

espnet2/asr/partially_AR_model.py line 160

_ Functional Correctness_ | _ Critical_ | ⚡ Quick win

Wrong loop variable: every mask in a batch resolves to the same position.

The inner loop iterates m, but line 160 passes the outer batch index i to _get_mask_idx. Ruff also reports m as unused. With the default max_mask_parallel=-1 the outer loop runs a single iteration with i=0, so all registered masks use the position of the first mask token. prev_tokens, next_token, and the primer are then identical for every mask, and forward fills all masked slots from the same context.

Pass the per-mask index instead.

espnet2/asr/partially_AR_model.py line 172

_ Functional Correctness_ | _ Major_ | ⚡ Quick win

Pass a scalar EOS value on both branches.

When mask_idx is last, [self.eos] makes torch.LongTensor([m[1]]) a (1, 1) tensor. init_hyp assigns it to the scalar eoses[i], which can fail because in-place assignment cannot expand the target.

                 next_token = (
-                    yseq_with_mask[0, mask_idx + 1].tolist()
+                    int(yseq_with_mask[0, mask_idx + 1])
                     if mask_idx < len(yseq_with_mask[0]) - 1
-                    else [self.eos]
+                    else self.eos
                 )

Source: Linters/SAST tools

Generated with Claude Code