How does Atlas find all callers of a PyTorch `nn.Module` method?

Atlas uses the `lsp` tool's `findReferences` operation. After identifying the `nn.Module` method with `documentSymbol`, Atlas queries the Language Server Protocol to locate every instance where that method is invoked across your entire PyTorch codebase, providing a comprehensive list of callers.

Can Atlas refactor PyTorch code without internet access?

Yes, Atlas can build its code index with local Ollama embeddings, keeping your PyTorch code off third-party servers. This allows for secure and offline refactoring, as core AI capabilities run locally, ensuring your sensitive `nn.Module` definitions remain private.

How does Atlas prevent breaking PyTorch tests during refactoring?

Atlas prevents test breakage by first running your `pytest` suite with its `bash` tool to establish a 'green' baseline. After each incremental change applied via `apply_patch`, Atlas re-runs `pytest` to immediately detect any regressions, especially those related to `torch.Tensor` shapes or `autograd` behavior.

What PyTorch files and configurations does Atlas read?

Atlas reads your `pyproject.toml` for project configuration, `nn.Module` definitions, training loops, and every `.to(device)` call. It understands common PyTorch idioms and file structures to provide context-aware assistance during refactoring.

How does Atlas ensure device placement changes don't break PyTorch?

Atlas is designed to understand PyTorch's device placement. It can read every `.to(device)` call and, during refactoring, will highlight or help modify these to ensure consistency. The continuous `pytest` runs after each change are crucial for catching any device-related runtime errors or performance regressions.

Can Atlas commit refactored PyTorch code on my behalf?

Yes, Atlas reads git branches, status, and diffs. After you approve the final changes, Atlas can stage the modified files and create commits on your behalf, integrating direct into your existing Git workflow for PyTorch development.

How does Atlas handle PyTorch package management during refactoring?

Atlas integrates with your project's package manager, `uv`. While its primary role is code refactoring, it understands the context provided by `uv` and your `pyproject.toml` to ensure that any code changes are consistent with your declared PyTorch dependencies.

Refactor a legacy PyTorch module with Atlas in 2026

To restructure an old PyTorch module without altering its behavior or breaking existing callers, Atlas provides a terminal-native AI coding agent that integrates directly with your PyTorch toolchain. It leverages `lsp` for comprehensive code analysis, `pytest` for rigorous behavior pinning, and `apply_patch` for precise, context-aware code modifications, ensuring your `nn.Module` definitions and training loops remain robust. You'll use `uv` for package management and `ruff format` for consistent code style.

How does Atlas map PyTorch module dependencies for refactoring?

Atlas maps a PyTorch module's public surface and its 100% of callers by leveraging the `lsp` tool's `documentSymbol` and `findReferences` operations. This ensures that every `nn.Module` definition, method, or function exported by the module is identified, along with all its usage sites across your codebase.

Before any refactoring begins, Atlas uses its `lsp` tool to perform a deep analysis of your PyTorch project. First, it executes `documentSymbol` on the target module to enumerate all public symbols, such as `nn.Module` subclasses, their methods, or standalone utility functions. For each identified symbol, Atlas then runs `findReferences` to locate every single callsite. This critical step prevents silent breakage by creating a complete inventory of dependencies. Atlas can also read your `pyproject.toml` to understand project structure and dependencies, ensuring it has a full context of your PyTorch environment. This comprehensive mapping is crucial for understanding the blast radius of any proposed change, especially when dealing with complex `torch.Tensor` operations or device placement logic that might be implicitly coupled across files.

How to pin PyTorch module behavior before refactoring?

Pinning the existing behavior of a PyTorch module is the first concrete step in a safe refactoring process, ensuring no regressions are introduced. Atlas achieves this by running your project's `pytest` suite with its `bash` tool, recording a green baseline of 0 failures before any code changes are made.

To guarantee that a refactor does not inadvertently alter the module's functionality, Atlas first establishes a behavioral baseline. Using the `bash` tool, Atlas executes your project's `pytest` command, for example: `bash -c "pytest tests/my_legacy_module_tests.py"`. This command runs the existing test suite for the module. Atlas then records the output, specifically looking for a 'green' run, meaning all tests passed. This baseline serves as the immutable contract for the module's behavior. If the module lacks sufficient test coverage, Atlas can be prompted to help write new `pytest` tests, focusing on critical aspects like `nn.Module` forward passes, loss calculations, or specific `torch.Tensor` manipulations, especially those involving `device` placement or `DataLoader` interactions, which are common sources of bugs in PyTorch.

