Exporters and the model zoo — rml internals

In short: every framework has two ways in. The direct exporters walk the framework's own traced graph (an ATen ExportedProgram, a jaxpr, or a frozen GraphDef); the StableHLO importer reads the StableHLO module each framework's compiler path produces (torch-mlir, jax.export, XLA's StableHLO stage). Both call one shared Builder, which computes a NumPy example value for every op, so each lowering is shape-checked node by node and ops on constants fold away. Clean-up passes then cancel transposes, drop broadcasts, merge repeated work and fuse written-out means, softmaxes, layer norms and GELUs back into single ops, and hoist an unrolled loop's per-step input projection into one product. A fixture script exports all 19 zoo models by all six routes, and C++ tests require rml to match every framework and all twelve results per model to agree.

Everything here runs offline in a virtualenv pinned by tools/requirements-export.txt (NumPy 2.5.4, PyTorch 2.14.1, JAX 0.12.0, TensorFlow 2.21.0 and torch-mlir 20261001), created with Python 3.12 or newer; rml itself never links or needs Python. The format the exporters write is described on The rml-graph IR, and where this fits in the whole stack on The inference stack.

Both routes end in the same builder and passes, so all six produce the same kind of file and are checked the same way, and against each other.

The shared builder

common.py holds the Builder that every exporter calls once per lowered op, and a NumPy reference implementation of every IR op (REFERENCE, L163). Carrying a concrete example value for every graph value is the central idea: it gives each value its recorded shape and dtype, it lets an exporter check its lowering against the framework node by node, and it folds constants.

Piece	Behavior	Source
`canonical`	Maps arrays onto the IR's two dtypes: floats and booleans to f32, integers to i32. An integer constant outside the i32 range is an `ExportError`.	L28
`Builder.op`	Rejects an op rml-graph does not have, computes the output's example with the NumPy reference, and returns a constant instead of a node when every input is constant. Ops that need f32 reject i32 inputs.	L279
`Builder.op_multi`	The same for an op with several outputs (`top_k` returns values and indices); the reference must produce exactly `count` examples.	L293
`Builder.output`	An output that is an input or a constant is routed through an identity reshape, because graph outputs must be computed.	L306
`Builder.evaluate`	Runs the whole graph on new inputs with the NumPy reference ops, so a written graph can be checked against the framework on any batch.	L318
`Builder.write`	Runs the passes, renumbers values densely, inlines small finite unnamed constants (at most 16 elements) in the JSON and stores the rest in safetensors. JSON has no infinities, so non-finite constants always go in the weights file.	L330
`upsample_nearest`	Nearest-neighbour upsampling of NCHW by an integer factor as reshape + `broadcast_to` + reshape, shared by the exporters that meet it (the U-Net).	L384
`rml_graph_run_binary`	Finds the `--config=opt` `rml-graph-run` by asking Bazel, because `.bazel/bin` points at whichever configuration built last; `RML_GRAPH_RUN` overrides it.	L392
`dump_json`	One node or value per line, for readable diffs and small files.	L421
`write_safetensors`	Writes F32 and I32 tensors in order, with every data offset 64-byte aligned.	L439
`write_golden`, `check_close`	Save the framework's inputs and outputs as `input.i` / `output.i`, and compare arrays within atol 1e-5 + rtol 1e-4.	L474

The direct exporters

Each direct exporter starts from the graph its framework already produces, so it sees what the framework would run, not a re-trace of Python code. Each one lowers op by op, checks every lowered value's shape against the framework's recorded shape, and raises ExportError naming the framework op when it has no lowering. There is no silent fallback. They keep the framework's high-level ops where the IR has them (graph_ir.md L10–12).

Exporter	Starts from	Notable lowerings	Source
PyTorch	`torch.export.export`, then the default decompositions to core ATen. Scaled dot-product attention and `upsample_nearest2d` are removed from the decomposition table (`_KEEP`, L17), so attention stays one `sdpa` op instead of a masked-softmax chain.	Parameters and buffers become named constants. A node whose inputs are all constant runs in torch itself (`_fold`, L278–291), which handles masks and aranges.	from_torch.py L318
JAX	`jax.make_jaxpr(fn)(params, *inputs)`. Nested calls (`jit`, `pjit`, `custom_jvp_call`, `custom_vjp_call`, `remat` and others) are inlined (L28).	`dot_general` becomes `transpose` + `reshape` + `matmul` + `reshape` (L62–74); convolutions and `reduce_window` map onto NCHW `conv2d` and pools. A `scan` is unrolled, its body inlined once per step, since shapes are static and the trip count is known (L264–293). Parameter leaves are named by their pytree path.	from_jax.py L303
TensorFlow	`tf.function` traced at fixed input shapes, then `convert_variables_to_constants_v2` freezes variables into `Const` ops, walked in topological order.	NHWC `Conv2D` and pools become `transpose` + NCHW op + `transpose`; SAME padding must be symmetric, and EXPLICIT padding equal on every side (L57–104).	from_tf.py L280

The StableHLO route

StableHLO is the MLIR dialect all three frameworks lower to, so one importer serves all of them (from_stablehlo.py L1–14). It parses the module's text with the MLIR Python bindings that ship in jaxlib (parse, L554–564) and walks func.func @main op by op. The leading arguments of @main become stored constants and the rest graph inputs (export, L567–586). The contraction, convolution and pooling lowerings are borrowed from the JAX exporter's Lowering, since StableHLO's ops carry the same dimension numbers as jax's primitives.

Producer	How it gets StableHLO	Source
`jax_module`	`jax.export.export(jax.jit(fn))`; parameters are passed as named constants, and an export that dropped an unused model input is an error.	L592
`tf_module`	Freezes the variables, wraps the frozen function in `tf.function(jit_compile=True)` and asks `experimental_get_compiler_ir(...)(stage="stablehlo")`.	L608
`torch_module`	Runs `torch.export` with the decompositions torch-mlir needs (including `upsample_nearest2d`), saves the program, and runs torch-mlir's `fx.export_and_import(..., output_type='stablehlo')` in a subprocess, because its MLIR runtime and those in jaxlib and TensorFlow crash in one process.	L620

Construct	What the importer does	Source
Text clean-up	torch-mlir's `dense_resource` weights are rewritten as hex `dense` literals, and TensorFlow's bit-packed `i1` literals are unpacked to one byte per element, which is what jaxlib's parser reads.	L83–122
Op dispatch	Unary, binary, compare and reduce ops map through tables; ops rml-graph lacks (`floor`, `ceil`, `round_nearest_even`, `sign`, `remainder`) are allowed only on constants, where they fold. Anything else, including an unknown `custom_call`, raises `ExportError` naming the op.	L217–329
`func.call`, `stablehlo.composite`	Calls are inlined. A composite is lowered directly when rml-graph has it (`chlo.top_k` becomes `top_k`), otherwise its decomposition is inlined.	L331–337
`stablehlo.while`	Unrolled: the condition runs in the builder, where the loop counter is a constant, and the body is inlined trip by trip until it is false. A condition not known at export time, or more than `MAX_LOOP_TRIPS` = 4096 trips, is an error.	L372–382
`stablehlo.sort`	Only a descending sort along the last axis, optionally carrying an iota, is accepted; it becomes `top_k` of every element, which puts ties in index order.	L450–471
`stablehlo.gather`	Whole slices along adjacent indexed axes: the operand is flattened over those axes and read with one rml-graph `gather`. Constant indices are clamped as StableHLO specifies; run-time indices (an embedding lookup) may index only one axis.	L502–551

Clean-up passes

Lowering leaves debris: identity reshapes, stacked transposes, the NHWC ↔ NCHW pairs TensorFlow's convolutions produce, the explicit broadcasts StableHLO spells out, and layer norms, softmaxes and GELUs written out as chains of elementwise ops, and the per-step input projections JAX and TensorFlow leave inside an unrolled recurrent loop. passes.py removes it before the graph is written, running the passes until none changes anything and dropping dead code after each round (run_passes, L586–591).

Pass	Rewrite	Source
`simplify_shape_ops`	Drops identity reshapes and transposes, collapses reshape-of-reshape, composes transposes (transpose(transpose(y, p), q) = transpose(y, p[q])), and turns a reshape of a reduction back to its input's rank into the reduction with `keep_dims`.	L41–80
`sink_transposes`	Rewrites f(transpose(x)) as transpose(f(x)) for elementwise f, permuting a constant other operand to match; only single-reader transposes move. Sinking below a conv's bias add and ReLU brings it next to the next layer's inverse transpose, where the two cancel.	L83–123
`drop_implicit_broadcasts`	Removes `broadcast_to`, and reshapes that only prepend size-1 axes, when every reader is a binary op that broadcasts anyway.	L131–162
`shrink_broadcast_constants`	A constant folded from a broadcast repeats along some axes; it is shrunk back to the slice it repeats.	L165–199
`unbatch_matmul_weights`	A constant matmul weight of rank > 2 whose batch slices are all equal (PyTorch's LSTM input projection) becomes one 2-D slice, so the executor runs the batch as one matrix product.	L202–227
`hoist_step_projections`	Products `reshape(slice(X, axis, [t, t+1])) @ W` that share one X, one axis and one constant 2-D W (an unrolled RNN or LSTM's input projection, step by step) become a single `X @ W` inserted before the first of them, each step reading its slice of the result. PyTorch's export already has that shape; this gives JAX and TensorFlow the same, taking their `lstm_b64` graphs from 128 to 65 matmuls. TensorFlow's slices also name the axes they keep whole, which the match skips.	L230–288
`fuse_reductions`	`reduce_sum(x) / N` becomes `reduce_mean`; exp(x − max x) / sum exp(x − max x) over one axis becomes `softmax`.	L306–343
`eliminate_common_subexpressions`	Merges nodes applying the same op with the same attributes to the same inputs.	L346–364
`fuse_layer_norm`	(x − mean x) · rsqrt(mean (x − mean x)² + eps) over the last axis becomes `layer_norm`, folding a following constant · weight + bias.	L431–470
`fuse_gelu`	0.5 x (1 + tanh(√(2/π)(x + 0.044715 x³))) and 0.5 x (1 + erf(x/√2)), multiplied in any order, become `gelu` with the matching `approximate`.	L527–546
`share_identical_constants`	Points every reader at one copy of each distinct constant array, preferring a named one. Unrolled loops fold the same weight once per step, and sharing keeps it in cache.	L549–569
`eliminate_dead_code`	Drops nodes and constants nothing reads.	L572–583

The fusions are why the three frameworks now export the tiny transformer block to similar sizes: 32 nodes from PyTorch, 42 from JAX and 29 from TensorFlow, where JAX and TensorFlow hand over layer norm, softmax and GELU decomposed. rml's own docs/progress.md checks off export-time fusion of layer norm and GELU (L163). Fusion inside rml, below the graph, is on Lowering and fast CPU.

The zoo

tools/rml_export/zoo/ defines each model once per framework (torch_zoo.py, jax_zoo.py, tf_zoo.py). All three load their parameters from weights(), which draws them from a NumPy generator seeded per model (1234 plus the model's index), so the three definitions compute the same function from identical weights (weights.py L1–7, L98–202). Each model has a "tiny" size, committed as a test fixture, and a "bench" size, which is timed (CONFIGS, L19–76). encoder and tiny_gpt take int32 token ids (TOKEN_MODELS, L16). Recurrent models are unrolled: JAX's lax.scan by the exporter, PyTorch's nn.RNN / nn.LSTM by torch.export, and TensorFlow's by a Python loop over time steps. Unrolled JAX and TensorFlow graphs then multiply each step's input by the same weight; hoist_step_projections turns those products back into one (Clean-up passes).

The table lists the 19 models in MODELS (L11–13) with their tiny node counts per route, read from the committed fixtures in rml/graph/testdata.

Model	Tiny	Bench	Direct (torch · jax · tf)	StableHLO (torch · jax · tf)
`mlp`	12 → 16 → 16 → 4	784 → 512 → 256 → 10	8 · 8 · 8	11 · 8 · 9
`cnn`	8×8, channels 1, 3, 4, 5 classes	28×28, channels 1, 16, 32, 10 classes	10 · 19 · 12	18 · 19 · 18
`transformer_block`	seq 5, dim 8, 2 heads, hidden 16	seq 64, dim 256, 4 heads, hidden 1024	32 · 42 · 29	50 · 42 · 37
`layernorm_gelu`	seq 3, dim 8, hidden 12	seq 64, dim 512, hidden 2048	6 · 6 · 4	7 · 6 · 7
`vit`	8×8, patch 4, dim 8, 2 layers	32×32×3, patch 4, dim 192, 4 layers	77 · 95 · 68	113 · 95 · 84
`encoder`	vocab 11, seq 5, dim 8, 2 layers	vocab 1024, seq 64, dim 256, 4 layers	71 · 96 · 61	106 · 96 · 78
`tiny_gpt`	vocab 11, seq 5, dim 8, 2 layers	vocab 1024, seq 64, dim 256, 4 layers	71 · 98 · 63	108 · 98 · 80
`rnn`	seq 5, input 3, hidden 4	seq 64, input 64, hidden 256	46 · 42 · 41	41 · 42 · 42
`lstm`	seq 5, input 3, hidden 4	seq 64, input 64, hidden 256	106 · 102 · 101	101 · 147 · 102
`ssm`	seq 5, input 3, hidden 4	seq 256, input 64, hidden 256	27 · 34 · 28	28 · 45 · 33
`autoencoder`	data 12, hidden 8, latent 2	data 784, hidden 512, latent 32	14 · 14 · 14	19 · 17 · 15
`vae`	data 12, hidden 8, latent 2	data 784, hidden 512, latent 32	20 · 20 · 20	26 · 23 · 21
`gan_generator`	data 12, hidden 8, latent 3	data 784, hidden 512, latent 64	9 · 9 · 9	12 · 9 · 10
`diffusion_mlp`	data 2, 4 freqs, hidden 8	data 64, 64 freqs, hidden 512	14 · 14 · 15	17 · 14 · 16
`unet`	8×8, channels 2, 4	32×32×3, channels 32, 64	23 · 29 · 34	27 · 29 · 39
`gcn`	6 nodes, 3 features, hidden 4	256 nodes, 64 features, hidden 128	14 · 11 · 7	13 · 11 · 12
`two_tower`	user 5, item 4, embed 3	user 128, item 128, embed 64	12 · 12 · 12	16 · 12 · 13
`clip`	8×8×2 image, vocab 11, seq 4	28×28×16 image, vocab 1024, seq 16	21 · 30 · 22	25 · 30 · 23
`moe`	seq 3, dim 4, 4 experts, top 2	seq 64, dim 256, 8 experts, top 2	57 · 63 · 43	68 · 64 · 63

Which /ml-models page each model covers is mapped in rml's docs/progress.md.

Fixtures and golden tests

make_graph_fixtures.py exports every tiny zoo model by six routes: torch, jax, tf, and stablehlo_torch, stablehlo_jax, stablehlo_tf, with PyTorch first because torch.export loads triton, which crashes if TensorFlow's or jaxlib's LLVM loaded before it (L27–30). For each pair it draws a batch of 2 inputs, exports the graph and runs the framework, writes <route>_<model>.json and .safetensors, evaluates the written graph with the NumPy reference ops and requires it to match the framework, and finally writes the framework's inputs and outputs to .golden.safetensors (L101–115). That is 114 committed graphs.

golden_test.cc closes the loop in C++ with no Python. It loads each fixture, runs it through rml's executor on the golden inputs, and requires every output element within 1e-5 of the framework's (L20–54). Its cases are every .json in the testdata directory, so a new fixture is tested without editing the test, and a second test fails if any direct framework and model pair is missing (L56–81).

Cross-framework agreement

Matching each framework separately does not prove the three definitions are the same model. Two more tests check that.

cross_framework_test.cc runs in bazel test //... with no Python. For every model with a fixture from all six routes, it first requires every route's golden inputs to be identical, then runs all six graphs and compares all twelve results (each route's framework output and its rml output) pairwise within 1e-5 (L80–136). A second test fails if any model is missing a route, or if there are fewer than 19 (L141–153).
tools/cross_framework_test.py runs the frameworks live and lowers each by both routes into rml-graph-run, then asserts that all nine results (three framework outputs, six rml runs) agree within 1e-5 + 1e-5 |b| with no NaNs (L184–211). Its nine cases per model are batches 1, 2, 5 and 7, inputs scaled by 4 and 10, all zeros, and the smallest and largest token ids (L53–65); RML_CROSS_BENCH=1 adds a bench-size case at 1e-4. It also checks that rml repeats itself bit for bit on one thread and agrees across 1 and 4 threads (L214–224), and that an op rml-graph lacks fails the export by name in every framework and in the StableHLO importer (L227–270). HoistStepProjectionsTest exports the tiny JAX and TensorFlow rnn and lstm, runs the passes, and checks that steps + 1 matmuls remain and that the NumPy evaluation is unchanged bit for bit (L273–296).

The suite found one real framework difference: with all-zero input every MoE router probability ties, and torch.topk on CPU picks different tied indices than JAX, TensorFlow and rml, which take the lowest. That one comparison is excluded in KNOWN_DIVERGENCE, and rml run on the PyTorch export must still agree (L67–75).

tools/rml_infer_example.py is the same idea as a short end-to-end demonstration: one small classifier (dense, layer norm, tanh-GELU, dense, softmax) written the ordinary way in PyTorch, JAX and TensorFlow on shared weights, lowered by the direct exporters, run by rml-graph-run, and printed with a table of pairwise maximum differences that must all be within 1e-5 (L148–179).

Exact tiny examples

Two examples run framework-defined models from C++ with no tolerance at all. Each model is written once per framework on the same non-zero dummy weights, exported by the direct exporters, run by the executor, and compared bit for bit, fused and unfused. That works because the dummy values make every op exact in f32, so the output has one value whatever the summation order, fusion, threading or exp and tanh approximation (tiny_models.h L1–10).

Example	Model	Why it is exact	Nodes (torch · jax · tf)
`tiny_net`	x[2, 4] → Dense(4, 8) → ReLU → Dense(8, 3)	Weights, biases and input are non-zero multiples of 1/8, so every product and partial sum is exact (tiny_net.h L1–8).	5 · 5 · 5
`tiny_models`	CNN: x[2, 1, 6, 6] → conv 3×3 → ReLU → max pool 2 → dense → y[2, 3]	Multiples of 1/8 again.	6 · 7 · 7
	LSTM: x[2, 12, 3] → LSTM (hidden 2) → last h → dense → y[2, 3]	Every gate pre-activation is exactly 0 or far enough into saturation (20 for sigmoid, 10 for tanh) that f32 sigmoid rounds to exactly 1 and tanh to ±1; the input bias of three gates is solved so their pre-activation on the last step of row 0 is exactly 0.	246 · 242 · 241
	Transformer: x[2, 3, 8] → pre-norm causal 2-head block → y[2, 3, 8]	Layer-norm rows have variance exactly 1 or 4, the head dim of 4 makes the 1/√4 scale exact, and one score per attention row is at least 40 above the rest, so softmax is one-hot.	32 · 42 · 29

The generators, make_tiny_net.py and make_tiny_models.py, check each framework's output and each written graph (evaluated in NumPy) against a NumPy reference with np.array_equal, then print the reference as C++ literals for the tests (make_tiny_net.py L103–119). The C++ side repeats the input formulas, and tiny_net also has a plain-loop ReferenceForward (tiny_net.cc L51–71). Building them found one framework detail: torch.nn.LayerNorm's one-pass CPU kernel rounds the variance by an ulp once a row's mean exceeds about 1/6, so the PyTorch block uses a two-pass layer norm like JAX and TensorFlow (make_tiny_models.py L338–351, progress.md L175).

bazel run //examples/tiny_net:tiny-net         # three exports, exact check
bazel run //examples/tiny_models:tiny-models   # nine exports, exact check

Tests

Check	Compares	Tolerance	Source
Per op, while exporting	Each lowered value's shape against the framework's (or StableHLO's) recorded shape	Exact	from_stablehlo.py L197
`make_graph_fixtures.py`	The written graph, evaluated in NumPy, against the framework's outputs	1e-5 + 1e-4 \|x\|	L109
`golden_test.cc`	rml's executor against the framework's outputs, every fixture	1e-5 absolute	L25
`cross_framework_test.cc`	All twelve results per model, pairwise	1e-5 absolute	L80
`cross_framework_test.py`	Three live frameworks and six rml runs, nine input cases, determinism, unsupported ops, the step-projection hoist	1e-5 + 1e-5 \|b\| (bench 1e-4)	L184
`tiny_net_test.cc`, `tiny_models_test.cc`	Three and nine framework exports run by the executor, fused and unfused, against the exact reference output	Bit for bit	L38
`rml_infer_example.py`	One classifier in three frameworks and in rml	1e-5 absolute	L148

Benchmarks reuse the same exporters at the bench sizes; how they are timed is on Benchmarks against the frameworks.