Just-in-Time Inspect Guide

Just-in-time inspection captures module hierarchy and input/output signatures during real model execution, without tuning or compilation. Use it to understand what JIT tuning would see before you commit to a tuning strategy.

Overview

Just-in-time inspection provides:

• Module Discovery: Capture modules that execute at runtime
• Hierarchy View: See parent/child relationships and depth
• Input/Output Signatures: Record tensor shapes, dtypes, and argument structure
• HTML Report: Generate a shareable inspection report

Quick Start

Enable inspection

Add a single import at the top of your script to enable inspection mode:

1
import aitune.torch.jit.enable_inspection as inspection  # noqa: F401

Run a real workload and save a report

The example below mirrors the Stable Diffusion inspection workflow used in tests/functional/pytorch/jit/006_jit_sd15_inspect.py:

1
import os
2
from logging import INFO, basicConfig
3
from pathlib import Path
4
5
import torch
6
from diffusers import StableDiffusionPipeline
7
8
import aitune.torch.jit.enable_inspection as inspection  # noqa: F401
9
10
11
def create_model():
12
    pipe = StableDiffusionPipeline.from_pretrained(
13
        "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16
14
    )
15
    pipe.to("cuda")
16
    return pipe
17
18
19
def main():
20
    basicConfig(level=INFO)
21
    prompt = (
22
        "A fluffy, orange tabby cat with bright green eyes is captured mid-air, "
23
        "pouncing playfully on a vibrant red ball of yarn"
24
    )
25
    pipe = create_model()
26
27
    def batch():
28
        with torch.no_grad():
29
            pipe([prompt] * 1, num_inference_steps=1)
30
            pipe([prompt] * 2, num_inference_steps=1)
31
32
    for _ in range(5):
33
        batch()
34
35
    html_path = "inspect_sd15.html"
36
    inspection.save_report(html_path, "SD15")
37
38
39
if __name__ == "__main__":
40
    main()

Open the report

$
firefox inspect_sd15.html
$
# or
$
google-chrome inspect_sd15.html

What the report shows

• Module summary: Total modules and basic statistics
• Module hierarchy: Tree view of the executed module structure
• Module details: Name, type, depth, parameter counts, call counts, execution times
• Inputs/outputs: Shapes, dtypes, and detected batch axes for each executed module

Example screenshots

The main view gives a top-level summary of discovered modules and quick navigation across the report sections.

If you click on a particular module, the detailed view shows execution stats and its recorded input/output signatures.

Notes

• Inspection data is collected only for modules that actually execute.
• Running multiple batches with different shapes helps capture dynamic behavior.
• The report is generated from in-memory inspection data at save_report time.

Next Steps

• Learn about Just-in-Time Tuning to actually tune models
• Compare with AOT Inspect for more control