camera_rope

def prope_qkv(
    q: torch.Tensor,  # (batch, num_heads, seqlen, head_dim)
    k: torch.Tensor,  # (batch, num_heads, seqlen, head_dim)
    v: torch.Tensor,  # (batch, num_heads, seqlen, head_dim)
    *,
    viewmats: torch.Tensor,  # (batch, cameras, 4, 4)
    Ks: Optional[torch.Tensor],  # (batch, cameras, 3, 3)
    patches_x: int = None,  # How many patches wide is each image?
    patches_y: int = None,  # How many patches tall is each image?
    image_width: int = None,  # Width of the image. Used to normalize intrinsics.
    image_height: int = None,  # Height of the image. Used to normalize intrinsics.
    coeffs_x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    coeffs_y: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    mask: Optional[torch.Tensor] = None,
    kv_cache=None,
    is_cache: bool = False,
    **kwargs,
) -> torch.Tensor:
    """Similar to torch.nn.functional.scaled_dot_product_attention, but applies PRoPE-style
    positional encoding.
Currently, we assume that the sequence length is equal to:

    cameras * patches_x * patches_y

And token ordering allows the `(seqlen,)` axis to be reshaped into
`(cameras, patches_x, patches_y)`.
"""
# We're going to assume self-attention: all inputs are the same shape.
(batch, num_heads, seqlen, head_dim) = q.shape
cameras = viewmats.shape[1]
assert q.shape == k.shape == v.shape
assert viewmats.shape == (batch, cameras, 4, 4)
assert Ks is None or Ks.shape == (batch, cameras, 3, 3)
# assert seqlen == cameras * patches_x * patches_y

apply_fn_q, apply_fn_kv, apply_fn_o = _prepare_apply_fns_all_dim(
    head_dim=head_dim,
    viewmats=viewmats,
    Ks=Ks,
    patches_x=patches_x,
    patches_y=patches_y,
    image_width=image_width,
    image_height=image_height,
    coeffs_x=coeffs_x,
    coeffs_y=coeffs_y,
)

query = apply_fn_q(q)
key = apply_fn_kv(k)
value = apply_fn_kv(v)

return query, key, value, apply_fn_o</code></pre></details>