Initial work for T5 model

frost-beta · Oct 3, 2024 · 1344ddc · 1344ddc
1 parent 3c3f148
commit 1344ddc
Show file tree

Hide file tree

Showing 2 changed files with 416 additions and 1 deletion.
diff --git a/src/models/llama.ts b/src/models/llama.ts
@@ -160,7 +160,7 @@ class Attention extends nn.Module {
     this.rope = initializeRoPE(args);
   }
 
-  forward(x: mx.array, mask: mx.array, cache: BaseKVCache) {
+  forward(x: mx.array, mask?: mx.array, cache?: BaseKVCache) {
     const [ B, L, D ] = x.shape;
 
     let queries = this.qProj.forward(x);