Struct Graph 

pub struct Graph { /* private fields */ }

An executable computation graph. Implements Module for composability.

Built via FlowBuilder. Supports parallel execution of independent nodes, observation of tagged outputs, profiling, and DOT/SVG visualization.

let g = FlowBuilder::from(Linear::new(4, 8)?)
    .through(GELU)
    .through(Linear::new(8, 2)?)
    .build()?;

// Forward pass (graph implements Module)
let y = g.forward(&x)?;

// Observation
g.end_step();
g.end_epoch();
let loss_trend = g.trend("loss");

// Visualization
let dot = g.dot();
g.svg(Some("graph.svg"))?;

Implementations

impl Graph

pub fn tagged(&self, tag: &str) -> Option<Variable>

Get the output of a tagged node from the last forward pass.

pub fn tag_names(&self) -> Vec<String>

Get all tag names defined in this graph.

pub fn collect(&self, tags: &[&str]) -> Result<()>

Snapshot current scalar values of tagged nodes into the batch buffer. Returns an error if any tag has a non-scalar output — use collect_with() with an explicit reduction for non-scalar tags.

pub fn collect_with(&self, tags: &[&str], reduce: Reduce) -> Result<()>

Snapshot tagged node values into the batch buffer using a reduction. Tag group names are automatically expanded to their members. Each tag’s output is reduced to a scalar and recorded individually.

pub fn record(&self, tag: &str, values: &[f64])

Inject external scalar values into the batch buffer.

Recorded values accumulate per step and are averaged on flush(). Use trend() to read epoch history for training decisions (early stopping, LR scheduling).

For human-facing output (terminal, live dashboard), use Monitor::log() instead.

pub fn record_scalar(&self, tag: &str, value: f64)

Record a single scalar value. Convenience wrapper around record.

pub fn latest_metrics(&self) -> Vec<(String, f64)>

Return the latest epoch value for every tag in the epoch history.

Tree-aware: automatically collects from labeled child subgraphs with dotted prefixes (e.g. a child labeled "subscan" with tag "ce" appears as "subscan.ce"). Parent metrics come first, then children in registration order.

Useful for bridging graph observation into Monitor::log(). Returns an empty vec if no epochs have been flushed yet.

Use latest_metrics_local() if you only want this graph’s own metrics.

pub fn latest_metrics_local(&self) -> Vec<(String, f64)>

Return latest epoch values for this graph only, without child metrics.

Use this when you need only the local metrics (e.g. when children report on a different cadence). See latest_metrics() for the tree-recursive version.

pub fn collected(&self, tag: &str) -> Vec<f64>

Read raw batch buffer for a tag (all values since last flush).

pub fn flush(&self, tags: &[&str])

Compute batch means, append to epoch history, clear batch buffer. Call once per epoch. If tags is empty, flushes all buffered tags.

Tree-aware: automatically recurses into labeled child subgraphs, so a single parent.flush(&[]) flushes the entire tree. Child buffers that are already empty (e.g. flushed separately) are skipped safely.

If you need different flush cadences per subgraph (e.g. flushing a child every 10 parent epochs), use flush_local() on both the parent and the child to manage them independently:

// Every epoch: flush parent only
parent.flush_local(&[]);
// Every 10 epochs: flush the child
if epoch % 10 == 0 {
    parent.child_graph("slow_child").unwrap().flush_local(&[]);
}

pub fn flush_local(&self, tags: &[&str])

Flush only this graph’s own batch buffer, without recursing into children.

Use this when you need independent flush cadences per subgraph. See flush() for the tree-recursive version.

pub fn flush_count(&self) -> usize

Number of flush calls that produced data.

pub fn trend(&self, tag: &str) -> Trend

Get epoch-level trend for a tag.

pub fn trends(&self, tags: &[&str]) -> TrendGroup

Get trends for multiple tags. Tag group names are automatically expanded to their member tags.

pub fn reset_trend(&self, tags: &[&str])

Clear epoch history. If tags is empty, clears all. Tag group names are automatically expanded.

pub fn traces(&self, tag: &str) -> Option<Vec<Variable>>

Get per-iteration trace outputs from loop nodes. Returns the trace buffer for the loop node associated with the given tag. The tag should be set on a node after the loop (the loop output flows to it). Returns None if no loop node with a trace buffer is found.

pub fn traces_by_node(&self, node_id: &str) -> Option<Vec<Variable>>

Get trace buffer directly from a loop node by node ID.

