3.10. Graph Models and Merge Layers

Sequential gives every layer exactly 1 input, and that input is the output of the layer in front of it. A model that splits, rejoins, reads 2 tensors, or answers 2 questions needs more than a chain. Graph is that model, and GraphBuilder is the only way to reach one.

This page covers the builder and the 5 calls that assemble a topology. It also covers the refusals that build gives before any data moves, and the 7 merge layers. Both types live in rustyml::neural_network::graph, and the prelude exports Graph, GraphBuilder, NodeId, and every merge layer.

Sequential is unchanged. When a chain cannot express your model, read this page. When it can, stay with Sequential. A chain built as a graph gives the same numbers, bit for bit, and it costs a little more to write.

3.10.1. Why a chain is not enough

A chain holds 1 shape of model. Each layer reads the layer in front of it, the first layer reads the input, and the last layer feeds the loss. 4 common architectures do not fit in that shape:

1. A residual connection. The input of a block reaches the output of that block a second time, so 1 tensor feeds 2 positions.
2. Weight sharing. 2 branches read the same encoder, so 1 set of arrays serves 2 positions and takes the sum of the 2 gradients.
3. Several inlets. A model reads a numeric table and a category table, and it joins them after a first transform of each.
4. Several outlets. A model predicts an amount and a class from 1 trunk, and each head takes its own loss.

The graph model expresses all 4. It holds a layer arena and a node list. An arena entry is a layer, with its arrays. A node is 1 call of 1 layer on the outputs of other nodes. The 2 lists are separate, and that is the whole idea:

// A chain: 1 call per layer, and each call reads the call before it.
let model = SequentialBuilder::new().add(a).add(b).add(c).build(&shape)?;

// A graph: the arena holds the layers, and the nodes say what reads what.
let mut builder = GraphBuilder::new();
let x = builder.input(shape);           // an inlet, which calls no layer
let one = builder.add(a, &[x]);         // 1 layer, put in the arena and called once
let two = builder.apply(shared, &[one]); // 1 more call of a layer already in the arena
let model = builder.build(&[two])?;     // the outlets, in the order the model gives them

Several nodes that name 1 arena entry are exactly what weight sharing is. The nodes each hold their own cache, because each has its own input and its own output. The arrays and the gradient belong to the arena entry, so the update reads the sum over every node that called it.

3.10.2. The builder: input, layer, apply, add, and build

GraphBuilder collects the arena and the node list, and build turns them into a Graph. The first 4 calls never fail, so a whole model reads as a run of let bindings. build is the 1 fallible call:

pub fn new() -> GraphBuilder;
pub fn new_with_seed(seed: u64) -> GraphBuilder;

pub fn input(&mut self, shape: Shape) -> NodeId;
pub fn layer<L: 'static + Layer>(&mut self, layer: L) -> LayerId;
pub fn apply(&mut self, layer: LayerId, inputs: &[NodeId]) -> NodeId;
pub fn add<L: 'static + Layer>(&mut self, layer: L, inputs: &[NodeId]) -> NodeId;

pub fn build(self, outputs: &[NodeId]) -> Result<Graph, Error>;

input records an inlet and gives back the node that carries it. An inlet calls no layer, and the caller supplies its tensor. The inlets keep registration order, and that order is the order that fit, train_batch, predict, and evaluate take their tensors in.

layer puts 1 layer in the arena and gives back its address. apply calls a layer of the arena on the outputs of other nodes, and gives back the node that holds the result. Whenever 2 positions of the model must share 1 set of arrays, use the pair.

add is layer and apply in 1 call, and it is the common case. Use it for every layer that the model reads exactly once. NodeId and LayerId are both usize, so keep them in named bindings and never mix them up by hand.

build takes the outlets, in the order the model gives its tensors. It consumes the builder and returns Result<Graph, Error>. The walk it runs is the subject of 3.10.6.

new_with_seed records the seed of the mini-batch shuffle, exactly as the sequential builder does. The seed governs the shuffle alone. Each layer seeds its own weight draw through with_random_state (see 7.1).

The 2 models differ on 5 points, and agree everywhere else:

3.10.3. A residual connection

