def lin_step2

bendtorch-gpu@1.3.0.0 · main.bend:278

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

278def lin_step2(+i: Nat, +o: Nat, +op: T.Opt, p: T.Lin<i, o>, g: T.Lin<i, o>, st: T.Two<T.OptState<TA.Mat<o, i>>, T.OptState<TA.Mat<1n, o>>>) -> T.Step<T.Lin<i, o>>:
279  match p g st:
280    case T.Lin{w, b} T.Lin{gw, gb} T.Two{sw, sb}:
281      T.lin_s2(i, o, opt_mat(o, i, op, w, gw, sw), opt_mat(1n, o, op, b, gb, sb))