Masked Scatter/Gather operations

[superposition]

Description

Efficient masking operations for MAE-style masked prediction.

Operations

# Gather visible tokens (for encoder)
visible = masked_gather(tokens, mask)  # (B, T, D) → (B, T_visible, D)

# Scatter predictions back to full sequence
full = masked_scatter(visible, predictions, mask)  # → (B, T, D)

Implementation

• Fused gather with mask indexing
• Avoid materialized index tensors
• Support arbitrary mask patterns (random, block, etc.)

Tests

• Gather produces correct subset
• Scatter places values at correct positions
• Round-trip: scatter(gather(x, m), m) preserves masked positions
• Gradient flows through both operations
• Works with variable mask ratios (25%, 50%, 75%)
• Handles batch dimension correctly

Used In

MAE, JEPA, BEiT, data2vec, masked world models