A residual block adds the input of the block to its own output. The gradient then reaches the input of the block by 2 paths, and the model trains a correction instead of a whole transform. The pattern is 3 lines in a graph: 1 node for the block, 1 Add node, and 1 head.

use ndarray::Array;
use rustyml::prelude::*;

fn main() {
    // 6 samples of 8 features, and a 2-value target for each.
    let x = Array::ones((6, 8)).into_dyn();
    let y = Array::zeros((6, 2)).into_dyn();

    let mut builder = GraphBuilder::new();
    let input = builder.input(Shape::known(&[6, 8]));

    // The block keeps the width of its input, so the sum below is legal.
    let hidden = builder.add(Dense::new(8, Activation::ReLU).unwrap(), &[input]);

    // The inlet reaches the sum, and it also reaches the block. 1 tensor, 2 readers.
    let sum = builder.add(Add::new(), &[input, hidden]);
    let head = builder.add(Dense::new(2, Activation::Linear).unwrap(), &[sum]);
    let mut model = builder.build(&[head]).unwrap();

    model.compile(
        SGD::new(0.01, 0.0, false, 0.0).unwrap(),
        MeanSquaredError::new(),
    );
    model.summary();

    let history = model.fit(&[&x], &[&y], 20).unwrap();
    assert!(*history.loss().last().unwrap() < history.loss()[0]);

    // 1 tensor per outlet, so a model with 1 outlet gives a vector of length 1.
    assert_eq!(model.predict(&[&x]).unwrap()[0].shape(), &[6, 2]);
}

summary() prints 1 row per node, and the last column names the nodes that each node reads:

Model: "graph"
Node   Layer (type)             Output Shape              Param #  Reads
----------------------------------------------------------------------------------------
0      input                    (6, 8)                             
1      dense (Dense)            (6, 8)                         72  [0]
2      add (Add)                (6, 8)                          0  [0, 1]
3      dense_1 (Dense)          (6, 2)                         18  [2]
----------------------------------------------------------------------------------------
Total params: 90
Trainable params: 90
Non-trainable params: 0

Node 2 reads [0, 1], and that pair is the residual connection. The Add layer holds no array, so it reports 0 parameters. The shape column comes from the same pure walk that build ran. A model that has never seen a tensor still prints the true shape of every node.

The backward pass sums the 2 paths into node 0 before it moves past that node. The walk runs the node positions in reverse, and every reader of a node runs before the node itself. The consumers of 1 node therefore contribute in descending node position, which pins the order of a sum that f32 addition does not commute.

3.10.4. Weight sharing: 1 layer and 2 nodes

layer and apply are the pair that shares weights. layer puts 1 encoder in the arena. Each apply adds a node that calls it. The model below reads 2 tensors, runs both through 1 tower, and scores the gap between the 2 codes.

use ndarray::Array;
use rustyml::prelude::*;

