| 53 | |
| 54 | |
| 55 | class FreqEncoder(nn.Module): |
| 56 | def __init__(self, input_dim=3, degree=4): |
| 57 | super().__init__() |
| 58 | |
| 59 | self.input_dim = input_dim |
| 60 | self.degree = degree |
| 61 | self.output_dim = input_dim + input_dim * 2 * degree |
| 62 | |
| 63 | def __repr__(self): |
| 64 | return f"FreqEncoder: input_dim={self.input_dim} degree={self.degree} output_dim={self.output_dim}" |
| 65 | |
| 66 | def forward(self, inputs, **kwargs): |
| 67 | # inputs: [..., input_dim] |
| 68 | # return: [..., ] |
| 69 | |
| 70 | prefix_shape = list(inputs.shape[:-1]) |
| 71 | inputs = inputs.reshape(-1, self.input_dim) |
| 72 | |
| 73 | outputs = freq_encode(inputs, self.degree, self.output_dim) |
| 74 | |
| 75 | outputs = outputs.reshape(prefix_shape + [self.output_dim]) |
| 76 | |
| 77 | return outputs |