Masked Mean Pooling

Average token embeddings over the valid (non-padding) positions of each sequence. embeddings has shape (batch, time, dim) and mask has shape (batch, time) with True for real tokens. A sequence with no valid token gets the zero vector.

Implement solve(embeddings, mask).

Returns. Return an array of shape (batch, dim).

Examples

Example 1

Input:

solve([[[1.0, 1.0], [3.0, 3.0]]], [[True, True]])

Output:

[[2.0, 2.0]]

Example 2

Input:

solve([[[1.0, 1.0], [3.0, 3.0]]], [[True, False]])

Output:

[[1.0, 1.0]]

Example 3

Input:

solve([[[5.0, 5.0]]], [[False]])

Output:

[[0.0, 0.0]]