Module Nx.Ptree

Parameter trees: structures with tensor leaves. Ptree.S is the traversal interface shared across the Raven ecosystem; Ptree.t is the stock dynamic instance.

Parameter trees: structures with tensor leaves.

S is the traversal interface the Raven ecosystem shares: autodiff transformations, structural optimizers and checkpointing all operate on any structure implementing it — typically a record of tensors with hand-written one-line traversals. The concrete t is the stock instance for structures only known at runtime.

module type S = sig ... end

Structures whose tensor leaves can be traversed.

type tensor =
  1. | P : ('a, 'b) Nx_effect.t -> tensor
    (*

    A packed tensor. The wrapper hides dtype and layout parameters.

    *)
type t =
  1. | Tensor of tensor
    (*

    A tensor leaf.

    *)
  2. | List of t list
    (*

    An ordered list branch.

    *)
  3. | Dict of (string * t) list
    (*

    A dict branch. Keys are strings.

    *)

The stock dynamic parameter tree; itself satisfies S.

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

tensor x is Tensor (P x).

val list : t list -> t

list ts is List ts.

val dict : (string * t) list -> t

dict kvs is Dict kvs.

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 differ in structure, dict keys, or leaf dtype.

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

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