apple / apple/coremltools

Map MIL ops back to their originating nn.Module for debugging and profiling

Open
#2,751 3 comments 0 reactions 0 assignees View on GitHub
feature request
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.