apply method

Tensor apply(Tensor qOrK, { int startPos = 0, })

Apply RoPE to a [N, headDim] query or key tensor whose rows correspond to absolute positions [startPos, startPos + N).

Returns a new tensor of the same shape.

Implementation

Tensor apply(Tensor qOrK, {int startPos = 0}) {
  if (qOrK.shape.length != 2 || qOrK.shape[1] != headDim) {
    throw ArgumentError(
      'RopeCache.apply: expected shape [N, $headDim], got '
      '${qOrK.shape}',
    );
  }
  final n = qOrK.shape[0];
  if (startPos + n > maxCtx) {
    throw ArgumentError(
      'RopeCache.apply: window [$startPos, ${startPos + n}) exceeds '
      'maxCtx=$maxCtx (rebuild with a larger maxCtx)',
    );
  }
  // Assemble cos/sin windows of shape [N, headDim] by stacking rows.
  final cosRows = <Tensor>[];
  final sinRows = <Tensor>[];
  for (int i = 0; i < n; i++) {
    cosRows.add(_cosPerPos[startPos + i]);
    sinRows.add(_sinPerPos[startPos + i]);
  }
  final cosBlock = cosRows.length == 1
      ? cosRows.first
      : TensorConcat.concat(cosRows, axis: 0);
  final sinBlock = sinRows.length == 1
      ? sinRows.first
      : TensorConcat.concat(sinRows, axis: 0);

  final rotated = qOrK.matmul(_rotateHalfP); // rotate_half via matmul
  return (qOrK * cosBlock) + (rotated * sinBlock);
}