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

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

# LAW (gqa_head): with g = 1 + i + k query heads per key/value head, query head j·g + i reads key/value head
# (j·g + i) / g = j, the head of its own group: T.Attention's forward pass picks head hi / g, and its backward pass
# sums the gradients of heads j·g + 0 .. j·g + g - 1 into key/value head j, so both see the same groups.
law gqa_head:
  for +j: Nat
  for +i: Nat
  for +k: Nat
  {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}

Uses (11):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Equal.sym
  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.div
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat.divmod.go
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat.mul
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Pair
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Pair.fst
  https://jend.leeeet.dev/definition.txt?path=bend-ml-nat-lemmas@0.1.2.0/main.bend&name=add_zero
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/PROOF.bend&name=div_blocks
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/PROOF.bend&name=div_rest