def Opt.mat

bendtorch@1.3.0.0 · main.bend:432

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

432# A mismatched state is initialized only at the training boundary; Opt.mat requires the matching variant.
433# The fallback leaves the values untouched, so it cannot fabricate a moment buffer.
434def Opt.mat(+r: Nat, +c: Nat, +o: Opt, w: TA.Mat<r, c>, g: TA.Mat<r, c>, st: OptState<TA.Mat<r, c>>) -> MatStep<r, c>:
435  match o w g st:
436    case SGD{lr} TA.Mat{wd} TA.Mat{gd} _:
437      Opt.sgd_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, wd, 0), Array.get(F32, gd, 0)))
438    case Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} OneMom{TA.Mat{md}}:
439      Opt.mom_fin(r, c, mom_l(Nat.mul(r, c), 0, lr, mu, Array.get(F32, wd, 0), Array.get(F32, gd, 0), Array.get(F32, md, 0)))
440    case Adam{lr, b1, b2, eps, wd2, dec, step, bc2, c1, c2} TA.Mat{wd} TA.Mat{gd} TwoMom{TA.Mat{md}, TA.Mat{vd}}:
441      Opt.fin(r, c, adam_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd2, dec, step, bc2, Array.get(F32, wd, 0), Array.get(F32, gd, 0), Array.get(F32, md, 0), Array.get(F32, vd, 0)))
442    case Adam{lr, b1, b2, eps, wd2, dec, step, bc2, c1, c2} TA.Mat{wd} TA.Mat{gd} Packed{t}:
443      adamh_mat(r, c, o, wd, gd, t)
444    case _ w g st:
445      MatStep{w, g, st}