def groups

bend-ml-llama@0.2.1.0 · main.bend:589

import Base

589# =====================================================================================================
590# Grouped-query attention
591# =====================================================================================================
592# the query heads in n groups of g, in order: group j holds query heads j·g .. j·g + g - 1
593def groups(n: Nat, +g: Nat, +hs: List<&2, Head>) -> List<&2, List<&2, Head>>:
594  match n:
595    case 0n:
596      Nil{}
597    case 1n+p:
598      List.take(&2, Head, hs, g) <> groups(p, g, List.drop(&2, Head, hs, g))