Embedding

Lookup table, weight[indices]

Easy Fundamentals

Problem Description

Implement an embedding lookup table from scratch.

Signature

class MyEmbedding(nn.Module):
    def __init__(self, num_embeddings: int, embedding_dim: int): ...
    def forward(self, indices: Tensor) -> Tensor: ...

Rules

• self.weight: nn.Parameter of shape (num_embeddings, embedding_dim)

• Forward: index into weight matrix — weight[indices]

• Do NOT use nn.Embedding

Template

Implement the function below. Use only basic PyTorch operations.

# ✏️ YOUR IMPLEMENTATION HERE

class MyEmbedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim):
        super().__init__()
        pass

    def forward(self, indices):
        pass

Test Your Implementation

Use this code to debug before submitting.

# 🧪 Debug
emb = MyEmbedding(10, 4)
idx = torch.tensor([0, 3, 7])
print('Output shape:', emb(idx).shape)
print('Matches manual:', torch.equal(emb(idx)[0], emb.weight[0]))

Reference Solution

Try solving it yourself first! Click below to reveal the solution.

# ✅ SOLUTION

class MyEmbedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(num_embeddings, embedding_dim))

    def forward(self, indices):
        return self.weight[indices]

Tips

Run Locally

For interactive practice with auto-grading, run TorchCode locally:
pip install torch-judge then use check("embedding")

Key Concepts

Lookup table, weight[indices]

Dropout GELU

Embedding

Description Template Test Solution Tips