def adamh_leaf

bendtorch@1.3.0.0 · main.bend:394

import Base
import bend-ml-tensor-array@0.1.18.0/main.bend as TA

394# one matrix's step on its leaf; anything else leaves the values untouched, as Opt.mat's fallback
395def adamh_leaf(+r: Nat, +c: Nat, +o: Opt, wa: Array<F32>, ga: Array<F32>, t: PTree) -> MatStep<r, c>:
396  match o t:
397    case Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} PLeaf{+n, ma, va}:
398      ph_fin(r, c, n, adamh_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, wa, 0), Array.get(F32, ga, 0), Array.get(U32, ma, 0), Array.get(U32, va, 0)))
399    case Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} PLeafV{+n, ma, va}:
400      phv_fin(r, c, n, adamhv_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, wa, 0), Array.get(F32, ga, 0), Array.get(U32, ma, 0), Array.get(F32, va, 0)))
401    case _ t:
402      MatStep{TA.Mat{wa}, TA.Mat{ga}, Packed{t}}