call method

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

Implementation

Tensor call(Tensor x, {int startPos = 0}) {
  if (x.shape.isEmpty || x.shape.last != embedDim) {
    throw ArgumentError(
      'LearnedPositionalEmbedding: expected [..., $embedDim]; '
      'got ${x.shape}',
    );
  }
  if (x.shape.length == 3) {
    final b = x.shape[0];
    final s = x.shape[1];
    if (startPos + s > maxLen) {
      throw ArgumentError(
        'LearnedPositionalEmbedding: startPos+seqLen ${startPos + s} '
        'exceeds maxLen $maxLen',
      );
    }
    // Same positions across every batch element => shape [B, S].
    final positions = Tensor.fromList(
      [b, s],
      List<double>.generate(b * s, (i) => ((i % s) + startPos).toDouble()),
      device: x.device,
    );
    return x + table(positions);
  }
  if (x.shape.length != 2) {
    throw ArgumentError(
      'LearnedPositionalEmbedding: expected rank 2 or 3; got ${x.shape}',
    );
  }
  final n = x.shape[0];
  if (startPos + n > maxLen) {
    throw ArgumentError(
      'LearnedPositionalEmbedding: startPos+seqLen ${startPos + n} '
      'exceeds maxLen $maxLen',
    );
  }
  final positions = Tensor.fromList(
    [n],
    List<double>.generate(n, (i) => (i + startPos).toDouble()),
    device: x.device,
  );
  return x + table(positions);
}