def adamh_l

bendtorch@1.3.0.0 · main.bend:352

import Base

352# m and v packed. rm and rv hold the word of element t (index t >> 1), read after the previous element's write.
353def adamh_l(nleft: Nat, +t: U32, +lr: F32, +c1: F32, +b2: F32, +c2: F32, +eps: F32, +wd: F32, +dec: Bool, +step: F32, +bc2: F32, rw: Array<F32> & F32, rg: Array<F32> & F32, rm: Array<U32> & U32, rv: Array<U32> & U32) -> PQ:
354  match nleft rw rg rm rv:
355    case 0n Tuple{w, x} Tuple{g, y} Tuple{m, a} Tuple{v, b}:
356      PQ{w, g, m, v}
357    case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}:
358      +odd = U32.is_eq((t .&. 1 : U32), 1)
359      +gg = adam_grad(dec, wd, x, y)
360      +hm = ph_of(adam_m(c1, ph_to(ph_half(odd, a)), gg))
361      +hv = ph_of(adam_v(b2, c2, ph_to(ph_half(odd, b)), gg))
362      +nw = adam_w(step, bc2, eps, adam_decay(dec, lr, wd, x), ph_to(hm), ph_to(hv))
363      +u = (t + 1 : U32)
364      +k = U32.shrn(t, 1n)
365      +ku = U32.shrn(u, 1n)
366      adamh_l(p, u, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, t, nw), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(U32, Array.set(U32, m, k, ph_put(odd, a, hm)), ku), Array.get(U32, Array.set(U32, v, k, ph_put(odd, b, hv)), ku))