(n: int, div: int)
| 204 | replace_mask = new_max > gumbel_max |
| 205 | gumbel_max = torch.where(replace_mask, new_max, gumbel_max) |
| 206 | gumbel_max_idx = torch.where(replace_mask, new_max_idx_global, gumbel_max_idx) |
| 207 | return gumbel_max_idx.T |
| 208 | |
| 209 | |
| 210 | def cdiv(n: int, div: int) -> int: |
no outgoing calls
no test coverage detected