Map MIL ops back to their originating nn.Module for debugging and profiling
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## 🌱 Describe your Feature Request
Add a new ScopeSource.TORCHSCRIPT_MODULE_PATH scope, populated from torch._C.Node.scopeName(), that preserves
nn.ModuleList / nn.ModuleDict container attribute names. The existing TORCHSCRIPT_MODULE_NAME is derived from
getModuleHierarchy(), which drops those container names — so any two ops emitted from sibling ModuleList submodules
(e.g. self.layers_a[0] and self.layers_b[0]) carry byte-identical scope tokens and cannot be reverse-mapped to their
original named_modules() path. The new source records the full dotted path tokens — equal to a valid named_modules()
key — alongside the existing scopes. The change is additive: existing scope outputs are unchanged.
How can this feature be used?
Reverse-mapping a converted MIL op back to its originating PyTorch submodule.
Describe alternatives you've considered
User-side reconstruction from TORCHSCRIPT_MODULE_NAME by treating the scope tokens as an ordered subsequence of a
named_modules() path. Works for unique paths but is structurally insufficient when two sibling ModuleLists hold
same-shaped children — both produce identical scope tokens, so no tiebreaker on MODULE_NAME alone can recover the right
path. Measured ~24% wrong picks on a real model.
Replacing TORCHSCRIPT_MODULE_NAME's source with scopeName() directly. Would change observable behavior for existing
consumers. Chose additive instead.
Additional context
Implementation is additive. TORCHSCRIPT_MODULE_NAME, TORCHSCRIPT_MODULE_TYPE,
and _trim_scopename_for_weight outputs are byte-identical to before; coremltools/mil/tests/test_programs.py and
mil/passes/tests/test_passes.py scope tests pass unchanged. The only externally observable change is
InternalTorchIRNode.get_scope_info() now returning a 3-tuple (scope_name, scope_type, module_path) instead of a 2-tuple
— all in-tree callers updated.
Reopened #2749
Contributor guide
Research direction
Start at ScopeSource and InternalTorchIRNode.get_scope_info(), then inspect the existing scope tests in coremltools/mil/tests/test_programs.py and mil/passes/tests/test_passes.py. Trace how TorchScript node scopes become MIL scope data, preserving existing outputs while adding the module-path value and updating in-tree callers for the 3-tuple. Done means sibling ModuleList or ModuleDict paths can be distinguished and the named scope tests remain unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100