meanPoolAndNormalize function

List<double> meanPoolAndNormalize(
  1. ForwardResult result, {
  2. List<int>? attentionMask,
})

Mean-pools a token-level ForwardResult over the sequence axis and L2-normalizes it into a single embedding vector of length dim.

Input MUST be rank-3 [1, seq, dim] (per-token hidden states, e.g. an ONNX BERT/MiniLM output). When attentionMask is supplied, positions where the mask is 0 are excluded from the mean (the standard BERT-style masked mean-pool over real tokens only, ignoring padding).

This function is ONLY for the EmbeddingOutputContract.tokenLevel path. A rank-2 [1, dim] result is ALREADY the final embedding (e.g. LiteRT's compiled-model output, which bakes pooling — and possibly normalization — into the graph); it must be copied verbatim via EmbeddingOutputContract.pooledFinal, never routed here. Passing a rank-2 result throws rather than silently re-normalizing it — the D5 double-normalize regression this seam exists to prevent.

Throws ArgumentError for a non-rank-3 shape, a batch size other than 1, or an attentionMask whose length doesn't match the sequence length. Throws StateError if every token is masked out (nothing to average).

Implementation

List<double> meanPoolAndNormalize(
  ForwardResult result, {
  List<int>? attentionMask,
}) {
  final shape = result.shape;
  if (shape.length != 3) {
    throw ArgumentError(
      'meanPoolAndNormalize handles token-level `[1, seq, dim]` results only; '
      'got $shape. A rank-2 `[1, dim]` result is already the final embedding — '
      'copy it verbatim via EmbeddingOutputContract.pooledFinal, never re-pool '
      'or re-normalize it (the D5 double-normalize trap).',
    );
  }
  return _l2Normalize(_meanPoolOverSeq(result, shape, attentionMask));
}