synor/sdk/js/src/tensor.ts

/**
 * Tensor utilities for Synor Compute SDK
 */

import type { TensorData, TensorShape, Precision } from './types';

/**
 * Tensor wrapper for compute operations.
 */
export class Tensor {
  readonly data: Float32Array | Float64Array | Int32Array | Int8Array | Uint8Array;
  readonly shape: number[];
  readonly dtype: Precision;

  private constructor(
    data: Float32Array | Float64Array | Int32Array | Int8Array | Uint8Array,
    shape: number[],
    dtype: Precision
  ) {
    this.data = data;
    this.shape = shape;
    this.dtype = dtype;
  }

  /**
   * Create a tensor from various input formats.
   */
  static from(data: TensorData, dtype: Precision = 'fp32'): Tensor {
    if (data instanceof Float32Array) {
      return new Tensor(data, [data.length], 'fp32');
    }
    if (data instanceof Float64Array) {
      return new Tensor(data, [data.length], 'fp64');
    }
    if (data instanceof Int32Array) {
      return new Tensor(data, [data.length], 'int8');
    }
    if (data instanceof Int8Array) {
      return new Tensor(data, [data.length], 'int8');
    }
    if (data instanceof Uint8Array) {
      return new Tensor(data, [data.length], 'int8');
    }

    // Handle nested arrays
    const shape = Tensor.inferShape(data);
    const flat = Tensor.flatten(data);

    const typedArray = dtype === 'fp64'
      ? new Float64Array(flat)
      : dtype === 'fp32'
        ? new Float32Array(flat)
        : dtype === 'int8' || dtype === 'int4'
          ? new Int8Array(flat)
          : new Float32Array(flat);

    return new Tensor(typedArray, shape, dtype);
  }

  /**
   * Create a tensor of zeros.
   */
  static zeros(shape: number[], dtype: Precision = 'fp32'): Tensor {
    const size = shape.reduce((a, b) => a * b, 1);
    const data = dtype === 'fp64'
      ? new Float64Array(size)
      : dtype === 'fp32'
        ? new Float32Array(size)
        : new Int8Array(size);
    return new Tensor(data, shape, dtype);
  }

  /**
   * Create a tensor of ones.
   */
  static ones(shape: number[], dtype: Precision = 'fp32'): Tensor {
    const size = shape.reduce((a, b) => a * b, 1);
    const data = dtype === 'fp64'
      ? new Float64Array(size).fill(1)
      : dtype === 'fp32'
        ? new Float32Array(size).fill(1)
        : new Int8Array(size).fill(1);
    return new Tensor(data, shape, dtype);
  }

  /**
   * Create a tensor with random values.
   */
  static random(shape: number[], dtype: Precision = 'fp32'): Tensor {
    const size = shape.reduce((a, b) => a * b, 1);
    const data = dtype === 'fp64'
      ? Float64Array.from({ length: size }, () => Math.random())
      : dtype === 'fp32'
        ? Float32Array.from({ length: size }, () => Math.random())
        : Int8Array.from({ length: size }, () => Math.floor(Math.random() * 256) - 128);
    return new Tensor(data, shape, dtype);
  }

  /**
   * Create a tensor with normal distribution.
   */
  static randn(shape: number[], dtype: Precision = 'fp32'): Tensor {
    const size = shape.reduce((a, b) => a * b, 1);
    const values: number[] = [];

    for (let i = 0; i < size; i++) {
      // Box-Muller transform
      const u1 = Math.random();
      const u2 = Math.random();
      values.push(Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2));
    }

    const data = dtype === 'fp64'
      ? new Float64Array(values)
      : dtype === 'fp32'
        ? new Float32Array(values)
        : new Int8Array(values.map(v => Math.max(-128, Math.min(127, Math.round(v * 50)))));

