def chunks

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

import Base

243# n plain heads (the values, which RoPE does not touch)
244def chunks(n: Nat, +w: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>:
245  match n:
246    case 0n:
247      Nil{}
248    case 1n+p:
249      List.take(&2, F32, xs, w) <> chunks(p, w, List.drop(&2, F32, xs, w))