Import Loops in PyTorch

Identifying and visualizing import loops in the PyTorch codebase

In this post, we will visualize all import loops in the PyTorch codebase, propose a fix for one potentially unstable case, and use Codegen to refactor that fix.

You can find the complete jupyter notebook in our examples repository.

Import loops (or circular dependencies) occur when two or more Python modules depend on each other, creating a cycle. For example:

# module_a.py
from module_b import function_b

# module_b.py
from module_a import function_a

While Python can handle some import cycles through its import machinery, they can lead to runtime errors, import deadlocks, or initialization order problems.

Debugging import cycle errors can be a challenge, especially when they occur in large codebases. However, Codegen allows us to identify these loops through our visualization tools and fix them very deterministically and at scale.

Import loop in pytorch/torchgen/model.py

Visualize Import Loops in PyTorch

Using Codegen, we discovered several import cycles in PyTorch’s codebase. The code to gather and visualize these loops is as follows:

G = nx.MultiDiGraph()

# Add all edges to the graph
for imp in codebase.imports:
    if imp.from_file and imp.to_file:
        edge_color = "red" if imp.is_dynamic else "black"
        edge_label = "dynamic" if imp.is_dynamic else "static"

        # Store the import statement and its metadata
        G.add_edge(
            imp.to_file.filepath,
            imp.from_file.filepath,
            color=edge_color,
            label=edge_label,
            is_dynamic=imp.is_dynamic,
            import_statement=imp,  # Store the whole import object
            key=id(imp.import_statement),
        )
# Find strongly connected components
cycles = [scc for scc in nx.strongly_connected_components(G) if len(scc) > 1]

print(f" Found {len(cycles)} import cycles:")
for i, cycle in enumerate(cycles, 1):
    print(f"\nCycle #{i}:")
    print(f"Size: {len(cycle)} files")

    # Create subgraph for this cycle to count edges
    cycle_subgraph = G.subgraph(cycle)

    # Count total edges
    total_edges = cycle_subgraph.number_of_edges()
    print(f"Total number of imports in cycle: {total_edges}")

    # Count dynamic and static imports separately
    dynamic_imports = sum(1 for u, v, data in cycle_subgraph.edges(data=True) if data.get("color") == "red")
    static_imports = sum(1 for u, v, data in cycle_subgraph.edges(data=True) if data.get("color") == "black")

    print(f"Number of dynamic imports: {dynamic_imports}")
    print(f"Number of static imports: {static_imports}")

Here is one example visualized ⤵️

Import loops in pytorch/torchgen/model.py

Not all import cycles are problematic! Some cycles using dynamic imports can work perfectly fine:

PyTorch prevents most circular import issues through dynamic imports which can be seen through the import_symbol.is_dynamic property. If any edge in a strongly connected component is dynamic, runtime conflicts are typically resolved.

However, we discovered an import loop worth investigating between flex_decoding.py and flex_attention.py:

flex_decoding.py imports flex_attention.py twice — once dynamically and once at top-level. This mixed static/dynamic import pattern from the same module creates potential runtime instability.

Thus, we propose the following refactoring using Codegen:

Move Shared Code to a Separate utils.py File

# Create new utils file
utils_file = codebase.create_file("torch/_inductor/kernel/flex_utils.py")

# Get the two files involved in the import cycle
decoding_file = codebase.get_file("torch/_inductor/kernel/flex_decoding.py")
attention_file = codebase.get_file("torch/_inductor/kernel/flex_attention.py")
attention_file_path = "torch/_inductor/kernel/flex_attention.py"
decoding_file_path = "torch/_inductor/kernel/flex_decoding.py"

# Track symbols to move
symbols_to_move = set()

# Find imports from flex_attention in flex_decoding
for imp in decoding_file.imports:
    if imp.from_file and imp.from_file.filepath == attention_file_path:
        # Get the actual symbol from flex_attention
        if imp.imported_symbol:
            symbols_to_move.add(imp.imported_symbol)

# Move identified symbols to utils file
for symbol in symbols_to_move:
    symbol.move_to_file(utils_file)

print(f" Moved {len(symbols_to_move)} symbols to flex_utils.py")
for symbol in symbols_to_move:
    print(symbol.name)

Running this codemod will move all the shared symbols to a separate utils.py as well as resolve the imports from both files to point to the newly created file solving this potential unpredictable error that could lead issues later on.

Conclusion

Import loops are a common challenge in large Python codebases. Using Codegen, no matter the repo size, you will gain some new insights into your codebase’s import structure and be able to perform deterministic manipulations saving developer hours and future runtime errors.

Want to try it yourself? Check out our complete example of fixing import loops using Codegen.