add an example of make_axes to readme

stanford-crfm · Sep 5, 2024 · 019fa2d · 019fa2d
1 parent 754c2bb
commit 019fa2d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -47,6 +47,9 @@ Head = hax.Axis("head", 8)  # number of attention heads
 Key = hax.Axis("key", 64)  # key size
 Embed = hax.Axis("embed", 512)  # embedding size
 
+# alternatively:
+#Pos, KPos, Head, Key, Embed = hax.make_axes(pos=1024, key_pos=1024, head=8, key=64, embed=512)
+
 
 def attention_scores(Key, KPos, query, key, mask):
     # how similar is each query to each key