From 36ac6a639f35ad6386c8a425806e27eab8986c84 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20K=C3=A4mmerer?= Date: Sun, 23 Aug 2026 01:30:01 +0200 Subject: [PATCH] Return a most-likely path from forced_align when predecessors tie The per-cell choice in forced_align_impl is if (x2 > x1 && x2 > x0) { result = x2; backPtr = 2; } else if (x1 > x0 && x1 > x2) { result = x1; backPtr = 1; } else { result = x0; backPtr = 0; } When x1 == x2 and both exceed x0, neither condition holds and the else branch takes x0 -- the strictly worst of the three. The returned path is then not a most-likely one, which is the entire contract of the function. Reaching the second branch already implies that x2 is not strictly greater than both others, so `x1 > x0` alone is the correct and complete condition there; `&& x1 > x2` only diverts exact ties into the wrong branch. Both kernels carry the same chain, so both are corrected. Exhaustively verified on 3243 randomly generated small cases (T <= 5, vocabulary <= 4, integer log-probabilities so ties are exact) by comparing against the best score over all enumerated valid paths: 92 sub-optimal before, 0 after, worst deficit 3.0 log-probability. The existing forced_align tests are unaffected, as are 43 recorded alignments from real wav2vec2 emissions -- on non-degenerate emissions exact ties are rare, which is presumably why this survived. The added test is asserted on the score rather than on the path: whenever this bug can fire, x1 == x2 means two distinct paths reach that cell with the same score, so the optimum is never unique and the choice among optimal paths is a convention the test should not pin. --- .../forced_align/cpu/compute.cpp | 2 +- src/libtorchaudio/forced_align/gpu/compute.cu | 2 +- .../functional/functional_impl.py | 29 +++++++++++++++++++ 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/src/libtorchaudio/forced_align/cpu/compute.cpp b/src/libtorchaudio/forced_align/cpu/compute.cpp index c14c919e19..332b284bc9 100644 --- a/src/libtorchaudio/forced_align/cpu/compute.cpp +++ b/src/libtorchaudio/forced_align/cpu/compute.cpp @@ -113,7 +113,7 @@ void forced_align_impl( if (x2 > x1 && x2 > x0) { result = x2; backPtr_a[t * S + i] = 2; // backPtr_a[t][i] = 2 - } else if (x1 > x0 && x1 > x2) { + } else if (x1 > x0) { result = x1; backPtr_a[t * S + i] = 1; // backPtr_a[t][i] = 1 } else { diff --git a/src/libtorchaudio/forced_align/gpu/compute.cu b/src/libtorchaudio/forced_align/gpu/compute.cu index 444a4f8f6d..b92f6a0814 100644 --- a/src/libtorchaudio/forced_align/gpu/compute.cu +++ b/src/libtorchaudio/forced_align/gpu/compute.cu @@ -90,7 +90,7 @@ __global__ void falign_cuda_step_kernel( if (x2 > x1 && x2 > x0) { result = x2; backPtrBuffer_a[backPtrBufferLen][i] = 2; - } else if (x1 > x0 && x1 > x2) { + } else if (x1 > x0) { result = x1; backPtrBuffer_a[backPtrBufferLen][i] = 1; } else { diff --git a/test/torchaudio_unittest/functional/functional_impl.py b/test/torchaudio_unittest/functional/functional_impl.py index e2e42f1fe7..8ea49023a4 100644 --- a/test/torchaudio_unittest/functional/functional_impl.py +++ b/test/torchaudio_unittest/functional/functional_impl.py @@ -1168,6 +1168,35 @@ def test_forced_align(self, targets, ref_path, targets_dtype): self.assertEqual(hyp_path, ref_path) self.assertEqual(hyp_scores, ref_scores) + @parameterized.expand([(torch.int32,), (torch.int64,)]) + def test_forced_align_optimal_on_ties(self, targets_dtype): + """The alignment returned must be a most-likely path, also when two + candidate predecessors of a cell score exactly the same. + + The search space here is only 3**3 paths, five of which spell the + target; enumerating them gives a best score of -1.0, reached by + [1, 0, 2] and by [0, 1, 2]. Which of the two is returned is a + tie-breaking convention and is deliberately not asserted -- the score + is the invariant. Before this case was fixed the kernel returned + [1, 2, 2], scoring -3.0. + """ + log_probs = torch.tensor( + [[[0.0, -1.0, -1.0], [0.0, -1.0, -2.0], [-2.0, 0.0, 0.0]]], + dtype=self.dtype, + device=self.device, + ) + targets = torch.tensor([[1, 2]], dtype=targets_dtype, device=self.device) + blank = 0 + input_lengths = torch.tensor([log_probs.shape[1]], device=self.device) + target_lengths = torch.tensor([targets.shape[1]], device=self.device) + hyp_path, hyp_scores = F.forced_align(log_probs, targets, input_lengths, target_lengths, blank) + # the path spells the target once repeats and blanks are collapsed ... + path = hyp_path[0].tolist() + collapsed = [t for i, t in enumerate(path) if t != blank and (i == 0 or t != path[i - 1])] + self.assertEqual(collapsed, targets[0].tolist()) + # ... and it is a most likely one + self.assertEqual(hyp_scores.sum().item(), -1.0) + @parameterized.expand([(torch.int32,), (torch.int64,)]) def test_forced_align_fail(self, targets_dtype): log_probs = torch.rand(1, 5, 6, dtype=self.dtype, device=self.device)