(self, cost, thing)
| 202 | |
| 203 | @torch.compile |
| 204 | def update(self, cost, thing): |
| 205 | state_err = (self.recons_state - |
| 206 | thing.unsqueeze(-1)).square().sum(dim=0) |
| 207 | cand_cost = torch.gather( |
| 208 | cost.unsqueeze(-2).expand(-1, self.state_cand.shape[1], -1), -1, |
| 209 | self.state_cand.expand(len(cost), -1, 2**(self.K * self.V))) |
| 210 | best = torch.min(cand_cost, dim=-1) |
| 211 | cost = state_err + best.values.unsqueeze(-1).expand( |
| 212 | -1, -1, 2**(self.K * self.V)).reshape(state_err.shape) |
| 213 | prev_state = torch.gather( |
| 214 | self.state_cand.expand(thing.shape[1], -1, -1), -1, |
| 215 | best.indices.unsqueeze(-1))[..., 0] |
| 216 | return prev_state, cost |
| 217 | |
| 218 | def viterbi(self, X, overlap=None): |
| 219 | T, B = X.shape |
no outgoing calls
no test coverage detected