docs menu

Module Vega.Sgd_state

The structure of SGD states. Its walk visits velocity at the field "velocity" as a position of the parameter, then step at "step" as a fixed int32 tensor. Use it with Nx.Ptree.Payload and Nx.Ptree.cast; transformations take sgd_ptree.

#type ('p, 'd) t = ('p, 'd) sgd_state

The type for values with parameter 'a and fixed arrays on 'd.

#val walk : ('a, 'b, 'd, 'e) Nx.Ptree.Walk.cursor -> ('a, 'd) t -> ('b, 'e) t

walk c x walks every part of x with c and is x rebuilt from what the walks return.