Module `Checkpoint.Ptree`

Dynamic parameter trees as a named structure. Leaves are named by their path from the root: dict keys and zero-based list positions joined with "." (e.g. "layers.0.w"). A bare root tensor has the empty path and is named by the prefix argument of of_params and to_params alone.

include Nx.Ptree.S with type t = Rune.Ptree.t

type t = Rune.Ptree.t

The structure type.

val map : ('a 'b. ('a, 'b) Nx_effect.t -> ('a, 'b) Nx_effect.t) -> t -> t

map f t is t with f applied to every tensor leaf.

val map2 : 
  ('a 'b. ('a, 'b) Nx_effect.t -> ('a, 'b) Nx_effect.t -> ('a, 'b) Nx_effect.t) ->
  t ->
  t ->
  t

map2 f a b is a and b combined leafwise with f.

Raises Invalid_argument if a and b are not structurally equal.

val iter : ('a 'b. ('a, 'b) Nx_effect.t -> unit) -> t -> unit

iter f t applies f to every tensor leaf of t.

val names : t -> string list

names x is the name of each tensor leaf of x, in traversal order.