pub fn last_trace(&self, tag: &str) -> Option<Variable>

Get the last trace output from the most recent loop iteration.

Convenience wrapper around traces() that returns only the final iteration’s trace. Useful for chaining loops where the last output of one (e.g. scan) feeds into the next (e.g. read).

Returns None if the tag has no associated loop or the body produced no traces.

pub fn eta(&self, total_epochs: usize) -> f64

Estimated time remaining based on average flush duration.

Returns seconds remaining. Returns 0.0 if no flushes have occurred yet.

impl Graph

pub fn enable_profiling(&self)

Turn on per-node and per-level timing for subsequent forward calls.

pub fn disable_profiling(&self)

Turn off timing. Subsequent forward calls have zero profiling overhead.

pub fn profiling(&self) -> bool

Whether profiling is currently enabled.

pub fn profile(&self) -> Option<Profile>

Timing data from the most recent forward call, or None.

pub fn timing(&self, tag: &str) -> Duration

Duration of a tagged node from the most recent forward call.

pub fn collect_timings(&self, tags: &[&str])

Snapshot tagged node durations into the timing batch buffer. If tags is empty, all tagged nodes with timing data are collected.

pub fn flush_timings(&self, tags: &[&str])

Compute batch mean, append to timing epoch history, clear buffer. If tags is empty, flushes all buffered tags.

pub fn timing_trend(&self, tag: &str) -> Trend

Epoch-level trend over the timing history of a tagged node. Values are mean execution times in seconds.

pub fn timing_trends(&self, tags: &[&str]) -> TrendGroup

TrendGroup for timing trends of the given tags (expands groups).

pub fn reset_timing_trend(&self, tags: &[&str])

Clear timing epoch history. If tags is empty, clears all.

impl Graph

pub fn dot(&self) -> String

Graphviz DOT representation of the graph structure.

Nodes are grouped by execution level, shaped by type (input/output/activation/norm), and annotated with parameter counts. Render with: dot -Tsvg graph.dot -o graph.svg

std::fs::write("graph.dot", g.dot()).unwrap();

pub fn dot_with_profile(&self) -> String

Timing-annotated DOT using the most recent Profile. Nodes are color-coded green → yellow → red by relative execution time, with per-node microsecond timings and per-level wall clock + parallelism.

pub fn svg(&self, path: Option<&str>) -> Result<Vec<u8>>

Render the structural graph as SVG via the Graphviz dot binary.

Requires Graphviz installed (apt install graphviz / brew install graphviz). Optionally writes to path if provided. Returns the raw SVG bytes.

g.svg(Some("graph.svg"))?;

pub fn svg_with_profile(&self, path: Option<&str>) -> Result<Vec<u8>>

Render the timing-annotated graph as SVG via the Graphviz dot binary.

Requires a prior profiled forward pass. Optionally writes to path.

impl Graph

pub fn plot_html(&self, path: &str, tags: &[&str]) -> Result<()>

Generate a self-contained HTML file with training curves from epoch history.

This plots data accumulated via record() / flush() — the graph’s observation system for metrics that feed back into training decisions.

For a full dashboard with resource graphs, epoch log, and graph SVG, use Monitor::save_html() instead.

Tag group names are expanded. If tags is empty, all epoch history is plotted. Uses inline Canvas JS — no external dependencies.

pub fn plot_timings_html(&self, path: &str, tags: &[&str]) -> Result<()>

Generate a self-contained HTML file with per-node timing trend curves.

Requires prior collect_timings/flush_timings calls. Tag groups are expanded. If tags is empty, all timing history is plotted.

pub fn export_trends(&self, path: &str, tags: &[&str]) -> Result<()>

Export epoch history to CSV (one column per tag, one row per epoch).

Tag groups are expanded. If tags is empty, all epoch history is exported.

pub fn export_timing_trends(&self, path: &str, tags: &[&str]) -> Result<()>

Export timing epoch history to CSV (one column per tag, one row per epoch).

Tag groups are expanded. If tags is empty, all timing history is exported.

pub fn write_log( &self, path: &str, total_epochs: usize, tags: &[&str], ) -> Result<()>

Write a human-readable training log with per-epoch metrics and ETA. total_epochs is used for ETA (0 to omit).

impl Graph

pub fn snapshot_cpu(&self) -> Result<ModelSnapshot>

Snapshot the graph’s current state onto CPU.

