parakeet.ts 1.0.0

package/dist/hub.js

@@ -0,0 +1,86 @@
+/**
+ * HuggingFace Hub downloads — backend-agnostic.
+ *
+ * Files are cached under the HF hub layout so both backends share one cache
+ * directory. Downloads are atomic (write to `.tmp`, then rename) so an
+ * interrupted fetch never leaves a truncated file that later looks cached.
+ */
+import fs from 'node:fs';
+import path from 'node:path';
+/** Default HF cache root, honouring HF_HOME / HF_HUB_CACHE like huggingface_hub. */
+export function defaultCacheDir() {
+    if (process.env['HF_HUB_CACHE'])
+        return process.env['HF_HUB_CACHE'];
+    if (process.env['HF_HOME'])
+        return path.join(process.env['HF_HOME'], 'hub');
+    return path.join(process.env['HOME'] ?? process.env['USERPROFILE'] ?? '/tmp', '.cache', 'huggingface', 'hub');
+}
+/** Local directory a repo's files are cached in. */
+export function repoDir(repoId, cacheDir) {
+    return path.join(cacheDir ?? defaultCacheDir(), repoId.replace(/\//g, '--'));
+}
+/**
+ * Fetch one file from a repo, returning its local path. No-op if already cached.
+ */
+export async function downloadFromHub(repoId, filename, cacheDir, onProgress) {
+    const modelDir = repoDir(repoId, cacheDir);
+    fs.mkdirSync(modelDir, { recursive: true });
+    const localPath = path.join(modelDir, filename);
+    if (fs.existsSync(localPath))
+        return localPath;
+    fs.mkdirSync(path.dirname(localPath), { recursive: true });
+    const url = `https://huggingface.co/${repoId}/resolve/main/${filename}`;
+    const response = await fetch(url, { headers: { 'User-Agent': 'parakeet.ts/1.0.0' } });
+    if (!response.ok) {
+        throw new Error(`Failed to download ${url}: ${response.status} ${response.statusText}`);
+    }
+    const total = parseInt(response.headers.get('content-length') ?? '0', 10);
+    const tmpPath = `${localPath}.tmp`;
+    if (response.body) {
+        const writeStream = fs.createWriteStream(tmpPath);
+        let downloaded = 0;
+        const reader = response.body.getReader();
+        try {
+            for (;;) {
+                const { done, value } = await reader.read();
+                if (done)
+                    break;
+                if (!writeStream.write(value)) {
+                    await new Promise(resolve => writeStream.once('drain', () => resolve()));
+                }
+                downloaded += value.byteLength;
+                if (onProgress)
+                    onProgress(downloaded, total);
+            }
+        }
+        finally {
+            reader.releaseLock();
+        }
+        await new Promise((resolve, reject) => {
+            writeStream.end((err) => (err ? reject(err) : resolve()));
+        });
+    }
+    else {
+        const arrayBuffer = await response.arrayBuffer();
+        fs.writeFileSync(tmpPath, Buffer.from(arrayBuffer));
+    }
+    fs.renameSync(tmpPath, localPath);
+    return localPath;
+}
+/**
+ * Fetch several files, skipping any marked optional that 404.
+ * Returns the repo's local directory.
+ */
+export async function downloadRepoFiles(repoId, files, cacheDir, onProgress) {
+    for (const { name, optional } of files) {
+        try {
+            await downloadFromHub(repoId, name, cacheDir, onProgress ? (d, t) => onProgress(name, d, t) : undefined);
+        }
+        catch (err) {
+            if (!optional)
+                throw err;
+        }
+    }
+    return repoDir(repoId, cacheDir);
+}
+//# sourceMappingURL=hub.js.map

package/dist/hub.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"hub.js","sourceRoot":"","sources":["../src/hub.ts"],"names":[],"mappings":"AAAA;;;;;;GAMG;AACH,OAAO,EAAE,MAAM,SAAS,CAAC;AACzB,OAAO,IAAI,MAAM,WAAW,CAAC;AAI7B,oFAAoF;AACpF,MAAM,UAAU,eAAe;IAC7B,IAAI,OAAO,CAAC,GAAG,CAAC,cAAc,CAAC;QAAE,OAAO,OAAO,CAAC,GAAG,CAAC,cAAc,CAAC,CAAC;IACpE,IAAI,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC;QAAE,OAAO,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC,EAAE,KAAK,CAAC,CAAC;IAC5E,OAAO,IAAI,CAAC,IAAI,CACd,OAAO,CAAC,GAAG,CAAC,MAAM,CAAC,IAAI,OAAO,CAAC,GAAG,CAAC,aAAa,CAAC,IAAI,MAAM,EAC3D,QAAQ,EAAE,aAAa,EAAE,KAAK,CAC/B,CAAC;AACJ,CAAC;AAED,oDAAoD;AACpD,MAAM,UAAU,OAAO,CAAC,MAAc,EAAE,QAAiB;IACvD,OAAO,IAAI,CAAC,IAAI,CAAC,QAAQ,IAAI,eAAe,EAAE,EAAE,MAAM,CAAC,OAAO,CAAC,KAAK,EAAE,IAAI,CAAC,CAAC,CAAC;AAC/E,CAAC;AAED;;GAEG;AACH,MAAM,CAAC,KAAK,UAAU,eAAe,CACnC,MAAc,EACd,QAAgB,EAChB,QAAiB,EACjB,UAAuB;IAEvB,MAAM,QAAQ,GAAG,OAAO,CAAC,MAAM,EAAE,QAAQ,CAAC,CAAC;IAC3C,EAAE,CAAC,SAAS,CAAC,QAAQ,EAAE,EAAE,SAAS,EAAE,IAAI,EAAE,CAAC,CAAC;IAE5C,MAAM,SAAS,GAAG,IAAI,CAAC,IAAI,CAAC,QAAQ,EAAE,QAAQ,CAAC,CAAC;IAChD,IAAI,EAAE,CAAC,UAAU,CAAC,SAAS,CAAC;QAAE,OAAO,SAAS,CAAC;IAE/C,EAAE,CAAC,SAAS,CAAC,IAAI,CAAC,OAAO,CAAC,SAAS,CAAC,EAAE,EAAE,SAAS,EAAE,IAAI,EAAE,CAAC,CAAC;IAE3D,MAAM,GAAG,GAAG,0BAA0B,MAAM,iBAAiB,QAAQ,EAAE,CAAC;IACxE,MAAM,QAAQ,GAAG,MAAM,KAAK,CAAC,GAAG,EAAE,EAAE,OAAO,EAAE,EAAE,YAAY,EAAE,mBAAmB,EAAE,EAAE,CAAC,CAAC;IAEtF,IAAI,CAAC,QAAQ,CAAC,EAAE,EAAE,CAAC;QACjB,MAAM,IAAI,KAAK,CAAC,sBAAsB,GAAG,KAAK,QAAQ,CAAC,MAAM,IAAI,QAAQ,CAAC,UAAU,EAAE,CAAC,CAAC;IAC1F,CAAC;IAED,MAAM,KAAK,GAAG,QAAQ,CAAC,QAAQ,CAAC,OAAO,CAAC,GAAG,CAAC,gBAAgB,CAAC,IAAI,GAAG,EAAE,EAAE,CAAC,CAAC;IAC1E,MAAM,OAAO,GAAG,GAAG,SAAS,MAAM,CAAC;IAEnC,IAAI,QAAQ,CAAC,IAAI,EAAE,CAAC;QAClB,MAAM,WAAW,GAAG,EAAE,CAAC,iBAAiB,CAAC,OAAO,CAAC,CAAC;QAClD,IAAI,UAAU,GAAG,CAAC,CAAC;QACnB,MAAM,MAAM,GAAG,QAAQ,CAAC,IAAI,CAAC,SAAS,EAAE,CAAC;QACzC,IAAI,CAAC;YACH,SAAS,CAAC;gBACR,MAAM,EAAE,IAAI,EAAE,KAAK,EAAE,GAAG,MAAM,MAAM,CAAC,IAAI,EAAE,CAAC;gBAC5C,IAAI,IAAI;oBAAE,MAAM;gBAChB,IAAI,CAAC,WAAW,CAAC,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC;oBAC9B,MAAM,IAAI,OAAO,CAAO,OAAO,CAAC,EAAE,CAAC,WAAW,CAAC,IAAI,CAAC,OAAO,EAAE,GAAG,EAAE,CAAC,OAAO,EAAE,CAAC,CAAC,CAAC;gBACjF,CAAC;gBACD,UAAU,IAAI,KAAK,CAAC,UAAU,CAAC;gBAC/B,IAAI,UAAU;oBAAE,UAAU,CAAC,UAAU,EAAE,KAAK,CAAC,CAAC;YAChD,CAAC;QACH,CAAC;gBAAS,CAAC;YACT,MAAM,CAAC,WAAW,EAAE,CAAC;QACvB,CAAC;QACD,MAAM,IAAI,OAAO,CAAO,CAAC,OAAO,EAAE,MAAM,EAAE,EAAE;YAC1C,WAAW,CAAC,GAAG,CAAC,CAAC,GAAkB,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,MAAM,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC,CAAC,CAAC;QAC3E,CAAC,CAAC,CAAC;IACL,CAAC;SAAM,CAAC;QACN,MAAM,WAAW,GAAG,MAAM,QAAQ,CAAC,WAAW,EAAE,CAAC;QACjD,EAAE,CAAC,aAAa,CAAC,OAAO,EAAE,MAAM,CAAC,IAAI,CAAC,WAAW,CAAC,CAAC,CAAC;IACtD,CAAC;IAED,EAAE,CAAC,UAAU,CAAC,OAAO,EAAE,SAAS,CAAC,CAAC;IAClC,OAAO,SAAS,CAAC;AACnB,CAAC;AAED;;;GAGG;AACH,MAAM,CAAC,KAAK,UAAU,iBAAiB,CACrC,MAAc,EACd,KAAkD,EAClD,QAAiB,EACjB,UAAsE;IAEtE,KAAK,MAAM,EAAE,IAAI,EAAE,QAAQ,EAAE,IAAI,KAAK,EAAE,CAAC;QACvC,IAAI,CAAC;YACH,MAAM,eAAe,CAAC,MAAM,EAAE,IAAI,EAAE,QAAQ,EAC1C,UAAU,CAAC,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,UAAU,CAAC,IAAI,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC;QAC/D,CAAC;QAAC,OAAO,GAAG,EAAE,CAAC;YACb,IAAI,CAAC,QAAQ;gBAAE,MAAM,GAAG,CAAC;QAC3B,CAAC;IACH,CAAC;IACD,OAAO,OAAO,CAAC,MAAM,EAAE,QAAQ,CAAC,CAAC;AACnC,CAAC"}

package/dist/index.d.ts

@@ -0,0 +1,13 @@
+export { load, detectBackend, DEFAULT_MODELS } from './load.js';
+export type { LoadOptions, BackendKind } from './load.js';
+export { downloadFromHub, downloadRepoFiles, defaultCacheDir, repoDir } from './hub.js';
+export { ParakeetModel, StreamingParakeet, consumePcmStream } from './model.js';
+export type { ParakeetModelArgs, TranscribeOptions, StreamOptions } from './model.js';
+export type { ParakeetBackend, EncoderOutput, EncoderLayout, DecodeStepResult, DecoderStateHandle, } from './backend.js';
+export { decodeTDTGreedy, decodeRNNTGreedy } from './decode.js';
+export type { DecoderState, GreedyOptions } from './decode.js';
+export { getLogMel, loadAudioRaw, getWindow, makePreprocessArgs, computeMelFilterbanks, computeMelFilterbanksInterpolated, } from './audio.js';
+export type { PreprocessArgs, LogMel } from './audio.js';
+export { AlignedToken, AlignedSentence, AlignedResult, SentenceConfig, makeAlignedToken, makeAlignedSentence, makeAlignedResult, tokensToSentences, sentencesToResult, } from './alignment.js';
+export { decode } from './tokenizer.js';
+//# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAYA,OAAO,EAAE,IAAI,EAAE,aAAa,EAAE,cAAc,EAAE,MAAM,WAAW,CAAC;AAChE,YAAY,EAAE,WAAW,EAAE,WAAW,EAAE,MAAM,WAAW,CAAC;AAE1D,OAAO,EAAE,eAAe,EAAE,iBAAiB,EAAE,eAAe,EAAE,OAAO,EAAE,MAAM,UAAU,CAAC;AAExF,OAAO,EAAE,aAAa,EAAE,iBAAiB,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAC;AAChF,YAAY,EAAE,iBAAiB,EAAE,iBAAiB,EAAE,aAAa,EAAE,MAAM,YAAY,CAAC;AAEtF,YAAY,EACV,eAAe,EACf,aAAa,EACb,aAAa,EACb,gBAAgB,EAChB,kBAAkB,GACnB,MAAM,cAAc,CAAC;AAEtB,OAAO,EAAE,eAAe,EAAE,gBAAgB,EAAE,MAAM,aAAa,CAAC;AAChE,YAAY,EAAE,YAAY,EAAE,aAAa,EAAE,MAAM,aAAa,CAAC;AAE/D,OAAO,EACL,SAAS,EACT,YAAY,EACZ,SAAS,EACT,kBAAkB,EAClB,qBAAqB,EACrB,iCAAiC,GAClC,MAAM,YAAY,CAAC;AACpB,YAAY,EAAE,cAAc,EAAE,MAAM,EAAE,MAAM,YAAY,CAAC;AAEzD,OAAO,EACL,YAAY,EACZ,eAAe,EACf,aAAa,EACb,cAAc,EACd,gBAAgB,EAChB,mBAAmB,EACnB,iBAAiB,EACjB,iBAAiB,EACjB,iBAAiB,GAClB,MAAM,gBAAgB,CAAC;AAExB,OAAO,EAAE,MAAM,EAAE,MAAM,gBAAgB,CAAC"}

package/dist/index.js

@@ -0,0 +1,19 @@
+// parakeet.ts — shared, backend-agnostic API
+//
+// The model class, audio front-end, decode loops, alignment and tokenizer are
+// all backend-independent. Pick a backend by importing its loader:
+//
+//   import { fromLocal } from 'parakeet.ts/onnx';   // ONNX Runtime (CUDA on Linux)
+//   import { fromLocal } from 'parakeet.ts/mlx';    // MLX (Apple Silicon)
+//
+// Both return the same `ParakeetModel`.
+// Zero-config entry point: picks the backend for this platform and fetches the
+// model on first use.
+export { load, detectBackend, DEFAULT_MODELS } from './load.js';
+export { downloadFromHub, downloadRepoFiles, defaultCacheDir, repoDir } from './hub.js';
+export { ParakeetModel, StreamingParakeet, consumePcmStream } from './model.js';
+export { decodeTDTGreedy, decodeRNNTGreedy } from './decode.js';
+export { getLogMel, loadAudioRaw, getWindow, makePreprocessArgs, computeMelFilterbanks, computeMelFilterbanksInterpolated, } from './audio.js';
+export { makeAlignedToken, makeAlignedSentence, makeAlignedResult, tokensToSentences, sentencesToResult, } from './alignment.js';
+export { decode } from './tokenizer.js';
+//# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,6CAA6C;AAC7C,EAAE;AACF,8EAA8E;AAC9E,mEAAmE;AACnE,EAAE;AACF,oFAAoF;AACpF,2EAA2E;AAC3E,EAAE;AACF,wCAAwC;AAExC,+EAA+E;AAC/E,sBAAsB;AACtB,OAAO,EAAE,IAAI,EAAE,aAAa,EAAE,cAAc,EAAE,MAAM,WAAW,CAAC;AAGhE,OAAO,EAAE,eAAe,EAAE,iBAAiB,EAAE,eAAe,EAAE,OAAO,EAAE,MAAM,UAAU,CAAC;AAExF,OAAO,EAAE,aAAa,EAAE,iBAAiB,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAC;AAWhF,OAAO,EAAE,eAAe,EAAE,gBAAgB,EAAE,MAAM,aAAa,CAAC;AAGhE,OAAO,EACL,SAAS,EACT,YAAY,EACZ,SAAS,EACT,kBAAkB,EAClB,qBAAqB,EACrB,iCAAiC,GAClC,MAAM,YAAY,CAAC;AAGpB,OAAO,EAKL,gBAAgB,EAChB,mBAAmB,EACnB,iBAAiB,EACjB,iBAAiB,EACjB,iBAAiB,GAClB,MAAM,gBAAgB,CAAC;AAExB,OAAO,EAAE,MAAM,EAAE,MAAM,gBAAgB,CAAC"}

package/dist/load.d.ts

@@ -0,0 +1,51 @@
+/**
+ * Platform-detecting loader — the entry point most callers should use.
+ *
+ *   import { load } from 'parakeet.ts';
+ *   const model = await load();          // picks the backend, fetches the model
+ *   const result = await model.transcribe('audio.wav');
+ *
+ * Backend selection:
+ *   - Apple Silicon  -> MLX / Metal
+ *   - everything else -> ONNX Runtime (CUDA execution provider on Linux/Nvidia,
+ *     CPU otherwise)
+ *
+ * Each backend has its own default checkpoint on the Hub, because the two need
+ * different asset formats (safetensors vs exported ONNX graphs). Both produce
+ * the same transcripts — see docs/cuda.md.
+ */
+import type { ParakeetModel } from './model.js';
+export type BackendKind = 'auto' | 'mlx' | 'onnx';
+export interface LoadOptions {
+    /** Force a backend. Default 'auto'. */
+    backend?: BackendKind;
+    /**
+     * HF repo id or local directory. Defaults to the chosen backend's standard
+     * checkpoint.
+     */
+    model?: string;
+    /** Override the HF cache root. */
+    cacheDir?: string;
+    /** Called as each file downloads. */
+    onProgress?: (file: string, downloaded: number, total: number) => void;
+    /**
+     * Mel filterbank. 'interpolated' (default) matches NVIDIA's reference
+     * preprocessor; 'floor' reproduces the legacy parakeet-mlx filterbank.
+     */
+    filterbank?: 'floor' | 'interpolated';
+    /** ONNX only: override the execution provider (default: CUDA where available). */
+    executionProvider?: 'cuda' | 'cpu' | 'tensorrt' | 'coreml' | 'dml';
+}
+/** Standard checkpoint per backend — different asset formats, same transcripts. */
+export declare const DEFAULT_MODELS: {
+    readonly mlx: "mlx-community/parakeet-tdt-0.6b-v3";
+    readonly onnx: "istupakov/parakeet-tdt-0.6b-v3-onnx";
+};
+/** Which backend `load()` would choose on this machine. */
+export declare function detectBackend(): 'mlx' | 'onnx';
+/**
+ * Load Parakeet using the backend native to this platform, downloading the
+ * model on first use.
+ */
+export declare function load(options?: LoadOptions): Promise<ParakeetModel>;
+//# sourceMappingURL=load.d.ts.map

package/dist/load.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"load.d.ts","sourceRoot":"","sources":["../src/load.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;GAeG;AACH,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,YAAY,CAAC;AAEhD,MAAM,MAAM,WAAW,GAAG,MAAM,GAAG,KAAK,GAAG,MAAM,CAAC;AAElD,MAAM,WAAW,WAAW;IAC1B,uCAAuC;IACvC,OAAO,CAAC,EAAE,WAAW,CAAC;IACtB;;;OAGG;IACH,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,kCAAkC;IAClC,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,qCAAqC;IACrC,UAAU,CAAC,EAAE,CAAC,IAAI,EAAE,MAAM,EAAE,UAAU,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,KAAK,IAAI,CAAC;IACvE;;;OAGG;IACH,UAAU,CAAC,EAAE,OAAO,GAAG,cAAc,CAAC;IACtC,kFAAkF;IAClF,iBAAiB,CAAC,EAAE,MAAM,GAAG,KAAK,GAAG,UAAU,GAAG,QAAQ,GAAG,KAAK,CAAC;CACpE;AAED,mFAAmF;AACnF,eAAO,MAAM,cAAc;;;CAGjB,CAAC;AAEX,2DAA2D;AAC3D,wBAAgB,aAAa,IAAI,KAAK,GAAG,MAAM,CAE9C;AAED;;;GAGG;AACH,wBAAsB,IAAI,CAAC,OAAO,GAAE,WAAgB,GAAG,OAAO,CAAC,aAAa,CAAC,CAqB5E"}

package/dist/load.js

@@ -0,0 +1,34 @@
+/** Standard checkpoint per backend — different asset formats, same transcripts. */
+export const DEFAULT_MODELS = {
+    mlx: 'mlx-community/parakeet-tdt-0.6b-v3',
+    onnx: 'istupakov/parakeet-tdt-0.6b-v3-onnx',
+};
+/** Which backend `load()` would choose on this machine. */
+export function detectBackend() {
+    return process.platform === 'darwin' && process.arch === 'arm64' ? 'mlx' : 'onnx';
+}
+/**
+ * Load Parakeet using the backend native to this platform, downloading the
+ * model on first use.
+ */
+export async function load(options = {}) {
+    const kind = !options.backend || options.backend === 'auto'
+        ? detectBackend()
+        : options.backend;
+    if (kind === 'mlx') {
+        const { fromPretrained } = await import('./mlx/load.js');
+        return fromPretrained(options.model ?? DEFAULT_MODELS.mlx, {
+            cacheDir: options.cacheDir,
+            onProgress: options.onProgress,
+            filterbank: options.filterbank,
+        });
+    }
+    const { fromPretrained } = await import('./onnx/parakeet.js');
+    return fromPretrained(options.model ?? DEFAULT_MODELS.onnx, {
+        cacheDir: options.cacheDir,
+        onProgress: options.onProgress,
+        filterbank: options.filterbank,
+        executionProvider: options.executionProvider,
+    });
+}
+//# sourceMappingURL=load.js.map

package/dist/load.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"load.js","sourceRoot":"","sources":["../src/load.ts"],"names":[],"mappings":"AAyCA,mFAAmF;AACnF,MAAM,CAAC,MAAM,cAAc,GAAG;IAC5B,GAAG,EAAE,oCAAoC;IACzC,IAAI,EAAE,qCAAqC;CACnC,CAAC;AAEX,2DAA2D;AAC3D,MAAM,UAAU,aAAa;IAC3B,OAAO,OAAO,CAAC,QAAQ,KAAK,QAAQ,IAAI,OAAO,CAAC,IAAI,KAAK,OAAO,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,MAAM,CAAC;AACpF,CAAC;AAED;;;GAGG;AACH,MAAM,CAAC,KAAK,UAAU,IAAI,CAAC,UAAuB,EAAE;IAClD,MAAM,IAAI,GAAG,CAAC,OAAO,CAAC,OAAO,IAAI,OAAO,CAAC,OAAO,KAAK,MAAM;QACzD,CAAC,CAAC,aAAa,EAAE;QACjB,CAAC,CAAC,OAAO,CAAC,OAAO,CAAC;IAEpB,IAAI,IAAI,KAAK,KAAK,EAAE,CAAC;QACnB,MAAM,EAAE,cAAc,EAAE,GAAG,MAAM,MAAM,CAAC,eAAe,CAAC,CAAC;QACzD,OAAO,cAAc,CAAC,OAAO,CAAC,KAAK,IAAI,cAAc,CAAC,GAAG,EAAE;YACzD,QAAQ,EAAE,OAAO,CAAC,QAAQ;YAC1B,UAAU,EAAE,OAAO,CAAC,UAAU;YAC9B,UAAU,EAAE,OAAO,CAAC,UAAU;SAC/B,CAAC,CAAC;IACL,CAAC;IAED,MAAM,EAAE,cAAc,EAAE,GAAG,MAAM,MAAM,CAAC,oBAAoB,CAAC,CAAC;IAC9D,OAAO,cAAc,CAAC,OAAO,CAAC,KAAK,IAAI,cAAc,CAAC,IAAI,EAAE;QAC1D,QAAQ,EAAE,OAAO,CAAC,QAAQ;QAC1B,UAAU,EAAE,OAAO,CAAC,UAAU;QAC9B,UAAU,EAAE,OAAO,CAAC,UAAU;QAC9B,iBAAiB,EAAE,OAAO,CAAC,iBAAiB;KAC7C,CAAC,CAAC;AACL,CAAC"}

