github ↗
docs menu

Module Vega.Adan_state

The structure of Adan states. Its walk visits mu, delta, nu then prev_grads at those fields as positions of the parameter, then step at "step" as a fixed int32 tensor. Transformations take adan_ptree.

#type ('p, 'd) t = ('p, 'd) adan_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.