~/jendhub

def mom_l

bendtorch@1.3.0.0 · main.bend:246

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

246# buf <- buf*mu + g; w <- w + (-lr)*buf, as PyTorch's SGD with momentum (no dampening, not Nesterov)
247def mom_l(nleft: Nat, +t: U32, +lr: F32, +mu: F32, rw: Array<F32> & F32, rg: Array<F32> & F32, rm: Array<F32> & F32) -> Q3:
248  match nleft rw rg rm:
249    case 0n Tuple{w, x} Tuple{g, y} Tuple{m, b}:
250      Q3{w, g, m}
251    case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +b}:
252      +nb = ((b * mu : F32) + y : F32)
253      +u = (t + 1 : U32)
254      mom_l(p, u, lr, mu, Array.get(F32, Array.set(F32, w, t, (x + (F32.neg(lr) * nb : F32) : F32)), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(F32, Array.set(F32, m, t, nb), u))

Uses (13)

Used by (2)