BendtorchGpu.opt_mat: def of bendtorch-gpu@1.3.0.0
Import: import bendtorch-gpu@1.3.0.0/main.bend as BendtorchGpu
Lines 262-273 of bendtorch-gpu@1.3.0.0/main.bend: https://jend.leeeet.dev/read.txt?path=bendtorch-gpu@1.3.0.0/main.bend&start_line=262&end_line=273

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

# one matrix's step, in place: SGD writes two arrays, momentum three, Adam four
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{wa} TA.Mat{ga} _:
      sgd_fin(r, c, TG.run2(~ESGD, ~lf_sgd, TG.depth(Nat.mul(r, c)), ESGD{U32.from_nat(Nat.mul(r, c)), lr}, nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), wa, ga))
    case T.Momentum{+lr, +mu} TA.Mat{wa} TA.Mat{ga} T.OneMom{TA.Mat{ma}}:
      mom_fin(r, c, TG.run3(~EMOM, ~lf_mom, TG.depth(Nat.mul(r, c)), EMOM{U32.from_nat(Nat.mul(r, c)), lr, mu}, nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), wa, ga, ma))
    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, TG.run4(~EADAM, ~lf_adam, TG.depth(Nat.mul(r, c)), EADAM{U32.from_nat(Nat.mul(r, c)), lr, c1, b2, c2, eps, wd, dec, step, bc2}, nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), wa, ga, ma, va))
    # packed BF16 moments (CAdamWBF16): main.bend's step, in one block
    case o w g T.Packed{t}: T.Opt.mat(r, c, o, w, g, T.Packed{t})
    case _ w g st: T.MatStep{w, g, st}

Uses (23):
  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=Base/base.bend&name=U32.from_nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Unit
  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-gpu@0.1.5.1/main.bend&name=depth
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-gpu@0.1.5.1/main.bend&name=run2
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-gpu@0.1.5.1/main.bend&name=run3
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-gpu@0.1.5.1/main.bend&name=run4
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=EADAM
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=EMOM
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=ESGD
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=lf_adam
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=lf_mom
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=lf_sgd
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=mom_fin
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=nn
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=opt_fin
  https://jend.leeeet.dev/definition.txt?path=bendtorch-gpu@1.3.0.0/main.bend&name=sgd_fin
  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=Opt.mat
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=OptState