fn main() {
    let left = Array::ones((4, 5)).into_dyn();
    let right = Array::zeros((4, 5)).into_dyn();
    let target = Array::zeros((4, 1)).into_dyn();

    let mut builder = GraphBuilder::new();
    let a = builder.input(Shape::known(&[4, 5]));
    let b = builder.input(Shape::known(&[4, 5]));

    // 1 entry in the arena, and 2 nodes that call it. Both towers read 1 kernel.
    let tower = builder.layer(Dense::new(3, Activation::Tanh).unwrap());
    let left_code = builder.apply(tower, &[a]);
    let right_code = builder.apply(tower, &[b]);

    let gap = builder.add(Subtract::new(), &[left_code, right_code]);
    let score = builder.add(Dense::new(1, Activation::Sigmoid).unwrap(), &[gap]);
    let mut model = builder.build(&[score]).unwrap();

    // The shared layer holds 1 set of arrays, so the checkpoint names it 1 time.
    assert_eq!(
        model.weight_paths(),
        vec!["0.kernel", "0.bias", "2.kernel", "2.bias"]
    );
    model.summary();

    model.compile(
        Adam::new(0.05, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        BinaryCrossEntropy::new(),
    );
    model.fit(&[&left, &right], &[&target], 10).unwrap();
    assert_eq!(model.predict(&[&left, &right]).unwrap()[0].shape(), &[4, 1]);
}

Model: "graph"
Node   Layer (type)             Output Shape              Param #  Reads
----------------------------------------------------------------------------------------
0      input                    (4, 5)                             
1      input                    (4, 5)                             
2      dense (Dense)            (4, 3)                         18  [0]
3      dense_1 (Dense)          (4, 3)                     shared  [1]
4      subtract (Subtract)      (4, 3)                          0  [2, 3]
5      dense_2 (Dense)          (4, 1)                          4  [4]
----------------------------------------------------------------------------------------
Total params: 22
Trainable params: 22
Non-trainable params: 0

Node 3 prints shared in place of a parameter count. A shared layer counts once in the total, which is the whole point of sharing. The 2 nodes still get their own row, because each holds its own input, its own output, and therefore its own cache.

The checkpoint paths tell the same story from the other side. The tower is arena entry 0, and the head is arena entry 2, because the Subtract layer takes arena entry 1 and holds no array. The tower contributes 0.kernel and 0.bias once, whatever number of nodes call it. See 3.9 for the format of a path.

An arena position is also the layer half of every parameter address, so the shared tower takes 1 update per step. Each node ran its own backward pass and produced its own gradient. The gradient store sums the 2 at 1 address, and the optimizer reads the total.

3.10.5. Several inlets and several outlets

A model with several outlets needs 1 loss per outlet. compile gives every outlet a copy of 1 loss. That is right for 1 outlet, and for several outlets of the same kind. compile_many gives each outlet its own loss:

pub fn compile<O, L>(&mut self, optimizer: O, loss: L) -> &mut Self
where O: 'static + Optimizer, L: 'static + Loss + Clone;

pub fn compile_many<O>(&mut self, optimizer: O, losses: Vec<Box<dyn Loss>>)
    -> Result<&mut Self, Error>
where O: 'static + Optimizer;

pub fn with_loss_weights(&mut self, weights: &[f32]) -> Result<&mut Self, Error>;

compile_many refuses a loss count that does not match the outlet count. with_loss_weights sets the weight of each outlet in the total. The reported loss is the weighted sum, and the gradient that seeds each outlet scales by the same weight. Every weight starts at 1.0, and with_loss_weights refuses a weight that is not finite.

The model below reads 2 tables and answers 2 questions. Concatenate joins the 2 inlets after the model has seen both, and 1 trunk feeds 2 heads.

use ndarray::Array;
use rustyml::neural_network::traits::Loss;
use rustyml::prelude::*;

fn main() {
    // 2 inlets: 6 numeric features, and 2 flags, for the same 4 samples.
    let numbers = Array::ones((4, 6)).into_dyn();
    let flags = Array::zeros((4, 2)).into_dyn();

    // 2 outlets: 1 amount, and 1 one-hot label over 3 classes.
    let value = Array::zeros((4, 1)).into_dyn();
    let label = Array::from_shape_vec(
        (4, 3),
        vec![1.0_f32, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0],
    )
    .unwrap()
    .into_dyn();

    let mut builder = GraphBuilder::new();
    let a = builder.input(Shape::known(&[4, 6]));
    let b = builder.input(Shape::known(&[4, 2]));

    let joined = builder.add(Concatenate::new(-1), &[a, b]);
    let trunk = builder.add(Dense::new(8, Activation::ReLU).unwrap(), &[joined]);
    let amount = builder.add(Dense::new(1, Activation::Linear).unwrap(), &[trunk]);
    let kind = builder.add(
        Dense::new(3, Activation::Softmax { axis: -1 }).unwrap(),
        &[trunk],
    );

    // The outlets keep the order the build received them.
    let mut model = builder.build(&[amount, kind]).unwrap();
    assert_eq!(model.output_shapes()[0].to_string(), "(4, 1)");
    assert_eq!(model.output_shapes()[1].to_string(), "(4, 3)");

    // 1 loss per outlet, and 1 weight per outlet in the reported total.
    let losses: Vec<Box<dyn Loss>> = vec![
        Box::new(MeanSquaredError::new()),
        Box::new(CategoricalCrossEntropy::new(false)),
    ];
    model
        .compile_many(Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(), losses)
        .unwrap()
        .with_loss_weights(&[1.0, 0.5])
        .unwrap();

    let history = model.fit(&[&numbers, &flags], &[&value, &label], 40).unwrap();
    assert!(*history.loss().last().unwrap() < history.loss()[0]);

    // 1 tensor per outlet, in the same order.
    let predicted = model.predict(&[&numbers, &flags]).unwrap();
    assert_eq!(predicted.len(), 2);
    assert_eq!(predicted[0].shape(), &[4, 1]);
    assert_eq!(predicted[1].shape(), &[4, 3]);
}

3 orders govern every call of a graph model, and the model infers none of them from a name. The inlet order is the order of the input calls. The outlet order is the order of the slice that build received. The loss order and the weight order follow the outlet order. input_shapes() and output_shapes() report the 2 ends, and neither needs a forward pass.

The 2 losses on this page normalize differently, so the raw values are not on 1 scale. with_loss_weights is where you put them on 1 scale. Read 3.3 for the divisor that each loss family uses.

3.10.6. What build refuses, and why

build walks the graph once, before any data moves. It checks the topology first, then it threads the shape of each node into the nodes that read it. It builds every layer on the first node that reaches it, and each later node checks its shapes against that build.

The topology checks run in this order. Each of them describes a model that could not run, or could run and waste work:

1. The builder holds at least 1 node.
2. The output list is not empty, and every id in it names a node.
3. Every id that a node reads names a node, and names a node at an earlier position.
4. The graph holds at least 1 inlet.
5. Every node reaches at least 1 outlet.

Check 3 is what makes the position order a topological order. A node reads only nodes that already exist, so build needs no sort and no cycle detection. The builder cannot express a cycle at all, because apply gives back a node id only after it records the node.

Check 5 refuses a node that no outlet reaches. Such a node looks harmless and is not. The forward pass would compute it and park a cache for it. No backward pass would take that cache back, and the layer would take an update that the loss never saw. Name the node as an outlet, or leave it out of the graph.

The last 2 rows share 1 message shape, because both come from the layer and not from the topology. Graph has no Debug, so read a refused build through a match instead of unwrap_err:

use rustyml::error::Error;
use rustyml::prelude::*;

// `Graph` has no `Debug`, so a refused build is read through a match.
fn refusal(built: Result<Graph, Error>) -> String {
    match built {
        Ok(_) => panic!("the build was expected to refuse this graph"),
        Err(error) => error.to_string(),
    }
}

fn main() {
    // A node that no outlet reaches.
    let mut builder = GraphBuilder::new();
    let input = builder.input(Shape::known(&[2, 3]));
    let live = builder.add(Dense::new(2, Activation::Linear).unwrap(), &[input]);
    let _dead = builder.add(Dense::new(2, Activation::Linear).unwrap(), &[input]);
    println!("{}", refusal(builder.build(&[live])));

    // A merge layer that cannot merge the 2 shapes that reach it.
    let mut builder = GraphBuilder::new();
    let wide = builder.input(Shape::known(&[2, 3]));
    let narrow = builder.input(Shape::known(&[2, 5]));
    let sum = builder.add(Add::new(), &[wide, narrow]);
    println!("{}", refusal(builder.build(&[sum])));

    // A merge layer with the wrong fan-in. `Subtract` takes exactly 2 inputs.
    let mut builder = GraphBuilder::new();
    let one = builder.input(Shape::known(&[2, 3]));
    let two = builder.input(Shape::known(&[2, 3]));
    let three = builder.input(Shape::known(&[2, 3]));
    let gap = builder.add(Subtract::new(), &[one, two, three]);
    println!("{}", refusal(builder.build(&[gap])));
}

invalid input: node 2 reaches no output of the model. Name it as an output, or leave it
out of the graph
invalid input: node 2 (layer 0, `Add`) refused what reaches it: invalid input: Add cannot
merge the input shape (2, 3) with the input shape (2, 5). Axis 1 holds 3 on 1 side and 5
on the other, and neither extent is 1
invalid input: node 3 (layer 0, `Subtract`) refused what reaches it: invalid input: layer
`Subtract` takes exactly 2 inputs, and it received 3

Read the 3 parts of a layer refusal. node 2 is the position in the node list. layer 0 is the arena entry, which is the address a checkpoint path uses. `Add` is the type at that entry. The 3 parts together point at 1 line of your model.

3.10.7. The 7 merge layers

A merge layer takes several inputs and gives 1 output. The family has 7 members, and none of them holds a trainable array. They live in rustyml::neural_network::layers and the prelude exports all 7.

Subtract is the 1 member with a fixed input count, because a difference of 3 terms has no single meaning. Every other member accepts 1 input or more, and 1 input alone comes back unchanged.

Maximum and Minimum route a tie to the first input that holds the winning value, and every later input receives 0 there. An even share over the tied inputs is equally correct math. This crate pins the first-input rule, and the max pooling layers route a tied window position the same way. A NaN wins its position and keeps it, so a diverged model shows itself instead of hiding.

Multiply builds the gradient of input i from the product of the other inputs, and it never divides the whole product by input i. A quotient would give 0 / 0 wherever an input holds a 0.

The shape rule of the 6 elementwise layers

Add, Subtract, Multiply, Average, Maximum, and Minimum share 1 shape rule. The rule has 2 clauses, and the 2 clauses differ:

1. The batch axis does not broadcast. Axis 0 holds the same extent on every input, or it stays free on 1 side. A shape (2, 3) with a shape (1, 3) is an error.
2. Every axis after the batch axis does broadcast. An extent of 1 takes the extent of the other side, and a free axis leaves that axis of the output free.

Rank alignment inserts every extra axis after the batch axis. A shape (2, 3, 4) with a shape (2, 4) therefore behaves as (2, 3, 4) with (2, 1, 4). It never behaves as (2, 3, 4) with (1, 2, 4), which a rule that aligned from the left would produce.

The rule is pure shape algebra, so it answers with no tensor and no build behind it:

use rustyml::neural_network::traits::Layer;
use rustyml::prelude::*;

fn main() {
    let add = Add::new();

    // Axis 1 broadcasts, and the extent of 1 takes the extent of the other side.
    let merged = add
        .compute_output_shape_many(&[Shape::known(&[4, 3]), Shape::known(&[4, 1])])
        .unwrap();
    println!("{merged}");

    // Rank alignment inserts the extra axis after the batch axis.
    let aligned = add
        .compute_output_shape_many(&[Shape::known(&[4, 3, 5]), Shape::known(&[4, 5])])
        .unwrap();
    println!("{aligned}");

    // The batch axis is the 1 axis that does not broadcast.
    let refused = add
        .compute_output_shape_many(&[Shape::known(&[4, 3]), Shape::known(&[1, 3])])
        .unwrap_err();
    println!("{refused}");
}

(4, 3)
(4, 3, 5)
invalid input: Add cannot merge the input shape (4, 3) with the input shape (1, 3),
because the batch axis does not broadcast. Axis 0 holds 4 on 1 side and 1 on the other

An input that broadcast in the forward pass reached several positions of the output. Its gradient is the sum over every position it reached, and it comes back at its own shape. A hand-driven pass shows the whole cycle. forward_many_mut builds the layer from the tensors and completes the pass. It is the entry point for a caller that drives 1 merge layer by hand. forward_mut is its form for a layer with 1 input:

use ndarray::Array;
use rustyml::neural_network::traits::Layer;
use rustyml::prelude::*;

fn main() {
    // 2 samples of 3 features, and 1 offset per sample.
    let features = Array::from_shape_vec((2, 3), vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0])
        .unwrap()
        .into_dyn();
    let offset = Array::from_shape_vec((2, 1), vec![10.0_f32, 20.0])
        .unwrap()
        .into_dyn();

    let mut layer = Add::new();
    let mut ctx = Ctx::training();
    let sum = layer
        .forward_many_mut(&[&features, &offset], &mut ctx)
        .unwrap();
    assert_eq!(sum.shape(), &[2, 3]);
    assert_eq!(sum[[0, 0]], 11.0);
    assert_eq!(sum[[1, 2]], 26.0);

    // 1 gradient per input, each at the shape of its own input. The offset reached 3
    // positions of axis 1, so its gradient sums over those 3 positions.
    let grads = layer
        .backward_many(&Array::ones(sum.raw_dim()), &mut ctx)
        .unwrap();
    assert_eq!(grads[0].shape(), &[2, 3]);
    assert_eq!(grads[1].shape(), &[2, 1]);
    assert_eq!(grads[1][[0, 0]], 3.0);
    assert_eq!(ctx.pending_caches(), 0);

    // An inference pass parks no cache at all.
    let mut ctx = Ctx::inference();
    layer.forward_many(&[&features, &offset], &mut ctx).unwrap();
    assert_eq!(ctx.pending_caches(), 0);
}

