poolTokens function

Tensor poolTokens(Tensor tokenFeatures, PoolingMode mode)

Pool a token feature matrix [seqLen, embedDim] to a sentence vector [1, embedDim].

• PoolingMode.mean — arithmetic mean over the sequence axis. Standard sentence-BERT default.
• PoolingMode.cls — the first token's features. Useful when the backbone was trained with a "CLS" token at index 0.

Implementation

Tensor poolTokens(Tensor tokenFeatures, PoolingMode mode) {
  if (tokenFeatures.shape.length != 2) {
    throw ArgumentError(
      'poolTokens: expected [seqLen, D]; got ${tokenFeatures.shape}',
    );
  }
  final s = tokenFeatures.shape[0];
  if (s == 0) {
    throw ArgumentError('poolTokens: cannot pool empty sequence');
  }
  switch (mode) {
    case PoolingMode.mean:
      // `[1, S] @ [S, D]` = `[1, D]`, differentiable through both operands.
      final ones = Tensor.fill([1, s], 1.0 / s, device: tokenFeatures.device);
      return ones.matmul(tokenFeatures);
    case PoolingMode.cls:
      return tokenFeatures.sliceRows(0, 1);
  }
}