Remove the heacks to handle dynamically sized tensors

newhouseb · Apr 23, 2023 · a49e51b · a49e51b
1 parent 6ff18df
commit a49e51b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 11 deletions.
diff --git a/main.ts b/main.ts
@@ -1,4 +1,4 @@
-import { addMatrix, causalMask, copy, gelu, getSlice, layerNorm, linear, mapInPlace, merge, multiplyMatrix, softmax, split, tensor, transposeMatrix, unsqueeze } from "./math";
+import { Var, addMatrix, causalMask, copy, gelu, getSlice, layerNorm, linear, mapInPlace, merge, multiplyMatrix, softmax, split, tensor, transposeMatrix, unsqueeze } from "./math";
 import * as fs from 'fs';
 import { inflate } from 'zlib';
 import { decode } from '@msgpack/msgpack';
@@ -125,10 +125,7 @@ async function main() {
 
   const gpt = await loadSmallGPT();
 
-  const inputs = tensor([gpt.SequenceLength, gpt.EmbeddingDimensions])
-
-  // Fake truncate things
-  inputs.shape = [tokens.length as any, gpt.EmbeddingDimensions]
+  const inputs = tensor([Var(tokens.length, 'Sequence Length'), gpt.EmbeddingDimensions])
 
   // Map each token into an embedding + position vector
   tokens.map((token, i) => {
@@ -183,10 +180,10 @@ async function main() {
     console.log("Splitting out k, q, and v tensors");
 
     // Next split out each of the heads
-    const kHeads = split(k, gpt.EmbeddingDimensions / gpt.AttentionHeads as 64);
-    const qHeads = split(q, gpt.EmbeddingDimensions / gpt.AttentionHeads as 64);
-    const vHeads = split(v, gpt.EmbeddingDimensions / gpt.AttentionHeads as 64);
-    const aHeads = [] as Tensor<readonly [typeof gpt.SequenceLength, 64]>[];
+    const kHeads = split(k, Var(gpt.EmbeddingDimensions / gpt.AttentionHeads, 'Head Width'));
+    const qHeads = split(q, Var(gpt.EmbeddingDimensions / gpt.AttentionHeads, 'Head Width'));
+    const vHeads = split(v, Var(gpt.EmbeddingDimensions / gpt.AttentionHeads, 'Head Width'));
+    const aHeads = [] as Tensor<readonly [Var<'Sequence Length'>, Var<'Head Width'>]>[];
 
     console.log("Performing self-attention");
 

diff --git a/math.ts b/math.ts
@@ -1,6 +1,6 @@
 // Used to define a dynamically sized (at runtime) dimension
-type Var<N extends string> = number & { label: N };
-const Var = <L extends string>(d: number, label: L) => { return d as Var<L> };
+export type Var<N extends string> = number & { label: N };
+export const Var = <L extends string>(d: number, label: L) => { return d as Var<L> };
 
 // Used to ensure that there's not ambiguity that worms its way through the type system via Union
 type UnionToIntersection<U> = (U extends any ? (k: U) => void : never) extends ((k: infer I) => void) ? I : never;