    return new Tensor(data, shape, dtype);
  }

  /**
   * Get tensor size (total elements).
   */
  get size(): number {
    return this.shape.reduce((a, b) => a * b, 1);
  }

  /**
   * Get number of dimensions.
   */
  get ndim(): number {
    return this.shape.length;
  }

  /**
   * Get byte size.
   */
  get byteSize(): number {
    const bytesPerElement = this.dtype === 'fp64' ? 8
      : this.dtype === 'fp32' ? 4
        : this.dtype === 'fp16' || this.dtype === 'bf16' ? 2
          : 1;
    return this.size * bytesPerElement;
  }

  /**
   * Reshape tensor.
   */
  reshape(newShape: number[]): Tensor {
    const newSize = newShape.reduce((a, b) => a * b, 1);
    if (newSize !== this.size) {
      throw new Error(`Cannot reshape tensor of size ${this.size} to shape [${newShape}]`);
    }
    return new Tensor(this.data, newShape, this.dtype);
  }

  /**
   * Convert to different precision.
   */
  to(dtype: Precision): Tensor {
    if (dtype === this.dtype) return this;

    const newData = dtype === 'fp64'
      ? Float64Array.from(this.data)
      : dtype === 'fp32'
        ? Float32Array.from(this.data)
        : Int8Array.from(this.data);

    return new Tensor(newData, this.shape, dtype);
  }

  /**
   * Get value at index.
   */
  get(...indices: number[]): number {
    const idx = this.flatIndex(indices);
    return this.data[idx];
  }

  /**
   * Set value at index.
   */
  set(value: number, ...indices: number[]): void {
    const idx = this.flatIndex(indices);
    this.data[idx] = value;
  }

  /**
   * Convert to nested array.
   */
  toArray(): number[] | number[][] | number[][][] {
    if (this.ndim === 1) {
      return Array.from(this.data);
    }

    const result: any[] = [];
    const innerSize = this.shape.slice(1).reduce((a, b) => a * b, 1);

    for (let i = 0; i < this.shape[0]; i++) {
      const slice = new Tensor(
        this.data.slice(i * innerSize, (i + 1) * innerSize) as Float32Array,
        this.shape.slice(1),
        this.dtype
      );
      result.push(slice.toArray());
    }

    return result;
  }

  /**
   * Serialize tensor for transmission.
   */
  serialize(): { data: string; shape: number[]; dtype: Precision } {
    // Convert to base64 for efficient transmission
    const bytes = new Uint8Array(this.data.buffer);
    const base64 = btoa(String.fromCharCode(...bytes));

    return {
      data: base64,
      shape: this.shape,
      dtype: this.dtype,
    };
  }

  /**
   * Deserialize tensor from transmission format.
   */
  static deserialize(serialized: { data: string; shape: number[]; dtype: Precision }): Tensor {
    const binary = atob(serialized.data);
    const bytes = new Uint8Array(binary.length);
    for (let i = 0; i < binary.length; i++) {
      bytes[i] = binary.charCodeAt(i);
    }

    const data = serialized.dtype === 'fp64'
      ? new Float64Array(bytes.buffer)
      : serialized.dtype === 'fp32'
        ? new Float32Array(bytes.buffer)
        : new Int8Array(bytes.buffer);

    return new Tensor(data, serialized.shape, serialized.dtype);
  }

  /**
   * Calculate flat index from multi-dimensional indices.
   */
  private flatIndex(indices: number[]): number {
    if (indices.length !== this.ndim) {
      throw new Error(`Expected ${this.ndim} indices, got ${indices.length}`);
    }

    let idx = 0;
    let stride = 1;
    for (let i = this.ndim - 1; i >= 0; i--) {
      if (indices[i] < 0 || indices[i] >= this.shape[i]) {
        throw new Error(`Index ${indices[i]} out of bounds for axis ${i} with size ${this.shape[i]}`);
      }
      idx += indices[i] * stride;
      stride *= this.shape[i];
    }
    return idx;
  }

  /**
   * Infer shape from nested array.
   */
  private static inferShape(data: number[] | number[][]): number[] {
    if (!Array.isArray(data)) return [];
    if (!Array.isArray(data[0])) return [data.length];
    return [data.length, ...Tensor.inferShape(data[0] as number[] | number[][])];
  }

  /**
   * Flatten nested array.
   */
  private static flatten(data: number[] | number[][]): number[] {
    if (!Array.isArray(data)) return [data as unknown as number];
    return (data as any[]).flatMap(item =>
      Array.isArray(item) ? Tensor.flatten(item) : item
    );
  }
}