← Walkthroughsrc/speech_decoding/models/v14_converged_v3/masking.py
Selected training excerpt · lines 658–678 · source SHA-256 c61bd3660017
658 total = n_units * sum(lengths)659 # Per-cell unit id, in the SAME flat layout the band split below undoes: each band segment is660 # (n_units, length) row-major = contact-major, so a contact's `length` cells are contiguous.661 unit_of_cell = torch.cat([unit_of_contact.repeat_interleave(length) for length in lengths])662 n_unit = int(unit_of_contact.max()) + 1663 cells = torch.bincount(unit_of_cell, minlength=n_unit) # (U,) cells per attention unit664 cnt = torch.round(frac * cells.float()).long() # (U,) exact masked count per unit665 start = torch.cat([torch.zeros(1, dtype=torch.long, device=device), cells.cumsum(0)[:-1]])666 # Offsetting the key by the unit id makes the sort group by unit and order by key WITHIN unit,667 # so unit u occupies sorted positions [start[u], start[u]+cells[u]) and the global rank minus668 # start[u] IS the rank within the unit. One argsort pair, no python loop over units.669 key = torch.rand(n_rows, total, generator=generator, device=device)670 key = key + unit_of_cell[None, :].to(key.dtype) * 2.0671 rank = key.argsort(-1).argsort(-1) - start[unit_of_cell][None, :]672 flat = rank < cnt[unit_of_cell][None, :] # (R, total) exactly cnt[u] per (row, unit)673 out, off = [], 0674 for length in lengths:675 seg = n_units * length676 out.append(flat[:, off:off + seg].reshape(n_rows, n_units, length))677 off += seg678 return out[0], out[1], out[2]