LAWS.kv_cols: law of bendtorch@1.3.0.0
Import: import bendtorch@1.3.0.0/LAWS.bend as LAWS
Lines 248-260 of bendtorch@1.3.0.0/LAWS.bend: https://jend.leeeet.dev/read.txt?path=bendtorch@1.3.0.0/LAWS.bend&start_line=248&end_line=260

import Base
import bendtorch@1.3.0.0/main.bend as T

# LAW (kv_cols): the attention code takes the columns each query head reads from T.head_cols, in both passes (the head
# loops pass no head index further, so they cannot compute other columns). With g = 1 + i + k query heads per key/value
# head, query head j·g + i reads the key and value columns of head j: j·hd onwards.
law kv_cols:
  for +t: Nat
  for +hd: Nat
  for +w: Nat
  for +kvw: Nat
  for +sc: F32
  for +j: Nat
  for +i: Nat
  for +k: Nat
  {T.hc_kc(T.head_cols(T.Dm{t, 1n+Nat.add(i, k), hd, w, kvw, sc}, Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i))) == U32.from_nat(Nat.mul(j, hd)) : U32}

Uses (12):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Equal.cong
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat.add
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat.mul
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32.from_nat
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/LAWS.bend&name=kv_of_group
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=Dm
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=hc_kc
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=head_cols
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=kv_of