Neural Compiler

Compile a self-hosting subset of Scheme into differentiable PyTorch computation graphs - then run gradient descent through the program to its numeric constants.

The Neural Compiler translates Scheme programs into ordinary PyTorch compute graphs in which every numeric value is a differentiable tensor. Because the compiler is expressive enough to compile its own evaluator, it supports Differentiable Meta-Circular Interpretation (DMCI): a compiled Scheme interpreter executes programs represented as data while autograd propagates gradients from a task loss back to the continuous constants embedded in those programs - with no custom gradient code and no per-program recompilation.

This is the software behind two papers:

Paper 1 - The Neural Compiler: Program-to-Network Translation for Hybrid Scientific Machine Learning. The original system (arXiv:2605.22498).
Paper 2 - Compile Once, Differentiate Everywhere: A Differentiable Meta-Circular Interpreter. DMCI, the recursion/closure-capable differentiable interpreter and its experiments (arXiv:2606.09930).

This repository contains the code, experiments, and data. The manuscripts themselves are distributed separately.

Note on history. Earlier versions compiled Scheme to graph neural networks (via PyTorch Geometric). That approach has been retired - the current system compiles to plain differentiable PyTorch graphs with no torch-geometric dependency. A few legacy symbols (e.g. the SchemeGNN nn.Module wrapper) keep "GNN" in their names but are ordinary PyTorch modules.

Why
Installation
Quick start
Differentiating through programs
Command-line compiler (nncompile)
Differentiable Meta-Circular Interpretation (DMCI)
How it works
The Scheme subset
Batched, population, and GPU evaluation
Discrete structure search
Formal guarantees
Reproducing the papers
Repository layout
Testing
Public API reference
Citing this work
License

Why

Classical machine learning forces a choice: write a model as fixed code (fast, exact, but frozen - you cannot gradient-optimize its constants) or as a neural network (trainable, but an opaque approximation of whatever structure you actually know). The Neural Compiler removes that boundary. You write an executable scientific model as a small Scheme program; the compiler turns it into a differentiable module; and the program's numeric constants become trainable parameters fit by gradient descent - while the program's structure (its control flow, recursion, and composition) stays exactly as written.

Because programs are data (S-expressions), they can be generated by an LLM, swapped at runtime, or composed programmatically, and gradient correctness is guaranteed once - by the compilation of the interpreter itself - rather than re-verified per program.

Installation

git clone https://github.com/sheneman/dmci.git
cd dmci
pip install -e .

Requires Python ≥ 3.10 (developed/tested on 3.11) and PyTorch ≥ 2.0.

The editable install pulls the core runtime (torch, pytest). To run the experiments, figure scripts, and examples you also need the scientific stack:

pip install -r requirements.txt   # adds numpy, scipy
pip install matplotlib            # used by examples/ and paper figure scripts

Optional backends and tooling (commented in requirements.txt):

Extra	Install	Enables
JAX backend	`jax`, `jaxlib`	functional autograd via `backend="jax"`
CuPy backend	`cupy-cuda12x`	GPU-accelerated forward-only evaluation (no autograd)
LLM regeneration	`openai`	regenerating Experiment B's Scheme programs from an OpenAI-compatible endpoint

There is no torch-geometric dependency. GPU acceleration comes from running the default Torch backend on CUDA tensors.

