@thi.ng/tensors/tensor.js

0D/1D/2D/3D/4D tensors with extensible polymorphic operations and customizable storage

import { swizzle } from "@thi.ng/arrays/swizzle";
import { isNumber } from "@thi.ng/checks/is-number";
import { equiv, equivArrayLike } from "@thi.ng/equiv";
import { illegalArgs } from "@thi.ng/errors/illegal-arguments";
import { outOfBounds } from "@thi.ng/errors/out-of-bounds";
import { unsupported } from "@thi.ng/errors/unsupported";
import { dot2, dot3, dot4 } from "@thi.ng/vectors/dot";
import { eqDeltaS as _eqDelta } from "@thi.ng/vectors/eqdelta";
import { product, product2, product3, product4 } from "@thi.ng/vectors/product";
import { illegalShape } from "./errors.js";
import { format } from "./format.js";
import { STORAGE } from "./storage.js";
const { abs, ceil, min } = Math;
class ATensor {
  constructor(type, storage, data, shape, stride, offset = 0) {
    this.type = type;
    this.storage = storage;
    this.data = data;
    this.shape = shape;
    this.stride = stride;
    this.offset = offset;
  }
  get order() {
    return strideOrder(this.stride);
  }
  get orderedShape() {
    return swizzle(this.order)(this.shape);
  }
  get orderedStride() {
    return swizzle(this.order)(this.stride);
  }
  broadcast(shape, stride) {
    return new TENSOR_IMPLS[shape.length](
      this.type,
      this.storage,
      this.data,
      shape,
      stride,
      this.offset
    );
  }
  copy() {
    return new this.constructor(
      this.type,
      this.storage,
      this.data,
      this.shape.slice(),
      this.stride.slice(),
      this.offset
    );
  }
  empty(storage = this.storage) {
    return new this.constructor(
      this.type,
      storage,
      storage.alloc(this.length),
      this.shape.slice(),
      shapeToStride(this.shape)
    );
  }
  /**
   * Calls {@link ITensorStorage.release} with this tensor's data.
   */
  release() {
    return this.storage.release(this.data);
  }
  equiv(o) {
    return this === o || o instanceof ATensor && equiv(this.shape, o.shape) && equivArrayLike([...this], [...o]);
  }
  eqDelta(o, eps = 1e-6) {
    return this === o || equiv(this.shape, o.shape) && _eqDelta([...this], [...o], this.length, eps);
  }
  position(index) {
    const { order, stride } = this;
    index |= 0;
    index -= this.offset;
    const idx = order.map((o) => {
      const i = ~~(index / stride[o]);
      index -= i * stride[o];
      return i;
    });
    return swizzle(order)(idx);
  }
  hi(pos) {
    return new this.constructor(
      this.type,
      this.storage,
      this.data,
      __hi(pos, this),
      this.stride,
      this.offset
    );
  }
  lo(pos) {
    const { shape, offset } = __lo(pos, this);
    return new this.constructor(
      this.type,
      this.storage,
      this.data,
      shape,
      this.stride,
      offset
    );
  }
  crop(pos, size) {
    const { shape, offset } = __crop(pos, size, this);
    return new this.constructor(
      this.type,
      this.storage,
      this.data,
      shape,
      this.stride,
      offset
    );
  }
  step(select) {
    const { shape, stride, offset } = __step(select, this);
    return new this.constructor(
      this.type,
      this.storage,
      this.data,
      shape,
      stride,
      offset
    );
  }
  pick(select) {
    const { shape, stride, offset } = __pick(select, this);
    return tensor(this.type, shape, {
      data: this.data,
      storage: this.storage,
      copy: false,
      stride,
      offset
    });
  }
  pack(storage = this.storage) {
    return new this.constructor(
      this.type,
      storage,
      storage.from(this),
      this.shape.slice(),
      shapeToStride(this.shape)
    );
  }
  reshape(newShape, newStride) {
    const newLength = product(newShape);
    if (newLength !== this.length) illegalShape(newShape);
    return tensor(this.type, newShape, {
      storage: this.storage,
      data: this.data,
      copy: false,
      stride: newStride ?? shapeToStride(newShape),
      offset: this.offset
    });
  }
  transpose(order) {
    const reorder = swizzle(order);
    return new this.constructor(
      this.type,
      this.storage,
      this.data,
      reorder(this.shape),
      reorder(this.stride),
      this.offset
    );
  }
  toJSON() {
    return {
      buf: [...this],
      shape: this.shape,
      stride: shapeToStride(this.shape)
    };
  }
}
class Tensor0 extends ATensor {
  *[Symbol.iterator]() {
    yield this.data[this.offset];
  }
  get dim() {
    return 0;
  }
  get order() {
    return [0];
  }
  get length() {
    return 1;
  }
  index() {
    return this.offset;
  }
  position() {
    return [0];
  }
  get() {
    return this.data[this.offset];
  }
  set(_, v) {
    this.data[this.offset] = v;
    return this;
  }
  pick([x]) {
    if (x !== 0) outOfBounds(x);
    return new Tensor0(
      this.type,
      this.storage,
      this.data,
      [],
      [],
      this.offset
    );
  }
  resize(newShape, fill, storage = this.storage) {
    const newLength = product(newShape);
    const newData = storage.alloc(newLength);
    if (fill !== void 0) newData.fill(fill);
    newData[0] = this.get();
    return tensor(this.type, newShape, {
      storage,
      data: newData,
      copy: false
    });
  }
  transpose(_) {
    return unsupported();
  }
  toString() {
    return format(this.get());
  }
}
class Tensor1 extends ATensor {
  *[Symbol.iterator]() {
    let {
      data,
      shape: [sx],
      stride: [tx],
      offset
    } = this;
    for (; sx-- > 0; offset += tx) yield data[offset];
  }
  get dim() {
    return 1;
  }
  get order() {
    return [0];
  }
  get length() {
    return this.shape[0];
  }
  index([x]) {
    return this.offset + x * this.stride[0];
  }
  position(index) {
    return [~~(((index | 0) - this.offset) / this.stride[0])];
  }
  get([x]) {
    return this.data[this.offset + x * this.stride[0]];
  }
  set([x], v) {
    this.data[this.offset + x * this.stride[0]] = v;
    return this;
  }
  pick([x]) {
    if (x < 0 && x >= this.length) outOfBounds(x);
    return new Tensor0(
      this.type,
      this.storage,
      this.data,
      [],
      [],
      this.offset + x * this.stride[0]
    );
  }
  resize(newShape, fill, storage = this.storage) {
    const newLength = product(newShape);
    const newData = storage.alloc(newLength);
    if (fill !== void 0) newData.fill(fill);
    const {
      data,
      shape: [sx],
      stride: [tx]
    } = this;
    const n = min(sx, newLength);
    for (let i = this.offset, ii = 0, x = 0; x < sx && ii < n; x++, i += tx, ii++) {
      newData[ii] = data[i];
    }
    return tensor(this.type, newShape, {
      storage,
      data: newData,
      copy: false
    });
  }
  transpose(_) {
    return unsupported();
  }
  toString() {
    const res = [];
    for (let x of this) res.push(format(x));
    return res.join(" ");
  }
}
class Tensor2 extends ATensor {
  _n;
  *[Symbol.iterator]() {
    const {
      data,
      shape: [sx, sy],
      stride: [tx, ty]
    } = this;
    let ox, x, y;
    for (ox = this.offset, x = 0; x < sx; x++, ox += tx) {
      for (y = 0; y < sy; y++) {
        yield data[ox + y * ty];
      }
    }
  }
  get length() {
    return this._n || (this._n = product2(this.shape));
  }
  get dim() {
    return 2;
  }
  get order() {
    return abs(this.stride[1]) > abs(this.stride[0]) ? [1, 0] : [0, 1];
  }
  index(pos) {
    return this.offset + dot2(pos, this.stride);
  }
  get(pos) {
    return this.data[this.offset + dot2(pos, this.stride)];
  }
  set(pos, v) {
    this.data[this.offset + dot2(pos, this.stride)] = v;
    return this;
  }
  resize(newShape, fill, storage = this.storage) {
    const newLength = product(newShape);
    const newData = storage.alloc(newLength);
    if (fill !== void 0) newData.fill(fill);
    const {
      data,
      shape: [sx, sy],
      stride: [tx, ty]
    } = this;
    const n = min(this.length, newLength);
    let ox, x, y, i;
    for (ox = this.offset, i = 0, x = 0; x < sx; x++, ox += tx) {
      for (y = 0; y < sy && i < n; y++, i++) {
        newData[i] = data[ox + y * ty];
      }
    }
    return tensor(this.type, newShape, {
      storage,
      data: newData,
      copy: false
    });
  }
  toString() {
    const res = [];
    for (let i = 0; i < this.shape[0]; i++) {
      res.push(this.pick([i]).toString());
    }
    return res.join("\n");
  }
}
class Tensor3 extends ATensor {
  _n;
  *[Symbol.iterator]() {
    const {
      data,
      shape: [sx, sy, sz],
      stride: [tx, ty, tz]
    } = this;
    let ox, oy, x, y, z;
    for (ox = this.offset, x = 0; x < sx; x++, ox += tx) {
      for (oy = ox, y = 0; y < sy; y++, oy += ty) {
        for (z = 0; z < sz; z++) {
          yield data[oy + z * tz];
        }
      }
    }
  }
  get length() {
    return this._n || (this._n = product3(this.shape));
  }
  get dim() {
    return 3;
  }
  index(pos) {
    return this.offset + dot3(pos, this.stride);
  }
  get(pos) {
    return this.data[this.offset + dot3(pos, this.stride)];
  }
  set(pos, v) {
    this.data[this.offset + dot3(pos, this.stride)] = v;
    return this;
  }
  resize(newShape, fill, storage = this.storage) {
    const newLength = product(newShape);
    const newData = storage.alloc(newLength);
    if (fill !== void 0) newData.fill(fill);
    const {
      data,
      shape: [sx, sy, sz],
      stride: [tx, ty, tz]
    } = this;
    const n = min(this.length, newLength);
    let ox, oy, x, y, z, i;
    for (ox = this.offset, i = 0, x = 0; x < sx; x++, ox += tx) {
      for (oy = ox, y = 0; y < sy; y++, oy += ty) {
        for (z = 0; z < sz && i < n; z++, i++) {
          newData[i] = data[oy + z * tz];
        }
      }
    }
    return tensor(this.type, newShape, {
      storage,
      data: newData,
      copy: false
    });
  }
  toString() {
    const res = [];
    for (let i = 0; i < this.shape[0]; i++) {
      res.push(`--- ${i}: ---`, this.pick([i]).toString());
    }
    return res.join("\n");
  }
}
class Tensor4 extends ATensor {
  _n;
  *[Symbol.iterator]() {
    const {
      data,
      shape: [sx, sy, sz, sw],
      stride: [tx, ty, tz, tw],
      offset
    } = this;
    let ox, oy, oz, x, y, z, w;
    for (ox = offset, x = 0; x < sx; x++, ox += tx) {
      for (oy = ox, y = 0; y < sy; y++, oy += ty) {
        for (oz = oy, z = 0; z < sz; z++, oz += tz) {
          for (w = 0; w < sw; w++) {
            yield data[oz + w * tw];
          }
        }
      }
    }
  }
  get length() {
    return this._n || (this._n = product4(this.shape));
  }
  get dim() {
    return 4;
  }
  index(pos) {
    return this.offset + dot4(pos, this.stride);
  }
  get(pos) {
    return this.data[this.offset + dot4(pos, this.stride)];
  }
  set(pos, v) {
    this.data[this.offset + dot4(pos, this.stride)] = v;
    return this;
  }
  resize(newShape, fill, storage = this.storage) {
    const newLength = product(newShape);
    const newData = storage.alloc(newLength);
    if (fill !== void 0) newData.fill(fill);
    const {
      data,
      shape: [sx, sy, sz, sw],
      stride: [tx, ty, tz, tw]
    } = this;
    const n = min(this.length, newLength);
    let ox, oy, oz, x, y, z, w, i;
    for (ox = this.offset, i = 0, x = 0; x < sx; x++, ox += tx) {
      for (oy = ox, y = 0; y < sy; y++, oy += ty) {
        for (oz = oy, z = 0; z < sz; z++, oz += tz) {
          for (w = 0; w < sw && i < n; w++, i++) {
            newData[i] = data[oz + w * tw];
          }
        }
      }
    }
    return tensor(this.type, newShape, {
      storage,
      data: newData,
      copy: false
    });
  }
  toString() {
    const res = [];
    for (let i = 0; i < this.shape[0]; i++) {
      res.push(`--- cube ${i}: ---`, this.pick([i]).toString());
    }
    return res.join("\n");
  }
}
const TENSOR_IMPLS = [
  Tensor0,
  Tensor1,
  Tensor2,
  Tensor3,
  Tensor4
];
function tensor(...args) {
  if (Array.isArray(args[0])) return tensorFromArray(args[0], args[1]);
  if (args.length === 1) return __tensor0(args[0]);
  const type = args[0];
  const shape = args[1];
  const dim = shape.length;
  const opts = args[2];
  const storage = opts?.storage ?? STORAGE[type];
  const stride = opts?.stride ?? shapeToStride(shape);
  const offset = opts?.offset ?? computeOffset(shape, stride);
  let data;
  if (opts?.data) {
    if (opts?.copy === false) data = opts.data;
    else data = storage.from(opts.data);
  } else {
    data = storage.alloc(dim > 0 ? product(shape) : 1);
  }
  const ctor = TENSOR_IMPLS[dim];
  return ctor ? new ctor(type, storage, data, shape, stride, offset) : unsupported(`unsupported dimension: ${dim}`);
}
const __tensor0 = (x) => {
  const type = isNumber(x) ? "num" : "str";
  const storage = STORAGE[type];
  const data = storage.alloc(1);
  data[0] = x;
  return new Tensor0(type, storage, data, [], []);
};
const tensorImpl = (shape) => shape.length === 1 && shape[0] === 1 ? Tensor0 : TENSOR_IMPLS[shape.length];
function tensorFromArray(data, opts) {
  const shape = [data.length];
  let $data = data;
  while (Array.isArray($data[0])) {
    shape.push($data[0].length);
    $data = $data.flat();
  }
  const $type = opts?.type ?? (isNumber($data[0]) ? "num" : "str");
  if ($type === "str" && isNumber($data[0]))
    illegalArgs("mismatched data type");
  return tensor($type, shape, {
    data: $data,
    copy: $type !== "num" && $type !== "str",
    storage: opts?.storage
  });
}
const zeroes = (shape, type = "num", storage) => tensor(type, shape, { storage });
const ones = (shape, type = "num", storage) => constant(shape, 1, type, storage);
const constant = (shape, value, type, storage) => {
  const res = tensor(type, shape, { storage });
  res.data.fill(value);
  return res;
};
const shapeToStride = (shape) => {
  const n = shape.length;
  const stride = new Array(n);
  for (let i = n, s = 1; i-- > 0; s *= shape[i]) {
    stride[i] = s;
  }
  return stride;
};
const strideOrder = (strides) => strides.map((x, i) => [x, i]).sort((a, b) => abs(b[0]) - abs(a[0])).map((x) => x[1]);
const computeOffset = (shape, stride) => {
  let offset = 0;
  for (let i = 0; i < shape.length; i++) {
    if (stride[i] < 0) {
      offset -= (shape[i] - 1) * stride[i];
    }
  }
  return offset;
};
const __lo = (select, { shape, stride, offset }) => {
  const newShape = [];
  for (let i = 0, n = shape.length; i < n; i++) {
    const x = select[i];
    if (x > shape[i]) illegalShape(select);
    newShape.push(
      x >= 0 ? (offset += stride[i] * x, shape[i] - x) : shape[i]
    );
  }
  return { shape: newShape, offset };
};
const __hi = (select, { shape }) => {
  const newShape = [];
  for (let i = 0, n = shape.length; i < n; i++) {
    const x = select[i];
    if (x > shape[i]) illegalShape(select);
    newShape.push(x > 0 ? x : shape[i]);
  }
  return newShape;
};
const __crop = (lo, hi, src) => {
  const res = __lo(lo, src);
  const shape = __hi(hi, res);
  return { ...res, shape };
};
const __step = (select, { shape, stride, offset }) => {
  const newShape = shape.slice();
  const newStride = stride.slice();
  for (let i = 0, n = shape.length; i < n; i++) {
    const x = select[i];
    if (x) {
      if (x < 0) {
        offset += stride[i] * (shape[i] - 1);
        newShape[i] = ceil(-shape[i] / x);
      } else {
        newShape[i] = ceil(shape[i] / x);
      }
      newStride[i] *= x;
    }
  }
  return { shape: newShape, stride: newStride, offset };
};
const __pick = (select, { shape, stride, offset }) => {
  const newShape = [];
  const newStride = [];
  for (let i = 0, n = shape.length; i < n; i++) {
    const x = select[i];
    if (x >= 0) {
      offset += stride[i] * x;
    } else {
      newShape.push(shape[i]);
      newStride.push(stride[i]);
    }
  }
  return { shape: newShape, stride: newStride, offset };
};
export {
  ATensor,
  TENSOR_IMPLS,
  Tensor0,
  Tensor1,
  Tensor2,
  Tensor3,
  Tensor4,
  computeOffset,
  constant,
  ones,
  shapeToStride,
  strideOrder,
  tensor,
  tensorFromArray,
  tensorImpl,
  zeroes
};