How does Atlas apply structural changes to PyTorch code safely?

Atlas applies structural changes to your PyTorch code using the `apply_patch` tool, which is designed for precision and safety. This tool anchors on context lines and old lines, refusing to apply a hunk if the file has drifted, preventing unintended modifications to your `nn.Module` definitions or training loops. This ensures a 100% reliable application of changes.

Once the module's public surface is mapped and its behavior is pinned, Atlas proceeds with structural changes using the `apply_patch` tool. Unlike simple find-and-replace, `apply_patch` is robust: it takes a unified diff and attempts to apply each hunk. Crucially, it requires the context lines and the exact 'old_lines' of each hunk to match the current file content. If the file has been modified since the patch was generated, `apply_patch` will fail with a `Failed to find context` error, preventing a misapplication of changes. This is particularly vital in PyTorch, where subtle changes to `nn.Module` constructors, `forward` methods, or `torch.Tensor` shapes can lead to hard-to-debug runtime errors. Atlas computes a unified diff for every file edit and surfaces it for approval before writing, giving you full control over every change to your PyTorch codebase.

How does Atlas track and migrate PyTorch callsites during refactoring?

Atlas tracks remaining callsites in a `todowrite` list, ensuring that a partially migrated PyTorch module is never mistaken for a finished one. After each structural change and test re-run, Atlas updates this list, providing a clear, actionable roadmap for migrating all 100% of the module's callers.

Refactoring a PyTorch module often involves migrating numerous callsites across the codebase. Atlas manages this complex task by maintaining a `todowrite` list. After mapping all references with `lsp findReferences`, Atlas populates this list with every identified callsite. As structural changes are applied with `apply_patch` and corresponding callsites are updated, Atlas removes them from the `todowrite` list. This systematic approach ensures that no callsite is overlooked, preventing runtime errors that could arise from outdated invocations of your `nn.Module` or utility functions. The `todowrite` list acts as a persistent checklist, guiding the refactoring process until every single reference to the legacy module has been successfully migrated and verified, especially critical for ensuring correct `torch.Tensor` handling and `device` placement across the entire application.

What safety and review features does Atlas offer for PyTorch refactoring?

Atlas provides multiple layers of safety and review for PyTorch refactoring, including permission-gated tool calls, unified diffs for every edit, and robust git integration. Every Atlas tool call, such as `bash` or `apply_patch`, is permission-gated against allow, ask, and deny rules, giving you 100% control over agent actions.

Atlas is built with safety at its core, offering several features to ensure secure PyTorch refactoring. Before any tool like `bash` (for running `pytest`) or `apply_patch` is executed, Atlas checks against predefined permission rules (allow, ask, deny). This means you explicitly approve sensitive operations. For every proposed file edit, Atlas computes and displays a unified diff, allowing you to review and approve changes before they are written to disk. This is crucial for verifying modifications to `nn.Module` definitions, `torch.optim` configurations, or `DataLoader` setups. Furthermore, Atlas reads git branches, status, and diffs, and can stage and create commits on your behalf, ensuring that your refactoring work is properly version-controlled. It also snapshots file changes as git patches, enabling easy diffing and rolling back of edits, providing a robust safety net for complex PyTorch codebase transformations.

Step by step

01Initialize Atlas in your PyTorch project, ensuring it can read your `pyproject.toml` and `nn.Module` definitions. Atlas will build its code index using local Ollama embeddings, keeping your PyTorch code off third-party servers.
02Map the legacy PyTorch module's public surface: Ask Atlas to use the `lsp` tool's `documentSymbol` operation on your target module (e.g., `src/models/legacy_net.py`), then `findReferences` on each exported symbol to enumerate every callsite across your project.
03Pin existing PyTorch behavior: Instruct Atlas to run your project's `pytest` suite using the `bash` tool (e.g., `bash -c "pytest tests/legacy_module_tests.py"`) and record the green baseline before any code changes are made. This ensures `torch.Tensor` operations and `device` placements are stable.
04Restructure the PyTorch module: Guide Atlas to apply structural changes using `apply_patch`. Atlas will propose changes to `nn.Module` definitions or `forward` methods, anchoring on context lines and refusing to apply against a drifted file. Review each unified diff before approval.
05Re-run PyTorch tests incrementally: After each significant hunk lands via `apply_patch`, ask Atlas to re-run the relevant `pytest` tests with `bash` (e.g., `bash -c "pytest tests/legacy_module_tests.py"`). This immediate feedback loop catches regressions early, especially for `autograd` or `DataLoader` issues.
06Track remaining PyTorch callsites: Use the `todowrite` tool to keep a running list of unmigrated callsites. As Atlas updates callers of the refactored `nn.Module`, it will remove them from this list, ensuring a complete migration.
07Finalize and format: Once all callsites are migrated and tests pass, ask Atlas to run `ruff format` across the modified files to ensure consistent code style. Atlas will present the final diff for approval before writing.

