The rml-graph IR

In short: a model is <stem>.json plus <stem>.safetensors. The JSON lists values (a name, a dtype of f32 or i32, and a static shape), the graph's inputs and outputs, constants bound to stored tensors or inline data, and nodes in topological order. LoadGraph maps the weights and ParseGraph checks structure: every value is defined exactly once and before it is read, and every constant matches its recorded type. Op semantics come second: InferTypes compiles each node against its inputs' types and rejects any op, arity, attribute or recorded shape that disagrees with rml's own rules. A graph that passes both never aborts at run time.

The format sits in the middle of the inference stack: the exporters write it, and the executor runs it. The idea of a model as a dataflow graph is on computation graphs, which the header of graph.h links.

Two files

A model is two files that share a stem (graph_ir.md L27–30): <stem>.json holds the graph and <stem>.safetensors holds every weight and large constant. The format is inference only, with no backward pass (L3–6), and both export routes, direct and StableHLO, write the same format (L8–25).

LoadGraph(stem) reads the JSON into a string, opens the safetensors file through io::SafetensorsReader only if it exists, and hands both to ParseGraph, prefixing any error with the JSON path. A stored constant comes back from SafetensorsReader::Get as a tensor view onto the mapped bytes that shares ownership of the mapping, so the weights stay mapped for as long as the graph holds them, even after the reader itself is gone. The one exception is a tensor misaligned for its element type, which is copied rather than read with unaligned loads. How the mapping works is on Model files.

The JSON layout

The top-level object carries "format": "rml-graph", "version": 1, an optional producer string, and the arrays values, inputs, outputs, constants and nodes (graph_ir.md L34–54). The rules (L56–63):

• A value's id is its index in values, and every value is produced exactly once: as a graph input, a constant, or a node output.
• Shapes are static. Exporters trace at a fixed batch size, which is why each zoo model ships one fixture per size.
• dtype is "f32" or "i32". Exporters turn booleans into f32 masks of 0 and 1, and 64-bit integers into i32.
• Nodes are in topological order: a node may only read values defined before it.

In C++ the parsed graph is three small structs (graph.h L28–51):

Worked example: the tiny MLP

The committed PyTorch MLP fixture, testdata/torch_mlp.json, has 15 values, 6 constants and 8 nodes. Its input x is f32 [2, 12] (batch 2) and its output is value 14, f32 [2, 4]. The first layer reads:

{"name": "x",  "dtype": "f32", "shape": [2, 12]}       value 3, the graph input
{"name": "v7", "dtype": "f32", "shape": [12, 16]}      value 4, constant "const.4"
{"op": "matmul", "inputs": [3, 4], "outputs": [5]}     [2, 12] x [12, 16] -> [2, 16]
{"op": "add",    "inputs": [5, 0], "outputs": [6]}     + layers.0.bias [16], broadcast
{"op": "relu",   "inputs": [6],    "outputs": [7]}

The pattern repeats for the second and third layers, the last ending in an add with no activation. Each weight is stored in the safetensors file under a const.<id> name, and the biases keep their PyTorch parameter names. Shape inference re-derives every recorded shape here: matmul gives [2, 16] from NumPy rules, the bias add broadcasts [16] against it, and relu keeps its input's shape.

Loader validation

ParseGraph parses with nlohmann/json's non-throwing mode, so malformed text is a Status, not an exception. It then checks, in order:

Because nodes are read in file order and a read of an undefined value fails, topological order is enforced by the same pass that parses, with no separate sort. The loader does not look at op names: graph_ir.md lists unknown ops and bad attributes among the loader's rejections (L65–68), and in the code those checks run in the second pass, when an executor is built.

The op set

Binary ops broadcast with NumPy rules, axes may be negative, and unless stated otherwise inputs and outputs are f32 with one output per op (graph_ir.md L72–74). The full table, with every attribute, is graph_ir.md L76–104. In ops.cc the set is three tables, and IsKnownOp is membership in any of them:

Only the shape and data-movement ops take i32 values (gather takes them as its index); everything else is f32 only, which OtherOps records per op in an OpSpec next to its minimum and maximum input count (ops.cc L849–854). The set is small on purpose. Exporters lower framework ops onto it, so JAX's dot_general becomes transpose + reshape + matmul and TensorFlow's NHWC convolutions become transpose + conv2d + transpose, and an op an exporter cannot lower fails the export with the framework's op name (graph_ir.md L106–110).

Shape inference

InferTypes(graph) starts from the recorded types of every value, then walks the nodes in order. For each node it gathers its inputs' types, calls CompileOp, and compares every inferred output with the type the exporter recorded. The first disagreement fails with the node's index and op, for example "node 3 (matmul) produces f32 [2, 16] but v12 is recorded as f32 [2, 8]" (L27–31). Names are kept, so the returned vector is the graph's values with every type confirmed (shape_inference.h L13–17). Executor::Create calls it first, so no executor exists for a graph whose shapes do not check.

Shape rules and kernels live together. CompileOp returns a CompiledOp: the output types and the Kernel that will run the node. One function therefore decides both what an op produces and how it is computed, so they cannot drift apart. CompileAny dispatches on the op name: unary ops need one f32 input and keep its shape, binary ops need two and take the broadcast shape (Broadcast), and the rest go through their OpSpec, which checks arity and dtype before the op's own compiler runs. An unknown name is Unsupported. CompileOp then requires the node to list exactly as many outputs as the op produces (L922–928).

Static checking leaves one thing for run time: data the loader cannot see. A Kernel fails only on such data (ops.h L22–25), and the one case is a gather index outside its axis, which returns OutOfRange instead of reading out of bounds (ops.cc L689–699). The kernels themselves, the constant-weight packing CompileOp does for matmul, and the fast-path layouts belong to Executing a graph.

Tests