| 1 | # MIT License |
| 2 | # |
| 3 | # Copyright (c) Authors of |
| 4 | # "Cameras as Relative Positional Encoding" https://arxiv.org/pdf/2507.10496 |
| 5 | # |
| 6 | # Permission is hereby granted, free of charge, to any person obtaining a copy |
| 7 | # of this software and associated documentation files (the "Software"), to deal |
| 8 | # in the Software without restriction, including without limitation the rights |
| 9 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 10 | # copies of the Software, and to permit persons to whom the Software is |
| 11 | # furnished to do so, subject to the following conditions: |
| 12 | # |
| 13 | # The above copyright notice and this permission notice shall be included in all |
| 14 | # copies or substantial portions of the Software. |
| 15 | # |
| 16 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 17 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 18 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 19 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 20 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 21 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
| 22 | # SOFTWARE. |
| 23 | |
| 24 | # How to use PRoPE attention for self-attention: |
| 25 | # |
| 26 | # 1. Easiest way (fast): |
| 27 | # attn = PropeDotProductAttention(...) |
| 28 | # o = attn(q, k, v, viewmats, Ks) |
| 29 | # |
| 30 | # 2. More flexible way (fast): |
| 31 | # attn = PropeDotProductAttention(...) |
| 32 | # attn._precompute_and_cache_apply_fns(viewmats, Ks) |
| 33 | # q = attn._apply_to_q(q) |
| 34 | # k = attn._apply_to_kv(k) |
| 35 | # v = attn._apply_to_kv(v) |
| 36 | # o = F.scaled_dot_product_attention(q, k, v, **kwargs) |
| 37 | # o = attn._apply_to_o(o) |
| 38 | # |
| 39 | # 3. The most flexible way (but slower because repeated computation of RoPE coefficients): |
| 40 | # o = prope_dot_product_attention(q, k, v, ...) |
| 41 | # |
| 42 | # How to use PRoPE attention for cross-attention: |
| 43 | # |
| 44 | # attn_src = PropeDotProductAttention(...) |
| 45 | # attn_tgt = PropeDotProductAttention(...) |
| 46 | # attn_src._precompute_and_cache_apply_fns(viewmats_src, Ks_src) |
| 47 | # attn_tgt._precompute_and_cache_apply_fns(viewmats_tgt, Ks_tgt) |
| 48 | # q_src = attn_src._apply_to_q(q_src) |
| 49 | # k_tgt = attn_tgt._apply_to_kv(k_tgt) |
| 50 | # v_tgt = attn_tgt._apply_to_kv(v_tgt) |
| 51 | # o_src = F.scaled_dot_product_attention(q_src, k_tgt, v_tgt, **kwargs) |
| 52 | # o_src = attn_src._apply_to_o(o_src) |
| 53 | |
| 54 | from functools import partial |
| 55 | from typing import Callable, Optional, Tuple, List |
| 56 | |
| 57 | import torch |
| 58 | import torch.nn.functional as F |
| 59 | |
| 60 | |
| 61 | class PropeDotProductAttention(torch.nn.Module): |
| 62 | """PRoPE attention with precomputed RoPE coefficients.""" |
| 63 | |
| 64 | coeffs_x_0: torch.Tensor |
| 65 | coeffs_x_1: torch.Tensor |
| 66 | coeffs_y_0: torch.Tensor |
| 67 | coeffs_y_1: torch.Tensor |
| 68 | |
| 69 | def __init__( |
| 70 | self, |
| 71 | head_dim: int, |
| 72 | patches_x: int, |
| 73 | patches_y: int, |
| 74 | image_width: int, |
| 75 | image_height: int, |
| 76 | freq_base: float = 100.0, |
| 77 | freq_scale: float = 1.0, |
| 78 | precompute_coeffs: bool = True, |
| 79 | ): |
| 80 | super().__init__() |
| 81 | self.head_dim = head_dim |
| 82 | self.patches_x = patches_x |
| 83 | self.patches_y = patches_y |
| 84 | self.image_width = image_width |
| 85 | self.image_height = image_height |
| 86 | |
| 87 | if precompute_coeffs: |
| 88 | coeffs_x: Tuple[torch.Tensor, torch.Tensor] = _rope_precompute_coeffs( |
| 89 | torch.tile(torch.arange(patches_x), (patches_y,)), |
| 90 | freq_base=freq_base, |
| 91 | freq_scale=freq_scale, |
| 92 | feat_dim=head_dim // 4, |
| 93 | ) |
| 94 | coeffs_y: Tuple[torch.Tensor, torch.Tensor] = _rope_precompute_coeffs( |
| 95 | torch.repeat_interleave(torch.arange(patches_y), patches_x), |
| 96 | freq_base=freq_base, |
| 97 | freq_scale=freq_scale, |
| 98 | feat_dim=head_dim // 4, |
| 99 | ) |
| 100 | # Model builders instantiate modules on the meta device before |
| 101 | # loading weights. Meta buffers cannot later be materialized by a |
| 102 | # state dict because these coefficients are intentionally not |
| 103 | # persisted, so generate them lazily on the first real forward. |
| 104 | if coeffs_x[0].is_meta: |
| 105 | coeffs_x = coeffs_y = (None, None) |
| 106 | self.register_buffer("coeffs_x_0", coeffs_x[0], persistent=False) |
| 107 | self.register_buffer("coeffs_x_1", coeffs_x[1], persistent=False) |
| 108 | self.register_buffer("coeffs_y_0", coeffs_y[0], persistent=False) |
| 109 | self.register_buffer("coeffs_y_1", coeffs_y[1], persistent=False) |
| 110 | else: |
| 111 | self.coeffs_x_0 = None |
| 112 | self.coeffs_x_1 = None |
| 113 | self.coeffs_y_0 = None |
| 114 | self.coeffs_y_1 = None |
| 115 | |
| 116 | # override load_state_dict to not load coeffs if they exist (for backward compatibility) |
| 117 | def load_state_dict(self, state_dict, strict=True): |
| 118 | # remove coeffs from state_dict |
| 119 | state_dict.pop("coeffs_x_0", None) |
| 120 | state_dict.pop("coeffs_x_1", None) |
| 121 | state_dict.pop("coeffs_y_0", None) |
| 122 | state_dict.pop("coeffs_y_1", None) |
| 123 | super().load_state_dict(state_dict, strict) |
| 124 | |
| 125 | def forward( |
| 126 | self, |
| 127 | q: torch.Tensor, # (batch, num_heads, seqlen, head_dim) |
| 128 | k: torch.Tensor, # (batch, num_heads, seqlen, head_dim) |
| 129 | v: torch.Tensor, # (batch, num_heads, seqlen, head_dim) |
| 130 | viewmats: torch.Tensor, # (batch, cameras, 4, 4) |
| 131 | Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) |
| 132 | **kwargs, |
| 133 | ) -> torch.Tensor: |
| 134 | return prope_dot_product_attention( |
| 135 | q, |
| 136 | k, |
| 137 | v, |
| 138 | viewmats=viewmats, |
| 139 | Ks=Ks, |
| 140 | patches_x=self.patches_x, |
| 141 | patches_y=self.patches_y, |
| 142 | image_width=self.image_width, |
| 143 | image_height=self.image_height, |
| 144 | coeffs_x=(self.coeffs_x_0, self.coeffs_x_1), |
| 145 | coeffs_y=(self.coeffs_y_0, self.coeffs_y_1), |
| 146 | **kwargs, |
| 147 | ) |
| 148 | |
| 149 | def _precompute_and_cache_apply_fns( |
| 150 | self, viewmats: torch.Tensor, Ks: Optional[torch.Tensor], |
| 151 | coeffs_x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| 152 | coeffs_y: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| 153 | ): |
| 154 | (batch, cameras, _, _) = viewmats.shape |
| 155 | assert viewmats.shape == (batch, cameras, 4, 4) |
| 156 | assert Ks is None or Ks.shape == (batch, cameras, 3, 3) |
| 157 | # self.cameras = cameras |
| 158 | |
| 159 | cached_x = None if self.coeffs_x_0 is None else (self.coeffs_x_0, self.coeffs_x_1) |
| 160 | cached_y = None if self.coeffs_y_0 is None else (self.coeffs_y_0, self.coeffs_y_1) |
| 161 | self.apply_fn_q, self.apply_fn_kv, self.apply_fn_o = _prepare_apply_fns( |
| 162 | head_dim=self.head_dim, |
| 163 | viewmats=viewmats, |
| 164 | Ks=Ks, |
| 165 | patches_x=self.patches_x, |
| 166 | patches_y=self.patches_y, |
| 167 | image_width=self.image_width, |
| 168 | image_height=self.image_height, |
| 169 | coeffs_x=cached_x if coeffs_x is None else coeffs_x, |
| 170 | coeffs_y=cached_y if coeffs_y is None else coeffs_y, |
| 171 | ) |
| 172 | |
| 173 | def _apply_to_q(self, q: torch.Tensor) -> torch.Tensor: |
| 174 | (batch, num_heads, seqlen, head_dim) = q.shape |
| 175 | # assert seqlen == self.cameras * self.patches_x * self.patches_y |
| 176 | assert head_dim == self.head_dim |
| 177 | assert q.shape == (batch, num_heads, seqlen, head_dim) |
| 178 | assert self.apply_fn_q is not None |
| 179 | return self.apply_fn_q(q) |
| 180 | |
| 181 | def _apply_to_kv(self, kv: torch.Tensor) -> torch.Tensor: |
| 182 | (batch, num_heads, seqlen, head_dim) = kv.shape |
| 183 | # assert seqlen == self.cameras * self.patches_x * self.patches_y |
| 184 | assert head_dim == self.head_dim |
| 185 | assert kv.shape == (batch, num_heads, seqlen, head_dim) |
| 186 | assert self.apply_fn_kv is not None |
| 187 | return self.apply_fn_kv(kv) |
| 188 | |
| 189 | def _apply_to_o(self, o: torch.Tensor) -> torch.Tensor: |
| 190 | (batch, num_heads, seqlen, head_dim) = o.shape |
| 191 | # assert seqlen == self.cameras * self.patches_x * self.patches_y |
| 192 | assert head_dim == self.head_dim |
| 193 | assert o.shape == (batch, num_heads, seqlen, head_dim) |
| 194 | assert self.apply_fn_o is not None |
| 195 | return self.apply_fn_o(o) |
| 196 | |
| 197 | |
| 198 | def prope_dot_product_attention( |
| 199 | q: torch.Tensor, # (batch, num_heads, seqlen, head_dim) |
| 200 | k: torch.Tensor, # (batch, num_heads, seqlen, head_dim) |
| 201 | v: torch.Tensor, # (batch, num_heads, seqlen, head_dim) |
| 202 | *, |
| 203 | viewmats: torch.Tensor, # (batch, cameras, 4, 4) |
| 204 | Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) |
| 205 | patches_x: int, # How many patches wide is each image? |
| 206 | patches_y: int, # How many patches tall is each image? |
| 207 | image_width: int, # Width of the image. Used to normalize intrinsics. |
| 208 | image_height: int, # Height of the image. Used to normalize intrinsics. |
| 209 | coeffs_x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| 210 | coeffs_y: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| 211 | **kwargs, |
| 212 | ) -> torch.Tensor: |
| 213 | """Similar to torch.nn.functional.scaled_dot_product_attention, but applies PRoPE-style |
| 214 | positional encoding. |
| 215 | |
| 216 | Currently, we assume that the sequence length is equal to: |
| 217 | |
| 218 | cameras * patches_x * patches_y |
| 219 | |
| 220 | And token ordering allows the `(seqlen,)` axis to be reshaped into |
| 221 | `(cameras, patches_x, patches_y)`. |
| 222 | """ |
| 223 | # We're going to assume self-attention: all inputs are the same shape. |
| 224 | (batch, num_heads, seqlen, head_dim) = q.shape |
| 225 | cameras = viewmats.shape[1] |
| 226 | assert q.shape == k.shape == v.shape |
| 227 | assert viewmats.shape == (batch, cameras, 4, 4) |
| 228 | assert Ks is None or Ks.shape == (batch, cameras, 3, 3) |
| 229 | assert seqlen == cameras * patches_x * patches_y |
| 230 | |
| 231 | apply_fn_q, apply_fn_kv, apply_fn_o = _prepare_apply_fns( |
| 232 | head_dim=head_dim, |
| 233 | viewmats=viewmats, |
| 234 | Ks=Ks, |
| 235 | patches_x=patches_x, |
| 236 | patches_y=patches_y, |
| 237 | image_width=image_width, |
| 238 | image_height=image_height, |
| 239 | coeffs_x=coeffs_x, |
| 240 | coeffs_y=coeffs_y, |
| 241 | ) |
| 242 | |
| 243 | out = F.scaled_dot_product_attention( |
| 244 | query=apply_fn_q(q), |
| 245 | key=apply_fn_kv(k), |
| 246 | value=apply_fn_kv(v), |
| 247 | **kwargs, |
| 248 | ) |
| 249 | out = apply_fn_o(out) |
| 250 | assert out.shape == (batch, num_heads, seqlen, head_dim) |
| 251 | return out |
| 252 | |
| 253 | |
| 254 | def _prepare_apply_fns( |
| 255 | head_dim: int, # Q/K/V will have this last dimension |
| 256 | viewmats: torch.Tensor, # (batch, cameras, 4, 4) |
| 257 | Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) |
| 258 | patches_x: int, # How many patches wide is each image? |
| 259 | patches_y: int, # How many patches tall is each image? |
| 260 | image_width: int, # Width of the image. Used to normalize intrinsics. |
| 261 | image_height: int, # Height of the image. Used to normalize intrinsics. |
| 262 | coeffs_x: Optional[torch.Tensor] = None, |
| 263 | coeffs_y: Optional[torch.Tensor] = None, |
| 264 | ) -> Tuple[ |
| 265 | Callable[[torch.Tensor], torch.Tensor], |
| 266 | Callable[[torch.Tensor], torch.Tensor], |
| 267 | Callable[[torch.Tensor], torch.Tensor], |
| 268 | ]: |
| 269 | """Prepare transforms for PRoPE-style positional encoding.""" |
| 270 | device = viewmats.device |
| 271 | (batch, cameras, _, _) = viewmats.shape |
| 272 | dtype = viewmats.dtype |
| 273 | |
| 274 | # Normalize camera intrinsics. |
| 275 | if Ks is not None: |
| 276 | Ks_norm = torch.zeros_like(Ks) |
| 277 | Ks_norm[..., 0, 0] = Ks[..., 0, 0] / image_width |
| 278 | Ks_norm[..., 1, 1] = Ks[..., 1, 1] / image_height |
| 279 | Ks_norm[..., 0, 2] = Ks[..., 0, 2] / image_width - 0.5 |
| 280 | Ks_norm[..., 1, 2] = Ks[..., 1, 2] / image_height - 0.5 |
| 281 | Ks_norm[..., 2, 2] = 1.0 |
| 282 | del Ks |
| 283 | |
| 284 | # Compute the camera projection matrices we use in PRoPE. |
| 285 | # - K is an `image<-camera` transform. |
| 286 | # - viewmats is a `camera<-world` transform. |
| 287 | # - P = lift(K) @ viewmats is an `image<-world` transform. |
| 288 | P = torch.einsum("...ij,...jk->...ik", _lift_K(Ks_norm), viewmats) |
| 289 | P_T = P.transpose(-1, -2) |
| 290 | P_inv = torch.einsum( |
| 291 | "...ij,...jk->...ik", |
| 292 | _invert_SE3(viewmats), |
| 293 | _lift_K(_invert_K(Ks_norm)), |
| 294 | ) |
| 295 | |
| 296 | else: |
| 297 | # GTA formula. P is `camera<-world` transform. |
| 298 | P = viewmats |
| 299 | P_T = P.transpose(-1, -2) |
| 300 | P_inv = _invert_SE3(viewmats) |
| 301 | |
| 302 | assert P.shape == P_inv.shape == (batch, cameras, 4, 4) |
| 303 | |
| 304 | # Precompute cos/sin terms for RoPE. We use tiles/repeats for 'row-major' |
| 305 | # broadcasting. |
| 306 | if coeffs_x is None: |
| 307 | coeffs_x = _rope_precompute_coeffs( |
| 308 | torch.tile(torch.arange(patches_x, device=device), (patches_y * cameras,)), |
| 309 | freq_base=100.0, |
| 310 | freq_scale=1.0, |
| 311 | feat_dim=head_dim // 4, |
| 312 | dtype=dtype, |
| 313 | ) |
| 314 | if coeffs_y is None: |
| 315 | coeffs_y = _rope_precompute_coeffs( |
| 316 | torch.tile( |
| 317 | torch.repeat_interleave( |
| 318 | torch.arange(patches_y, device=device), patches_x |
| 319 | ), |
| 320 | (cameras,), |
| 321 | ), |
| 322 | freq_base=100.0, |
| 323 | freq_scale=1.0, |
| 324 | feat_dim=head_dim // 4, |
| 325 | dtype=dtype, |
| 326 | ) |
| 327 | |
| 328 | # Block-diagonal transforms to the inputs and outputs of the attention operator. |
| 329 | assert head_dim % 4 == 0 |
| 330 | transforms_q = [ |
| 331 | (partial(_apply_tiled_projmat, matrix=P_T), head_dim // 2), |
| 332 | (partial(_rope_apply_coeffs, coeffs=coeffs_x), head_dim // 4), |
| 333 | (partial(_rope_apply_coeffs, coeffs=coeffs_y), head_dim // 4), |
| 334 | ] |
| 335 | transforms_kv = [ |
| 336 | (partial(_apply_tiled_projmat, matrix=P_inv), head_dim // 2), |
| 337 | (partial(_rope_apply_coeffs, coeffs=coeffs_x), head_dim // 4), |
| 338 | (partial(_rope_apply_coeffs, coeffs=coeffs_y), head_dim // 4), |
| 339 | ] |
| 340 | transforms_o = [ |
| 341 | (partial(_apply_tiled_projmat, matrix=P), head_dim // 2), |
| 342 | (partial(_rope_apply_coeffs, coeffs=coeffs_x, inverse=True), head_dim // 4), |
| 343 | (partial(_rope_apply_coeffs, coeffs=coeffs_y, inverse=True), head_dim // 4), |
| 344 | ] |
| 345 | |
| 346 | apply_fn_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q) |
| 347 | apply_fn_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv) |
| 348 | apply_fn_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o) |
| 349 | return apply_fn_q, apply_fn_kv, apply_fn_o |
| 350 | |
| 351 | |
| 352 | def _apply_tiled_projmat( |
| 353 | feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) |
| 354 | matrix: torch.Tensor, # (batch, cameras, D, D) or (batch, seqlen, D, D) |
| 355 | ) -> torch.Tensor: |
| 356 | """Apply projection matrix to features.""" |
| 357 | # - seqlen => (cameras, patches_x * patches_y) |
| 358 | # - feat_dim => (feat_dim // 4, 4) |
| 359 | (batch, num_heads, seqlen, feat_dim) = feats.shape |
| 360 | matrix = matrix.to(device=feats.device, dtype=feats.dtype) |
| 361 | D = matrix.shape[-1] |
| 362 | assert feat_dim % D == 0, f"feat_dim={feat_dim} must be divisible by D={D}" |
| 363 | |
| 364 | if matrix.shape[1] == seqlen: |
| 365 | # Per-ray projection: matrix shape [B, seqlen, D, D] |
| 366 | feats_ = feats.view(batch, num_heads, seqlen, feat_dim // D, D) |
| 367 | out = torch.einsum("btij,bntpj->bntpi", matrix, feats_) |
| 368 | return out.reshape(feats.shape) |
| 369 | |
| 370 | # Per-camera projection (original implementation) |
| 371 | cameras = matrix.shape[1] |
| 372 | assert seqlen > cameras and seqlen % cameras == 0 |
| 373 | assert matrix.shape == (batch, cameras, D, D) |
| 374 | assert feat_dim % D == 0 |
| 375 | return torch.einsum( |
| 376 | "bcij,bncpkj->bncpki", |
| 377 | matrix, |
| 378 | feats.reshape((batch, num_heads, cameras, -1, feat_dim // D, D)), |
| 379 | ).reshape(feats.shape) |
| 380 | |
| 381 | |
| 382 | def _rope_precompute_coeffs( |
| 383 | positions: torch.Tensor, # (seqlen,) |
| 384 | freq_base: float, |
| 385 | freq_scale: float, |
| 386 | feat_dim: int, |
| 387 | dtype: torch.dtype = torch.float32, |
| 388 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 389 | """Precompute RoPE coefficients.""" |
| 390 | assert len(positions.shape) == 1 |
| 391 | assert feat_dim % 2 == 0 |
| 392 | num_freqs = feat_dim // 2 |
| 393 | freqs = freq_scale * ( |
| 394 | freq_base |
| 395 | ** ( |
| 396 | -torch.arange(num_freqs, device=positions.device)[None, None, None, :] |
| 397 | / num_freqs |
| 398 | ) |
| 399 | ) |
| 400 | angles = positions[None, None, :, None] * freqs |
| 401 | # Shape should be: `(batch, num_heads, seqlen, num_freqs)`; we're |
| 402 | # broadcasting across `batch` and `num_heads`. |
| 403 | assert angles.shape == (1, 1, positions.shape[0], num_freqs) |
| 404 | return torch.cos(angles).to(dtype), torch.sin(angles).to(dtype) |
| 405 | |
| 406 | |
| 407 | def _rope_apply_coeffs( |
| 408 | feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) |
| 409 | coeffs: Tuple[torch.Tensor, torch.Tensor], |
| 410 | inverse: bool = False, |
| 411 | ) -> torch.Tensor: |
| 412 | """Apply RoPE coefficients to features. We adopt a 'split' ordering |
| 413 | convention. (in contrast to 'interleaved')""" |
| 414 | cos, sin = coeffs |
| 415 | # We allow (cos, sin) to be either with shape (1, 1, seqlen, feat_dim // 2), |
| 416 | # or (1, 1, seqlen_per_image, feat_dim // 2) and we repeat it to |
| 417 | # match the shape of feats. |
| 418 | if cos.shape[2] != feats.shape[2]: |
| 419 | n_repeats = feats.shape[2] // cos.shape[2] |
| 420 | cos = cos.repeat(1, 1, n_repeats, 1) |
| 421 | sin = sin.repeat(1, 1, n_repeats, 1) |
| 422 | assert len(feats.shape) == len(cos.shape) == len(sin.shape) == 4 |
| 423 | assert cos.shape[-1] == sin.shape[-1] == feats.shape[-1] // 2 |
| 424 | x_in = feats[..., : feats.shape[-1] // 2] |
| 425 | y_in = feats[..., feats.shape[-1] // 2 :] |
| 426 | return torch.cat( |
| 427 | ( |
| 428 | [cos * x_in + sin * y_in, -sin * x_in + cos * y_in] |
| 429 | if not inverse |
| 430 | else [cos * x_in - sin * y_in, sin * x_in + cos * y_in] |
| 431 | ), |
| 432 | dim=-1, |
| 433 | ) |
| 434 | |
| 435 | |
| 436 | def _apply_block_diagonal( |
| 437 | feats: torch.Tensor, # (..., dim) |
| 438 | func_size_pairs: List[Tuple[Callable[[torch.Tensor], torch.Tensor], int]], |
| 439 | ) -> torch.Tensor: |
| 440 | """Apply a block-diagonal function to an input array. |
| 441 | |
| 442 | Each function is specified as a tuple with form: |
| 443 | |
| 444 | ((Tensor) -> Tensor, int) |
| 445 | |
| 446 | Where the integer is the size of the input to the function. |
| 447 | """ |
| 448 | funcs, block_sizes = zip(*func_size_pairs) |
| 449 | assert feats.shape[-1] == sum(block_sizes) |
| 450 | x_blocks = torch.split(feats, block_sizes, dim=-1) |
| 451 | out = torch.cat( |
| 452 | [f(x_block) for f, x_block in zip(funcs, x_blocks)], |
| 453 | dim=-1, |
| 454 | ) |
| 455 | assert out.shape == feats.shape, "Input/output shapes should match." |
| 456 | return out |
| 457 | |
| 458 | |
| 459 | def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: |
| 460 | """Invert a 4x4 SE(3) matrix.""" |
| 461 | assert transforms.shape[-2:] == (4, 4) |
| 462 | Rinv = transforms[..., :3, :3].transpose(-1, -2) |
| 463 | out = torch.zeros_like(transforms) |
| 464 | out[..., :3, :3] = Rinv |
| 465 | out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3]) |
| 466 | out[..., 3, 3] = 1.0 |
| 467 | return out |
| 468 | |
| 469 | |
| 470 | def _lift_K(Ks: torch.Tensor) -> torch.Tensor: |
| 471 | """Lift 3x3 matrices to homogeneous 4x4 matrices.""" |
| 472 | assert Ks.shape[-2:] == (3, 3) |
| 473 | out = torch.zeros(Ks.shape[:-2] + (4, 4), device=Ks.device, dtype=Ks.dtype) |
| 474 | out[..., :3, :3] = Ks |
| 475 | out[..., 3, 3] = 1.0 |
| 476 | return out |
| 477 | |
| 478 | |
| 479 | def _invert_K(Ks: torch.Tensor) -> torch.Tensor: |
| 480 | """Invert 3x3 intrinsics matrices. Assumes no skew.""" |
| 481 | assert Ks.shape[-2:] == (3, 3) |
| 482 | out = torch.zeros_like(Ks) |
| 483 | out[..., 0, 0] = 1.0 / Ks[..., 0, 0] |
| 484 | out[..., 1, 1] = 1.0 / Ks[..., 1, 1] |
| 485 | out[..., 0, 2] = -Ks[..., 0, 2] / Ks[..., 0, 0] |
| 486 | out[..., 1, 2] = -Ks[..., 1, 2] / Ks[..., 1, 1] |
| 487 | out[..., 2, 2] = 1.0 |
| 488 | return out |
| 489 |