BendMlLlama.groups: def of bend-ml-llama@0.2.1.0
Import: import bend-ml-llama@0.2.1.0/main.bend as BendMlLlama
Lines 589-598 of bend-ml-llama@0.2.1.0/main.bend: https://jend.leeeet.dev/read.txt?path=bend-ml-llama@0.2.1.0/main.bend&start_line=589&end_line=598

import Base

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

Uses (5):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=List
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=List.drop
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=List.take
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Head