def ln_step2

bendtorch-gpu@1.3.0.0 · main.bend:556

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

556# the gain and the bias take this file's optimizer step, as Linear's two matrices do
557def ln_step2(+d: Nat, +op: T.Opt, p: T.LN<d>, g: T.LN<d>, st: T.Two<T.OptState<TA.Mat<1n, d>>, T.OptState<TA.Mat<1n, d>>>) -> T.Step<T.LN<d>>:
558  match p g st:
559    case T.LN{w, b} T.LN{gw, gb} T.Two{sw, sb}:
560      T.ln_s2(d, opt_mat(1n, d, op, w, gw, sw), opt_mat(1n, d, op, b, gb, sb))