The schema you pass to Model(...) becomes a tree. The generated root branch encodes one processed observation, nested Branch(...) nodes define shared child contexts, and leaf tensorfields read values from input records.
Internally, the tree is backed by anytree and pydantic, but the public model is easiest to read as root, branch, and leaf nodes.
Each node has a stable slash-delimited address (jv.Address):
order root branch context encoder, embed
|-- source_channel category leaf, inactive
|-- line_items branch context encoder, embed
| |-- sku category leaf
| |-- quantity number leaf
| `-- price number leaf, masked reconstruction
`-- returned category leaf, supervised target
order is the root context. jv.Address("order", "line_items") is the branch address rendered as order/line_items. The leaves are the typed requests that read values from input records.
Thus, an order is defined by a customer ID, a list of items purchased, and whether or not the customer returned any of the items.
Node Kinds
Node kind
Created by
Runtime role
Root branch
Model(..., name=...)
Encodes the whole processed observation.
Branch
Branch(..., name=..., length=...)
Defines a child context, optionally applies local attention, and pools the representation to its parent.
Leaf tensorfield
Number, Category, Set, Vector, and other tensorfields
Reads source values, stores typed tensors, embeds visible input, and decodes predictions when trained or requested.
Note
Tensorfields represent an extensible typing system. There are many built-in data types, with more in development.
A branch names the coordinate space shared by its child leaves. Its encoder can apply transformer attention within that context, then cross-attention pooling rolls the context up to the parent. Leaf nodes are the only nodes that bind directly to values with a request-level query. See Query Paths for how leaf queries are inferred.
Note
Under the hood, the root branch uses the same machinery as other branches, but it is always a singleton context with length=1. Use nested Branch(length=...) for repeated data.
Nodes As N-Dimensional Arrays
Every schema node corresponds to an N-dimensional tensor shape. Branch nodes add dimensions. Leaf tensorfields store typed arrays at the shape defined by their branch ancestors.
For the order/line_items/price leaf above, the inherited shape is:
For user-declared branches, length defines the retained dimension and overflow controls whether overlong query results keep the head, keep the tail, or raise an error.
At runtime, a scalar tensorfield such as price has arrays like:
state (batch, 1, 32) - if the content is valued, padded, masked, or null
content (batch, 1, 32) - the actual content payload (numerical value, token indices, etc.)
trainable (batch, 1, 32) - boolean mask to determine whether the content may contribute to loss
Some field types add content dimensions. For example, Vector content includes the configured vector width, and Set content includes vocabulary dimensions. The shared rule is that state carries the node’s value-state array and content carries the type-specific value array.
The state array distinguishes observed values, explicit nulls, padded slots, and training-time masks. See Built-In Data Types for the shared state vocabulary.
Embedders convert leaf arrays into d_model vectors:
A model forward pass follows the tree from leaves upward, then decodes selected leaves from their available context.
Raw records are preprocessed and encoded into active leaf tensorfields.
Training-time masking and pruning mark hidden positions as masked, save the original values in targets, and set trainable=True where reconstruction loss should be computed.
Active, non-target leaves run their tensorfield embedders and send Parcel objects to their parent branches.
Branch encoders run from deepest branches back to the root. Each branch concatenates child parcels, applies its configured attention layers, pools the child context, and sends one parcel to its parent.
Branches configured with embed=True emit embedding predictions from their pooled payloads.
Leaves that are trainable, targets, or configured with embed=True run their decoders.
Losses are computed for decoded leaf predictions that have trainable targets. Pure embedding outputs are emitted only at inference time.
Pooling And Heritage
Pooling happens in two places.
Branch pooling compresses a branch context before passing it upward to its parent context. After a branch concatenates its child payloads and runs attention layers, learned-query cross-attention produces the vector payload for that branch node.
Decoder pooling builds a prediction context for a leaf. Each leaf has a heritage: the addresses along its path from the root through the leaf itself. For order/line_items/price, the heritage is:
order
order/line_items
order/line_items/price
The decoder gathers the available outgoing parcels from that heritage, concatenates them, then pools them into the target leaf shape. The default pooling="query" uses learned-query cross-attention. pooling="mean" repeats the mean of the heritage context for each target slot.
This is why a decoded leaf can use information from its own branch, ancestor contexts, and visible sibling fields. If the leaf itself was masked, its own parcel still carries masked-state information, while the original value is kept only in targets for loss computation.
Roles And Mutations
Schema roles change how nodes participate in training and prediction:
Setting
Effect
p_mask=0.15
Randomly hides individual leaf values for reconstruction.
p_prune=0.15
Randomly hides whole leaf instances for reconstruction.
target=True
Always hides a leaf from input and decodes it as a supervised target; shorthand for p_prune=1.0.
mask=jv.Mask(...)
Hides selected coordinates inside a Branch and projects that selection onto descendant leaves.
embed=True
Emits an embedding for root, branch, or leaf nodes during prediction.
active=False
Keeps a leaf in the schema but removes it from encoding, forward passes, losses, and prediction until it is reactivated.
Model mutations edit the same tree:
Method
Use
model.update(predicate, **values)
Change selected node attributes such as weight, target, p_mask, embed, or active.
model.extend(predicate, new_field, ...)
Add new children under one selected branch node.
model.delete(predicate)
Permanently remove selected nodes and their descendants.
model.reset(predicate)
Reinitialize selected runtime modules while keeping schema values.
model.override(predicate, **values)
Temporarily change nodes inside a context manager, then restore them.
Mutations rebuild the runtime graph and reload compatible state where shapes still match. They are blocked while Lightning owns an active training, validation, test, or prediction loop.
Use predicates such as jv.where("name") == "price" or jv.where("type") == "number" to select nodes. The same predicate style is used across schema mutation workflows.