| 29 | |
| 30 | @staticmethod |
| 31 | def scaled_l2(x, codewords, scale): |
| 32 | num_codes, channels = codewords.size() |
| 33 | batch_size = x.size(0) |
| 34 | reshaped_scale = scale.view((1, 1, num_codes)) |
| 35 | expanded_x = x.unsqueeze(2).expand((batch_size, x.size(1), num_codes, channels)) |
| 36 | reshaped_codewords = codewords.view((1, 1, num_codes, channels)) |
| 37 | |
| 38 | scaled_l2_norm = reshaped_scale * (expanded_x - reshaped_codewords).pow(2).sum(dim=3) |
| 39 | return scaled_l2_norm |
| 40 | |
| 41 | @staticmethod |
| 42 | def aggregate(assignment_weights, x, codewords): |