src/speech_decoding/models/v14_converged_v3/pack_r4.py

Selected training excerpt · lines 201–208 · source SHA-256 10e2eebf5adf

201    B, total = masked.shape202    device = masked.device203    visible = ~masked204    if m_vis is None:205        m_vis = int(visible.sum(1).max().item())206    ar = torch.arange(total, device=device)207    key = torch.where(visible, ar[None, :].expand(B, total), total)208    idx = key.sort(dim=1).values[:, :m_vis]  # (B, M_vis) visible columns, grid order