# facebookresearch/neural-rewriter architecture

> Architecture diagram of the GitHub repository facebookresearch/neural-rewriter, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-22). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/facebookresearch/neural-rewriter
- Repository: https://github.com/facebookresearch/neural-rewriter

## Overview

1. **What this repository is**
   - This is a PyTorch-based research codebase for the “Neural Rewriter” framework from the NeurIPS 2019 paper.
   - It implements one shared pattern: take a combinatorial instance, encode its current state, choose a local rewrite/action, and iteratively improve the solution.

2. **Top-level task split**
   - The repo is organized around three problem domains:
     - **Halide expression simplification** (`run_Halide.py`, `Halide_search.py`)
     - **Job scheduling** (`run_jsp.py`, `jsp_nonNN_baselines.py`)
     - **Vehicle routing** (`run_vrp.py`)
   - Each task has its own model, input encoder, and rewriter, but they share the same training/evaluation scaffolding.

3. **Shared model architecture**
   - `src/models/BaseModel.py` defines the common model interface.
   - Task-specific models extend it:
     - `HalideModel.py`
     - `jspModel.py`
     - `vrpModel.py`
   - The shared pattern is: **parse instance → encode state → score candidate rewrites/actions → select and apply a rewrite**.

4. **Input representation layer**
   - `src/models/data_utils/` is the preprocessing boundary between raw datasets and neural models.
   - It includes different structured representations:
     - `Tree.py` for expression trees
     - `Dag.py` for DAG-like structures
     - `Seq.py` for sequential representations
     - `parser.py`, `data_utils.py`, `utils.py` for loading and transforming examples
   - This layer is important because the three tasks do not share the same raw state format.

5. **Task-specific encoders**
   - `src/models/modules/` contains the neural feature extractors:
     - `HalideInputEncoder.py`
     - `jspInputEncoder.py`
     - `vrpInputEncoder.py`
     - `mlp.py` as a generic building block
   - These encoders convert structured problem states into fixed-size embeddings used by the policy/model head.

6. **Rewrite/action selection**
   - `src/models/rewriter/` holds the logic that maps model outputs to actual local edits:
     - `HalideRewriter.py`
     - `jspRewriter.py`
     - `vrpRewriter.py`
   - This is the key architectural boundary: the neural network does not directly solve the whole problem, it proposes local transformations that the task-specific rewriter applies.

7. **Training and experiment control**
   - `src/models/model_utils/` contains shared experiment infrastructure:
     - `supervisor.py` for training/evaluation orchestration
     - `logger.py` for logging
   - `src/arguments.py` centralizes CLI options such as dataset paths, model directory, evaluation mode, epochs, and checkpointing cadence.
   - The run scripts wire together arguments, data loading, model construction, and supervisor logic.

8. **Halide pipeline**
   - `run_Halide.py` is the main training/evaluation entry point for expression simplification.
   - `Halide_search.py` provides a non-neural search baseline over Halide rewrite rules.
   - This task depends on expression-tree parsing and on the Halide rule set conceptually, but the repo’s code focuses on learning and search around those rewrite traces.

9. **Job scheduling pipeline**
   - `run_jsp.py` trains/evaluates the neural rewriter for online scheduling with a bounded pending queue.
   - `jsp_nonNN_baselines.py` runs heuristic baselines such as SJF, EJF, and random.
   - The data flow here is: generated scheduling trace → structured state → encoder → rewrite/scheduling decision → updated schedule state.

10. **Vehicle routing pipeline**
    - `run_vrp.py` is the entry point for routing experiments.
    - The vehicle-routing model works over customer/depot state and capacity constraints, using the VRP-specific encoder and rewriter.
    - External baselines mentioned in the README include OR-Tools and reinforcement-learning solvers, but the repo’s core path is the neural rewrite loop.

11. **Dataset generation**
    - `data_generator/jspDatagen.py` and `data_generator/vrpDatagen.py` synthesize training/evaluation data for the scheduling and routing tasks.
    - These generators are part of the offline data pipeline, not the model runtime.
    - For Halide, the dataset comes from generated expressions and rewrite traces derived from the Halide ecosystem.

12. **External dependencies and boundaries**
    - **PyTorch** is the core runtime for all neural components.
    - **Halide**, **Z3**, and **OR-Tools** appear as external systems for dataset creation or baseline comparisons, not as internal subsystems.
    - The architectural boundary is clear: external solvers/heuristics provide data or reference performance, while the repo’s own code implements the neural rewrite policy and its task-specific application layer.

