PIXELBANKv9.1.0
Menu

Many-to-One Output Selector

Problem Statement

Select the final relevant output from a batch of variable-length sequences.

Background

In "Many-to-One" architectures (e.g., sentiment analysis), sequences in a batch often have different lengths. We pad shorter sequences but need to extract the hidden state at the last valid timestep for each sequence.

Your Task

Write a function get_final_states(hidden_states, sequence_lengths) that extracts the final hidden state for each sequence in the batch.

Input Format

  • hidden_states: numpy array of shape (batch_size, max_seq_len, hidden_dim)
  • sequence_lengths: list of integers, length of each sequence in the batch

Output Format

Return a numpy array of shape (batch_size, hidden_dim) containing the last valid hidden state for each sequence.

Example:

Input:
hidden_states shape (2, 3, 4), sequence_lengths = [2, 3]
Output:
Shape (2, 4) with hidden_states[0,1,:] and hidden_states[1,2,:]
Reasoning:

For batch 0, take index 1 (length 2, 0-indexed). For batch 1, take index 2 (length 3).

Constraints:

  • 1 <= batch_size <= 64
  • 1 <= max_seq_len <= 512
  • 1 <= hidden_dim <= 512
  • sequence_lengths[i] >= 1 for all i
🔒

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.
Many-to-One Output Selector - Medium | PixelBank