Parameters are detached from autograd so no grad metadata crosses the thread boundary. Buffers have no grad and are copied as-is.

let snap = graph.snapshot_cpu()?;
worker.submit(move || {
    snap.save_file("epoch_10.fdl.gz").unwrap();
});

impl Graph

pub fn tree_children(&self) -> HashMap<&str, &Graph>

Direct children: label -> child graph.

pub fn child_graph(&self, label: &str) -> Option<&Graph>

Get a direct child graph by label (one level only).

pub fn subgraph(&self, path: &str) -> Result<&Graph>

Get a subgraph at any depth via dot-path.

pub fn is_composed(&self) -> bool

Whether this graph has been composed into a parent graph.

pub fn internal_tags(&self) -> &HashSet<String>

Tags marked as internal (hidden from parent resolution).

pub fn validate_path(&self, path: &str) -> Result<PathKind>

Validate that a path resolves, returning what it resolves to.

pub fn parameters_at(&self, path: &str) -> Result<Vec<Parameter>>

All parameters at a label path.

pub fn named_parameters_at( &self, path: &str, ) -> Result<Vec<(String, Parameter)>>

Named parameters at a label path, using the target’s own namespace. For subgraphs: delegates to the child graph’s named_parameters(). For tags: qualifies with the tag name as prefix.

pub fn named_buffers_at(&self, path: &str) -> Result<Vec<(String, Buffer)>>

Named buffers at a label path, using the target’s own namespace.

pub fn freeze(&self, path: &str) -> Result<()>

Freeze all parameters at the given label path.

pub fn thaw(&self, path: &str) -> Result<()>

Thaw (unfreeze) all parameters at the given label path.

pub fn is_frozen(&self, path: &str) -> Result<bool>

Check if all parameters at the path are frozen. Returns true only if there are parameters and ALL are frozen.

pub fn load_subgraph_checkpoint( &self, path: &str, file: &str, ) -> Result<LoadReport>

Load a checkpoint into a specific subgraph.

The checkpoint’s structural hash is validated against the target subgraph’s hash. Named parameters/buffers are matched within the subgraph’s own namespace.

pub fn set_training_at(&self, path: &str, training: bool) -> Result<()>

Set training mode on a specific subgraph or tagged module.

pub fn tagged_at(&self, path: &str) -> Result<Option<Variable>>

Get a tagged output by label path. Returns Err if the path doesn’t exist (null – wiring bug). Returns Ok(None) if the path exists but hasn’t been computed yet (nil). Returns Ok(Some(v)) if the value is available.

pub fn collect_at(&self, paths: &[&str]) -> Result<()>

Collect metrics from label paths into observation buffers. Each path must resolve to a tag (not a subgraph). Metrics are stored in the target graph’s batch buffer.

pub fn record_at(&self, path: &str, value: f64) -> Result<()>

Record a scalar metric at a label path. For dotted paths, the metric is stored in the target graph’s buffer under the final segment name.

pub fn trend_at(&self, path: &str) -> Result<Trend>

Get trend for a label-path metric. For dotted paths, reads from the target graph’s epoch history.

impl Graph

pub fn tree_summary(&self) -> String

Generate a tree summary showing subgraph structure, tags, and parameter counts.

pub fn param_summary(&self) -> String

Per-subgraph parameter count breakdown.

impl Graph

pub fn reset_state(&self)

Clear all forward-reference state buffers to None. Call when starting inference on a new sequence.

pub fn detach_state(&self)

Break gradient chain on forward-reference state buffers and module state. Call between training steps to prevent unbounded graph growth.

pub fn has_state(&self) -> bool

Returns true if this graph has forward-reference state.

pub fn end_step(&self)

End-of-step housekeeping: detach state (cut gradient chain but preserve values for the next forward), collect timings, increment step counter.

For recurrent models this implements truncated BPTT — state carries over between steps but gradients don’t flow across step boundaries. Call end_sequence to fully wipe state when starting a new independent sequence.

for token in sequence {
    let y = graph.forward(&token)?;
    // ... backward, optimize ...
    graph.end_step();       // keep state, cut gradients
}
graph.end_sequence();       // wipe state for next sequence

pub fn end_sequence(&self)

End-of-sequence housekeeping: fully reset state buffers to None. Call between independent sequences so the model starts fresh.

For non-recurrent graphs (no forward refs) this is a no-op.

pub fn end_epoch(&self)

End-of-epoch housekeeping: flush all observation and timing buffers, increment epoch counter.

