Skip the Multiplication

A colleague's pipeline one-hot encodes every token before it reaches the embedding layer. A sentence of $n$ tokens arrives as an $n \times V$ matrix $X$ — one row per token, where $V$ is the vocabulary size — and the layer multiplies it by the $V \times d$ embedding table $E$ :

(XE)_{ij} = \sum_{k=0}^{V-1} X_{ik}\, E_{kj}

With $V$ in the tens of thousands, that sum is enormous for every single entry of the result, and almost every term in it is zero.

Task: write embed(one_hot_rows, table) that returns exactly the $n \times d$ matrix $XE$ — without carrying out the multiplication.

one_hot_rows is $X$ , as a list of rows. A token's row has a single 1 and zeros everywhere else.
Padding slots arrive as a row of all zeros (that is how this pipeline encodes <PAD>). Your output for them must still be whatever $XE$ gives.
table is $E$ , as a list of $V$ rows — one row per vocabulary entry, each of length $d$ .

Return a list of $n$ lists, with every entry rounded to 4 decimal places.

Skip the Multiplication

Examples

Hints