| """Continuous positional embedding helpers for voxel and query tokens.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| class ContinuousFourierPositionEmbed(nn.Module): |
| """Encode continuous coordinates into a learned token-space positional embedding.""" |
|
|
| def __init__( |
| self, |
| *, |
| input_dim: int, |
| embed_dim: int, |
| num_bands: int = 6, |
| hidden_dim: int | None = None, |
| include_input: bool = True, |
| ) -> None: |
| super().__init__() |
| self.input_dim = int(input_dim) |
| self.num_bands = int(num_bands) |
| self.include_input = bool(include_input) |
| hidden_dim = int(hidden_dim) if hidden_dim is not None else max(int(embed_dim), 64) |
| self.register_buffer( |
| "freq_bands", |
| 2.0 ** torch.arange(self.num_bands, dtype=torch.float32) * math.pi, |
| persistent=False, |
| ) |
| encoded_dim = (self.input_dim if self.include_input else 0) + self.input_dim * self.num_bands * 2 |
| self.net = nn.Sequential( |
| nn.Linear(encoded_dim, hidden_dim), |
| nn.GELU(), |
| nn.Linear(hidden_dim, embed_dim), |
| ) |
|
|
| def _encode(self, x: torch.Tensor) -> torch.Tensor: |
| pieces = [x] if self.include_input else [] |
| angles = x.unsqueeze(-1) * self.freq_bands |
| pieces.append(torch.sin(angles).flatten(-2)) |
| pieces.append(torch.cos(angles).flatten(-2)) |
| return torch.cat(pieces, dim=-1) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.net(self._encode(x)) |
|
|
|
|
| __all__ = ["ContinuousFourierPositionEmbed"] |
|
|
