github ↗
docs menu

Module Ptree.Walk

The functions a walk is written with, one per kind of part.

Each walks one part at the cursor's path. A sub-structure is walked by its own walk, and a variant names its case before its other parts:

ocaml
type ('a, 'd) weight =
  | Float of 'a
  | Mxfp4 of { blocks : 'd uint8_t; scales : 'd uint8_t }

let walk_weight c =
  let open Nx.Ptree.Walk in
  function
  | Float w ->
      case c "float";
      Float (leaf c w)
  | Mxfp4 { blocks; scales } ->
      case c "mxfp4";
      let blocks = field c "blocks" tensor blocks in
      let scales = field c "scales" tensor scales in
      Mxfp4 { blocks; scales }
#type ('a, 'b, 'd, 'e) cursor

The type for cursors of a walk that turns the parameter's positions from 'a into 'b and the fixed arrays from the set 'd to 'e. A cursor stands at a path of the value walked.

#val field : ('a, 'b, 'd, 'e) cursor -> string -> (('a, 'b, 'd, 'e) cursor -> 'u -> 'w) -> 'u -> 'w

field c name walk x is walk c' x, c' at c's path then Field name.

#val index : ('a, 'b, 'd, 'e) cursor -> int -> (('a, 'b, 'd, 'e) cursor -> 'u -> 'w) -> 'u -> 'w

index c i walk x is walk c' x, c' at c's path then Index i. Tuples and arrays number their parts with it; an array also reports its length with int.

#val leaf : ('a, 'b, 'd, 'e) cursor -> 'a -> 'b

leaf c x walks x, a position of the parameter, at c's path.

#val tensor : ('a, 'b, 'd, 'e) cursor -> ('v, 's, 'd) tensor -> ('v, 's, 'e) tensor

tensor c x walks x, a fixed array, at c's path. A description treats it as any other array; Payload's maps and cast keep it, and place places it.

#val int : ('a, 'b, 'd, 'e) cursor -> int -> int

int c n reports n at c's path and is n. A bool is reported as 0 or 1.

#val case : ('a, 'b, 'd, 'e) cursor -> string -> unit

case c tag reports the variant case tag at c's path.

#val option : (('a, 'b, 'd, 'e) cursor -> 'u -> 'w) -> ('a, 'b, 'd, 'e) cursor -> 'u option -> 'w option

option walk c x reports whether x is present at c's path and, if it is, walks its content with walk at the same path.

#val list : (('a, 'b, 'd, 'e) cursor -> 'u -> 'w) -> ('a, 'b, 'd, 'e) cursor -> 'u list -> 'w list

list walk c l reports l's length at c's path, then walks its i-th element with walk at c's path then Index i, from the first.

#val under : (('p, 'q, 'd, 'e) cursor -> 'u -> 'w) -> (('a, 'b, 'd, 'e) cursor -> 'p -> 'q) -> ('a, 'b, 'd, 'e) cursor -> 'u -> 'w

under outer inner c x walks x with outer at c's path, and each position of outer's parameter with inner at that position's path, inner's paths after the position's as nest makes them. Its visits are outer's in walk order, with inner's visits of each position where outer reaches it. The fixed arrays of outer and of inner go from 'd to 'e alike, so under U.walk V.walk walks (('a, 'd) V.t, 'd) U.t as a structure's walk, and place moves every array it walks.

A frozen layer is Linear (Ptree) with a fixed array at each position:

ocaml
type ('a, 'd) tuned = {
  head : 'a linear;
  base : 'd Nx.float32_t linear;
}

module Tuned = struct
  type ('a, 'd) t = ('a, 'd) tuned

  let walk c x =
    let open Nx.Ptree.Walk in
    let head = field c "head" Linear.walk x.head in
    let base = field c "base" (under Linear.walk tensor) x.base in
    { head; base }
end

Nx.Ptree.place (module Tuned) Gpu.on x places head and base, and Nx.Ptree.cast (module Tuned) dt x casts head alone.

#val structure : ('s, 'd) t -> ('a, 'b, 'd, 'd) cursor -> 's -> 's

structure s c x walks x, a value of the description s, at c's path: each of its arrays as tensor walks one, each of its reports as s makes it, at s's paths after c's. tensor c x is structure Nx.Ptree.tensor c x.

A description is at one set, so a walk that calls structure is too: it makes a description with v, never a structure's module, and place cannot move what it walks. A structure at the positions of another's parameter is walked by under, and a part of a fixed type that moves between sets by its own walk, as Rng.walk walks a key:

ocaml
type ('a, 'd) state = {
  params : 'a linear;
  key : 'd Nx.Rng.key;
  step : 'd Nx.int32_t;
}

module State = struct
  type ('a, 'd) t = ('a, 'd) state

  let walk c s =
    let open Nx.Ptree.Walk in
    let params = field c "params" Linear.walk s.params in
    let key = field c "key" Nx.Rng.walk s.key in
    let step = field c "step" tensor s.step in
    { params; key; step }
end

Nx.Ptree.place (module State) Gpu.on s places the parameters, the key and the step.