Skip to content

rotary_embedding_3d

3D Rotary Position Embedding (RoPE) for video transformers.

Reference: https://arxiv.org/pdf/2104.09864.pdf

Classes

fastvideo.layers.rotary_embedding_3d.RotaryPositionalEmbedding3D 

RotaryPositionalEmbedding3D(head_dim: int, base: float = 10000.0)

Bases: Module

3D Rotary Positional Embedding for video transformers.

Splits the head dimension across temporal, height, and width dimensions, computing separate rotary embeddings for each and concatenating them.

Parameters:

Name Type Description Default
head_dim int

Dimension of each attention head

 required
base float

Base value for exponential frequency

 10000.0
Source code in fastvideo/layers/rotary_embedding_3d.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(
    self,
    head_dim: int,
    base: float = 10000.0,
):
    """
    Args:
        head_dim: Dimension of each attention head
        base: Base value for exponential frequency
    """
    super().__init__()
    self.head_dim = head_dim
    assert self.head_dim % 8 == 0, "head_dim must be a multiple of 8 for 3D RoPE"
    self.base = base

    # Cache for precomputed frequencies
    self.freqs_dict: dict[tuple, torch.Tensor] = {}

Functions

fastvideo.layers.rotary_embedding_3d.RotaryPositionalEmbedding3D.forward 
forward(q: Tensor, k: Tensor, grid_size: tuple[int, int, int]) -> tuple[Tensor, Tensor]

Apply 3D rotary positional embedding to queries and keys.

Parameters:

Name Type Description Default
q Tensor

Query tensor [B, num_heads, seq_len, head_dim]

 required
k Tensor

Key tensor [B, num_heads, seq_len, head_dim]

 required
grid_size tuple[int, int, int]

(T, H, W) tuple of grid dimensions

 required

Returns:

Type Description
(q_rotated, k_rotated)

