123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209
import {  Array1D,  Array2D,  Array3D,  Array4D,  Array5D,  Array6D,  Rank,  Shape,  Shape1D,  Shape2D,  Shape3D,  Shape4D,  Shape5D,  Shape6D,} from "../api/shape.ts";import { inferShape, length } from "./util.ts";
/** * A generic N-dimensional tensor. */export class Tensor<R extends Rank> {  shape: Shape[R];  data: Float32Array;
  constructor(data: Float32Array, shape: Shape[R]) {    this.shape = shape;    this.data = data;  }
  /**   * Creates an empty tensor.   */  static zeroes<R extends Rank>(shape: Shape[R]): Tensor<R> {    return new Tensor(new Float32Array(length(shape)), shape);  }
  /**   * Serialise a tensor into JSON.   */  toJSON() {    const data = new Array(this.data.length).fill(1);    this.data.forEach((value, i) => data[i] = value);    return { data, shape: this.shape };  }}
/** * Create an nth rank tensor from the given nthD array and shape. * ```ts * tensor([1, 2, 3, 4], [2, 2]); * ``` */export function tensor<R extends Rank>(values: Float32Array, shape: Shape[R]) {  return new Tensor(values, shape);}
/** * Create a 1D tensor from the given 1D array. * * ```ts * tensor1D([1, 2, 3, 4]); * ``` */export function tensor1D(values: Array1D) {  const shape = inferShape(values) as Shape1D;  return new Tensor(new Float32Array(values), shape);}
/** * Create a 2D tensor from the given 2D array. * * ```ts * tensor2D([ *  [1, 2, 3, 4], *  [5, 6, 7, 8], * ]); * ``` */export function tensor2D(values: Array2D) {  const shape = inferShape(values) as Shape2D;  return new Tensor(new Float32Array(values.flat(1)), shape);}
/** * Create a 3D tensor from the given 3D array. * * ```ts * tensor3D([ *  [ *    [1, 2, 3, 4], *    [5, 6, 7, 8], *  ], *  [ *    [1, 2, 3, 4], *    [5, 6, 7, 8], *  ], * ]); * ``` */export function tensor3D(values: Array3D) {  const shape = inferShape(values) as Shape3D;  return new Tensor(new Float32Array(values.flat(2)), shape);}
/** * Create a 4D tensor from the given 4D array. * * ```ts * tensor4D([ *  [ *    [ *      [1, 2, 3], *      [4, 5, 6], *    ], *    [ *      [1, 2, 3], *      [4, 5, 6], *    ], *  ], *  [ *    [ *      [1, 2, 3], *      [4, 5, 6], *    ], *    [ *      [1, 2, 3], *      [4, 5, 6], *    ] *  ], * ]); * ``` */export function tensor4D(values: Array4D) {  const shape = inferShape(values) as Shape4D;  return new Tensor(new Float32Array(values.flat(3)), shape);}
/** * Create a 5D tensor from the given 5D array. * * ```ts * tensor5D([ *   [ *     [ *       [ *         [1, 2, 3], *         [4, 5, 6], *       ], *       [ *         [1, 2, 3], *         [4, 5, 6], *       ], *     ], *     [ *       [ *         [1, 2, 3], *         [4, 5, 6], *       ], *       [ *         [1, 2, 3], *         [4, 5, 6], *       ], *     ], *   ], * ]); * ``` */export function tensor5D(values: Array5D) {  const shape = inferShape(values) as Shape5D;  return new Tensor(new Float32Array(values.flat(4)), shape);}
/** * Create a 6D tensor from the given 6D array. * ```ts * tensor6D([ *   [ *     [ *       [ *         [ *           [1, 2, 3], *           [4, 5, 6], *         ], *         [ *           [1, 2, 3], *           [4, 5, 6], *         ], *       ], *       [ *         [ *           [1, 2, 3], *           [4, 5, 6], *         ], *         [ *           [1, 2, 3], *           [4, 5, 6], *         ], *       ], *     ] *   ] * ]); * ``` */export function tensor6D(values: Array6D) {  const shape = inferShape(values) as Shape6D;  return new Tensor(new Float32Array(values.flat(5)), shape);}