RoPE positional embedding with no mixing of coordinates (axial) and no learnable weights. Supports two parametrizations of the rope parameters: either using `base` or `min_period` and `max_period`.
| 84 | |
| 85 | |
| 86 | class RopePositionEmbedding(nn.Module): |
| 87 | """RoPE positional embedding with no mixing of coordinates (axial) and no learnable weights. |
| 88 | |
| 89 | Supports two parametrizations of the rope parameters: either using `base` or `min_period` and `max_period`. |
| 90 | """ |
| 91 | |
| 92 | def __init__( |
| 93 | self, |
| 94 | embed_dim: int, |
| 95 | *, |
| 96 | num_heads: int, |
| 97 | base: Optional[float] = 100.0, |
| 98 | min_period: Optional[float] = None, |
| 99 | max_period: Optional[float] = None, |
| 100 | normalize_coords: Literal["min", "max", "separate"] = "separate", |
| 101 | shift_coords: Optional[float] = None, |
| 102 | jitter_coords: Optional[float] = None, |
| 103 | rescale_coords: Optional[float] = None, |
| 104 | dtype: Optional[torch.dtype] = None, |
| 105 | device: Optional[torch.device] = None, |
| 106 | ): |
| 107 | super().__init__() |
| 108 | assert embed_dim % (4 * num_heads) == 0 |
| 109 | both_periods = min_period is not None and max_period is not None |
| 110 | if (base is None and not both_periods) or (base is not None and both_periods): |
| 111 | raise ValueError("Either `base` or `min_period`+`max_period` must be provided.") |
| 112 | |
| 113 | D_head = embed_dim // num_heads |
| 114 | self.base = base |
| 115 | self.min_period = min_period |
| 116 | self.max_period = max_period |
| 117 | self.D_head = D_head |
| 118 | self.normalize_coords = normalize_coords |
| 119 | self.shift_coords = shift_coords |
| 120 | self.jitter_coords = jitter_coords |
| 121 | self.rescale_coords = rescale_coords |
| 122 | |
| 123 | self.dtype = dtype |
| 124 | self.register_buffer( |
| 125 | "periods", |
| 126 | torch.empty(D_head // 4, device=device, dtype=dtype), |
| 127 | persistent=True, |
| 128 | ) |
| 129 | self._init_weights() |
| 130 | |
| 131 | def forward(self, *, H: int, W: int) -> tuple[Tensor, Tensor]: |
| 132 | device = self.periods.device |
| 133 | dtype = self.dtype |
| 134 | dd = {"device": device, "dtype": dtype} |
| 135 | |
| 136 | # Prepare coords in range [-1, +1] |
| 137 | if self.normalize_coords == "max": |
| 138 | max_HW = max(H, W) |
| 139 | coords_h = torch.arange(0.5, H, **dd) / max_HW |
| 140 | coords_w = torch.arange(0.5, W, **dd) / max_HW |
| 141 | elif self.normalize_coords == "min": |
| 142 | min_HW = min(H, W) |
| 143 | coords_h = torch.arange(0.5, H, **dd) / min_HW |