Rotated query and key tensors
Source code in fastvideo/layers/rotary_embedding_3d.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
def forward(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    grid_size: tuple[int, int, int],
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Apply 3D rotary positional embedding to queries and keys.

    Args:
        q: Query tensor [B, num_heads, seq_len, head_dim]
        k: Key tensor [B, num_heads, seq_len, head_dim]
        grid_size: (T, H, W) tuple of grid dimensions

    Returns:
        (q_rotated, k_rotated): Rotated query and key tensors
    """
    # Register grid size if not cached
    if grid_size not in self.freqs_dict:
        self.register_grid_size(grid_size)

    # Get cached frequencies
    freqs_cis = self.freqs_dict[grid_size].to(q.device)

    # Cast to float32 for precision
    q_, k_ = q.float(), k.float()
    freqs_cis = freqs_cis.float()

    # Compute cos and sin
    cos = freqs_cis.cos()
    sin = freqs_cis.sin()

    # Reshape for broadcasting: [1, 1, seq_len, head_dim]
    cos = rearrange(cos, "n d -> 1 1 n d")
    sin = rearrange(sin, "n d -> 1 1 n d")

    # Apply rotation
    q_ = (q_ * cos) + (rotate_half(q_) * sin)
    k_ = (k_ * cos) + (rotate_half(k_) * sin)

    # Cast back to original dtype
    return q_.type_as(q), k_.type_as(k)
fastvideo.layers.rotary_embedding_3d.RotaryPositionalEmbedding3D.precompute_freqs_3d 
precompute_freqs_3d(grid_size: tuple[int, int, int]) -> Tensor

Precompute 3D rotary frequencies.

Parameters:

Name Type Description Default
grid_size tuple[int, int, int]

(num_frames, height, width)

 required

Returns:

Name Type Description
freqs Tensor

[THW, head_dim] tensor of frequencies
Source code in fastvideo/layers/rotary_embedding_3d.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def precompute_freqs_3d(self, grid_size: tuple[int, int,
                                               int]) -> torch.Tensor:
    """
    Precompute 3D rotary frequencies.

    Args:
        grid_size: (num_frames, height, width)

    Returns:
        freqs: [T*H*W, head_dim] tensor of frequencies
    """
    num_frames, height, width = grid_size

    # Split head_dim across 3 dimensions
    # Temporal gets the remainder to ensure exact division
    dim_t = self.head_dim - 4 * (self.head_dim // 6)
    dim_h = 2 * (self.head_dim // 6)
    dim_w = 2 * (self.head_dim // 6)

    # Compute frequency bands for each dimension
    freqs_t = 1.0 / (self.base**(
        torch.arange(0, dim_t, 2)[:(dim_t // 2)].float() / dim_t))
    freqs_h = 1.0 / (self.base**(
        torch.arange(0, dim_h, 2)[:(dim_h // 2)].float() / dim_h))
    freqs_w = 1.0 / (self.base**(
        torch.arange(0, dim_w, 2)[:(dim_w // 2)].float() / dim_w))

    # Create position grids
    grid_t = torch.arange(num_frames, dtype=torch.float32)
    grid_h = torch.arange(height, dtype=torch.float32)
    grid_w = torch.arange(width, dtype=torch.float32)

    # Compute frequencies for each position
    freqs_t = torch.einsum("..., f -> ... f", grid_t, freqs_t)
    freqs_h = torch.einsum("..., f -> ... f", grid_h, freqs_h)
    freqs_w = torch.einsum("..., f -> ... f", grid_w, freqs_w)

    # Duplicate for complex pair representation
    freqs_t = repeat(freqs_t, "... n -> ... (n r)", r=2)
    freqs_h = repeat(freqs_h, "... n -> ... (n r)", r=2)
    freqs_w = repeat(freqs_w, "... n -> ... (n r)", r=2)

    # Broadcast and concatenate across all 3 dimensions
    freqs = broadcast(
        [
            freqs_t[:, None, None, :],  # [T, 1, 1, dim_t]
            freqs_h[None, :, None, :],  # [1, H, 1, dim_h]
            freqs_w[None, None, :, :],  # [1, 1, W, dim_w]
        ],
        dim=-1,
    )

    # Flatten spatial dimensions: [T, H, W, head_dim] -> [T*H*W, head_dim]
    freqs = rearrange(freqs, "T H W D -> (T H W) D")

    return freqs
fastvideo.layers.rotary_embedding_3d.RotaryPositionalEmbedding3D.register_grid_size 
register_grid_size(grid_size: tuple[int, int, int]) -> None

Precompute and register frequencies for a given grid size.

Parameters:

Name Type Description Default
grid_size tuple[int, int, int]

(T, H, W) tuple of grid dimensions

 required
Source code in fastvideo/layers/rotary_embedding_3d.py 
80
81
82
83
84
85
86
87
88
def register_grid_size(self, grid_size: tuple[int, int, int]) -> None:
    """
    Precompute and register frequencies for a given grid size.

    Args:
        grid_size: (T, H, W) tuple of grid dimensions
    """
    if grid_size not in self.freqs_dict:
        self.freqs_dict[grid_size] = self.precompute_freqs_3d(grid_size)

Functions

fastvideo.layers.rotary_embedding_3d.apply_rotary_emb_3d 

apply_rotary_emb_3d(q: Tensor, k: Tensor, rope_module: RotaryPositionalEmbedding3D, grid_size: tuple[int, int, int]) -> tuple[Tensor, Tensor]

Convenience function to apply 3D RoPE.

Parameters:

Name Type Description Default
q Tensor

Query tensor [B, num_heads, seq_len, head_dim]

 required
k Tensor

Key tensor [B, num_heads, seq_len, head_dim]

 required
rope_module RotaryPositionalEmbedding3D

RotaryPositionalEmbedding3D module

 required
grid_size tuple[int, int, int]

(T, H, W) grid dimensions

 required

Returns:

Type Description
(q_rotated, k_rotated)

Rotated tensors
Source code in fastvideo/layers/rotary_embedding_3d.py 
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def apply_rotary_emb_3d(
    q: torch.Tensor,
    k: torch.Tensor,
    rope_module: RotaryPositionalEmbedding3D,
    grid_size: tuple[int, int, int],
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Convenience function to apply 3D RoPE.

    Args:
        q: Query tensor [B, num_heads, seq_len, head_dim]
        k: Key tensor [B, num_heads, seq_len, head_dim]
        rope_module: RotaryPositionalEmbedding3D module
        grid_size: (T, H, W) grid dimensions

    Returns:
        (q_rotated, k_rotated): Rotated tensors
    """
    return rope_module(q, k, grid_size)

fastvideo.layers.rotary_embedding_3d.broadcast 

broadcast(tensors: list[Tensor], dim: int = -1) -> Tensor

Broadcast and concatenate tensors along a dimension.

Source code in fastvideo/layers/rotary_embedding_3d.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def broadcast(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor:
    """
    Broadcast and concatenate tensors along a dimension.
    """
    num_tensors = len(tensors)
    shape_lens = set(len(t.shape) for t in tensors)
    assert len(
        shape_lens) == 1, "tensors must all have the same number of dimensions"

    shape_len = list(shape_lens)[0]
    dim = (dim + shape_len) if dim < 0 else dim

    dims = list(zip(*[list(t.shape) for t in tensors], strict=False))
    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]

    assert all(
        len(set(t[1])) <= 2 for t in
        expandable_dims), "invalid dimensions for broadcastable concatenation"

    max_dims = [(t[0], max(t[1])) for t in expandable_dims]
    expanded_dims = [(t[0], (t[1], ) * num_tensors) for t in max_dims]
    expanded_dims.insert(dim, (dim, dims[dim]))
    expandable_shapes = list(zip(*[t[1] for t in expanded_dims], strict=False))
    tensors = [
        t[0].expand(*t[1])
        for t in zip(tensors, expandable_shapes, strict=False)
    ]

    return torch.cat(tensors, dim=dim)

fastvideo.layers.rotary_embedding_3d.rotate_half 

rotate_half(x: Tensor) -> Tensor

Rotate half the hidden dims of the input.

Source code in fastvideo/layers/rotary_embedding_3d.py 
44
45
46
47
48
49
50
51
def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """
    Rotate half the hidden dims of the input.
    """
    x = rearrange(x, "... (d r) -> ... d r", r=2)
    x1, x2 = x.unbind(dim=-1)
    x = torch.stack((-x2, x1), dim=-1)
    return rearrange(x, "... d r -> ... (d r)")