Skip to content

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

Description

@RKey-PairApple

🌱 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    feature requestFunctionality does not currently exist, would need to be created as a new feature (type)

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions