Graph is Meganeura’s representation of your model’s forward pass. You build it imperatively by calling methods that append nodes and return NodeId handles. When you pass the finished graph to build_session, Meganeura runs autodiff, e-graph optimization, and GPU compilation — all from the same declarative description.
Creating a graph
CallGraph::new() to create an empty graph:
g.input(...), g.parameter(...), g.matmul(...), and so on appends a node to the graph and returns a NodeId.
NodeId
NodeId is a u32 alias that acts as a stable handle to a node. You pass it as an argument to any operation that consumes that tensor:
Adding inputs
Useg.input for floating-point feature tensors and g.input_u32 for integer token indices:
"x", "tokens") is what you use later when feeding data to the session:
Adding parameters
Parameters are learnable weights. The optimizer updates them each step.HuggingFace naming convention
When loading pretrained weights from safetensors files, Meganeura matches parameter names exactly. Follow the standard HuggingFace naming scheme so weights load automatically:| Layer type | Parameter name pattern |
|---|---|
| Linear weight | layer_name.weight |
| Linear bias | layer_name.bias |
| RmsNorm scale | layer_name.weight |
| Attention Q | layer_name.self_attn.q_proj.weight |
nn:: layer constructors follow this convention automatically when you pass a name matching the checkpoint’s key prefix.
Adding constants
Useg.constant for fixed tensors embedded in the graph, and g.scalar for single float values:
Setting outputs
Callg.set_outputs with the loss node (or any other terminal nodes) before passing the graph to build_session:
set_outputs before building a session. The autodiff engine traces backward from output nodes to compute gradients.
Calling toposort before autodiff
The e-graph optimizer can append nodes out of insertion order. Calltoposort() before running autodiff to ensure every node’s inputs have lower IDs than the node itself:
build_session calls this for you internally, but if you are calling autodiff::differentiate directly you must sort first.
toposort() also strips Nop nodes that were marked dead by the optimizer. The returned graph has consecutive IDs with no gaps.Full example: two-layer MLP for MNIST
The following is taken directly fromexamples/mnist.rs:
Define inputs and parameters
Call
g.input and g.parameter to create leaf nodes. Each returns a NodeId.Build the forward pass
Chain operations —
matmul, bias_add, relu — by passing NodeIds as arguments. Each call appends a node and returns its NodeId.Attach a loss function
Pass the logits and labels to
g.cross_entropy_loss (or another loss). The output is a scalar [1] node.Available loss functions
| Method | Description | ||
|---|---|---|---|
g.cross_entropy_loss(logits, labels) | Categorical cross-entropy; both inputs [N, C]. | ||
g.bce_loss(pred, labels) | Binary cross-entropy; pred should be in (0, 1). | ||
g.mse_loss(pred, target) | Mean squared error: mean((pred - target)²). | ||
g.l1_loss(pred, target) | Mean absolute error: `mean( | pred - target | )`. |
[1] node suitable as the single output to set_outputs.