@ohos-ports/sillytavern-transformers 2.17.2-beta.1

package/src/utils/tensor.js

@@ -0,0 +1,1250 @@
+/**
+ * @file Helper module for `Tensor` processing.
+ * 
+ * These functions and classes are only used internally, 
+ * meaning an end-user shouldn't need to access anything here.
+ * 
+ * @module utils/tensor
+ */
+
+import { ONNX } from '../backends/onnx.js';
+
+import {
+    interpolate_data,
+    permute_data
+} from './maths.js';
+
+
+const DataTypeMap = Object.freeze({
+    float32: Float32Array,
+    float64: Float64Array,
+    string: Array, // string[]
+    int8: Int8Array,
+    uint8: Uint8Array,
+    int16: Int16Array,
+    uint16: Uint16Array,
+    int32: Int32Array,
+    uint32: Uint32Array,
+    int64: BigInt64Array,
+    uint64: BigUint64Array,
+    bool: Uint8Array,
+});
+
+/**
+ * @typedef {keyof typeof DataTypeMap} DataType
+ * @typedef {import('./maths.js').AnyTypedArray | any[]} DataArray
+ */
+
+const ONNXTensor = ONNX.Tensor;
+
+export class Tensor {
+    /** @type {number[]} Dimensions of the tensor. */
+    dims;
+
+    /** @type {DataType} Type of the tensor. */
+    type;
+
+    /** @type {DataArray} The data stored in the tensor. */
+    data;
+
+    /** @type {number} The number of elements in the tensor. */
+    size;
+
+    /**
+     * Create a new Tensor or copy an existing Tensor.
+     * @param {[DataType, DataArray, number[]]|[import('onnxruntime-common').Tensor]} args
+     */
+    constructor(...args) {
+        // NOTE: onnxruntime >= 1.17 (including 1.27 used here) exposes `data` and
+        // `location` as prototype getters backed by internal fields (`cpuData`,
+        // `dataLocation`), so `Object.assign(this, tensor)` no longer copies `data`
+        // and the native binding rejects session inputs lacking a string `location`.
+        // Read the properties explicitly via getters - this works with both ORT 1.14
+        // (own properties) and ORT >= 1.17 (getters).
+        if (!(args[0] instanceof ONNXTensor)) {
+            // Create new tensor
+            args = [new ONNXTensor(
+                /** @type {DataType} */(args[0]),
+                /** @type {Exclude<import('./maths.js').AnyTypedArray, Uint8ClampedArray>} */(args[1]),
+                args[2]
+            )];
+        }
+        // Create shallow copy of the ONNX tensor
+        const ortTensor = args[0];
+        Object.assign(this, {
+            type: ortTensor.type,
+            data: ortTensor.data,
+            dims: ortTensor.dims,
+            size: ortTensor.size,
+            location: ortTensor.location ?? 'cpu',
+        });
+
+        return new Proxy(this, {
+            get: (obj, key) => {
+                if (typeof key === 'string') {
+                    let index = Number(key);
+                    if (Number.isInteger(index)) {
+                        // key is an integer (i.e., index)
+                        return obj._getitem(index);
+                    }
+                }
+                // @ts-ignore
+                return obj[key];
+            },
+            set: (obj, key, value) => {
+                // TODO allow setting of data
+
+                // @ts-ignore
+                return obj[key] = value;
+            }
+        });
+    }
+
+    /**
+     * Returns an iterator object for iterating over the tensor data in row-major order.
+     * If the tensor has more than one dimension, the iterator will yield subarrays.
+     * @returns {Iterator} An iterator object for iterating over the tensor data in row-major order.
+     */
+    *[Symbol.iterator]() {
+        const [iterLength, ...iterDims] = this.dims;
+
+        if (iterDims.length > 0) {
+            const iterSize = iterDims.reduce((a, b) => a * b);
+            for (let i = 0; i < iterLength; ++i) {
+                yield this._subarray(i, iterSize, iterDims);
+            }
+        } else {
+            yield* this.data
+        }
+
+    }
+
+    /**
+     * Index into a Tensor object.
+     * @param {number} index The index to access.
+     * @returns {Tensor} The data at the specified index.
+     */
+    _getitem(index) {
+        const [iterLength, ...iterDims] = this.dims;
+
+        index = safeIndex(index, iterLength);
+
+        if (iterDims.length > 0) {
+            const iterSize = iterDims.reduce((a, b) => a * b);
+            return this._subarray(index, iterSize, iterDims);
+        } else {
+            return new Tensor(this.type, [this.data[index]], iterDims);
+        }
+    }
+
+    /**
+     * @param {number|bigint} item The item to search for in the tensor
+     * @returns {number} The index of the first occurrence of item in the tensor data.
+     */
+    indexOf(item) {
+        for (let index = 0; index < this.data.length; ++index) {
+            // Note: == instead of === so we can match Ints with BigInts
+            if (this.data[index] == item) {
+                return index;
+            }
+        }
+        return -1;
+    }
+
+    /**
+     * @param {number} index 
+     * @param {number} iterSize 
+     * @param {any} iterDims 
+     * @returns {Tensor}
+     */
+    _subarray(index, iterSize, iterDims) {
+        const o1 = index * iterSize;
+        const o2 = (index + 1) * iterSize;
+
+        // We use subarray if available (typed array), otherwise we use slice (normal array)
+        const data =
+            ('subarray' in this.data)
+                ? this.data.subarray(o1, o2)
+                : this.data.slice(o1, o2);
+        return new Tensor(this.type, data, iterDims);
+    }
+
+    /**
+     * Returns the value of this tensor as a standard JavaScript Number. This only works
+     * for tensors with one element. For other cases, see `Tensor.tolist()`.
+     * @returns {number|bigint} The value of this tensor as a standard JavaScript Number.
+     * @throws {Error} If the tensor has more than one element.
+     */
+    item() {
+        if (this.data.length !== 1) {
+            throw new Error(`a Tensor with ${this.data.length} elements cannot be converted to Scalar`);
+        }
+        return this.data[0];
+    }
+
+    /**
+     * Convert tensor data to a n-dimensional JS list
+     * @returns {Array}
+     */
+    tolist() {
+        return reshape(this.data, this.dims)
+    }
+
+    /**
+     * Return a new Tensor with the sigmoid function applied to each element.
+     * @returns {Tensor} The tensor with the sigmoid function applied.
+     */
+    sigmoid() {
+        return this.clone().sigmoid_();
+    }
+
+    /**
+     * Applies the sigmoid function to the tensor in place.
+     * @returns {Tensor} Returns `this`.
+     */
+    sigmoid_() {
+        for (let i = 0; i < this.data.length; ++i) {
+            this.data[i] = 1 / (1 + Math.exp(-this.data[i]));
+        }
+        return this;
+    }
+
+    /**
+     * Return a new Tensor with every element multiplied by a constant.
+     * @param {number} val The value to multiply by.
+     * @returns {Tensor} The new tensor.
+     */
+    mul(val) {
+        return this.clone().mul_(val);
+    }
+
+    /**
+     * Multiply the tensor by a constant in place.
+     * @param {number} val The value to multiply by.
+     * @returns {Tensor} Returns `this`.
+     */
+    mul_(val) {
+        for (let i = 0; i < this.data.length; ++i) {
+            this.data[i] *= val;
+        }
+        return this;
+    }
+
+
+    /**
+     * Return a new Tensor with every element added by a constant.
+     * @param {number} val The value to add by.
+     * @returns {Tensor} The new tensor.
+     */
+    add(val) {
+        return this.clone().add_(val);
+    }
+
+    /**
+     * Add the tensor by a constant in place.
+     * @param {number} val The value to add by.
+     * @returns {Tensor} Returns `this`.
+     */
+    add_(val) {
+        for (let i = 0; i < this.data.length; ++i) {
+            this.data[i] += val;
+        }
+        return this;
+    }
+    clone() {
+        return new Tensor(this.type, this.data.slice(), this.dims.slice());
+    }
+
+    slice(...slices) {
+        // This allows for slicing with ranges and numbers
+        let newTensorDims = [];
+        let newOffsets = [];
+
+        // slices is an array of numbers or arrays of numbers
+        // e.g., slices = [0, [1, 3], null, [0, 3]]
+        for (let sliceIndex = 0; sliceIndex < this.dims.length; ++sliceIndex) {
+            let slice = slices[sliceIndex];
+
+            if (slice === null || slice === undefined) {
+                // null or undefined means take the whole dimension
+                newOffsets.push([0, this.dims[sliceIndex]]);
+                newTensorDims.push(this.dims[sliceIndex]);
+
+            } else if (typeof slice === 'number') {
+                slice = safeIndex(slice, this.dims[sliceIndex], sliceIndex);
+
+                // A number means take a single element
+                newOffsets.push([slice, slice + 1]);
+
+            } else if (Array.isArray(slice) && slice.length === 2) {
+                // An array of length 2 means take a range of elements
+
+                if (slice[0] > slice[1]) {
+                    throw new Error(`Invalid slice: ${slice}`);
+                }
+
+                let offsets = [
+                    Math.max(slice[0], 0),
+                    Math.min(slice[1], this.dims[sliceIndex])
+                ];
+
+                newOffsets.push(offsets);
+                newTensorDims.push(offsets[1] - offsets[0]);
+
+            } else {
+                throw new Error(`Invalid slice: ${slice}`);
+            }
+        }
+
+        let newDims = newOffsets.map(([start, end]) => end - start);
+        let newBufferSize = newDims.reduce((a, b) => a * b);
+
+        // Allocate memory
+        // @ts-ignore
+        let data = new this.data.constructor(newBufferSize);
+
+        // Precompute strides
+        const stride = this.stride();
+
+        for (let i = 0; i < newBufferSize; ++i) {
+            let originalIndex = 0;
+            for (let j = newDims.length - 1, num = i; j >= 0; --j) {
+                const size = newDims[j];
+                originalIndex += ((num % size) + newOffsets[j][0]) * stride[j];
+                num = Math.floor(num / size);
+            }
+            data[i] = this.data[originalIndex];
+        }
+        return new Tensor(this.type, data, newTensorDims);
+
+    }
+
+    /**
+     * Return a permuted version of this Tensor, according to the provided dimensions.
+     * @param  {...number} dims Dimensions to permute.
+     * @returns {Tensor} The permuted tensor.
+     */
+    permute(...dims) {
+        return permute(this, dims);
+    }
+
+    // TODO: implement transpose. For now (backwards compatibility), it's just an alias for permute()
+    transpose(...dims) {
+        return this.permute(...dims);
+    }
+
+    // TODO add .max() and .min() methods
+
+    /**
+     * Returns the sum of each row of the input tensor in the given dimension dim.
+     * 
+     * @param {number} [dim=null] The dimension or dimensions to reduce. If `null`, all dimensions are reduced.
+     * @param {boolean} keepdim Whether the output tensor has `dim` retained or not.
+     * @returns The summed tensor
+     */
+    sum(dim = null, keepdim = false) {
+        return this.norm(1, dim, keepdim);
+    }
+
+    /**
+     * Returns the matrix norm or vector norm of a given tensor.
+     * @param {number|string} [p='fro'] The order of norm
+     * @param {number} [dim=null] Specifies which dimension of the tensor to calculate the norm across.
+     * If dim is None, the norm will be calculated across all dimensions of input.
+     * @param {boolean} [keepdim=false] Whether the output tensors have dim retained or not.
+     * @returns {Tensor} The norm of the tensor.
+     */
+    norm(p = 'fro', dim = null, keepdim = false) {
+        if (p === 'fro') {
+            // NOTE: Since we only support integer dims, Frobenius norm produces the same result as p=2.
+            p = 2;
+        } else if (typeof p === 'string') {
+            throw Error(`Unsupported norm: ${p}`);
+        }
+
+        if (dim === null) {
+            // @ts-ignore
+            let val = this.data.reduce((a, b) => a + (b ** p), 0) ** (1 / p);
+            return new Tensor(this.type, [val], []);
+        }
+
+        // Negative indexing
+        dim = safeIndex(dim, this.dims.length);
+
+        // Calculate the shape of the resulting array after summation
+        const resultDims = this.dims.slice(); // Copy the original dimensions
+        resultDims[dim] = 1; // Remove the specified axis
+
+        // Create a new array to store the accumulated values
+        // @ts-ignore
+        const result = new this.data.constructor(this.data.length / this.dims[dim]);
+
+        // Iterate over the data array
+        for (let i = 0; i < this.data.length; ++i) {
+
+            // Calculate the index in the resulting array
+            let resultIndex = 0;
+
+            for (let j = this.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
+                const size = this.dims[j];
+                if (j !== dim) {
+                    const index = num % size;
+                    resultIndex += index * resultMultiplier;
+                    resultMultiplier *= resultDims[j];
+                }
+                num = Math.floor(num / size);
+            }
+
+            // Accumulate the value at the current index
+            result[resultIndex] += (this.data[i]) ** p;
+        }
+
+        if (p !== 1) {
+            for (let i = 0; i < result.length; ++i) {
+                result[i] = result[i] ** (1 / p);
+            }
+        }
+
+        if (!keepdim) {
+            resultDims.splice(dim, 1);
+        }
+
+        return new Tensor(this.type, result, resultDims);
+    }
+
+    /**
+     * Performs `L_p` normalization of inputs over specified dimension. Operates in place.
+     * @param {number} [p=2] The exponent value in the norm formulation
+     * @param {number} [dim=1] The dimension to reduce
+     * @returns {Tensor} `this` for operation chaining.
+     */
+    normalize_(p = 2.0, dim = 1) {
+        dim = safeIndex(dim, this.dims.length);
+
+        const norm = this.norm(p, dim, true);
+
+        for (let i = 0; i < this.data.length; ++i) {
+
+            // Calculate the index in the resulting array
+            let resultIndex = 0;
+
+            for (let j = this.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
+                const size = this.dims[j];
+                if (j !== dim) {
+                    const index = num % size;
+                    resultIndex += index * resultMultiplier;
+                    resultMultiplier *= this.dims[j];
+                }
+                num = Math.floor(num / size);
+            }
+
+            // Divide by normalized value
+            this.data[i] /= norm.data[resultIndex];
+        }
+
+        return this;
+    }
+
+    /**
+     * Performs `L_p` normalization of inputs over specified dimension.
+     * @param {number} [p=2] The exponent value in the norm formulation
+     * @param {number} [dim=1] The dimension to reduce
+     * @returns {Tensor} The normalized tensor.
+     */
+    normalize(p = 2.0, dim = 1) {
+        return this.clone().normalize_(p, dim);
+    }
+
+    /**
+     * Compute and return the stride of this tensor.
+     * Stride is the jump necessary to go from one element to the next one in the specified dimension dim.
+     * @returns {number[]} The stride of this tensor.
+     */
+    stride() {
+        return dimsToStride(this.dims);
+    }
+
+    /**
+     * Returns a tensor with all specified dimensions of input of size 1 removed.
+     * 
+     * NOTE: The returned tensor shares the storage with the input tensor, so changing the contents of one will change the contents of the other.
+     * If you would like a copy, use `tensor.clone()` before squeezing.
+     * 
+     * @param {number} [dim=null] If given, the input will be squeezed only in the specified dimensions.
+     * @returns The squeezed tensor
+     */
+    squeeze(dim = null) {
+        return new Tensor(
+            this.type,
+            this.data,
+            calc_squeeze_dims(this.dims, dim)
+        )
+    }
+
+    /**
+     * In-place version of @see {@link Tensor.squeeze}
+     */
+    squeeze_(dim = null) {
+        this.dims = calc_squeeze_dims(this.dims, dim);
+        return this;
+    }
+
+    /**
+     * Returns a new tensor with a dimension of size one inserted at the specified position.
+     * 
+     * NOTE: The returned tensor shares the same underlying data with this tensor.
+     * 
+     * @param {number} dim The index at which to insert the singleton dimension
+     * @returns The unsqueezed tensor
+     */
+    unsqueeze(dim = null) {
+        return new Tensor(
+            this.type,
+            this.data,
+            calc_unsqueeze_dims(this.dims, dim)
+        );
+    }
+
+    /**
+     * In-place version of @see {@link Tensor.unsqueeze}
+     */
+    unsqueeze_(dim = null) {
+        this.dims = calc_unsqueeze_dims(this.dims, dim);
+        return this;
+    }
+
+    /**
+     * In-place version of @see {@link Tensor.flatten}
+     */
+    flatten_(start_dim = 0, end_dim = -1) {
+        // TODO validate inputs
+        end_dim = (end_dim + this.dims.length) % this.dims.length;
+
+        let dimsToKeepBefore = this.dims.slice(0, start_dim);
+        let dimsToFlatten = this.dims.slice(start_dim, end_dim + 1);
+        let dimsToKeepAfter = this.dims.slice(end_dim + 1);
+
+        this.dims = [...dimsToKeepBefore, dimsToFlatten.reduce((a, b) => a * b, 1), ...dimsToKeepAfter]
+        return this;
+    }
+
+    /**
+     * Flattens input by reshaping it into a one-dimensional tensor.
+     * If `start_dim` or `end_dim` are passed, only dimensions starting with `start_dim`
+     * and ending with `end_dim` are flattened. The order of elements in input is unchanged.
+     * @param {number} start_dim the first dim to flatten
+     * @param {number} end_dim the last dim to flatten
+     * @returns The flattened tensor.
+     */
+    flatten(start_dim = 0, end_dim = -1) {
+        return this.clone().flatten_(start_dim, end_dim);
+    }
+
+    /**
+     * Returns a new tensor with the same data as the `self` tensor but of a different `shape`.
+     * @param  {...number} dims the desired size
+     * @returns {Tensor} The tensor with the same data but different shape
+     */
+    view(...dims) {
+        // TODO: validate dims
+        let inferredIndex = -1;
+        for (let i = 0; i < dims.length; ++i) {
+            if (dims[i] === -1) {
+                if (inferredIndex !== -1) {
+                    throw new Error("Only one dimension can be inferred");
+                }
+                inferredIndex = i;
+            }
+        }
+
+        if (inferredIndex !== -1) {
+            // Some dimension must be inferred
+            const productOther = dims.reduce((product, curr, index) => {
+                return index !== inferredIndex ? product * curr : product
+            }, 1);
+
+            dims[inferredIndex] = this.data.length / productOther;
+        }
+        return new Tensor(this.type, this.data, dims); // NOTE: uses same underlying storage
+    }
+
+    neg_() {
+        for (let i = 0; i < this.data.length; ++i) {
+            this.data[i] = -this.data[i];
+        }
+        return this;
+    }
+    neg() {
+        return this.clone().neg_();
+    }
+
+    /**
+     * In-place version of @see {@link Tensor.clamp}
+     */
+    clamp_(min, max) {
+        for (let i = 0; i < this.data.length; ++i) {
+            this.data[i] = Math.min(Math.max(this.data[i], min), max);
+        }
+        return this;
+    }
+
+    /**
+     * Clamps all elements in input into the range [ min, max ]
+     * @param {number} min lower-bound of the range to be clamped to
+     * @param {number} max upper-bound of the range to be clamped to
+     * @returns the output tensor.
+     */
+    clamp(min, max) {
+        return this.clone().clamp_(min, max);
+    }
+
+    /**
+     * In-place version of @see {@link Tensor.round}
+     */
+    round_() {
+        for (let i = 0; i < this.data.length; ++i) {
+            this.data[i] = Math.round(this.data[i]);
+        }
+        return this;
+    }
+
+    /**
+     * Rounds elements of input to the nearest integer.
+     * @returns the output tensor.
+     */
+    round() {
+        return this.clone().round_();
+    }
+
+    /**
+     * Performs Tensor dtype conversion.
+     * @param {DataType} type The desired data type.
+     * @returns {Tensor} The converted tensor.
+     */
+    to(type) {
+        // If the self Tensor already has the correct dtype, then self is returned.
+        if (this.type === type) return this;
+
+        // Otherwise, the returned tensor is a copy of self with the desired dtype.
+        if (!DataTypeMap.hasOwnProperty(type)) {
+            throw new Error(`Unsupported type: ${type}`);
+        }
+        // @ts-ignore
+        return new Tensor(type, DataTypeMap[type].from(this.data), this.dims);
+    }
+}
+
+/**
+ * This creates a nested array of a given type and depth (see examples).
+ * 
+ * @example
+ *   NestArray<string, 1>; // string[]
+ * @example
+ *   NestArray<number, 2>; // number[][]
+ * @example
+ *   NestArray<string, 3>; // string[][][] etc.
+ * @template T
+ * @template {number} Depth
+ * @template {never[]} [Acc=[]]
+ * @typedef {Acc['length'] extends Depth ? T : NestArray<T[], Depth, [...Acc, never]>} NestArray
+ */
+
+/**
+ * Reshapes a 1-dimensional array into an n-dimensional array, according to the provided dimensions.
+ *
+ * @example
+ *   reshape([10                    ], [1      ]); // Type: number[]      Value: [10]
+ *   reshape([1, 2, 3, 4            ], [2, 2   ]); // Type: number[][]    Value: [[1, 2], [3, 4]]
+ *   reshape([1, 2, 3, 4, 5, 6, 7, 8], [2, 2, 2]); // Type: number[][][]  Value: [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
+ *   reshape([1, 2, 3, 4, 5, 6, 7, 8], [4, 2   ]); // Type: number[][]    Value: [[1, 2], [3, 4], [5, 6], [7, 8]]
+ * @param {T[]|DataArray} data The input array to reshape.
+ * @param {DIM} dimensions The target shape/dimensions.
+ * @template T
+ * @template {[number]|number[]} DIM
+ * @returns {NestArray<T, DIM["length"]>} The reshaped array.
+ */
+function reshape(data, dimensions) {
+
+    const totalElements = data.length;
+    const dimensionSize = dimensions.reduce((a, b) => a * b);
+
+    if (totalElements !== dimensionSize) {
+        throw Error(`cannot reshape array of size ${totalElements} into shape (${dimensions})`);
+    }
+
+    /** @type {any} */
+    let reshapedArray = data;
+
+    for (let i = dimensions.length - 1; i >= 0; i--) {
+        reshapedArray = reshapedArray.reduce((acc, val) => {
+            let lastArray = acc[acc.length - 1];
+
+            if (lastArray.length < dimensions[i]) {
+                lastArray.push(val);
+            } else {
+                acc.push([val]);
+            }
+
+            return acc;
+        }, [[]]);
+    }
+
+    return reshapedArray[0];
+}
+
+/**
+ * Permutes a tensor according to the provided axes.
+ * @param {any} tensor The input tensor to permute.
+ * @param {Array} axes The axes to permute the tensor along.
+ * @returns {Tensor} The permuted tensor.
+ */
+export function permute(tensor, axes) {
+    const [permutedData, shape] = permute_data(tensor.data, tensor.dims, axes);
+    return new Tensor(tensor.type, permutedData, shape);
+}
+
+
+/**
+ * Interpolates an Tensor to the given size.
+ * @param {Tensor} input The input tensor to interpolate. Data must be channel-first (i.e., [c, h, w])
+ * @param {number[]} size The output size of the image
+ * @param {string} mode The interpolation mode
+ * @param {boolean} align_corners Whether to align corners.
+ * @returns {Tensor} The interpolated tensor.
+ */
+export function interpolate(input, [out_height, out_width], mode = 'bilinear', align_corners = false) {
+
+    // Input image dimensions
+    const in_channels = input.dims.at(-3) ?? 1;
+    const in_height = input.dims.at(-2);
+    const in_width = input.dims.at(-1);
+
+    let output = interpolate_data(
+        /** @type {import('./maths.js').TypedArray}*/(input.data),
+        [in_channels, in_height, in_width],
+        [out_height, out_width],
+        mode,
+        align_corners
+    );
+    return new Tensor(input.type, output, [in_channels, out_height, out_width]);
+}
+
+/**
+ * Perform mean pooling of the last hidden state followed by a normalization step.
+ * @param {Tensor} last_hidden_state Tensor of shape [batchSize, seqLength, embedDim]
+ * @param {Tensor} attention_mask Tensor of shape [batchSize, seqLength]
+ * @returns {Tensor} Returns a new Tensor of shape [batchSize, embedDim].
+ */
+export function mean_pooling(last_hidden_state, attention_mask) {
+    // last_hidden_state: [batchSize, seqLength, embedDim]
+    // attention_mask:    [batchSize, seqLength]
+
+    let shape = [last_hidden_state.dims[0], last_hidden_state.dims[2]];
+    // @ts-ignore
+    let returnedData = new last_hidden_state.data.constructor(shape[0] * shape[1]);
+    let [batchSize, seqLength, embedDim] = last_hidden_state.dims;
+
+    let outIndex = 0;
+    for (let i = 0; i < batchSize; ++i) {
+        let offset = i * embedDim * seqLength;
+
+        for (let k = 0; k < embedDim; ++k) {
+            let sum = 0;
+            let count = 0;
+
+            let attnMaskOffset = i * seqLength;
+            let offset2 = offset + k;
+            // Pool over all words in sequence
+            for (let j = 0; j < seqLength; ++j) {
+                // index into attention mask
+                let attn = Number(attention_mask.data[attnMaskOffset + j]);
+
+                count += attn;
+                sum += last_hidden_state.data[offset2 + j * embedDim] * attn;
+            }
+
+            let avg = sum / count;
+            returnedData[outIndex++] = avg;
+        }
+    }
+
+    return new Tensor(
+        last_hidden_state.type,
+        returnedData,
+        shape
+    )
+}
+
+/**
+ * Apply Layer Normalization for last certain number of dimensions.
+ * @param {Tensor} input The input tensor
+ * @param {number[]} normalized_shape input shape from an expected input of size
+ * @param {Object} options The options for the layer normalization
+ * @param {number} [options.eps=1e-5] A value added to the denominator for numerical stability.
+ * @returns {Tensor} The normalized tensor.
+ */
+export function layer_norm(input, normalized_shape, {
+    eps = 1e-5,
+} = {}) {
+    if (input.dims.length !== 2) {
+        throw new Error('`layer_norm` currently only supports 2D input.');
+    }
+
+    const [batchSize, featureDim] = input.dims;
+
+    if (normalized_shape.length !== 1 && normalized_shape[0] !== featureDim) {
+        throw new Error('`normalized_shape` must be a 1D array with shape `[input.dims[1]]`.');
+    }
+
+    const [std, mean] = std_mean(input, 1, 0, true);
+
+    // @ts-ignore
+    const returnedData = new input.data.constructor(input.data.length);
+
+    for (let i = 0; i < batchSize; ++i) {
+        const offset = i * featureDim;
+        for (let j = 0; j < featureDim; ++j) {
+            const offset2 = offset + j;
+            returnedData[offset2] = (input.data[offset2] - mean.data[i]) / (std.data[i] + eps);
+        }
+    }
+    return new Tensor(input.type, returnedData, input.dims);
+}
+
+/**
+ * Helper function to calculate new dimensions when performing a squeeze operation.
+ * @param {number[]} dims The dimensions of the tensor.
+ * @param {number|number[]|null} dim The dimension(s) to squeeze.
+ * @returns The new dimensions.
+ * @private
+ */
+function calc_squeeze_dims(dims, dim) {
+    dims = dims.slice();
+    if (dim === null) {
+        dims = dims.filter((d) => d !== 1);
+    } else if (typeof dim === 'number') {
+        if (dims[dim] === 1) {
+            dims.splice(dim, 1);
+        }
+    } else if (Array.isArray(dim)) {
+        dims = dims.filter((x, i) => {
+            return x !== 1 || !dim.includes(i);
+        });
+    }
+    return dims;
+}
+
+/**
+ * Helper function to calculate new dimensions when performing an unsqueeze operation.
+ * @param {number[]} dims The dimensions of the tensor.
+ * @param {number} dim The dimension to unsqueeze.
+ * @returns The new dimensions.
+ * @private
+ */
+function calc_unsqueeze_dims(dims, dim) {
+    // Dimension out of range (e.g., "expected to be in range of [-4, 3], but got 4")
+    // + 1 since we allow inserting at the end (i.e. dim = -1)
+    dim = safeIndex(dim, dims.length + 1);
+    dims = dims.slice();
+    // Insert 1 into specified dimension
+    dims.splice(dim, 0, 1);
+    return dims;
+}
+
+/**
+ * Safely calculate the index for an array of a given size, allowing negative indexing.
+ * @param {number} index The index that will be used.
+ * @param {number} size The size of the array.
+ * @param {number} [dimension=null] The dimension that the index is for (optional).
+ * @returns {number} The index, guaranteed to be non-negative and less than `arrayLength`.
+ * 
+ * @throws {Error} If the index is out of range.
+ * @private
+ */
+function safeIndex(index, size, dimension = null) {
+    if (index < -size || index >= size) {
+        throw new Error(`IndexError: index ${index} is out of bounds for dimension${dimension === null ? '' : ' ' + dimension} with size ${size}`);
+    }
+
+    if (index < 0) {
+        // Negative indexing, ensuring positive index
+        index = ((index % size) + size) % size;
+    }
+    return index;
+}
+
+/**
+ * Concatenates an array of tensors along a specified dimension.
+ * @param {Tensor[]} tensors The array of tensors to concatenate.
+ * @param {number} dim The dimension to concatenate along.
+ * @returns {Tensor} The concatenated tensor.
+ */
+export function cat(tensors, dim = 0) {
+    dim = safeIndex(dim, tensors[0].dims.length);
+
+    // TODO do validation of shapes
+
+    const resultDims = tensors[0].dims.slice();
+    resultDims[dim] = tensors.reduce((a, b) => a + b.dims[dim], 0);
+
+    // Create a new array to store the accumulated values
+    const resultSize = resultDims.reduce((a, b) => a * b, 1);
+    // @ts-ignore
+    const result = new tensors[0].data.constructor(resultSize);
+
+    // Create output tensor of same type as first
+    const resultType = tensors[0].type;
+
+    if (dim === 0) {
+        // Handle special case for performance reasons
+
+        let offset = 0;
+        for (let t of tensors) {
+            result.set(t.data, offset);
+            offset += t.data.length;
+        }
+
+    } else {
+
+        let currentDim = 0;
+
+        for (let t = 0; t < tensors.length; ++t) {
+            let tensor = tensors[t];
+
+            // Iterate over the data array
+            for (let i = 0; i < tensor.data.length; ++i) {
+                // Calculate the index in the resulting array
+                let resultIndex = 0;
+
+                for (let j = tensor.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
+                    const size = tensor.dims[j];
+                    let index = num % size;
+                    if (j === dim) {
+                        index += currentDim;
+                    }
+                    resultIndex += index * resultMultiplier;
+                    resultMultiplier *= resultDims[j];
+                    num = Math.floor(num / size);
+                }
+                // Accumulate the value at the current index
+                result[resultIndex] = tensor.data[i];
+            }
+
+            currentDim += tensor.dims[dim];
+        }
+    }
+    return new Tensor(resultType, result, resultDims);
+}
+
+/**
+ * Stack an array of tensors along a specified dimension.
+ * @param {Tensor[]} tensors The array of tensors to stack.
+ * @param {number} dim The dimension to stack along.
+ * @returns {Tensor} The stacked tensor.
+ */
+export function stack(tensors, dim = 0) {
+    // TODO do validation of shapes
+    // NOTE: stack expects each tensor to be equal size
+    return cat(tensors.map(t => t.unsqueeze(dim)), dim);
+}
+
+
+/**
+ * Calculates the standard deviation and mean over the dimensions specified by dim. dim can be a single dimension or `null` to reduce over all dimensions.
+ * @param {Tensor} input the input tenso
+ * @param {number|null} dim the dimension to reduce. If None, all dimensions are reduced.
+ * @param {number} correction difference between the sample size and sample degrees of freedom. Defaults to Bessel's correction, correction=1.
+ * @param {boolean} keepdim whether the output tensor has dim retained or not.
+ * @returns {Tensor[]} A tuple of (std, mean) tensors.
+ */
+export function std_mean(input, dim = null, correction = 1, keepdim = false) {
+
+    if (dim === null) {
+        // None to reduce over all dimensions.
+        // @ts-ignore
+        const sum = input.data.reduce((a, b) => a + b, 0);
+        const mean = sum / input.data.length;
+        // @ts-ignore
+        const std = Math.sqrt(input.data.reduce((a, b) => a + (b - mean) ** 2, 0) / (input.data.length - correction));
+
+        const meanTensor = new Tensor(input.type, [mean], [/* scalar */]);
+        const stdTensor = new Tensor(input.type, [std], [/* scalar */]);
+
+        return [stdTensor, meanTensor];
+    }
+
+    // Negative indexing
+    dim = safeIndex(dim, input.dims.length);
+
+    const meanTensor = mean(input, dim, keepdim);
+
+    // Calculate the shape of the resulting array after summation
+    const resultDims = input.dims.slice(); // Copy the original dimensions
+    resultDims[dim] = 1; // Remove the specified axis
+
+    // Create a new array to store the accumulated values
+    // @ts-ignore
+    const result = new input.data.constructor(input.data.length / input.dims[dim]);
+
+    // Iterate over the data array
+    for (let i = 0; i < input.data.length; ++i) {
+
+        // Calculate the index in the resulting array
+        let resultIndex = 0;
+
+        for (let j = input.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
+            const size = input.dims[j];
+            if (j !== dim) {
+                const index = num % size;
+                resultIndex += index * resultMultiplier;
+                resultMultiplier *= resultDims[j];
+            }
+            num = Math.floor(num / size);
+        }
+
+        // Accumulate the value at the current index
+        result[resultIndex] += (input.data[i] - meanTensor.data[resultIndex]) ** 2;
+    }
+
+    for (let i = 0; i < result.length; ++i) {
+        result[i] = Math.sqrt(result[i] / (input.dims[dim] - correction));
+    }
+
+    if (!keepdim) {
+        resultDims.splice(dim, 1);
+    }
+
+    const stdTensor = new Tensor(input.type, result, resultDims);
+
+    return [stdTensor, meanTensor];
+}
+
+
+/**
+ * Returns the mean value of each row of the input tensor in the given dimension dim.
+ * @param {Tensor} input the input tensor.
+ * @param {number|null} dim the dimension to reduce.
+ * @param {boolean} keepdim whether the output tensor has dim retained or not.
+ * @returns A new tensor with means taken along the specified dimension.
+ */
+export function mean(input, dim = null, keepdim = false) {
+
+    if (dim === null) {
+        // None to reduce over all dimensions.
+        // @ts-ignore
+        let val = input.data.reduce((a, b) => a + b, 0);
+        return new Tensor(input.type, [val / input.data.length], [/* scalar */]);
+    }
+
+    // Negative indexing
+    dim = safeIndex(dim, input.dims.length);
+
+    // Calculate the shape of the resulting array after summation
+    const resultDims = input.dims.slice(); // Copy the original dimensions
+    resultDims[dim] = 1; // Remove the specified axis
+
+    // Create a new array to store the accumulated values
+    // @ts-ignore
+    const result = new input.data.constructor(input.data.length / input.dims[dim]);
+
+    // Iterate over the data array
+    for (let i = 0; i < input.data.length; ++i) {
+
+        // Calculate the index in the resulting array
+        let resultIndex = 0;
+
+        for (let j = input.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
+            const size = input.dims[j];
+            if (j !== dim) {
+                const index = num % size;
+                resultIndex += index * resultMultiplier;
+                resultMultiplier *= resultDims[j];
+            }
+            num = Math.floor(num / size);
+        }
+
+        // Accumulate the value at the current index
+        result[resultIndex] += input.data[i];
+    }
+
+    if (input.dims[dim] !== 1) {
+        for (let i = 0; i < result.length; ++i) {
+            result[i] = result[i] / input.dims[dim];
+        }
+    }
+
+    if (!keepdim) {
+        resultDims.splice(dim, 1);
+    }
+
+    return new Tensor(input.type, result, resultDims);
+}
+
+
+/**
+ *
+ * Measures similarity between two temporal sequences (e.g., input audio and output tokens
+ * to generate token-level timestamps).
+ * @param {Tensor} matrix 
+ * @returns {number[][]}
+ */
+export function dynamicTimeWarping(matrix) {
+    const [output_length, input_length] = matrix.dims;
+
+    const outputShape = [output_length + 1, input_length + 1];
+
+    const cost = new Tensor(
+        'float32',
+        new Float32Array(outputShape[0] * outputShape[1]).fill(Infinity),
+        outputShape
+    );
+
+    const trace = new Tensor(
+        'float32',
+        new Float32Array(outputShape[0] * outputShape[1]).fill(-1),
+        outputShape
+    )
+
+    // same as `cost[0][0] = 0`;
+    cost[0].data[0] = 0;
+
+    for (let j = 1; j < input_length + 1; ++j) {
+        for (let i = 1; i < output_length + 1; ++i) {
+
+            const c0 = cost[i - 1][j - 1].item();
+            const c1 = cost[i - 1][j].item();
+            const c2 = cost[i][j - 1].item();
+
+            let c, t;
+            if (c0 < c1 && c0 < c2) {
+                c = c0;
+                t = 0;
+            } else if (c1 < c0 && c1 < c2) {
+                c = c1;
+                t = 1;
+            } else {
+                c = c2;
+                t = 2;
+            }
+
+            cost[i].data[j] = matrix[i - 1][j - 1].item() + c;
+            trace[i].data[j] = t;
+        }
+    }
+
+    // backtrace
+    let i = output_length;
+    let j = input_length;
+
+    // @ts-ignore
+    trace.data.fill(2, 0, outputShape[1]) // trace[0, :] = 2
+    for (let i = 0; i < outputShape[0]; ++i) { // trace[:, 0] = 1
+        trace[i].data[0] = 1;
+    }
+
+    let text_indices = [];
+    let time_indices = [];
+
+    while (i > 0 || j > 0) {
+        text_indices.push(i - 1);
+        time_indices.push(j - 1);
+
+        const t = trace[i][j].item();
+        switch (t) {
+            case 0:
+                --i; --j;
+                break;
+            case 1:
+                --i;
+                break;
+            case 2:
+                --j;
+                break;
+            default:
+                throw new Error(
+                    `Internal error in dynamic time warping. Unexpected trace[${i}, ${j}]. Please file a bug report.`
+                )
+        }
+    }
+
+    text_indices.reverse();
+    time_indices.reverse();
+
+    return [text_indices, time_indices];
+
+}
+
+function dimsToStride(dims) {
+    const stride = new Array(dims.length);
+    for (let i = dims.length - 1, s2 = 1; i >= 0; --i) {
+        stride[i] = s2;
+        s2 *= dims[i];
+    }
+    return stride;
+}
+
+/**
+ * Returns a tensor filled with the scalar value 1, with the shape defined by the variable argument size.
+ * @param {number[]} size A sequence of integers defining the shape of the output tensor.
+ */
+export function ones(size) {
+    const numElements = size.reduce((a, b) => a * b, 1);
+    return new Tensor(
+        'int64',
+        new BigInt64Array(numElements).fill(1n),
+        size
+    )
+}
+
+/**
+ * Returns a tensor filled with the scalar value 1, with the same size as input.
+ * @param {Tensor} tensor The size of input will determine size of the output tensor.
+ * @returns The ones tensor.
+ */
+export function ones_like(tensor) {
+    return ones(tensor.dims);
+}
+
+/**
+ * Quantizes the embeddings tensor to binary or unsigned binary precision.
+ * @param {Tensor} tensor The tensor to quantize.
+ * @param {'binary'|'ubinary'} precision The precision to use for quantization.
+ * @returns {Tensor} The quantized tensor.
+ */
+export function quantize_embeddings(tensor, precision) {
+    if (tensor.dims.length !== 2) {
+        throw new Error("The tensor must have 2 dimensions");
+    }
+    if (tensor.dims.at(-1) % 8 !== 0) {
+        throw new Error("The last dimension of the tensor must be a multiple of 8");
+    }
+    if (!['binary', 'ubinary'].includes(precision)) {
+        throw new Error("The precision must be either 'binary' or 'ubinary'");
+    }
+
+    const signed = precision === 'binary';
+    const dtype = signed ? 'int8' : 'uint8';
+
+    // Create a typed array to store the packed bits
+    const cls = signed ? Int8Array : Uint8Array;
+    const inputData = tensor.data;
+    const outputData = new cls(inputData.length / 8);
+
+    // Iterate over each number in the array
+    for (let i = 0; i < inputData.length; ++i) {
+        // Determine if the number is greater than 0
+        const bit = inputData[i] > 0 ? 1 : 0;
+
+        // Calculate the index in the typed array and the position within the byte
+        const arrayIndex = Math.floor(i / 8);
+        const bitPosition = i % 8;
+
+        // Pack the bit into the typed array
+        outputData[arrayIndex] |= bit << (7 - bitPosition);
+        if (signed && bitPosition === 0) {
+            outputData[arrayIndex] -= 128;
+        }
+    };
+
+    return new Tensor(dtype, outputData, [tensor.dims[0], tensor.dims[1] / 8]);
+}