def gnb_sums

bendtorch@1.3.0.0 · main.bend:1945

import Base

1945# the backward pass is LayerNorm's with the block in place of the row: the two means (of dxh and of dxh·xh) are taken
1946# over the whole block, then each cell's dx, dg[ch] += dy·xh and db[ch] += dy, in order
1947def gnb_sums(left: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +mu: F32, +r: F32, q: R3SS) -> R3SS:
1948  match left q:
1949    case 0n t:
1950      t
1951    case 1n+l R3SS{x, dy, g, +sa, +sb}:
1952      gnb_sums(l, (off + c : U32), j, c, cg, mu, r, lnb_sums(Nat.sub(cg, 1n), j, (off + c : U32), mu, r, sa, sb, Array.get(F32, x, ((off + c : U32) + j : U32)), Array.get(F32, dy, ((off + c : U32) + j : U32)), Array.get(F32, g, j)))