def sgd_leaf
import bendtorch-fast@1.3.0.0/main.bend as BendtorchFastimport Base
import bend-ml-tensor-fork@0.1.5.0/main.bend as TF
import bendtorch@1.3.0.0/main.bend as T
50def sgd_leaf(+lr: F32, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
51 match r:
52 case TF.B2{w, g}:
53 p2_b2(T.sgd_l(cnt, i0, lr, Array.get(F32, w, i0), Array.get(F32, g, i0)))