src/speech_decoding/models/v14_converged_v3/shaft_batch.py

Selected training excerpt · lines 127–133 · source SHA-256 b06c18319c16

127        self._subj_w = pool.subject_n_shafts.to(torch.float64) ** self.alpha128129    def draw(self, generator: torch.Generator) -> ShaftRef:130        u = int(torch.multinomial(self._subj_w, 1, generator=generator).item())131        idxs = self.pool.subject_ref_idxs[u]132        j = int(torch.randint(len(idxs), (1,), generator=generator).item())133        return self.pool.refs[idxs[j]]