Ctx is the per-pass context. Ctx::training() and Ctx::inference() pick the mode, and the context carries the caches, the gradients, and the non-trainable state that a training pass changes. A layer takes &self in both passes, so it holds no cache and no gradient of its own. A graph model builds its own context on every call, and a caller reaches Ctx only to drive 1 layer by hand.

Concatenate and the full-rank axis

Concatenate shares none of the rule above. It runs no batch check, and it broadcasts nothing. Every input holds the same rank, and every axis except the joined axis holds the same extent. The joined axis of the output holds the sum of the input extents.

Concatenate::new(axis: i32) is infallible, and the axis counts against the full rank. The batch axis is part of that rank, so the axis 0 is legal and joins the batches of the inputs. A negative axis counts back from the end, so the axis -1 is the last axis. The layer keeps the axis as given, and it resolves the axis against the rank of the inputs on each call. A layer that later meets another rank therefore still joins the axis you named.

The 2 rules invert each other on the same pair of shapes. An elementwise layer refuses a shape (2, 3) with a shape (1, 3), and Concatenate gives (3, 3) on the axis 0 here. An elementwise layer gives (2, 3) for a shape (2, 3) with a shape (2, 1). Concatenate refuses that same pair on the axis 0, because an extent of 1 stays an extent of 1.

