def opt_mat

bendtorch-fast@1.3.0.0 · main.bend:158

import Base
import bend-ml-tensor-array@0.1.18.0/main.bend as TA
import bend-ml-tensor-fork@0.1.5.0/main.bend as TF
import bendtorch@1.3.0.0/main.bend as T

158# SGD splits two arrays, momentum three, Adam four. No unused parameter-sized buffers.
159def opt_mat(+r: Nat, +c: Nat, +o: T.Opt, w: TA.Mat<r, c>, g: TA.Mat<r, c>, st: T.OptState<TA.Mat<r, c>>) -> T.MatStep<r, c>:
160  match o w g st:
161    case T.SGD{+lr} TA.Mat{wd} TA.Mat{gd} _:
162      sgd_fin(r, c, TF.par2(~F32, ~F32, ~F32, ~sgd_leaf, lr, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wd, gd))
163    case T.Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} T.OneMom{TA.Mat{md}}:
164      mom_fin(r, c, TF.par3(~T.Opt, ~F32, ~F32, ~F32, ~mom_leaf, o, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wd, gd, md))
165    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} TA.Mat{wa} TA.Mat{ga} T.TwoMom{TA.Mat{ma}, TA.Mat{va}}:
166      opt_fin(r, c, TF.par4(~T.Opt, ~F32, ~F32, ~F32, ~F32, ~opt_leaf, o, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wa, ga, ma, va))
167    # packed BF16 moments (CAdamWBF16): blocks of whole words
168    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} TA.Mat{wa} TA.Mat{ga} T.Packed{t}:
169      ph_mat(r, c, o, wa, ga, t)
170    case _ w g st: T.MatStep{w, g, st}