src/speech_decoding/models/v14_converged_v3/objective.py

Selected training excerpt · lines 850–867 · source SHA-256 5718458b537d

850        B, N = bands[0].shape[0], bands[0].shape[1]851        T = bands[0].shape[-1]852        f_max = max(nb for nb, _ in PER_BAND_SPECS)853        padded: list[Tensor] = []854        for b, (nb, _) in enumerate(PER_BAND_SPECS):855            xb = bands[b].transpose(-1, -2)  # (B, N, T, F_b)856            if nb < f_max:857                xb = F.pad(xb, (0, f_max - nb))  # zero-pad freq to F_MAX858            padded.append(xb)859        stack = torch.stack(padded, dim=1)  # (B, N_BANDS, N, T, F_MAX)860        flat = stack.reshape(B, len(PER_BAND_SPECS) * N * T, f_max)861        lin = (grid.band * N + grid.contact) * T + grid.time_pos  # (total,)862        target = flat[:, lin, :]  # (B, total, F_MAX)863        device = grid.band.device864        f_by_band = torch.tensor([nb for nb, _ in PER_BAND_SPECS], device=device)865        feat_count = f_by_band[grid.band]  # (total,)866        feat_valid = torch.arange(f_max, device=device)[None, :] < feat_count[:, None]867        return target, feat_valid, feat_count