JIT / TorchScript

Active contributors: zou3519, eellison

Purpose

The "JIT" or "TorchScript" stack is PyTorch's original graph compilation system. It introduced torch.jit.trace, torch.jit.script, and the .pt (TorchScript) serialization format. While torch.compile and torch.export are now the preferred path forward, JIT/TorchScript is still widely used in production: most pre-2.x deployed models are TorchScript, mobile/edge inference still uses it, and major chunks of the codebase (ONNX export, parts of distributed) build on JIT IR.

The package lives at torch/jit/ (Python) and torch/csrc/jit/ (C++).

Directory layout

Path	Contents
torch/jit/	Python frontend
torch/jit/__init__.py	script, trace, save, load
torch/jit/_script.py	ScriptModule, ScriptFunction
torch/jit/_trace.py	Tracing implementation
torch/jit/_freeze.py	Module freezing (constant folding for inference)
torch/jit/mobile/	Mobile lite-interpreter glue
torch/csrc/jit/	C++ implementation
torch/csrc/jit/frontend/	TorchScript parser, type inference
torch/csrc/jit/ir/	The JIT IR (Graph, Node, Value)
torch/csrc/jit/passes/	Optimization passes (constant prop, peephole, freeze, ONNX prep, …)
torch/csrc/jit/runtime/	The interpreter and the JIT compiler
torch/csrc/jit/python/	Python <-> JIT bindings
torch/csrc/jit/serialization/	.pt file format read/write
torch/csrc/jit/codegen/	Older codegen (NNC, fuser); largely superseded by Inductor
torch/csrc/jit/api/	C++-side API for libtorch users

Key abstractions

Type	File	Purpose
torch::jit::Graph	torch/csrc/jit/ir/ir.h	The JIT IR (different from torch.fx.Graph)
torch::jit::Node	torch/csrc/jit/ir/ir.h	An op with inputs (Value*s) and outputs
torch::jit::Value	torch/csrc/jit/ir/ir.h	A typed value flowing along edges
ScriptModule	torch/jit/_script.py	A compiled nn.Module
Method	torch/csrc/jit/api/method.h	One callable inside a ScriptModule
OperatorRegistry	torch/csrc/jit/runtime/operator.h	The JIT-side operator catalogue (overlaps with the dispatcher)
Pickler / Unpickler	torch/csrc/jit/serialization/pickler.h	The TorchScript pickle format

How it works

Tracing (torch.jit.trace)

trace runs a function with example inputs while a special tracing dispatch mode is active, recording each ATen op call as a JIT IR node. The result is a ScriptFunction whose IR Graph is exactly the sequence of ATen ops the example took. Tracing specializes to the control flow seen on the example; data-dependent control flow is silently lost.

Scripting (torch.jit.script)

script parses Python source (using a TorchScript-aware parser in torch/csrc/jit/frontend/) and compiles a typed subset of Python to JIT IR. Supported features include if/while/for, simple classes, tuples/lists/dicts, optionals, and a curated stdlib subset. Anything dynamic that doesn't fit the type system is rejected at compile time.

Optimization passes

torch/csrc/jit/passes/ contains dozens of passes that mutate the IR: constant propagation, peephole, dead-code elimination, common-subexpression elimination, batch matmul fusion, freezing (folding constants from a ScriptModule's parameters), ONNX-prep, shape analysis, etc. These run when torch.jit.optimize_for_inference or torch.jit.freeze is called.

Runtime

There are two runtimes:

• Interpreter (torch/csrc/jit/runtime/interpreter.cpp) — a stack-based VM that walks the IR.
• Static runtime (torch/csrc/jit/runtime/static/) — a faster Meta-internal runtime that pre-allocates tensor outputs and uses a custom out-of-place op interface; used in production inference at Meta.

The fuser stack (NNC, TensorExpr) under torch/csrc/jit/codegen/ was the precursor to Inductor; it remains for legacy ScriptModule users but is no longer the focus of new compiler work.

Serialization

.pt files are zip archives containing pickled IValues, IR (as TorchScript source), and tensor data; readers and writers live in torch/csrc/jit/serialization/. The "weights only" torch.load path in torch/_weights_only_unpickler.py is layered on top of this format. See Serialization.

Mobile

torch.jit.mobile (and the C++ torch::jit::mobile namespace) provides a stripped-down interpreter that runs .ptl (lite-interpreter) files. Used by PyTorch Mobile and its successor ExecuTorch.

Integration points

• Dispatcher. JIT ops are dispatcher ops; they share the same registration table.
• ONNX export (legacy path). The torch.onnx.export(...) legacy path traces through JIT. See ONNX.
• Mobile. PyTorch Mobile runs lite-interpreter on JIT-serialized models.
• Distributed. torch.distributed.rpc historically used JIT IR for serialization of remote calls.

Entry points for modification

• New TorchScript builtin → add to torch/csrc/jit/runtime/register_*.cpp.
• New optimization pass → drop a file in torch/csrc/jit/passes/, register from torch/csrc/jit/passes/optimize_for_inference.cpp or wherever appropriate.
• Most modern PyTorch development happens in torch.compile/torch.export rather than JIT. New high-level features should usually go there.

Key source files

File	Purpose
torch/jit/_script.py	ScriptModule, ScriptFunction
torch/jit/_trace.py	Tracing
torch/csrc/jit/ir/ir.h	JIT IR (Graph, Node, Value)
torch/csrc/jit/frontend/parser.cpp	TorchScript parser
torch/csrc/jit/runtime/interpreter.cpp	Bytecode interpreter
torch/csrc/jit/passes/freeze_module.cpp	Freezing for inference
torch/csrc/jit/serialization/pickler.cpp	.pt writer
torch/csrc/jit/serialization/unpickler.cpp	.pt reader