BendtorchFast.opt_mat: def of bendtorch-fast@1.3.0.0
Import: import bendtorch-fast@1.3.0.0/main.bend as BendtorchFast
Lines 158-170 of bendtorch-fast@1.3.0.0/main.bend: https://jend.leeeet.dev/read.txt?path=bendtorch-fast@1.3.0.0/main.bend&start_line=158&end_line=170

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

# SGD splits two arrays, momentum three, Adam four. No unused parameter-sized buffers.
def 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>:
  match o w g st:
    case T.SGD{+lr} TA.Mat{wd} TA.Mat{gd} _:
      sgd_fin(r, c, TF.par2(~F32, ~F32, ~F32, ~sgd_leaf, lr, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wd, gd))
    case T.Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} T.OneMom{TA.Mat{md}}:
      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))
    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}}:
      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))
    # packed BF16 moments (CAdamWBF16): blocks of whole words
    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} TA.Mat{wa} TA.Mat{ga} T.Packed{t}:
      ph_mat(r, c, o, wa, ga, t)
    case _ w g st: T.MatStep{w, g, st}

Uses (18):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat.mul
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-array@0.1.19.0/main.bend&name=Mat
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-fork@0.1.6.0/main.bend&name=par2
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-fork@0.1.6.0/main.bend&name=par3
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-fork@0.1.6.0/main.bend&name=par4
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=ew_depth
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=mom_fin
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=mom_leaf
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=opt_fin
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=opt_leaf
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=ph_mat
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=sgd_fin
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=sgd_leaf
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=MatStep
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=Opt
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=OptState