onnxruntime/js/common/lib/tensor.ts
Yulong Wang ecca11340a
[js/common] allow creating (u)int64 tensors in 2 ways (#16541)
### Description
allow creating (u)int64 tensors from either a number array or a bigint
array.

before:

```js
// TypeScript think is good, but actually does not work
// runtime error: Uncaught TypeError: Cannot convert 1 to a BigInt
const myTensor1 = new Tensor('int64', [1, 2, 3, 4], [2, 2]);

// runtime good, but TypeScript thinks myTensor2 is a string tensor
const myTensor2 = new Tensor('int64', [1n, 2n, 3n, 4n], [2, 2]);
```

after:
```js
// both work at runtime and TypeScript populates the correct types
const myTensor1 = new Tensor('int64', [1, 2, 3, 4], [2, 2]);
const myTensor2 = new Tensor('int64', [1n, 2n, 3n, 4n], [2, 2]);
```
2023-07-11 21:07:36 -07:00

251 lines
8.4 KiB
TypeScript

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
import {TensorFactory} from './tensor-factory.js';
import {Tensor as TensorImpl} from './tensor-impl.js';
import {TypedTensorUtils} from './tensor-utils.js';
/* eslint-disable @typescript-eslint/no-redeclare */
/**
* represent a basic tensor with specified dimensions and data type.
*/
interface TypedTensorBase<T extends Tensor.Type> {
/**
* Get the dimensions of the tensor.
*/
readonly dims: readonly number[];
/**
* Get the data type of the tensor.
*/
readonly type: T;
/**
* Get the buffer data of the tensor.
*/
readonly data: Tensor.DataTypeMap[T];
}
export declare namespace Tensor {
interface DataTypeMap {
float32: Float32Array;
uint8: Uint8Array;
int8: Int8Array;
uint16: Uint16Array;
int16: Int16Array;
int32: Int32Array;
int64: BigInt64Array;
string: string[];
bool: Uint8Array;
float16: Uint16Array; // Keep using Uint16Array until we have a concrete solution for float 16.
float64: Float64Array;
uint32: Uint32Array;
uint64: BigUint64Array;
// complex64: never;
// complex128: never;
// bfloat16: never;
}
interface ElementTypeMap {
float32: number;
uint8: number;
int8: number;
uint16: number;
int16: number;
int32: number;
int64: bigint;
string: string;
bool: boolean;
float16: number; // Keep using Uint16Array until we have a concrete solution for float 16.
float64: number;
uint32: number;
uint64: bigint;
// complex64: never;
// complex128: never;
// bfloat16: never;
}
type DataType = DataTypeMap[Type];
type ElementType = ElementTypeMap[Type];
/**
* represent the data type of a tensor
*/
export type Type = keyof DataTypeMap;
}
/**
* Represent multi-dimensional arrays to feed to or fetch from model inferencing.
*/
export interface TypedTensor<T extends Tensor.Type> extends TypedTensorBase<T>, TypedTensorUtils<T> {}
/**
* Represent multi-dimensional arrays to feed to or fetch from model inferencing.
*/
export interface Tensor extends TypedTensorBase<Tensor.Type>, TypedTensorUtils<Tensor.Type> {}
export interface TensorConstructor {
// #region specify element type
/**
* Construct a new string tensor object from the given type, data and dims.
*
* @param type - Specify the element type.
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(type: 'string', data: Tensor.DataTypeMap['string']|readonly string[],
dims?: readonly number[]): TypedTensor<'string'>;
/**
* Construct a new bool tensor object from the given type, data and dims.
*
* @param type - Specify the element type.
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(type: 'bool', data: Tensor.DataTypeMap['bool']|readonly boolean[], dims?: readonly number[]): TypedTensor<'bool'>;
/**
* Construct a new 64-bit integer typed tensor object from the given type, data and dims.
*
* @param type - Specify the element type.
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new<T extends 'uint64'|'int64'>(
type: T, data: Tensor.DataTypeMap[T]|readonly bigint[]|readonly number[],
dims?: readonly number[]): TypedTensor<T>;
/**
* Construct a new numeric tensor object from the given type, data and dims.
*
* @param type - Specify the element type.
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new<T extends Exclude<Tensor.Type, 'string'|'bool'|'uint64'|'int64'>>(
type: T, data: Tensor.DataTypeMap[T]|readonly number[], dims?: readonly number[]): TypedTensor<T>;
// #endregion
// #region infer element types
/**
* Construct a new float32 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Float32Array, dims?: readonly number[]): TypedTensor<'float32'>;
/**
* Construct a new int8 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Int8Array, dims?: readonly number[]): TypedTensor<'int8'>;
/**
* Construct a new uint8 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Uint8Array, dims?: readonly number[]): TypedTensor<'uint8'>;
/**
* Construct a new uint16 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Uint16Array, dims?: readonly number[]): TypedTensor<'uint16'>;
/**
* Construct a new int16 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Int16Array, dims?: readonly number[]): TypedTensor<'int16'>;
/**
* Construct a new int32 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Int32Array, dims?: readonly number[]): TypedTensor<'int32'>;
/**
* Construct a new int64 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: BigInt64Array, dims?: readonly number[]): TypedTensor<'int64'>;
/**
* Construct a new string tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: readonly string[], dims?: readonly number[]): TypedTensor<'string'>;
/**
* Construct a new bool tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: readonly boolean[], dims?: readonly number[]): TypedTensor<'bool'>;
/**
* Construct a new float64 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Float64Array, dims?: readonly number[]): TypedTensor<'float64'>;
/**
* Construct a new uint32 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Uint32Array, dims?: readonly number[]): TypedTensor<'uint32'>;
/**
* Construct a new uint64 tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: BigUint64Array, dims?: readonly number[]): TypedTensor<'uint64'>;
// #endregion
// #region fall back to non-generic tensor type declaration
/**
* Construct a new tensor object from the given type, data and dims.
*
* @param type - Specify the element type.
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(type: Tensor.Type, data: Tensor.DataType|readonly number[]|readonly string[]|readonly bigint[]|readonly boolean[],
dims?: readonly number[]): Tensor;
/**
* Construct a new tensor object from the given data and dims.
*
* @param data - Specify the tensor data.
* @param dims - Specify the dimension of the tensor. If omitted, a 1-D tensor is assumed.
*/
new(data: Tensor.DataType, dims?: readonly number[]): Tensor;
// #endregion
}
// eslint-disable-next-line @typescript-eslint/naming-convention
export const Tensor = TensorImpl as (TensorConstructor & TensorFactory);