apply_rope_2d_blh

apply_rope_2d_blh(x, rope_2d_cache)

Apply 2D RoPE to a tensor laid out as [batch_size, H, W, hidden_dim].

The channel dimension C is split into two equal parts: C_y and C_x. RoPE is applied independently along Y (to C_y) and X (to C_x). For pairwise rotations, C must be divisible by 4 so that each half is even.

Parameters:

• x (Tensor) – Input tensor of shape [batch_size, H, W, hidden_dim] with hidden_dim % 4 == 0.

• rope_2d_cache (tuple[Tensor, Tensor, Tensor, Tensor]) – tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] - The cache of 2D RoPE cos/sin for y and x axes, organized as (cos_y, sin_y, cos_x, sin_x).

Returns:

Tensor with the same shape as x.

Return type:

Tensor

Broadcasting:

• cos_y/sin_y are reshaped to [1, hidden_dim/2, H, 1] for the first half.

• cos_x/sin_x are reshaped to [1, hidden_dim/2, 1, W] for the second half.