law both_sum

bendtorch@1.3.0.0 · LAWS.bend:61

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

61# LAW (both_sum): T.Both's own code, running two linear layers side by side on the two inputs of a pair, gives a backward
62# pass that is the adjoint of its forward pass: dy1·y1 + dy2·y2 == dx1·x1 + dx2·x2. Each layer gets its own input's
63# gradient and nothing crosses between the two sides. Attention applies RoPE to q and k this way (Both(RoPE, Both(RoPE,
64# Id))).
65law both_sum:
66  for +w1: List<&2, List<&2, Nat>>
67  for +w2: List<&2, List<&2, Nat>>
68  for +x1: List<&2, Nat>
69  for +x2: List<&2, Nat>
70  for +dy1: List<&2, Nat>
71  for +dy2: List<&2, Nat>
72  for g1: List<&2, List<&2, Nat>>
73  for g2: List<&2, List<&2, Nat>>
74  {Spec.pdot(T.Two{dy1, dy2}, Spec.fy(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Layer.fwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.BothL(), T.Two{w1, w2}, T.Two{x1, x2}))) == Spec.pdot(Spec.bdx(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Layer.bwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.BothL(), T.Two{w1, w2}, Spec.fsv(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Layer.fwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.BothL(), T.Two{w1, w2}, T.Two{x1, x2})), T.Two{dy1, dy2}, T.Two{g1, g2})), T.Two{x1, x2}) : Nat}