call method

Tensor call(
  1. Tensor x, {
  2. int startPos = 0,
})

Adds sinusoidal PE to x. Accepts [seqLen, embedDim] or a batched [batch, seqLen, embedDim]. startPos is the position of the first row of x (used by autoregressive / cached inference so that a single new token gets the encoding for its true position rather than position 0).

Implementation

Tensor call(Tensor x, {int startPos = 0}) {
  if (x.shape.isEmpty || x.shape.last != embedDim) {
    throw ArgumentError(
      'SinusoidalPositionalEncoding: expected [..., $embedDim]; '
      'got ${x.shape}',
    );
  }
  if (x.shape.length == 3) {
    final b = x.shape[0];
    final s = x.shape[1];
    final peBS = _computePEBatched(b, s, startPos, x.device);
    return (x.reshape([b * s, embedDim]) + peBS).reshape([b, s, embedDim]);
  }
  if (x.shape.length != 2) {
    throw ArgumentError(
      'SinusoidalPositionalEncoding: expected rank 2 or 3; got ${x.shape}',
    );
  }
  final n = x.shape[0];
  final pe = _computePE(n, startPos, x.device);
  return x + pe;
}