package/dist/mlx/attention.d.ts

@@ -0,0 +1,61 @@
+import { MxArray } from '@mlx-node/core';
+import { Module, Linear, WeightMap } from './nn.js';
+import { ConformerCache } from './cache.js';
+export declare class MultiHeadAttention extends Module {
+    linearQ: Linear;
+    linearK: Linear;
+    linearV: Linear;
+    linearOut: Linear;
+    readonly nHead: number;
+    readonly headDim: number;
+    readonly scale: number;
+    constructor(nHead: number, nFeat: number, bias?: boolean);
+    forward(q: MxArray, k: MxArray, v: MxArray, posEmb: MxArray | null, mask: MxArray | null, cache: ConformerCache | null): MxArray;
+    loadWeights(weights: WeightMap, prefix: string): void;
+}
+export declare class RelPositionMultiHeadAttention extends Module {
+    linearQ: Linear;
+    linearK: Linear;
+    linearV: Linear;
+    linearOut: Linear;
+    linearPos: Linear;
+    posBiasU: MxArray;
+    posBiasV: MxArray;
+    readonly nHead: number;
+    readonly headDim: number;
+    readonly scale: number;
+    constructor(nHead: number, nFeat: number, bias?: boolean);
+    private relShift;
+    forward(q: MxArray, k: MxArray, v: MxArray, posEmb: MxArray | null, mask: MxArray | null, cache: ConformerCache | null): MxArray;
+    loadWeights(weights: WeightMap, prefix: string): void;
+}
+export declare class RelPositionMultiHeadLocalAttention extends RelPositionMultiHeadAttention {
+    readonly contextSize: [number, number];
+    constructor(nHead: number, nFeat: number, bias?: boolean, contextSize?: [number, number]);
+    forward(q: MxArray, k: MxArray, v: MxArray, posEmb: MxArray | null, mask: MxArray | null, cache: ConformerCache | null): MxArray;
+}
+export declare class RelPositionalEncoding extends Module {
+    readonly dModel: number;
+    maxLen: number;
+    readonly scaleInput: boolean;
+    readonly scale: number;
+    pe: MxArray;
+    constructor(dModel: number, maxLen?: number, scaleInput?: boolean);
+    calculatePE(): void;
+    forward(x: MxArray, offset?: number): [MxArray, MxArray];
+    loadWeights(_weights: WeightMap, _prefix: string): void;
+}
+export declare class LocalRelPositionalEncoding extends Module {
+    readonly dModel: number;
+    readonly maxLen: number;
+    readonly scaleInput: boolean;
+    readonly scale: number;
+    readonly leftContext: number;
+    readonly rightContext: number;
+    pe: MxArray;
+    constructor(dModel: number, maxLen?: number, scaleInput?: boolean, contextSize?: [number, number]);
+    calculatePE(): void;
+    forward(x: MxArray, _offset?: number): [MxArray, MxArray];
+    loadWeights(_weights: WeightMap, _prefix: string): void;
+}
+//# sourceMappingURL=attention.d.ts.map