pub fn step_count(&self) -> usize

Number of completed training steps.

pub fn epoch_count(&self) -> usize

Number of completed training epochs.

pub fn tag_group(&self, name: &str) -> Option<&[String]>

Get member tags of a tag group, or None if not registered.

pub fn forward_multi(&self, inputs: &[Variable]) -> Result<Variable>

Forward with multiple inputs (for graphs with Input ports). Inputs are in declaration order: From entry first, then each Input.

pub fn set_device(&self, device: Device)

Move all parameters, state buffers, and module buffers to a device.

pub fn named_parameters(&self) -> Vec<(String, Parameter)>

Return parameters with qualified names: "prefix/param_name".

The prefix is the tag name if the node is tagged, otherwise the node ID (e.g. "linear_1"). When a node has multiple parameters with the same name, suffixes _0, _1, … are appended to disambiguate.

pub fn named_buffers(&self) -> Vec<(String, Buffer)>

Return buffers with qualified names, using the same prefix logic as named_parameters().

pub fn label(&self) -> Option<&str>

Human-readable label set via FlowBuilder::label().

pub fn structural_hash(&self) -> &str

Full 64-character hex structural hash (computed lazily, cached).

pub fn short_hash(&self) -> &str

First 8 characters of the structural hash.

pub fn save_checkpoint(&self, path: &str) -> Result<()>

Save all parameters and buffers to a checkpoint file.

Embeds the structural hash for architecture validation on load. Supports .gz extension for gzip compression.

pub fn load_checkpoint(&self, path: &str) -> Result<LoadReport>

Load parameters and buffers from a checkpoint file.

Validates the structural hash against this graph’s architecture. Returns a LoadReport describing what was loaded, skipped, or missing.

Trait Implementations

impl Metrics for &Graph

Graph only: &model — reads latest epoch history.

fn into_metrics(self) -> Vec<(String, f64)>

Convert into owned (name, value) pairs for recording.

impl Module for Graph

fn name(&self) -> &str

Human-readable type name used as node ID prefix in graph visualization. Override to return a lowercase identifier (e.g., “linear”, “gelu”).

fn as_graph(&self) -> Option<&Graph>

Upcast to Graph for hierarchical tree composition. Override in Graph to enable subgraph nesting with label-path addressing.

fn structural_hash(&self) -> Option<String>

SHA-256 hex hash of module architecture for checkpoint validation. Override in composite modules (Graph) that compute a deterministic hash from their topology and parameter shapes.

fn forward(&self, input: &Variable) -> Result<Variable>

Run the forward pass on input and return the result.

fn parameters(&self) -> Vec<Parameter>

Return this module’s learnable parameters. Default: recursively collects from sub_modules() with pointer dedup. Leaf modules should override to return their own parameters.

fn set_training(&self, training: bool)

Set training/eval mode. Affects Dropout, BatchNorm, etc. Override in modules with mode-dependent behavior.

fn move_to_device(&self, device: Device)

Move all parameters and buffers to the given device. Override in modules like BatchNorm that hold non-parameter state.

fn buffers(&self) -> Vec<Buffer>

Return this module’s non-learnable persistent buffers (e.g., running stats). Default: recursively collects from sub_modules() with pointer dedup. Leaf modules should override to return their own buffers.

fn sub_modules(&self) -> Vec<Rc<dyn Module>>

Return direct child modules for recursive tree walks. Override in composite modules (loops, switches, gates).

fn train(&self)

Set training mode. Shorthand for set_training(true).

fn eval(&self)

Set eval mode. Shorthand for set_training(false).

fn trace(&self) -> Option<Variable>

Return per-iteration side output for loop tracing. Override in loop body modules that capture trajectory data (e.g., attention fixation points). Returns None by default. When Some, the loop executor collects traces accessible via Graph::traces().

fn as_named_input(&self) -> Option<&dyn NamedInputModule>

Upcast to NamedInputModule for multi-input graphs. Override in types that implement NamedInputModule to enable receiving additional named inputs via graph using().

fn reset(&self)

Reset internal state (e.g. recurrent hidden state) between sequences. Called by loops before iterating to clear stale tensors whose grad_fns may reference freed saved tensors. Override in stateful modules.

fn detach_state(&self)

Detach internal state from the computation graph (for truncated BPTT). Called between training steps to break gradient chains on state carried across forward passes (e.g., recurrent hidden state). Override in stateful modules.