🎓 A1Learner
🔍

PyTorch TorchScript Scripting

Scripting compiles a model's Python code into a standalone program. The compiled program keeps every branch written in the code.

Scripting the branching model

The scripted copy takes the negative branch on its own. No example input steered it there.

import torch

class Branchy(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * 3

scripted = torch.jit.script(Branchy())
print(scripted(torch.tensor([-2.0])).item())

Both branches in one program

One scripted program answers both signs correctly. The branch test runs again on every call.

import torch

class Branchy(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * 3

scripted = torch.jit.script(Branchy())
print(scripted(torch.tensor([2.0])).item())
print(scripted(torch.tensor([-2.0])).item())

A loop survives scripting

A loop in the code compiles along with the rest of forward. The scripted sum matches a hand count.

import torch

class Summer(torch.nn.Module):
    def forward(self, x):
        total = torch.zeros(())
        for i in range(x.shape[0]):
            total = total + x[i]
        return total

scripted = torch.jit.script(Summer())
print(scripted(torch.tensor([1.0, 2.0, 3.0])).item())

Example

import torch

class Branchy(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * 3

torch.manual_seed(0)
model = Branchy()
scripted = torch.jit.script(model)

print("Input | Eager | Scripted | Match")
for value in [2.0, -2.0, 0.5, -0.5]:
    x = torch.tensor([value])
    eager = model(x).item()
    got = scripted(x).item()
    print(f"{value} | {eager:.1f} | {got:.1f} | {eager == got}")

Key points

  • Use torch.jit.script(model) to compile a model from its code
  • Scripting reads the code itself, and no example input is needed
  • Each input chooses its branch at run time, just like the original
  • Scripting keeps the code's logic, and tracing keeps one run's path

In this example, the same Branchy model from the tracing lesson is scripted instead. Four inputs test both branches, two positive and two negative. Every scripted answer matches the eager model. The negative rows are the proof that matters. Scripting compiled the branch test itself, so each input chooses its own branch.

Note

Torch 2.14 marks jit.script as deprecated, and its warning points to torch.export

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does scripting compile?