package/dist/mlx/attention.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"attention.d.ts","sourceRoot":"","sources":["../../src/mlx/attention.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,OAAO,EAAE,MAAM,gBAAgB,CAAC;AACzC,OAAO,EAAE,MAAM,EAAE,MAAM,EAAE,SAAS,EAAoB,MAAM,SAAS,CAAC;AACtE,OAAO,EAAE,cAAc,EAAE,MAAM,YAAY,CAAC;AAiC5C,qBAAa,kBAAmB,SAAQ,MAAM;IAC5C,OAAO,EAAE,MAAM,CAAC;IAChB,OAAO,EAAE,MAAM,CAAC;IAChB,OAAO,EAAE,MAAM,CAAC;IAChB,SAAS,EAAE,MAAM,CAAC;IAElB,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAC;IACvB,QAAQ,CAAC,OAAO,EAAE,MAAM,CAAC;IACzB,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAC;gBAEX,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,IAAI,UAAO;IAYrD,OAAO,CACL,CAAC,EAAE,OAAO,EACV,CAAC,EAAE,OAAO,EACV,CAAC,EAAE,OAAO,EACV,MAAM,EAAE,OAAO,GAAG,IAAI,EACtB,IAAI,EAAE,OAAO,GAAG,IAAI,EACpB,KAAK,EAAE,cAAc,GAAG,IAAI,GAC3B,OAAO;IAuBV,WAAW,CAAC,OAAO,EAAE,SAAS,EAAE,MAAM,EAAE,MAAM,GAAG,IAAI;CAMtD;AAMD,qBAAa,6BAA8B,SAAQ,MAAM;IACvD,OAAO,EAAE,MAAM,CAAC;IAChB,OAAO,EAAE,MAAM,CAAC;IAChB,OAAO,EAAE,MAAM,CAAC;IAChB,SAAS,EAAE,MAAM,CAAC;IAClB,SAAS,EAAE,MAAM,CAAC;IAElB,QAAQ,EAAG,OAAO,CAAC;IACnB,QAAQ,EAAG,OAAO,CAAC;IAEnB,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAC;IACvB,QAAQ,CAAC,OAAO,EAAE,MAAM,CAAC;IACzB,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAC;gBAEX,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,IAAI,UAAO;IAarD,OAAO,CAAC,QAAQ;IAkBhB,OAAO,CACL,CAAC,EAAE,OAAO,EACV,CAAC,EAAE,OAAO,EACV,CAAC,EAAE,OAAO,EACV,MAAM,EAAE,OAAO,GAAG,IAAI,EACtB,IAAI,EAAE,OAAO,GAAG,IAAI,EACpB,KAAK,EAAE,cAAc,GAAG,IAAI,GAC3B,OAAO;IA0DV,WAAW,CAAC,OAAO,EAAE,SAAS,EAAE,MAAM,EAAE,MAAM,GAAG,IAAI;CAoBtD;AAOD,qBAAa,kCAAmC,SAAQ,6BAA6B;IACnF,QAAQ,CAAC,WAAW,EAAE,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;gBAGrC,KAAK,EAAE,MAAM,EACb,KAAK,EAAE,MAAM,EACb,IAAI,UAAO,EACX,WAAW,GAAE,CAAC,MAAM,EAAE,MAAM,CAAc;IAM5C,OAAO,CACL,CAAC,EAAE,OAAO,EACV,CAAC,EAAE,OAAO,EACV,CAAC,EAAE,OAAO,EACV,MAAM,EAAE,OAAO,GAAG,IAAI,EACtB,IAAI,EAAE,OAAO,GAAG,IAAI,EACpB,KAAK,EAAE,cAAc,GAAG,IAAI,GAC3B,OAAO;CAoEX;AAMD,qBAAa,qBAAsB,SAAQ,MAAM;IAC/C,QAAQ,CAAC,MAAM,EAAE,MAAM,CAAC;IACxB,MAAM,EAAE,MAAM,CAAC;IACf,QAAQ,CAAC,UAAU,EAAE,OAAO,CAAC;IAC7B,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAC;IACvB,EAAE,EAAG,OAAO,CAAC;gBAED,MAAM,EAAE,MAAM,EAAE,MAAM,SAAO,EAAE,UAAU,UAAO;IAS5D,WAAW,IAAI,IAAI;IAkBnB,OAAO,CAAC,CAAC,EAAE,OAAO,EAAE,MAAM,SAAI,GAAG,CAAC,OAAO,EAAE,OAAO,CAAC;IAkBnD,WAAW,CAAC,QAAQ,EAAE,SAAS,EAAE,OAAO,EAAE,MAAM,GAAG,IAAI;CAGxD;AAED,qBAAa,0BAA2B,SAAQ,MAAM;IACpD,QAAQ,CAAC,MAAM,EAAE,MAAM,CAAC;IACxB,QAAQ,CAAC,MAAM,EAAE,MAAM,CAAC;IACxB,QAAQ,CAAC,UAAU,EAAE,OAAO,CAAC;IAC7B,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAC;IACvB,QAAQ,CAAC,WAAW,EAAE,MAAM,CAAC;IAC7B,QAAQ,CAAC,YAAY,EAAE,MAAM,CAAC;IAC9B,EAAE,EAAG,OAAO,CAAC;gBAGX,MAAM,EAAE,MAAM,EACd,MAAM,SAAO,EACb,UAAU,UAAO,EACjB,WAAW,GAAE,CAAC,MAAM,EAAE,MAAM,CAAc;IAW5C,WAAW,IAAI,IAAI;IAkBnB,OAAO,CAAC,CAAC,EAAE,OAAO,EAAE,OAAO,SAAI,GAAG,CAAC,OAAO,EAAE,OAAO,CAAC;IAKpD,WAAW,CAAC,QAAQ,EAAE,SAAS,EAAE,OAAO,EAAE,MAAM,GAAG,IAAI;CACxD"}

