idea.ts

let implementMe: any, math: any;

const compute = Symbol("compute");
const source = Symbol("source");

interface Computable {
  [compute](): Promise<TensorBuffer>;
}

class TensorBuffer implements Computable {
  asArrayBuffer(): ArrayBuffer { return implementMe; }
  async [compute]() { return this; }
}

class AddOp implements Computable {
  inputs: Tensor[];

  constructor(...inputs: Tensor[]) {
    this.inputs = inputs;
  }

  async [compute]() {
    const tbufs: TensorBuffer[] = await Promise.all(
        this.inputs.map(t => t[compute]()));
    return math.add(...tbufs);
  }
}

class MatMulOp implements Computable {
  constructor(private input1: Tensor, private input2: Tensor) {}

  async [compute]() {
    const [tb1, tb2]: TensorBuffer[] = await Promise.all([
      this.input1[compute](),
      this.input2[compute]()
    ]);
    return math.matmul(tb1, tb2);
  }
}

class Tensor {
  [source]: Computable;

  constructor(value: Computable) {
    this[source] = value;
  }

  add(...addend: Tensor[]): Tensor {
    return new Tensor(new AddOp(this, ...addend));
  }

  matmul(that: Tensor): Tensor {
    return new Tensor(new MatMulOp(this, that));
  }

  async data(): Promise<ArrayBuffer> {
    return (await this[compute]()).asArrayBuffer();
  }

  async [compute](): Promise<TensorBuffer> {
    if (!(this[source] instanceof TensorBuffer)) {
      this[source] = await this[source][compute]();
    }
    return this[source] as TensorBuffer;
  }
}