Git LFS (required for the paper's raw data)

Raw experiment results and figure data are stored with Git LFS. Install LFS before cloning so these files materialize as real content instead of pointer stubs:

git lfs install
git clone https://github.com/sheneman/dmci.git
# already cloned? fetch the real files:
git lfs pull

LFS tracks all committed data and binary artifacts (*.csv, *.json, *.dat, *.pt, *.png, *.pdf, … - see .gitattributes), including the per-run results under experiments/*/results*/. Without git lfs pull, the aggregation scripts will read pointer stubs rather than data.

Quick start

run_scheme compiles and evaluates a single expression. Scalar programs return a Python float:

from neural_compiler.compiler import run_scheme

run_scheme("(+ (* 3 4) 5)")                              # 17.0

# External inputs are passed by name:
run_scheme("(+ (* 3 x) (- y 1))", {"x": 4.0, "y": 7.0}) # 18.0

# Conditionals, let-bindings, iteration, and recursion all work:
run_scheme("(if (> x 0) 1 -1)", {"x": 5.0})                                            # 1.0
run_scheme("(let ((a 2) (b 3)) (* a b))")                                              # 6.0
run_scheme("(loop ((i 1) (acc 0)) (if (> i n) acc (recur (+ i 1) (+ acc i))))", {"n": 5.0})  # 15.0
run_scheme("(letrec ((fact (lambda (n) (if (<= n 1) 1 (* n (fact (- n 1))))))) (fact 5))")    # 120.0

Multi-form programs with top-level (define ...) use run_program (and, optionally, a standard-library prelude of list utilities such as map/filter/fold-left):

from neural_compiler.compiler import run_program

run_program("(define (sq x) (* x x)) (sq 6)", {})              # 36.0
run_program("(map (lambda (v) (* v v)) (list 1 2 3))", {}, prelude=True)

compile_scheme / compile_program return a reusable ComputeGraph you can inspect and evaluate repeatedly:

from neural_compiler.compiler import compile_scheme
from neural_compiler.evaluator import evaluate

g = compile_scheme("(+ (* a b) (- c d))", inputs={"a": None, "b": None, "c": None, "d": None})
print(g.input_names, g.depth(), len(g.nodes))    # ['a', 'b', 'c', 'd'] <depth> <#nodes>

evaluate(g, {"a": 3.0, "b": 4.0, "c": 10.0, "d": 3.0})   # 19.0

In compile_scheme(source, inputs=...), inputs is a declaration (keys are input names, values are None placeholders). Concrete values are supplied later to evaluate/run_scheme.

Differentiating through programs

This is the point of the project. Compile a program once, then let autograd flow gradients through it.

Gradients w.r.t. inputs - evaluate_batched preserves the autograd graph:

import torch
from neural_compiler.compiler import compile_scheme
from neural_compiler.evaluator import evaluate_batched

g = compile_scheme("(* x x)", inputs={"x": None})
x = torch.tensor([1., 2., 3.], requires_grad=True)

out = evaluate_batched(g, {"x": x})   # tensor([1., 4., 9.])
out.sum().backward()
print(x.grad)                          # tensor([2., 4., 6.])  =  d(x^2)/dx

Fitting a program's constants by gradient descent - make the constant a torch.nn.Parameter, evaluate through the compiled program with tagged values, and run a standard optimizer loop:

import torch
from neural_compiler.compiler import compile_program
from neural_compiler.evaluator import evaluate
from neural_compiler.runtime.tagged_value import make_float, unwrap_number

# Recover the constant a in  y = a * x  from data, by descending through the compiled program.
graph = compile_program("(* a x)", inputs={"a": None, "x": None})

a   = torch.nn.Parameter(torch.tensor(0.0))
opt = torch.optim.Adam([a], lr=0.1)
xs  = torch.linspace(0.5, 3.0, 8)
ys  = 3.0 * xs                                  # ground truth a = 3.0

for _ in range(200):
    loss = torch.tensor(0.0)
    for x, y in zip(xs, ys):
        pred = unwrap_number(evaluate(graph, {"a": make_float(a), "x": make_float(x)}))
        loss = loss + (pred - y) ** 2
    opt.zero_grad(); loss.backward(); opt.step()

print(a.item())   # ≈ 3.0

Autograd footgun. evaluate() returns a detached Python float for a scalar, non-tagged result (it calls .item()). To keep gradients you must either (a) use the tagged path - wrap inputs with make_float(...) and read the result with unwrap_number(...), as above - or (b) use evaluate_batched, which always returns a live tensor. The experiment training loops (experiments/exp_*/baselines.py) use the tagged path.

This same loop, scaled up, is exactly how the paper fits LLM-generated scientific models, recursive filters, ODE integrators, and so on.

Command-line compiler (`nncompile`)

pip install -e . registers an nncompile console tool with four subcommands. It compiles a Scheme .scm file to a portable, backend-agnostic .ncg artifact, or emits a ready-to-import differentiable torch.nn.Module. Input variables are auto-detected - you don't declare them.

nncompile compile model.scm -o model.ncg          # Scheme -> portable compiled artifact
nncompile emit    model.scm --params a,b -o m.py  # -> standalone torch.nn.Module
nncompile run     model.ncg --inputs '{"x": 4.0}' # evaluate on a backend
nncompile info    model.scm                        # inputs, structure, backends

A .ncg is the compiled program as data: the backend (PyTorch / JAX / NumPy / CuPy) is chosen at run time, not baked into the file - compile once, differentiate everywhere. One artifact runs on every backend:

nncompile run model.ncg --inputs '{"x": 4.0}'                   # PyTorch (default; differentiable)
nncompile run model.ncg --inputs '{"x": 4.0}' --backend numpy  # forward-only

Gradient backends are torch and jax, but jax covers only the scalar straight-line direct-compile path; the differentiable meta-circular interpreter (heap/tagged programs) is torch-only (see "Why PyTorch for the differentiable interpreter"). numpy and cupy are forward-only, with numpy serving as the reference oracle.

Emit a `torch.nn.Module` for optimization / search

emit writes a self-contained .py exposing a CompiledModel(torch.nn.Module): the inputs you pass to --params become learnable nn.Parameters, and the rest become forward() data arguments. Drop it into any training loop:

# model.scm:  (/ (* k (* q1 q2)) (* r r))      # Coulomb's law; k is learnable
# $ nncompile emit model.scm --params k -o coulomb.py
import torch
from coulomb import CompiledModel

model = CompiledModel(k=0.1)                       # initial guess
opt = torch.optim.Adam(model.parameters(), lr=0.2)
for _ in range(400):
    pred = torch.stack([model(q1=3.0, q2=4.0, r=float(r)) for r in rs])
    loss = ((pred - targets) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()
print(model.params["k"].item())                   # -> recovers the true k

The emitted module embeds the compiled graph and depends only on neural_compiler + torch.

--backend jax emits a functional JAX module instead (an apply(params, **data) you optimize with jax.grad and jax.vmap).
--dmci (on any subcommand) compiles via the meta-circular interpreter - your program runs as quoted data through the compiled self-hosted evaluator (the Paper 2 method); same value, heap-backed graph.

evaluate_batched also runs loop/recur with data-dependent bounds per batch element (masked padded iteration), so batched training of iterative programs is correct. Full details and the .ncg format are in docs/cli.md.

Differentiable Meta-Circular Interpretation (DMCI)

The compiler is expressive enough to compile its own Scheme evaluator (bootstrap/compiler.scm, a self-hosted interpreter written in Scheme). That compiled evaluator is itself a differentiable PyTorch graph. To run a target program, you hand it to the evaluator as quoted data:

;; conceptually: compile the evaluator once, then interpret a program-as-data
<contents of bootstrap/compiler.scm>
(scheme-eval '(<your target program>) <environment>)

Because the interpreter is compiled once and the program is data, gradients flow from a loss, through the compiled interpreter (environment lookup, dispatch, heap operations, arithmetic primitives), all the way to the numeric constants of the interpreted program - without recompiling for each new program. New programs (e.g. ones generated by an LLM) are just new data fed to the same differentiable evaluator. experiments/exp_b does exactly this for 15 LLM-generated models.

What makes this work end-to-end:

Tagged values. Every runtime value is a fixed 14-dimensional tensor: a 10-way one-hot type tag (nil, bool, int, float, char, symbol, pair, string, closure, vector) concatenated with a 4-float payload. The numeric value lives in payload slot 0, so autograd flows through the payload while the tag merely routes dispatch (neural_compiler/runtime/tagged_value.py).
An autograd-preserving heap. Pairs, lists, and closures live on a dictionary-backed heap (neural_compiler/runtime/heap.py) that stores the exact tensor written and never mutates in place, so the autograd chain survives arbitrarily many cons/car/cdr operations.
Soft control flow. if is evaluated as a differentiable MUX in straight-line code (and lazily inside recursion to bound depth); branch decisions are made by program structure, not by learnable constants.

Why PyTorch for the differentiable interpreter

The repo has a real backend abstraction (neural_compiler/backend/: torch / numpy / jax / cupy), but it covers the scalar, straight-line direct-compile path only. The differentiable meta-circular interpreter - dictionary heap, tagged values, and the variable-length trampolined recur loop - is torch-only by design, not by oversight.

DMCI's differentiability needs gradients to flow through data-dependent control flow: which eval-apply clause fires, the (ref obs k) index, and the loop-termination test are all decided by reading structural ints/bools off tagged values with .item(). That .item() call is a deliberate non-differentiable boundary separating structural control (discrete, program-determined) from the numeric dataflow carried on-tape as tensors. Neither torch nor JAX differentiates the discrete branch choice itself (a step discontinuity); both give the gradient along the executed path. PyTorch's define-by-run (tape-based) autograd records the ops that actually execute, so it differentiates the realized trajectory - including the data-dependent-length loop - with no restructuring.

JAX traces to a functional jaxpr: it forbids data-dependent Python branching on traced values (concretization error → use lax.cond), and lax.while_loop is not reverse-mode differentiable. So the variable-length trampoline would have to be rewritten as a fixed-length lax.scan(max_iter) + masking. JAX is therefore future re-architecture (payoff: jit fusion of the per-step overhead + vmap over the population path), not a drop-in swap; backend/jax_backend.py's jax_grad is intentionally limited to the scalar path above.

NumPy and CuPy have no autograd → forward-only, so they cannot host a differentiable interpreter. NumPy already plays its correct role: the forward-only reference oracle the validation gate checks every DMCI result against. (CuPy is redundant with torch.cuda and loses gradients; the single-filter interpreter walk is latency/dispatch-bound, so GPU does not help it regardless.) The batched tensor/matrix ops and logdet determinant are likewise torch-only. See neural_compiler/backend/README.md.

Distribute the interpreter as a pre-compiled artifact

compile_dmci(program, ...) bakes one program into the graph (as a quote_const), so each program serializes to its own .ncg. For the compile the interpreter once, then run any program through it deployment, compile the bare interpreter - the evaluator with the program and environment as runtime inputs - and hand it programs as data:

from neural_compiler import compile_interpreter, evaluate_program, save_compiled, load_compiled
from neural_compiler.runtime.tagged_value import unwrap_number

interp = compile_interpreter()             # compile the evaluator ONCE
save_compiled(interp, "interpreter.ncg")   # ~285 KB portable JSON - ship it anywhere

interp = load_compiled("interpreter.ncg")  # consumer side: no Scheme toolchain, no recompilation
y = evaluate_program(interp, "(* a (exp (* (- 0 b) x)))", {"x": 1.5, "a": 2.5, "b": 0.8})
z = evaluate_program(interp, "(+ (* a x) (* b (* x x)))", {"x": 2.0, "a": 3.0, "b": 1.5})  # different program, same graph
print(unwrap_number(y).item(), unwrap_number(z).item())

Bindings may be Python numbers, scalar tensors (learnable parameters - gradients flow back through the interpreter to them), or batched [N] tensors (the whole batch runs in one interpreter walk). The program string is parsed and materialized onto the eval heap and bound to the interpreter's program/env inputs at run time; multi-form programs (with defines) work too. This is the literal form of compile once, differentiate everywhere: one backend-agnostic artifact, any program - no per-program recompilation. (A program whose branch decision depends on a batched input is the one case that cannot be vectorized; see Batched DMCI.)

How it works

Scheme source
   │  parse  (tokenizer + recursive-descent parser; desugars cond/let*/when/unless/quasiquote)
   ▼
AST            (frozen dataclasses: Const, Var, If, Lambda, Let, App, Loop, Recur, Letrec, Quote, Begin, Define, SoftChoice)
   │  to_anf  (A-Normal Form: every argument and every `if` test is trivial; compound subexpressions hoisted to lets)
   ▼
ANF
   │  optimize_tco  (self- and mutual-tail-recursive letrec → loop/recur for O(1) stack)
   ▼
ANF (optimized)
   │  build_graph
   ▼
ComputeGraph   (dataflow DAG of GraphNodes; loop bodies and recursive functions become nested subgraphs)
   │  evaluate / evaluate_batched   (PyTorch; autograd-enabled)
   ▼
Result (tensor)

Parser / AST - neural_compiler/parser/. A hand-written tokenizer and recursive-descent parser produce an immutable AST; surface forms like cond, let*, when, unless, and quasiquote are desugared to core nodes.
A-Normal Form - neural_compiler/anf/transform.py. Names every intermediate subexpression ("SSA for functional languages") so each maps 1:1 to a graph node.
Tail-call optimization - neural_compiler/anf/tco.py. Self-tail-recursive functions become loop/recur; mutually-tail-recursive sets become a single dispatch loop. Non-tail (tree) recursion such as naive Fibonacci is left as lazy recursion (call nodes + a function-body subgraph).
Graph builder - neural_compiler/graph/builder.py. Lowers ANF to a ComputeGraph of integer-keyed GraphNodes. Node op_types include const, input, if, loop, recur, call, dynamic_call, make_closure, soft_choice, and primitive operators. Loop bodies (LoopBody) run iteratively; recursive functions (FunctionBody) are invoked lazily.
Evaluator - neural_compiler/evaluator/engine.py. Walks the graph in PyTorch. Programs that touch the heap (cons/lists/closures/quote) run through the tagged-value evaluator; plain arithmetic programs use a lighter path. The tagged evaluator trampolines tail calls - including the self-hosted interpreter's own scheme-eval/eval-apply loop - so tail recursion (and recursive DMCI) runs in constant Python stack; only non-tail (tree) recursion nests.

The Scheme subset

Feature	Syntax	Example
Integer / float literals	`42`, `3.14`	`42`
Booleans	`#t`, `#f`	`#t`
Variables	`x`, `my-var`	`x`
Arithmetic	`+ - * /`, `modulo`, `remainder`	`(+ 1 2)`
Comparison	`= < > <= >=`	`(> x 0)`
Logic	`and`, `or`, `not`	`(and a b)`
Math	`abs min max sin cos exp log sqrt pow`	`(sqrt x)`
Conditionals	`if`, `cond`, `when`, `unless`	`(if t a b)`
Local bindings	`let`, `let*`	`(let ((x 1)) (+ x 2))`
Lambda / closures	`lambda`	`(lambda (x) (* x x))`
Iteration	`loop` / `recur`	`(loop ((n 5) (acc 1)) (if (= n 0) acc (recur (- n 1) (* acc n))))`
Recursion	`letrec`	`(letrec ((f (lambda (n) ...))) (f 5))`
Definitions	`define` (multi-form, via `compile_program`)	`(define (sq x) (* x x))`
Data	`quote`, `cons`, `car`, `cdr`, `list`, pair/list predicates	`(car (list 1 2 3))`
Vectors / matrices	`vec dot cross norm matmul det inv …`	`(dot a b)`
Differentiable choice	`soft-choice`	`(soft-choice (a b c) weights)`

The prelude (prelude=True) adds higher-order list utilities: map, filter, fold-left, fold-right, for-each, assoc, member, and the c[ad]+r accessors.

MUX vs. lazy if. In straight-line (combinational) code, if evaluates both branches and blends them (a differentiable multiplexer). Inside recursive functions, if is evaluated lazily (taken branch only) to prevent infinite descent in base cases.

Batched, population, and GPU evaluation

evaluate_batched(graph, inputs) walks the graph once with batched tensors, so throughput scales by amortizing the Python-level graph-walking cost across all inputs:

import torch
from neural_compiler.compiler import compile_scheme
from neural_compiler.evaluator import evaluate_batched

g = compile_scheme("(+ (* a x) b)", inputs={"a": None, "x": None, "b": None})

# Feature-dimension batching: evaluate one program on many inputs at once.
out = evaluate_batched(g, {"a": torch.tensor(2.0),
                           "x": torch.tensor([1., 2., 3.]),
                           "b": torch.tensor(1.0)})        # tensor([3., 5., 7.])

Population batching is the same mechanism via broadcasting: shape parameters (M, 1) and data (N,) to evaluate M parameterizations against N points at once, with independent gradients per population member. This is what makes population-based optimization practical (experiments/exp_h).
GPU - there is no device flag; place input tensors on cuda and the (default) Torch backend follows.
compile_batched(graph) wraps the batched forward pass with torch.compile.

Batching the heap (incl. DMCI). evaluate_batched has two paths. Pure arithmetic/comparison/logic and loop programs take a fast heap-free walk (in batched loops the recur-vs-terminate decision is read from one batch element, so all elements share the iteration schedule). Programs that build data on the heap - cons/car/cdr/list/ quote - including the meta-circular interpreter (DMCI) - are routed to the heap-backed evaluator, which batches natively: pass tagged inputs of shape (N, VALUE_DIM) and it runs one heap-backed walk in which structural values (pairs, symbols, the quoted AST, loop counters) stay scalar and data-independent while only numeric leaves carry the batch. This fits a whole dataset through the compiled interpreter in a single walk per epoch (≈ one interpreter walk instead of N). The requirement is that control flow and heap structure be identical across the batch - true when fitting the constants of a fixed program; a branch whose decision depends on the batched data cannot be vectorized this way and raises a clear error rather than silently miscomputing (a data-dependent branch that nonetheless resolves uniformly across the batch still evaluates).

Backends are selected at evaluation time - evaluate(graph, inputs, backend=...):

Backend	Autograd	Notes
`torch` (default)	✅	the only path with batching, `compile_batched`, and the full evaluator
`jax`	✅ (functional)	`jax_grad` / `jax_value_and_grad`; straight-line control flow w.r.t. the differentiated variable
`numpy`	❌	reference forward evaluation
`cupy`	❌	GPU-accelerated forward-only

Throughput figures reported in Paper 2 (e.g. ~875× at batch size 1024, and a 3,848× end-to-end population-optimization speedup) are benchmark results measured on specific hardware (single CPU core / NVIDIA A100) and specific models, not guarantees. Notably, for those models CPU outperformed the A100 at all tested batch sizes due to CUDA kernel-launch overhead. See Reproducing the papers.

Discrete structure search

Beyond fitting continuous constants, the interpreter exposes a differentiable soft-choice construct for searching over discrete program structure (e.g. which operator to use). Options are blended by a Gumbel-Softmax (or straight-through softmax) over learnable logits, annealed from soft to near-discrete:

(soft-choice (e1 e2 e3) weights)   ; differentiable selection among ≥2 option expressions

Global controls (neural_compiler.evaluator.set_soft_choice_tau / set_soft_choice_gumbel / set_soft_choice_hard) set the temperature and relaxation mode. experiments/exp_e uses this to recover unknown operators in a small symbolic task. This is an exploratory capability: as with prior differentiable interpreters (TerpreT), gradient-based search over discrete program structure succeeds on only a modest fraction of restarts - the project's strength is continuous parameter optimization within a known program structure.

Formal guarantees

Paper 2 proves three results about the compilation (§Theoretical Guarantees; proofs in the appendix). They are stated over a core language L_DMCI (x | c | λx.e | (e₁ e₂) | if | letrec | cons/car/cdr | op, with the 13 core primitives + - * / sin cos exp log sqrt abs pow min max):

Compilation correctness. Every supported program, run by the compiled graph, produces the same value as the source semantics (to floating-point precision). Proof by structural induction.
Gradient correctness almost everywhere. For learnable constants θ, the gradient computed through the compiled graph equals the gradient through the source semantics for almost every θ (a full-Lebesgue-measure set).
Composition preservation. Composing two a.e.-gradient-correct programs yields an a.e.-gradient-correct program.

The honest caveat - "almost everywhere" is load-bearing. Gradients are guaranteed on the open, full-measure trace-constant region Θ_tc where the discrete execution trace (branch decisions, tag dispatches) does not change with θ. The complement is a measure-zero union of analytic hypersurfaces (branch boundaries), where the compiled program merely inherits the source program's non-differentiability rather than introducing new discontinuities. Concretely, a learnable parameter inside a branch condition (e.g. (< x α)) receives zero gradient there.

These guarantees rest on the autograd-preserving heap (a proposition in the paper, and verified in neural_compiler/runtime/heap.py: reads return the identical tensor object written, with no in-place mutation). Empirically, trajectory equivalence holds to a maximum final-loss difference < 7×10⁻⁷ across 171 (program, seed) pairs in Experiments A–C - corroborating Theorem 2, though not part of the formal statement.

The implementation's full operator table extends the core with vector/matrix ops (vec, dot, cross, matmul, det, inv, …). The theorems are stated over the 13-primitive core and are not claimed to cover those extensions. sqrt/log are clamped at 1e-8, creating a zero-gradient region near zero (technically correct, but uninformative); none of the experiments evaluate them there.

Reproducing the papers

All experiments run from the repo root as modules and write per-run JSON/CSV to experiments/<exp>/results/ (Git LFS). Six experiments back Paper 2 (A, B, C, D, E, H). Run git lfs pull first so the committed raw data is real content.

Exp	What it shows	Run
A	Constant learning across 5 methods (DMCI, direct compile, hand-coded, finite-diff, evolution strategy) × 6 programs × 10 seeds	`python -m experiments.exp_a.run_all --output-dir experiments/exp_a/results --skip-existing`
B	15 LLM-generated scientific models (Qwen 3.6-35B); the headline DMCI-vs-baselines comparison	`python -m experiments.exp_b.run_all --output-dir experiments/exp_b/results --skip-existing`
C	8 recursive/iterative scientific models (Lotka–Volterra, SIR, IIR filter, cascaded EMA, …)	`python -m experiments.exp_c.run_all --output-dir experiments/exp_c/results --skip-existing`
D	Compile-vs-train cost tradeoff inside a genetic-programming loop	`python -m experiments.exp_d.exp_d --method both --seed 0`
E	Discrete operator recovery via `soft-choice` (64 operator combinations)	`python -m experiments.exp_e.exp_e1 --task all --output-dir experiments/exp_e/results`
H	Batched/GPU throughput and population-batching speedups	`python -m experiments.exp_h.exp_h --part all --device cuda`

One-command integrity checks. Four aggregators reconstruct the paper's tables from the committed per-run data (JSONs + the SLURM .out logs for Exp H Parts A/B) and self-check against the manuscript values, so any drift is caught:

python -m experiments.exp_b.aggregate_table   # Table 6:  65/65 cells match (+ 60/60 figure coords)
python -m experiments.exp_c.aggregate_table   # Table 8:  43/43 cells match
python -m experiments.exp_d.aggregate_table   # Table 9:  10/10 cells match
python -m experiments.exp_h.aggregate_table   # Exp H tables: 275/275 cells match
python -m experiments.check_manifests         # completeness: each results/ dir matches its committed MANIFEST.txt

Experiment B can also be re-run on the actual cached LLM programs (rather than reference implementations) via a uniform evaluator wrapper:

python -m experiments.exp_b.run_all --use-llm-cache --output-dir experiments/exp_b/results_llm --skip-existing

Figure data. The aggregated .dat/.json series that the manuscript figures are drawn from are committed under the per-experiment results directories (experiments/*/results*/), produced by each experiment's aggregation scripts.

Data-key note. A few committed result keys predate model renames - notably M05_hookes_spring is reported as "M05 Damped oscillator" in the paper. The aggregators (aggregate_table.py) map keys to paper labels.

LLM regeneration (optional). Experiment B's Scheme programs are cached in experiments/exp_b/llm_cache/, so training is fully reproducible offline. To regenerate them you need an OpenAI-compatible endpoint configured via MINDROUTER_BASE_URL, MINDROUTER_API_KEY, and MINDROUTER_MODEL (the paper uses qwen/qwen3.6-35b): python -m experiments.exp_b.run_all --generate-only --api mindrouter.

Scale & environment. Full grids are large and were run on a CPU/GPU cluster (hundreds of runs per experiment; some recursive models take hours each). The SLURM scripts (experiments/*/slurm_*.sh) target the authors' environment and are illustrative, not universal requirements. experiments/exp_f and exp_g (LLM-in-the-loop model discovery and runtime compositional modeling) are in-progress prototypes - not part of either paper and not yet reproducible.

Repository layout

Path	Contents
`neural_compiler/`	the compiler and runtime (see below)
`neural_compiler/parser/`	tokenizer, recursive-descent parser, AST nodes
`neural_compiler/anf/`	A-Normal Form transform + tail-call optimization
`neural_compiler/graph/`	`ComputeGraph` / graph builder
`neural_compiler/evaluator/`	PyTorch evaluator (`engine.py`), batched engine, `DirectModule`, `SchemeGNN`
`neural_compiler/ops/`	scalar/vector/matrix primitives and tagged (heap) ops
`neural_compiler/runtime/`	tagged-value representation, the autograd-preserving heap, symbols
`neural_compiler/backend/`	torch / numpy / jax / cupy backends
`neural_compiler/compiler.py`	top-level entry points (`compile_scheme`, `run_scheme`, …)
`bootstrap/compiler.scm`	the self-hosted Scheme evaluator compiled for DMCI
`experiments/`	experiments A–H (configs, runners, SLURM scripts, LFS results)
`examples/`	runnable demos (e.g. coefficient learning, damped pendulum, compositional generalization)
`large_examples/`	larger Scheme models (`diffesm_s.scm`, `diffsoc_s.scm`)
`benchmarks/`	microbenchmark drivers and plots
`tests/`	unit + integration test suites
`docs/`	design notes (some predate Paper 2)

Testing

pytest                # full suite
pytest tests/unit     # fast unit tests
pytest -k self_hosting # e.g. the meta-circular evaluator compiles and runs

The suite has 823 test functions across 30 modules (15 unit + 15 integration; 363 + 460 functions) covering the parser, ANF/TCO, graph builder, primitives, tagged values and heap, the sequential / batched / multi-backend evaluators, loops, recursion, composition, GPU batching, self-hosting (compiling bootstrap/compiler.scm), and a fixed regression suite.

Public API reference

Common entry points are re-exported at the top level; everything else lives in the submodules:

# convenience re-exports
from neural_compiler import (
    compile_scheme, compile_program, run_scheme, run_program,
    compile_dmci,         # (program, input_names=None, prelude=False) -> DMCI graph (program embedded)
    compile_interpreter,  # (prelude=True) -> bare interpreter graph (program/env are runtime inputs)
    evaluate_program,     # (interp_graph, program, bindings) -> TaggedValue (run any program through it)
    save_compiled,    # (graph, path, source=None) -> writes a portable .ncg artifact
    load_compiled,    # (path) -> ComputeGraph  (run with evaluate(..., backend=...))
)

# compile / run (full signatures)
from neural_compiler.compiler import (
    compile_scheme,   # (source, inputs=None) -> ComputeGraph         (single expression)
    compile_program,  # (source, inputs=None, prelude=False) -> ComputeGraph   (multi-form, with defines)
    run_scheme,       # (source, inputs=None, backend=None) -> float          (compile + evaluate)
    run_program,      # (source, inputs=None, prelude=False, backend=None) -> float
)

# evaluate
from neural_compiler.evaluator import (
    evaluate,             # (graph, inputs=None, backend=None, **kw)   -> float | Tensor
    evaluate_batched,     # (graph, inputs)                            -> Tensor (autograd-preserving)
    compile_batched,      # (graph, input_names=None)                  -> callable (torch.compile)
    set_soft_choice_tau, set_soft_choice_gumbel, set_soft_choice_hard,
    jax_grad, jax_value_and_grad,
    SchemeGNN, DirectModule,    # nn.Module wrappers (plain PyTorch; "GNN" is a legacy name)
)

# tagged values (needed for gradient-preserving evaluation of scalar programs)
from neural_compiler.runtime.tagged_value import make_float, make_int, unwrap_number

# serialize / emit (also driven by the `nncompile` CLI)
from neural_compiler.serialize import graph_to_dict, graph_from_dict, save_compiled, load_compiled
from neural_compiler.emit import emit_torch_module   # -> source of a standalone torch.nn.Module

ComputeGraph exposes .nodes, .input_names, .root_id, .loops, .functions, .depth(), .topological_order(), and the properties .has_loops / .has_functions.

Citing this work

Paper 2 (DMCI) is available at arXiv:2606.09930.

@article{sheneman_neural_compiler,
  title        = {The Neural Compiler: Program-to-Network Translation for Hybrid Scientific Machine Learning},
  author       = {Sheneman, Lucas},
  journal      = {arXiv preprint arXiv:2605.22498},
  year         = {2026},
  eprint       = {2605.22498},
  archivePrefix = {arXiv},
  url          = {https://arxiv.org/abs/2605.22498}
}

@article{sheneman_dmci,
  title        = {Compile Once, Differentiate Everywhere: A Differentiable Meta-Circular Interpreter},
  author       = {Sheneman, Lucas},
  journal      = {arXiv preprint arXiv:2606.09930},
  year         = {2026},
  eprint       = {2606.09930},
  archivePrefix = {arXiv},
  url          = {https://arxiv.org/abs/2606.09930}
}

Navigation Menu

Repository files navigation

Neural Compiler

Table of contents

Why

Installation

Git LFS (required for the paper's raw data)

Quick start

Differentiating through programs

Command-line compiler (`nncompile`)

Emit a `torch.nn.Module` for optimization / search

Differentiable Meta-Circular Interpretation (DMCI)

Why PyTorch for the differentiable interpreter

Distribute the interpreter as a pre-compiled artifact

How it works

The Scheme subset

Batched, population, and GPU evaluation

Discrete structure search

Formal guarantees

Reproducing the papers

Repository layout

Testing

Public API reference

Citing this work

License

About

Releases

Packages

Contributors

Languages

Name		Name	Last commit message	Last commit date
Latest commit History 4 Commits
benchmarks		benchmarks
bootstrap		bootstrap
docs		docs
examples		examples
experiments		experiments
large_examples		large_examples
neural_compiler		neural_compiler
tests		tests
tools		tools
.gitattributes		.gitattributes
.gitignore		.gitignore
LICENSE		LICENSE
README.md		README.md
REPRODUCIBILITY.md		REPRODUCIBILITY.md
dmci.png		dmci.png
pyproject.toml		pyproject.toml
reproduce.sh		reproduce.sh
requirements.txt		requirements.txt

Folders and files

Latest commit

History

Repository files navigation

Neural Compiler

Table of contents

Why

Installation

Git LFS (required for the paper's raw data)

Quick start

Differentiating through programs

Command-line compiler (nncompile)

Emit a torch.nn.Module for optimization / search

Differentiable Meta-Circular Interpretation (DMCI)

Why PyTorch for the differentiable interpreter

Distribute the interpreter as a pre-compiled artifact

How it works

The Scheme subset

Batched, population, and GPU evaluation

Discrete structure search

Formal guarantees

Reproducing the papers

Repository layout

Testing

Public API reference

Citing this work

License

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages

Command-line compiler (`nncompile`)

Emit a `torch.nn.Module` for optimization / search