def mom_fin
import bendtorch-fast@1.3.0.0/main.bend as BendtorchFastimport 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
79def mom_fin(-r: Nat, -c: Nat, q: TF.B3<F32, F32, F32>) -> T.MatStep<r, c>:
80 match q:
81 case TF.B3{w, g, m}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.OneMom{TA.Mat{m}}}