chromadb-default-embed
Version:
Chroma's fork of @xenova/transformers serving as our default embedding function
1,158 lines (976 loc) • 36.9 kB
JavaScript
/**
* @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,
transpose_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) {
if (args[0] instanceof ONNXTensor) {
// Create shallow copy
Object.assign(this, args[0]);
} else {
// Create new tensor
Object.assign(this, new ONNXTensor(
/** @type {DataType} */(args[0]),
/** @type {Exclude<import('./maths.js').AnyTypedArray, Uint8ClampedArray>} */(args[1]),
args[2]
));
}
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 transposed version of this Tensor, according to the provided dimensions.
* @param {...number} dims Dimensions to transpose.
* @returns {Tensor} The transposed tensor.
*/
transpose(...dims) {
return transpose(this, dims);
}
// TODO: rename transpose to permute
// TODO: implement transpose
// 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];
}
/**
* Transposes a tensor according to the provided axes.
* @param {any} tensor The input tensor to transpose.
* @param {Array} axes The axes to transpose the tensor along.
* @returns {Tensor} The transposed tensor.
*/
export function transpose(tensor, axes) {
const [transposedData, shape] = transpose_data(tensor.data, tensor.dims, axes);
return new Tensor(tensor.type, transposedData, 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
)
}
/**
* 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);
}