def mom_leaf

bendtorch-fast@1.3.0.0 · main.bend:59

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

59def mom_leaf(+op: T.Opt, +i0: U32, +cnt: Nat, r: TF.B3<F32, F32, F32>) -> TF.B3<F32, F32, F32>:
60  match op r:
61    case T.Momentum{+lr, +mu} TF.B3{w, g, m}:
62      mom_b(T.mom_l(cnt, i0, lr, mu, Array.get(F32, w, i0), Array.get(F32, g, i0), Array.get(F32, m, i0)))
63    case _ r: r