src/speech_decoding/models/v14_converged_v3/towers.py

Selected training excerpt · lines 326–341 · source SHA-256 5610fc7e8fc4

326        xf = x_flat.reshape(B * M, d)327        taps: dict[int, Tensor] = {}328        levels: list[Tensor] = []329        for i, blk in enumerate(self.blocks):330            # r4 is L1-only (Design B); an L2 in the layout has no flat path (guard).331            if not isinstance(blk, L1Block):332                raise TypeError("forward_flat is L1-only (Design B); found a non-L1 block")333            xf = blk.forward_flat(334                xf, depth_b, time_b, cu_b, cu_drop_b, max_seqlen, rope_cs=rope_cs,335            )336            if (i + 1) in tap_blocks:337                taps[i + 1] = xf.reshape(B, M, d)338            if self.deep_sup and (i + 1) in self.sup_taps:339                levels.append(self.norms_block[self.sup_taps.index(i + 1)](xf))340        out = torch.cat(levels, dim=-1) if self.deep_sup else self.norm_out(xf)341        out = out.reshape(B, M, out.shape[-1])