type PTree

bendtorch@1.3.0.0 · main.bend:50

import Base

50# AdamW's moments stored as packed BF16 (CAdamWBF16): a tree that follows the parameters' pairs as the step walks them,
51# with the buffers of each matrix at its leaf. A packed buffer has half the words of its matrix (two BF16 values a
52# U32, element 2k in the low half of word k, torch/dtype.bend's layout), so it cannot be a value of the parameters'
53# type P: the tree holds the words itself. A leaf is allocated on the matrix's first step (PNew marks a part not yet
54# stepped, and stays at the parameter-free layers), so the full-size buffers never exist. v16 picks v packed as well;
55# otherwise v is F32 (PLeafV). A leaf keeps its matrix's count n, for checkpoints.
56type PTree is Type:
57  PNew{v16: Bool}
58  PLeaf{n: U32, m: Array<U32>, v: Array<U32>}
59  PLeafV{n: U32, m: Array<U32>, v: Array<F32>}
60  PNode{a: PTree, b: PTree}