function mean_pooling(last_hidden_state, attention_mask) {
    // last_hidden_state: [batchSize, seqLength, embedDim]
    // attention_mask:    [batchSize, seqLength]
    const lastHiddenStateData = last_hidden_state.data;
    const attentionMaskData = attention_mask.data;

    const shape = [last_hidden_state.dims[0], last_hidden_state.dims[2]];

    // @ts-ignore
    const returnedData = new lastHiddenStateData.constructor(shape[0] * shape[1]);
    const [batchSize, seqLength, embedDim] = last_hidden_state.dims;

    let outIndex = 0;
    for (let i = 0; i < batchSize; ++i) {
        const offset = i * embedDim * seqLength;

        for (let k = 0; k < embedDim; ++k) {
            let sum = 0;
            let count = 0;

            const attnMaskOffset = i * seqLength;
            const offset2 = offset + k;
            // Pool over all words in sequence
            for (let j = 0; j < seqLength; ++j) {
                // index into attention mask
                const attn = Number(attentionMaskData[attnMaskOffset + j]);

                count += attn;
                sum += lastHiddenStateData[offset2 + j * embedDim] * attn;
            }

            const avg = sum / count;
            returnedData[outIndex++] = avg;
        }
    }

    return new Tensor(
        last_hidden_state.type,
        returnedData,
        shape
    )
}
Perform mean pooling of the last hidden state followed by a normalization step.
@param {Tensor} last_hidden_state Tensor of shape [batchSize, seqLength, embedDim]
@param {Tensor} attention_mask Tensor of shape [batchSize, seqLength]
@returns {Tensor} Returns a new Tensor of shape [batchSize, embedDim].