~/jendhub

def OptState.new

bendtorch@1.3.0.0 · main.bend:905

import bendtorch@1.3.0.0/main.bend as Bendtorch
import Base

905def OptState.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, cfg: OptCfg) -> OptState<P>:
906  match cfg:
907    case CSGD{lr}: NoMom{}
908    case CMomentum{lr, mu}: OneMom{Layer.zeros(P, X, Y, S, m)}
909    case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: Packed{PNew{v16}}
910    case _: TwoMom{Layer.zeros(P, X, Y, S, m), Layer.zeros(P, X, Y, S, m)}

Uses (5)

Used by (1)