use ndarray::Array;
use rustyml::neural_network::traits::Layer;
use rustyml::prelude::*;

fn main() {
    let left = Array::from_shape_vec((2, 2), vec![1.0_f32, 2.0, 3.0, 4.0])
        .unwrap()
        .into_dyn();
    let right = Array::from_shape_vec((2, 1), vec![5.0_f32, 6.0])
        .unwrap()
        .into_dyn();

    // The axis -1 is the last axis, so the join adds the feature counts.
    let mut features = Concatenate::new(-1);
    let mut ctx = Ctx::training();
    let joined = features
        .forward_many_mut(&[&left, &right], &mut ctx)
        .unwrap();
    assert_eq!(joined.shape(), &[2, 3]);
    assert_eq!(joined[[0, 2]], 5.0);

    // Each input takes the gradient of its own band, at its own shape.
    let grad = Array::from_shape_vec((2, 3), vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0])
        .unwrap()
        .into_dyn();
    let grads = features.backward_many(&grad, &mut ctx).unwrap();
    assert_eq!(grads[0].shape(), &[2, 2]);
    assert_eq!(grads[1].shape(), &[2, 1]);
    assert_eq!(grads[1][[1, 0]], 6.0);

    // The axis 0 is the batch axis, and the join stacks the samples.
    let mut samples = Concatenate::new(0);
    let mut ctx = Ctx::inference();
    let stacked = samples.forward_many_mut(&[&left, &left], &mut ctx).unwrap();
    assert_eq!(stacked.shape(), &[4, 2]);
}