Frequently asked questions

How does Atlas find all callers of a PyTorch `nn.Module` method?: Atlas uses the `lsp` tool's `findReferences` operation. After identifying the `nn.Module` method with `documentSymbol`, Atlas queries the Language Server Protocol to locate every instance where that method is invoked across your entire PyTorch codebase, providing a comprehensive list of callers.
Can Atlas refactor PyTorch code without internet access?: Yes, Atlas can build its code index with local Ollama embeddings, keeping your PyTorch code off third-party servers. This allows for secure and offline refactoring, as core AI capabilities run locally, ensuring your sensitive `nn.Module` definitions remain private.
How does Atlas prevent breaking PyTorch tests during refactoring?: Atlas prevents test breakage by first running your `pytest` suite with its `bash` tool to establish a 'green' baseline. After each incremental change applied via `apply_patch`, Atlas re-runs `pytest` to immediately detect any regressions, especially those related to `torch.Tensor` shapes or `autograd` behavior.
What PyTorch files and configurations does Atlas read?: Atlas reads your `pyproject.toml` for project configuration, `nn.Module` definitions, training loops, and every `.to(device)` call. It understands common PyTorch idioms and file structures to provide context-aware assistance during refactoring.
How does Atlas ensure device placement changes don't break PyTorch?: Atlas is designed to understand PyTorch's device placement. It can read every `.to(device)` call and, during refactoring, will highlight or help modify these to ensure consistency. The continuous `pytest` runs after each change are crucial for catching any device-related runtime errors or performance regressions.
Can Atlas commit refactored PyTorch code on my behalf?: Yes, Atlas reads git branches, status, and diffs. After you approve the final changes, Atlas can stage the modified files and create commits on your behalf, integrating direct into your existing Git workflow for PyTorch development.
How does Atlas handle PyTorch package management during refactoring?: Atlas integrates with your project's package manager, `uv`. While its primary role is code refactoring, it understands the context provided by `uv` and your `pyproject.toml` to ensure that any code changes are consistent with your declared PyTorch dependencies.

Try Atlas in your terminal

The terminal-native AI coding agent. Free core, single binary.

Install Atlas

Related guides

Refactor a Legacy Module with Atlas in 2026

How to refactor a legacy module with Atlas in 2026: findReferences maps every callsite, apply_patch refuses to apply against a drifted file, and bash proves behavior.

Atlas for PyTorch: Terminal-Native AI Coding for nn.Module, Devices, and Autograd in 2026

Atlas is a terminal-native AI coding agent for PyTorch in 2026, where device placement, autograd, and DataLoader worker counts cause most bugs and most slowness.

Locate Where a Behavior is Implemented in PyTorch with Atlas in 2026

Pinpoint the exact file and symbol responsible for a behavior in your PyTorch codebase using Atlas. Leverage semantic search, grep, and LSP tools for precise results.

Rename a symbol across the repo in PyTorch with Atlas in 2026

Effortlessly rename PyTorch functions, classes, or constants across your entire codebase in 2026 with Atlas. Leverage lsp, grep, and edit for precise, verified refactoring.

Document a PyTorch Module with a README using Atlas in 2026

In 2026, Atlas helps PyTorch developers generate accurate READMEs for modules by analyzing live code, ensuring documentation reflects current device placement, autograd, and DataLoader configurations.

Plan a Multi-File Change Before Editing in PyTorch with Atlas in 2026

Design and review complex, multi-file PyTorch changes with Atlas's plan agent before modifying a single line of code. Leverage real PyTorch tools like pytest and uv.

Write Unit Tests for Untested PyTorch Code with Atlas in 2026

Add real unit tests to untested PyTorch modules using Atlas. Leverage pytest, uv, and ruff format to match existing repo conventions and ensure code quality.

Trace a runtime bug from a stack trace in PyTorch with Atlas in 2026

Trace PyTorch runtime bugs from production stack traces with Atlas. Identify the responsible line and generate fixes without a debugger, streamlining your PyTorch development.

Browse this resource hub