PIXELBANKv9.1.0
Menu

2D Sine-Cosine Positional Embedding

Problem Statement

Build the fixed 2D sin-cos positional embedding used by many ViT/MAE implementations: concatenate an independent 1D sin-cos embedding for the row coordinate and the column coordinate.

Background

For a 1D position pos and an embedding width d (even), the sin-cos embedding uses d/2 frequencies

ωi=110000 2i/d,i=0,…,d/2−1\omega_i = \frac{1}{10000^{\,2i/d}}, \quad i = 0, \dots, d/2 - 1

and emits [sin(posw_0), ..., sin(posw_{d/2-1}), cos(posw_0), ..., cos(posw_{d/2-1})].

For a 2D grid, each patch has a (row, col) coordinate. Give half the width to the row embedding and half to the column embedding, then concatenate [row_emb, col_emb] into a length-dim vector. Patches are enumerated row-major.

Your Task

Implement:

def pos_embed_2d(grid, dim):

Return a nested list of shape (grid*grid, dim), every value rounded to 4 decimals.

Input Format

  • grid (int): side of the square patch grid.
  • dim (int): total embedding width, divisible by 4.

Output Format

  • A (grid*grid, dim) nested list, rounded to 4 decimals.

Sample

emb = pos_embed_2d(2, 4)
print(emb[0])

Output:

[0.0, 1.0, 0.0, 1.0]

Example:

Input:
emb = pos_embed_2d(2, 4)
print(emb[0])
Output:
[0.0, 1.0, 0.0, 1.0]
Reasoning:
  • Split the total dimension dim=4dim=4 into two halves for the row and column embeddings, so each 1D embedding has width daxis=4/2=2d_{axis} = 4 / 2 = 2.
  • Determine the coordinates for the first patch (index 0) in the 2×22 \times 2 grid: it is at row r=0r=0 and column c=0c=0.
  • Calculate the row embedding for position 00 with width 22. The single frequency is ω0=10000−2(0)/2=1\omega_0 = 10000^{-2(0)/2} = 1. The angles are 0⋅1=00 \cdot 1 = 0, yielding [sin⁡(0),cos⁡(0)]=[0.0,1.0][\sin(0), \cos(0)] = [0.0, 1.0].
  • Calculate the column embedding for position 00 with width 22. Similarly, the angle is 00, yielding [sin⁡(0),cos⁡(0)]=[0.0,1.0][\sin(0), \cos(0)] = [0.0, 1.0].
  • Concatenate the row and column embeddings to form the final vector: [0.0,1.0]⊕[0.0,1.0]=[0.0,1.0,0.0,1.0][0.0, 1.0] \oplus [0.0, 1.0] = [0.0, 1.0, 0.0, 1.0].
  • The final output is [0.0, 1.0, 0.0, 1.0]

Constraints:

  • dim is divisible by 4; each axis gets dim/2, each with dim/4 frequencies.
  • Position 0 gives sin=0, cos=1, so the first patch is [0,1,...,0,1].
  • Round every entry to 4 decimals; avoid emitting -0.0.
🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Test Results

0/0
Run code to see test results.
2D Sine-Cosine Positional Embedding - Medium | PixelBank