3.10.8. Average: 1 name that 2 categories give

Average is the 1 name that 2 categories of the crate give to an item. The metrics category gives the averaging mode of the classification scores, and the neural network category gives the merge layer. A glob of both cannot give the name 1 meaning, so the crate settles it once.

The root of the prelude keeps metrics::Average, the averaging mode. The merge layer stays reachable as rustyml::prelude::neural_network::Average and as rustyml::neural_network::layers::Average. An alias in the import list settles the name in your own file:

// The glob of the prelude root gives the averaging mode of the classification scores.
use rustyml::prelude::*;
// The merge layer keeps its own path, and an alias settles the name in 1 line.
use rustyml::neural_network::layers::Average as AverageLayer;

fn main() {
    let mode: Average = Average::Macro;
    println!("{mode:?}");

    let mut builder = GraphBuilder::new();
    let a = builder.input(Shape::known(&[2, 4]));
    let b = builder.input(Shape::known(&[2, 4]));
    let mean = builder.add(AverageLayer::new(), &[a, b]);
    let model = builder.build(&[mean]).unwrap();
    assert_eq!(model.output_shapes()[0].to_string(), "(2, 4)");
}

A file that reads use rustyml::prelude::neural_network::*; instead of the root glob gets the merge layer under the plain name, and no averaging mode. Pick the import that matches the file, and reach for the alias only where 1 file needs both.

