def adam_loop

import bendtorch-gpu@1.3.0.0/main.bend as BendtorchGpu

import Base
import bend-ml-tensor-gpu@0.1.5.0/main.bend as TG
import bendtorch@1.3.0.0/main.bend as T

235def adam_loop(k: Nat, +i: U32, +l: 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<F32> & F32, rv: Array<F32> & F32) -> TG.O4:
236  match k rw rg rm rv:
237    case 0n Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}: TG.O4{w, g, m, v}
238    case 1n+q Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}:
239      +gg = T.adam_grad(dec, wd, x, y)
240      +nm = T.adam_m(c1, a, gg)
241      +nv = T.adam_v(b2, c2, b, gg)
242      +nw = T.adam_w(step, bc2, eps, T.adam_decay(dec, lr, wd, x), nm, nv)
243      +u = (i + l : U32)
244      adam_loop(q, u, l, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, i, nw), u), Array.get(F32, Array.set(F32, g, i, 0.0), u), Array.get(F32, Array.set(F32, m, i, nm), u), Array.get(F32, Array.set(F32, v, i, nv), u))

Uses (16)

Used by (1)

bendtorch-gpu@1.3.0.0/main:lf_adam