package/dist/mlx/attention.js

@@ -0,0 +1,330 @@
+import { MxArray } from '@mlx-node/core';
+import { Module, Linear, softmax } from './nn.js';
+function s(...dims) {
+    return BigInt64Array.from(dims.map(BigInt));
+}
+// ---------------------------------------------------------------------------
+// Scaled dot-product attention (standard)
+// ---------------------------------------------------------------------------
+function scaledDotProductAttention(q, // [batch, heads, q_seq, head_dim]
+k, // [batch, heads, k_seq, head_dim]
+v, // [batch, heads, k_seq, head_dim]
+scale, mask) {
+    // scores: [batch, heads, q_seq, k_seq]
+    let scores = q.matmul(k.transpose(new Int32Array([0, 1, 3, 2]))).mulScalar(scale);
+    if (mask !== null) {
+        // mask is additive (0 or -inf) or boolean — add directly
+        scores = scores.add(mask);
+    }
+    const attn = softmax(scores, 3);
+    return attn.matmul(v); // [batch, heads, q_seq, head_dim]
+}
+// ---------------------------------------------------------------------------
+// Standard multi-head attention
+// ---------------------------------------------------------------------------
+export class MultiHeadAttention extends Module {
+    linearQ;
+    linearK;
+    linearV;
+    linearOut;
+    nHead;
+    headDim;
+    scale;
+    constructor(nHead, nFeat, bias = true) {
+        super();
+        this.nHead = nHead;
+        this.headDim = Math.floor(nFeat / nHead);
+        this.scale = Math.pow(this.headDim, -0.5);
+        this.linearQ = new Linear(nFeat, nFeat, bias);
+        this.linearK = new Linear(nFeat, nFeat, bias);
+        this.linearV = new Linear(nFeat, nFeat, bias);
+        this.linearOut = new Linear(nFeat, nFeat, bias);
+    }
+    forward(q, k, v, posEmb, mask, cache) {
+        const Qp = this.linearQ.forward(q);
+        const Kp = this.linearK.forward(k);
+        const Vp = this.linearV.forward(v);
+        const qShape = Qp.shape();
+        const batch = Number(qShape[0]);
+        const qSeq = Number(qShape[1]);
+        const kSeq = Number(Kp.shape()[1]);
+        const Qr = Qp.reshape(s(batch, qSeq, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        let Kr = Kp.reshape(s(batch, kSeq, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        let Vr = Vp.reshape(s(batch, kSeq, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        if (cache) {
+            [Kr, Vr] = cache.updateAndFetchKV(Kr, Vr);
+        }
+        const o = scaledDotProductAttention(Qr, Kr, Vr, this.scale, mask);
+        const out = o.transpose(new Int32Array([0, 2, 1, 3])).reshape(s(batch, qSeq, this.headDim * this.nHead));
+        return this.linearOut.forward(out);
+    }
+    loadWeights(weights, prefix) {
+        this.linearQ.loadWeights(weights, `${prefix}.linear_q`);
+        this.linearK.loadWeights(weights, `${prefix}.linear_k`);
+        this.linearV.loadWeights(weights, `${prefix}.linear_v`);
+        this.linearOut.loadWeights(weights, `${prefix}.linear_out`);
+    }
+}
+// ---------------------------------------------------------------------------
+// Relative-position multi-head attention
+// ---------------------------------------------------------------------------
+export class RelPositionMultiHeadAttention extends Module {
+    linearQ;
+    linearK;
+    linearV;
+    linearOut;
+    linearPos;
+    posBiasU; // [nHead, headDim]
+    posBiasV; // [nHead, headDim]
+    nHead;
+    headDim;
+    scale;
+    constructor(nHead, nFeat, bias = true) {
+        super();
+        this.nHead = nHead;
+        this.headDim = Math.floor(nFeat / nHead);
+        this.scale = Math.pow(this.headDim, -0.5);
+        this.linearQ = new Linear(nFeat, nFeat, bias);
+        this.linearK = new Linear(nFeat, nFeat, bias);
+        this.linearV = new Linear(nFeat, nFeat, bias);
+        this.linearOut = new Linear(nFeat, nFeat, bias);
+        this.linearPos = new Linear(nFeat, nFeat, false);
+    }
+    relShift(x) {
+        // x: [B, H, Tq, posLen]
+        const xShape = x.shape();
+        const B = Number(xShape[0]);
+        const H = Number(xShape[1]);
+        const Tq = Number(xShape[2]);
+        const posLen = Number(xShape[3]);
+        // Pad: [B, H, Tq, posLen+1]
+        const padded = x.pad(new Int32Array([0, 0, 0, 0, 0, 0, 1, 0]), 0.0);
+        // Reshape: [B, H, posLen+1, Tq]
+        const reshaped = padded.reshape(s(B, H, posLen + 1, Tq));
+        // Slice off first row: [B, H, posLen, Tq]
+        const sliced = reshaped.slice(s(0, 0, 1, 0), s(B, H, posLen + 1, Tq));
+        // Reshape back: [B, H, Tq, posLen]
+        return sliced.reshape(s(B, H, Tq, posLen));
+    }
+    forward(q, k, v, posEmb, mask, cache) {
+        if (posEmb === null)
+            throw new Error('posEmb is required for RelPositionMultiHeadAttention');
+        const Qp = this.linearQ.forward(q);
+        const Kp = this.linearK.forward(k);
+        const Vp = this.linearV.forward(v);
+        const P = this.linearPos.forward(posEmb);
+        const qShape = Qp.shape();
+        const batch = Number(qShape[0]);
+        const qSeq = Number(qShape[1]);
+        const kSeq = Number(Kp.shape()[1]);
+        const pBatch = Number(P.shape()[0]);
+        const posLen = Number(P.shape()[1]);
+        // posEmb broadcast if batch > 1
+        const Pb = pBatch === 1 && batch > 1
+            ? P.broadcastTo(s(batch, posLen, Number(P.shape()[2])))
+            : P;
+        const Qr = Qp.reshape(s(batch, qSeq, this.nHead, this.headDim));
+        const Qu = Qr.add(this.posBiasU).transpose(new Int32Array([0, 2, 1, 3]));
+        const Qv = Qr.add(this.posBiasV).transpose(new Int32Array([0, 2, 1, 3]));
+        let Kr = Kp.reshape(s(batch, kSeq, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        let Vr = Vp.reshape(s(batch, kSeq, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        const Pr = Pb.reshape(s(batch, posLen, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        if (cache) {
+            [Kr, Vr] = cache.updateAndFetchKV(Kr, Vr);
+        }
+        const kLen = Number(Kr.shape()[2]);
+        // Content-based scores: [batch, heads, qSeq, kLen]
+        const matrixAC = Qu.matmul(Kr.transpose(new Int32Array([0, 1, 3, 2]))).mulScalar(this.scale);
+        // Position-based scores: [batch, heads, qSeq, posLen] → rel_shift → [batch, heads, qSeq, kLen]
+        let matrixBD = Qv.matmul(Pr.transpose(new Int32Array([0, 1, 3, 2])));
+        matrixBD = this.relShift(matrixBD);
+        matrixBD = matrixBD
+            .slice(s(0, 0, 0, 0), s(batch, this.nHead, qSeq, kLen))
+            .mulScalar(this.scale);
+        let scores = matrixAC.add(matrixBD);
+        if (mask !== null) {
+            scores = scores.add(mask);
+        }
+        const attn = softmax(scores, 3);
+        const o = attn.matmul(Vr);
+        const out = o.transpose(new Int32Array([0, 2, 1, 3])).reshape(s(batch, qSeq, -1 /* nHead*headDim */));
+        // Note: -1 not supported by reshape; use actual size:
+        const outReshaped = o.transpose(new Int32Array([0, 2, 1, 3])).reshape(s(batch, qSeq, this.nHead * this.headDim));
+        return this.linearOut.forward(outReshaped);
+    }
+    loadWeights(weights, prefix) {
+        this.linearQ.loadWeights(weights, `${prefix}.linear_q`);
+        this.linearK.loadWeights(weights, `${prefix}.linear_k`);
+        this.linearV.loadWeights(weights, `${prefix}.linear_v`);
+        this.linearOut.loadWeights(weights, `${prefix}.linear_out`);
+        this.linearPos.loadWeights(weights, `${prefix}.linear_pos`);
+        const entry = weights.get(`${prefix}.pos_bias_u`);
+        if (entry) {
+            this.posBiasU = MxArray.fromFloat32(entry.data, BigInt64Array.from(entry.shape.map(BigInt)));
+        }
+        else {
+            this.posBiasU = MxArray.zeros(s(this.nHead, this.headDim), null);
+        }
+        const entryV = weights.get(`${prefix}.pos_bias_v`);
+        if (entryV) {
+            this.posBiasV = MxArray.fromFloat32(entryV.data, BigInt64Array.from(entryV.shape.map(BigInt)));
+        }
+        else {
+            this.posBiasV = MxArray.zeros(s(this.nHead, this.headDim), null);
+        }
+    }
+}
+// ---------------------------------------------------------------------------
+// Local relative-position attention
+// Uses standard attention with a local window mask (no custom Metal kernel)
+// ---------------------------------------------------------------------------
+export class RelPositionMultiHeadLocalAttention extends RelPositionMultiHeadAttention {
+    contextSize;
+    constructor(nHead, nFeat, bias = true, contextSize = [256, 256]) {
+        super(nHead, nFeat, bias);
+        this.contextSize = contextSize;
+    }
+    forward(q, k, v, posEmb, mask, cache) {
+        if (posEmb === null)
+            throw new Error('posEmb is required');
+        const Qp = this.linearQ.forward(q);
+        const Kp = this.linearK.forward(k);
+        const Vp = this.linearV.forward(v);
+        const P = this.linearPos.forward(posEmb);
+        const qShape = Qp.shape();
+        const batch = Number(qShape[0]);
+        const qSeq = Number(qShape[1]);
+        const kSeq = Number(Kp.shape()[1]);
+        const posLen = Number(P.shape()[1]);
+        const Pb = Number(P.shape()[0]) === 1 && batch > 1
+            ? P.broadcastTo(s(batch, posLen, Number(P.shape()[2])))
+            : P;
+        const Qr = Qp.reshape(s(batch, qSeq, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        let Kr = Kp.reshape(s(batch, kSeq, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        let Vr = Vp.reshape(s(batch, kSeq, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        const Pr = Pb.reshape(s(batch, posLen, this.nHead, this.headDim)).transpose(new Int32Array([0, 2, 1, 3]));
+        if (cache) {
+            [Kr, Vr] = cache.updateAndFetchKV(Kr, Vr);
+        }
+        const kLen = Number(Kr.shape()[2]);
+        const [leftCtx, rightCtx] = this.contextSize;
+        // Content scores
+        const Qu = Qr.add(this.posBiasU.expandDims(1));
+        const Qv = Qr.add(this.posBiasV.expandDims(1));
+        const matrixAC = Qu.matmul(Kr.transpose(new Int32Array([0, 1, 3, 2]))).mulScalar(this.scale);
+        // Position scores — simplified: use relative positions from posEmb
+        const matrixBD = Qv.matmul(Pr.transpose(new Int32Array([0, 1, 3, 2]))).mulScalar(this.scale);
+        let scores = matrixAC; // start with content scores
+        // Build local attention mask — positions outside [t-left, t+right] get -inf
+        // We build a [1, 1, qSeq, kLen] mask
+        const maskData = new Float32Array(qSeq * kLen).fill(-Infinity);
+        for (let qi = 0; qi < qSeq; qi++) {
+            // Align q to the right side of k (k might be longer due to cache)
+            const kOffset = kLen - qSeq;
+            const ki_start = Math.max(0, kOffset + qi - leftCtx);
+            const ki_end = Math.min(kLen, kOffset + qi + rightCtx + 1);
+            for (let ki = ki_start; ki < ki_end; ki++) {
+                maskData[qi * kLen + ki] = 0.0;
+            }
+        }
+        const localMask = MxArray.fromFloat32(maskData, s(1, 1, qSeq, kLen));
+        scores = scores.add(localMask);
+        if (mask !== null) {
+            scores = scores.add(mask);
+        }
+        // Add position bias where in range
+        scores = scores.add(matrixBD);
+        const attn = softmax(scores, 3);
+        const o = attn.matmul(Vr);
+        const outReshaped = o.transpose(new Int32Array([0, 2, 1, 3])).reshape(s(batch, qSeq, this.nHead * this.headDim));
+        return this.linearOut.forward(outReshaped);
+    }
+}
+// ---------------------------------------------------------------------------
+// Positional encodings
+// ---------------------------------------------------------------------------
+export class RelPositionalEncoding extends Module {
+    dModel;
+    maxLen;
+    scaleInput;
+    scale;
+    pe; // [1, 2*maxLen-1, dModel]
+    constructor(dModel, maxLen = 5000, scaleInput = true) {
+        super();
+        this.dModel = dModel;
+        this.maxLen = maxLen;
+        this.scaleInput = scaleInput;
+        this.scale = scaleInput ? Math.sqrt(dModel) : 1.0;
+        this.calculatePE();
+    }
+    calculatePE() {
+        const totalLen = 2 * this.maxLen - 1;
+        const pe = new Float32Array(totalLen * this.dModel);
+        for (let i = 0; i < totalLen; i++) {
+            const pos = this.maxLen - 1 - i; // from (maxLen-1) down to -(maxLen-1)
+            for (let j = 0; j < this.dModel; j += 2) {
+                const expFactor = Math.exp((-j * Math.log(10000.0)) / this.dModel);
+                pe[i * this.dModel + j] = Math.sin(pos * expFactor);
+                if (j + 1 < this.dModel) {
+                    pe[i * this.dModel + j + 1] = Math.cos(pos * expFactor);
+                }
+            }
+        }
+        this.pe = MxArray.fromFloat32(pe, BigInt64Array.from([1n, BigInt(totalLen), BigInt(this.dModel)]));
+    }
+    forward(x, offset = 0) {
+        const inputLen = Number(x.shape()[1]) + offset;
+        if (inputLen > this.maxLen) {
+            this.maxLen = inputLen + 1;
+            this.calculatePE();
+        }
+        const scaledX = x.mulScalar(this.scale);
+        const bufferLen = Number(this.pe.shape()[1]);
+        const startIdx = Math.floor(bufferLen / 2) - (inputLen - 1);
+        const endIdx = Math.floor(bufferLen / 2) + (inputLen - 1) + 1;
+        const posEmb = this.pe.slice(s(0, startIdx, 0), s(1, endIdx, this.dModel));
+        return [scaledX, posEmb];
+    }
+    loadWeights(_weights, _prefix) {
+        // PE is computed, not learned (though pos_bias_u/v in attention are learned)
+    }
+}
+export class LocalRelPositionalEncoding extends Module {
+    dModel;
+    maxLen;
+    scaleInput;
+    scale;
+    leftContext;
+    rightContext;
+    pe; // [1, leftCtx+rightCtx+1, dModel]
+    constructor(dModel, maxLen = 5000, scaleInput = true, contextSize = [256, 256]) {
+        super();
+        this.dModel = dModel;
+        this.maxLen = maxLen;
+        this.scaleInput = scaleInput;
+        this.scale = scaleInput ? Math.sqrt(dModel) : 1.0;
+        [this.leftContext, this.rightContext] = contextSize;
+        this.calculatePE();
+    }
+    calculatePE() {
+        const totalLen = this.leftContext + this.rightContext + 1;
+        const pe = new Float32Array(totalLen * this.dModel);
+        for (let i = 0; i < totalLen; i++) {
+            const pos = this.leftContext - i; // from leftCtx down to -rightCtx
+            for (let j = 0; j < this.dModel; j += 2) {
+                const expFactor = Math.exp((-j * Math.log(10000.0)) / this.dModel);
+                pe[i * this.dModel + j] = Math.sin(pos * expFactor);
+                if (j + 1 < this.dModel) {
+                    pe[i * this.dModel + j + 1] = Math.cos(pos * expFactor);
+                }
+            }
+        }
+        this.pe = MxArray.fromFloat32(pe, BigInt64Array.from([1n, BigInt(totalLen), BigInt(this.dModel)]));
+    }
+    forward(x, _offset = 0) {
+        const scaledX = x.mulScalar(this.scale);
+        return [scaledX, this.pe];
+    }
+    loadWeights(_weights, _prefix) { }
+}
+//# sourceMappingURL=attention.js.map