3.10.9. Training, inference, and persistence

The training surface of a graph model repeats the sequential surface, with 1 slice in place of each tensor:

pub fn train_batch(&mut self, xs: &[&Tensor], ys: &[&Tensor]) -> Result<f32, Error>;
pub fn fit(&mut self, xs: &[&Tensor], ys: &[&Tensor], epochs: u32) -> Result<History, Error>;
pub fn fit_with_batches(&mut self, xs: &[&Tensor], ys: &[&Tensor], epochs: u32,
    batch_size: usize) -> Result<History, Error>;

pub fn predict(&self, xs: &[&Tensor]) -> Result<Vec<Tensor>, Error>;
pub fn evaluate(&self, xs: &[&Tensor], ys: &[&Tensor]) -> Result<f32, Error>;

fit is full-batch, and each epoch runs exactly 1 gradient step over everything you give it. fit_with_batches reshuffles the sample order every epoch, trains on fixed-size chunks, and refuses a batch_size of 0. Every tensor of a batch call must agree on the sample count, and a disagreement gives Error::DimensionMismatch.

predict and evaluate both take &self and change nothing. A forward pass writes no state into any layer, so several threads can run inference against 1 built model with no lock. predict needs no compile, evaluate needs a loss, and every training call needs an optimizer and a loss.

weight_paths, weight, save_to_path, load_from_path, and load_partial_from_path behave exactly as they do on a sequential model, and the file format is the same. A path is <arena entry>.<array name>, so a shared layer contributes 1 set of paths. Read 3.9 for the format, the validation, and the partial load.

use ndarray::Array;
use rustyml::prelude::*;

// 1 definition of the architecture, which the live model and the reload target share.
fn make_arch() -> Graph {
    let mut builder = GraphBuilder::new_with_seed(7);
    let x = builder.input(Shape::with_free_batch(&[8, 4]));
    let hidden = builder.add(
        Dense::new(4, Activation::ReLU).unwrap().with_random_state(1),
        &[x],
    );
    let sum = builder.add(Add::new(), &[x, hidden]);
    let head = builder.add(
        Dense::new(1, Activation::Linear)
            .unwrap()
            .with_random_state(2),
        &[sum],
    );
    builder.build(&[head]).unwrap()
}

fn main() {
    let x = Array::from_shape_fn((8, 4), |(r, c)| (r + c) as f32 / 8.0).into_dyn();
    let y = Array::from_shape_fn((8, 1), |(r, _)| r as f32 / 8.0).into_dyn();

    let mut model = make_arch();
    model.compile(
        SGD::new(0.05, 0.9, false, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    // Mini-batches, reshuffled every epoch from the seed the builder recorded.
    let history = model.fit_with_batches(&[&x], &[&y], 30, 4).unwrap();
    assert!(*history.loss().last().unwrap() < history.loss()[0]);

    // `evaluate` scores the model you hold now, and it changes nothing.
    let scored = model.evaluate(&[&x], &[&y]).unwrap();

    let path = std::env::temp_dir().join("graph_demo.bin");
    model.save_to_path(&path).unwrap();

    // The file holds the arrays alone, so the same architecture is rebuilt first.
    let mut restored = make_arch();
    restored.load_from_path(&path).unwrap();
    restored.compile(
        SGD::new(0.05, 0.9, false, 0.0).unwrap(),
        MeanSquaredError::new(),
    );
    assert_eq!(restored.evaluate(&[&x], &[&y]).unwrap(), scored);
    std::fs::remove_file(&path).unwrap();
}

A graph that holds 1 chain agrees with the sequential model of that chain, bit for bit. The agreement covers the loss of every epoch and every trained array. The test a_chain_graph_agrees_with_the_sequential_model in src/neural_network/graph.rs pins that agreement. Use the graph model where the topology needs it, and keep Sequential everywhere else.