(
self,
embed_dim: int,
*,
num_heads: int,
base: Optional[float] = 100.0,
min_period: Optional[float] = None,
max_period: Optional[float] = None,
normalize_coords: Literal["min", "max", "separate"] = "separate",
shift_coords: Optional[float] = None,
jitter_coords: Optional[float] = None,
rescale_coords: Optional[float] = None,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
)
| 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 |
nothing calls this directly
no test coverage detected