call method

Implementation

({MlxArray scaledInput, MlxArray posEmb}) call(
  MlxArray input, {
  int offset = 0,
}) {
  final batch = input.shape[0];
  final inputLen = input.shape[1] + offset;
  _ensureCache(inputLen);
  final pos = _pe!;
  final bufferLen = pos.shape[1];
  final startIdx = bufferLen ~/ 2 - (inputLen - 1);
  final endIdx = bufferLen ~/ 2 + (inputLen - 1) + 1;
  final scaled = (() {
    if (!scaleInput) {
      return input;
    }
    _scaleTensor ??= MlxArray.full([], math.sqrt(dModel.toDouble()));
    try {
      return mx.multiply(input, _scaleTensor!);
    } finally {
    }
  })();
  final sliced = pos.slice(
    start: <int>[0, startIdx, 0],
    stop: <int>[1, endIdx, dModel],
  );
  final typed = sliced.astype(input.dtype);
  sliced.close();
  final posEmb = batch == 1
      ? typed
      : (() {
          final out = mx.broadcastTo(
            typed,
            <int>[batch, endIdx - startIdx, dModel],
          );
          typed.close();
          return out;
        })();
  return (scaledInput: scaled, posEmb: posEmb);
}