Main entry points: [`src/models/model_utils/supervisor.py`](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/model_utils/supervisor.py) (Supervisor), [`src/models/data_utils/data_utils.py`](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/data_utils.py) (Data utils), [`src/models/BaseModel.py`](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/BaseModel.py) (Base model), [`src/arguments.py`](https://github.com/facebookresearch/neural-rewriter/blob/master/src/arguments.py) (CLI args)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_shared["Shared core"]
  node_arguments["CLI args<br/>config<br/>[arguments.py]"]
  node_base_model["Base model<br/>model interface<br/>[BaseModel.py]"]
  node_supervisor["Supervisor<br/>train/eval orchestration<br/>[supervisor.py]"]
  node_logger["Logger<br/>logging<br/>[logger.py]"]
  node_data_utils["Data utils<br/>preprocessing<br/>[data_utils.py]"]
  node_parser["Parser<br/>instance parsing<br/>[parser.py]"]
  node_tree(("Tree<br/>[Tree.py]"))
  node_dag(("DAG<br/>graph state<br/>[Dag.py]"))
  node_seq(("Sequence<br/>[Seq.py]"))
  node_mlp["MLP<br/>neural block<br/>[mlp.py]"]
end

subgraph group_halide["Halide"]
  node_halide_model["Halide model<br/>task model<br/>[HalideModel.py]"]
  node_halide_encoder["Halide encoder<br/>input encoder"]
  node_halide_rewriter["Halide rewriter<br/>rewrite policy<br/>[HalideRewriter.py]"]
  node_run_halide["Halide run<br/>entry point<br/>[run_Halide.py]"]
  node_halide_search["Halide search<br/>baseline search<br/>[Halide_search.py]"]
end

subgraph group_jsp["Job scheduling"]
  node_jsp_model["JSP model<br/>task model<br/>[jspModel.py]"]
  node_jsp_encoder["JSP encoder<br/>input encoder<br/>[jspInputEncoder.py]"]
  node_jsp_rewriter["JSP rewriter<br/>rewrite policy<br/>[jspRewriter.py]"]
  node_run_jsp["JSP run<br/>entry point<br/>[run_jsp.py]"]
  node_jsp_baselines["JSP baselines<br/>heuristics"]
end

subgraph group_vrp["Vehicle routing"]
  node_vrp_model["VRP model<br/>task model<br/>[vrpModel.py]"]
  node_vrp_encoder["VRP encoder<br/>input encoder<br/>[vrpInputEncoder.py]"]
  node_vrp_rewriter["VRP rewriter<br/>rewrite policy<br/>[vrpRewriter.py]"]
  node_run_vrp["VRP run<br/>entry point<br/>[run_vrp.py]"]
end

subgraph group_data["Data generation"]
  node_jsp_datagen["JSP data gen<br/>synthetic data<br/>[jspDatagen.py]"]
  node_vrp_datagen["VRP data gen<br/>synthetic data<br/>[vrpDatagen.py]"]
end

node_arguments -->|"configures"| node_run_halide
node_arguments -->|"configures"| node_run_jsp
node_arguments -->|"configures"| node_run_vrp
node_run_halide -->|"drives"| node_supervisor
node_run_jsp -->|"drives"| node_supervisor
node_run_vrp -->|"drives"| node_supervisor
node_supervisor -->|"logs to"| node_logger
node_supervisor -->|"trains"| node_base_model
node_data_utils -->|"parses"| node_parser
node_data_utils -->|"represents"| node_tree
node_data_utils -->|"represents"| node_dag
node_data_utils -->|"represents"| node_seq
node_halide_model -->|"extends"| node_base_model
node_halide_model -->|"uses"| node_halide_encoder
node_halide_model -->|"produces actions"| node_halide_rewriter
node_halide_encoder -->|"builds on"| node_mlp
node_halide_rewriter -->|"applies to"| node_tree
node_halide_search -->|"searches"| node_tree
node_jsp_model -->|"extends"| node_base_model
node_jsp_model -->|"uses"| node_jsp_encoder
node_jsp_model -->|"produces actions"| node_jsp_rewriter
node_jsp_encoder -->|"builds on"| node_mlp
node_jsp_rewriter -->|"applies to"| node_seq
node_vrp_model -->|"extends"| node_base_model
node_vrp_model -->|"uses"| node_vrp_encoder
node_vrp_model -->|"produces actions"| node_vrp_rewriter
node_vrp_encoder -->|"builds on"| node_mlp
node_vrp_rewriter -->|"applies to"| node_dag
node_jsp_datagen -->|"feeds"| node_run_jsp
node_vrp_datagen -->|"feeds"| node_run_vrp

click node_arguments "https://github.com/facebookresearch/neural-rewriter/blob/master/src/arguments.py"
click node_base_model "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/BaseModel.py"
click node_supervisor "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/model_utils/supervisor.py"
click node_logger "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/model_utils/logger.py"
click node_data_utils "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/data_utils.py"
click node_parser "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/parser.py"
click node_tree "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/Tree.py"
click node_dag "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/Dag.py"
click node_seq "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/Seq.py"
click node_mlp "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/modules/mlp.py"
click node_halide_model "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/HalideModel.py"
click node_halide_encoder "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/modules/HalideInputEncoder.py"
click node_halide_rewriter "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/rewriter/HalideRewriter.py"
click node_run_halide "https://github.com/facebookresearch/neural-rewriter/blob/master/src/run_Halide.py"
click node_halide_search "https://github.com/facebookresearch/neural-rewriter/blob/master/src/Halide_search.py"
click node_jsp_model "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/jspModel.py"
click node_jsp_encoder "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/modules/jspInputEncoder.py"
click node_jsp_rewriter "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/rewriter/jspRewriter.py"
click node_run_jsp "https://github.com/facebookresearch/neural-rewriter/blob/master/src/run_jsp.py"
click node_jsp_baselines "https://github.com/facebookresearch/neural-rewriter/blob/master/src/jsp_nonNN_baselines.py"
click node_vrp_model "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/vrpModel.py"
click node_vrp_encoder "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/modules/vrpInputEncoder.py"
click node_vrp_rewriter "https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/rewriter/vrpRewriter.py"
click node_run_vrp "https://github.com/facebookresearch/neural-rewriter/blob/master/src/run_vrp.py"
click node_jsp_datagen "https://github.com/facebookresearch/neural-rewriter/blob/master/data_generator/jspDatagen.py"
click node_vrp_datagen "https://github.com/facebookresearch/neural-rewriter/blob/master/data_generator/vrpDatagen.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_arguments,node_base_model,node_supervisor,node_logger,node_data_utils,node_parser,node_tree,node_dag,node_seq,node_mlp toneBlue
class node_halide_model,node_halide_encoder,node_halide_rewriter,node_run_halide,node_halide_search toneAmber
class node_jsp_model,node_jsp_encoder,node_jsp_rewriter,node_run_jsp,node_jsp_baselines toneMint
class node_vrp_model,node_vrp_encoder,node_vrp_rewriter,node_run_vrp toneRose
class node_jsp_datagen,node_vrp_datagen toneIndigo
```

## Components

### Shared core

Common model, data, and training infrastructure used across tasks.

- [CLI args](https://github.com/facebookresearch/neural-rewriter/blob/master/src/arguments.py): config · Central command-line options for training, evaluation, and checkpoints. · `src/arguments.py`
- [Base model](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/BaseModel.py): model interface · Common model abstraction extended by task-specific models. · `src/models/BaseModel.py`
- [Supervisor](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/model_utils/supervisor.py): train/eval orchestration · Shared experiment loop for training, evaluation, and checkpoint handling. · `src/models/model_utils/supervisor.py`
- [Logger](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/model_utils/logger.py): logging · Shared experiment logging utilities. · `src/models/model_utils/logger.py`
- [Data utils](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/data_utils.py): preprocessing · Boundary layer for loading, parsing, and transforming raw problem instances. · `src/models/data_utils/data_utils.py`
- [Parser](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/parser.py): instance parsing · Parses task-specific raw inputs into structured representations. · `src/models/data_utils/parser.py`
- [Tree](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/Tree.py): expression tree · Structured tree representation for Halide expressions. · `src/models/data_utils/Tree.py`
- [DAG](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/Dag.py): graph state · DAG-like state representation used by the data layer. · `src/models/data_utils/Dag.py`
- [Sequence](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/data_utils/Seq.py): sequence state · Sequential state representation used by the data layer. · `src/models/data_utils/Seq.py`
- [MLP](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/modules/mlp.py): neural block · Generic feed-forward building block used inside task models. · `src/models/modules/mlp.py`

### Halide

Expression simplification pipeline and rewrite search.

- [Halide model](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/HalideModel.py): task model · Neural model for expression simplification. · `src/models/HalideModel.py`
- [Halide encoder](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/modules/HalideInputEncoder.py): input encoder · Encodes Halide expression state into model features. · `src/models/modules/HalideInputEncoder.py`
- [Halide rewriter](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/rewriter/HalideRewriter.py): rewrite policy · Maps model scores to concrete Halide rewrite actions. · `src/models/rewriter/HalideRewriter.py`
- [Halide run](https://github.com/facebookresearch/neural-rewriter/blob/master/src/run_Halide.py): entry point · Training and evaluation entry point for Halide experiments. · `src/run_Halide.py`
- [Halide search](https://github.com/facebookresearch/neural-rewriter/blob/master/src/Halide_search.py): baseline search · Non-neural search baseline over Halide rewrite rules. · `src/Halide_search.py`

### Job scheduling

Scheduling pipeline, model, and heuristics.

- [JSP model](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/jspModel.py): task model · Neural model for job scheduling decisions. · `src/models/jspModel.py`
- [JSP encoder](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/modules/jspInputEncoder.py): input encoder · Encodes scheduling state and pending jobs for the policy. · `src/models/modules/jspInputEncoder.py`
- [JSP rewriter](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/rewriter/jspRewriter.py): rewrite policy · Selects and applies scheduling actions to the current state. · `src/models/rewriter/jspRewriter.py`
- [JSP run](https://github.com/facebookresearch/neural-rewriter/blob/master/src/run_jsp.py): entry point · Training and evaluation entry point for job scheduling. · `src/run_jsp.py`
- [JSP baselines](https://github.com/facebookresearch/neural-rewriter/blob/master/src/jsp_nonNN_baselines.py): heuristics · Non-neural scheduling baselines such as SJF, EJF, and random. · `src/jsp_nonNN_baselines.py`

### Vehicle routing

Routing pipeline and model.

- [VRP model](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/vrpModel.py): task model · Neural model for vehicle routing decisions. · `src/models/vrpModel.py`
- [VRP encoder](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/modules/vrpInputEncoder.py): input encoder · Encodes routing state, customers, depot, and capacity features. · `src/models/modules/vrpInputEncoder.py`
- [VRP rewriter](https://github.com/facebookresearch/neural-rewriter/blob/master/src/models/rewriter/vrpRewriter.py): rewrite policy · Selects route construction or local routing actions. · `src/models/rewriter/vrpRewriter.py`
- [VRP run](https://github.com/facebookresearch/neural-rewriter/blob/master/src/run_vrp.py): entry point · Training and evaluation entry point for routing experiments. · `src/run_vrp.py`

### Data generation

Offline synthetic data generators for training and evaluation.

- [JSP data gen](https://github.com/facebookresearch/neural-rewriter/blob/master/data_generator/jspDatagen.py): synthetic data · Offline generator for scheduling traces and training data. · `data_generator/jspDatagen.py`
- [VRP data gen](https://github.com/facebookresearch/neural-rewriter/blob/master/data_generator/vrpDatagen.py): synthetic data · Offline generator for routing instances and training data. · `data_generator/vrpDatagen.py`

## Connections

- CLI args → Halide run: configures · CLI options flow into the Halide entry point.
- CLI args → JSP run: configures · CLI options flow into the JSP entry point.
- CLI args → VRP run: configures · CLI options flow into the VRP entry point.
- Halide run → Supervisor: drives · Halide experiments are executed through the shared supervisor.
- JSP run → Supervisor: drives · JSP experiments are executed through the shared supervisor.
- VRP run → Supervisor: drives · VRP experiments are executed through the shared supervisor.
- Supervisor → Logger: logs to · Training and evaluation metrics are recorded by the shared logger.
- Supervisor → Base model: trains · Supervisor operates on the common model interface.
- Data utils → Parser: parses · Raw inputs are turned into structured instances.
- Data utils → Tree: represents · Expression trees are one of the core structured forms.
- Data utils → DAG: represents · Graph-like state is normalized through DAG representations.
- Data utils → Sequence: represents · Sequential state is normalized through sequence representations.
- Halide model → Base model: extends · Halide-specific model implements the shared interface.
- Halide model → Halide encoder: uses · Halide model consumes encoded expression state.
- Halide model → Halide rewriter: produces actions · Model outputs are turned into rewrite operations.
- Halide encoder → MLP: builds on · Encoder uses generic neural blocks where needed.
- Halide rewriter → Tree: applies to · Rewrite actions operate on expression trees.
- Halide search → Tree: searches · Baseline search explores expression rewrite traces.
- JSP model → Base model: extends · JSP-specific model implements the shared interface.
- JSP model → JSP encoder: uses · JSP model consumes scheduling-state embeddings.
- JSP model → JSP rewriter: produces actions · Model outputs are mapped to scheduling decisions.
- JSP encoder → MLP: builds on · Encoder uses generic neural blocks where needed.
- JSP rewriter → Sequence: applies to · Scheduling actions update sequential state.
- VRP model → Base model: extends · VRP-specific model implements the shared interface.
- VRP model → VRP encoder: uses · VRP model consumes routing-state embeddings.
- VRP model → VRP rewriter: produces actions · Model outputs are mapped to routing decisions.
- VRP encoder → MLP: builds on · Encoder uses generic neural blocks where needed.
- VRP rewriter → DAG: applies to · Routing actions update graph-structured route state.
- JSP data gen → JSP run: feeds · Synthetic scheduling traces supply JSP experiments.
- VRP data gen → VRP run: feeds · Synthetic routing instances supply VRP experiments.