Module tree 

Trees of arrays: structured values (tuples, vectors, maps, options, user structs and enums) whose array leaves have paths and a fixed order, as MLX’s tree_flatten.

• Leaves is the read-only, object-safe half: anything whose leaves can be visited, including borrowed views such as &T, (&A, &B) and [T]. crate::eval takes it, so eval(&(&loss, &grads)) evaluates both without cloning.
• Tree is an owned tree that can be mutated in place and rebuilt with new leaves of the same structure (Tree::unflatten); transforms differentiate, batch and compile trees.

Paths and order (MLX’s tree_flatten names): struct fields in declaration order, sequences and tuples by index (0, 1, …), maps by sorted key, Option, Box and newtype enum variants transparent, joined with .; a root array has the empty path.

Leaf kinds (LeafKind): a plain Array (always differentiable), a Param (a module parameter, trainable unless frozen) or a buffer (an array field marked #[tree(buffer)]: saved and loaded, never trained).

Structure (TreeDef): the root type, the ordered (path, kind, trainable) of every leaf, the active enum variants, the mode flags (train/eval) and a hash of the static fields (#[tree(skip)] fields of a Static type, such as a layer’s eps). Two trees with equal TreeDefs have the same structure; a compiled function retraces when the structure changes.

User types implement the traits with #[derive(Tree)] (nuvec::Tree; code that uses nuvec_core directly names it with #[tree(crate = "nuvec_core")]):

use nuvec_core::tree::Tree;
use nuvec_core::{Array, array};

#[derive(Clone, nuvec_macros::Tree)]
#[tree(crate = "nuvec_core")]
struct Batch {
    x: Array,
    y: Array,
    #[tree(skip)]
    len: usize, // not a leaf: hashed into the structure
}

let b = Batch { x: array![1.0, 2.0], y: array![0.0], len: 2 };
assert_eq!(b.paths(), ["x", "y"]);
let doubled = b.map(|_, a| a.multiply(2.0));
assert_eq!(doubled.x.to_vec::<f32>().unwrap(), [2.0, 4.0]);
assert_eq!(doubled.len, 2);

Static fields and retracing: everything in a TreeDef is part of a compiled function’s cache key, so switching train/eval, freezing, a different enum variant or a changed hashed field retraces once. A #[tree(skip, static_unchecked)] field is not hashed: its value is baked into compiled traces.

Structs

Param

A module parameter: an array and a frozen flag (trainable unless frozen).

TreeDef

The structure of a tree: its root type, the ordered (path, kind, trainable) of every leaf, the active enum variants, the mode flags and a hash of the static fields. Equal structures have equal TreeDefs (see the module documentation).

Visit

The read-only visitor that Leaves::visit reports to: it builds paths and dispatches leaves, child trees, enum variants, mode flags and static fields to a Visitor.

VisitMut

The mutable visitor that Tree::visit_mut reports to.

Enums

LeafKind

What a leaf is (see the module documentation).

LoadMode

How Tree::load_flat treats names it does not find.

Traits

BufferLeaf

What a #[tree(buffer)] field may hold: an Array, or an Option<Array> (a leaf when present, as MLX’s optional buffers such as affine-only quantization biases).

Leaves

Anything whose array leaves can be visited in a fixed order with their paths: owned trees and borrowed views (&T, (&A, &B), [T]). Object-safe.

Static

A value hashed into a TreeDef: the type of a #[tree(skip)] field (a configuration, a hyperparameter). A different value is a different structure, so a compiled function retraces instead of replaying a stale constant.

Tree

An owned tree of arrays: rebuilt with new leaves of the same structure, mutated in place.

Visitor

Receives what a Visit reports, with full paths. Every method but Visitor::leaf has an empty default.

VisitorMut

Receives what a VisitMut reports, with full paths.

Functions

collect_leaves

The leaves of any Leaves value, in order.