Update model inspector documentation

docs/src/palettization/config.md

@@ -116,7 +116,7 @@ op_config = OpKMeansPalettizerConfig(
 
 ## Examples
 
-Several examples below configure specific module types or module names. To determine these for your model, see [How to get names + types](../quantization/config.md#how-to-get-names--types-for-modules-and-ops). Since palettization only supports eager execution mode, only the eager mode guidance in that section is relevant.
+Several examples below configure specific module types or module names. To determine these for your model, use {class}`~coreai_opt.inspection.ModelInspector` with `execution_mode="eager"` — see [Inspecting Model Structure](../utils/model_inspection.md). Palettization supports eager mode only.
 
 ### Apply 4-bit palettization globally, 8-bit to linear layers
 

docs/src/quantization/config.md

@@ -184,7 +184,7 @@ The defaults are:
 
 In [Quantization Overview](overview.md) we saw how to use the default `W_INT8_A_INT8` config. [Config classes and their defaults](#config-classes-and-their-defaults) described the default settings in `QuantizerConfig()`, `ModuleQuantizerConfig()`, and `OpQuantizerConfig()`. Let us now see how to configure quantization when non-default settings are desired.
 
-Several examples below configure specific module names, module types, op names, or op types. To determine these for your model, see [How to get names + types for modules and ops](#how-to-get-names-types-for-modules-and-ops) (eager mode) or [Inspecting Model Structure](../utils/model_inspection.md) (graph mode).
+Several examples below configure specific module names, module types, op names, or op types. To determine these for your model, see [Inspecting Model Structure](../utils/model_inspection.md).
 
 ### Example: `W_MXFP4_A_FP8` applied to all supported ops
 
@@ -857,7 +857,7 @@ classDiagram
 
 ## How to get names + types for modules and ops
 
-**Graph mode** (for `module_name_configs`, `module_type_configs`, `op_name_config`, `op_type_config`): use {class}`~coreai_opt.inspection.ModelInspector` to discover module names, module types, op names, and op types.
+Use {class}`~coreai_opt.inspection.ModelInspector` to discover module names, module types, op names, and op types for both graph and eager execution modes.
 
 ```python
 import torch
@@ -866,40 +866,12 @@ from coreai_opt.inspection import ModelInspector
 
 model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))
 inspector = ModelInspector(
-    model, example_inputs=(torch.randn(1, 10),), execution_mode="graph"
+    # Use execution_mode="eager" for eager mode inspection.
+    model,
+    example_inputs=(torch.randn(1, 10),),
+    execution_mode="graph",
 )
 print(inspector.format_summary())
 ```
 
-See [Inspecting Model Structure](../utils/model_inspection.md) for full usage and examples.
-
-**Eager mode**: module names (for `module_name_configs`) can be obtained by inspecting `model.named_modules()`. This includes all modules in the model (nested and leaf). The names align with the structure of modules defined in code.
-
-Op names can be constructed by referring to the parent module and the op in it. Example:
-
-```python
-class TwoAddModule(torch.nn.Module):
-    def forward(self, x):
-        a = x + x
-        b = a + a
-        return b
-
-
-class Model(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.submodule_a = TwoAddModule()
-        self.submodule_b = TwoAddModule()
-
-    def forward(self, x):
-        x = self.submodule_a(x)
-        x = self.submodule_b(x)
-        return x
-
-
-# the names of ops in the model will be :
-# - submodule_a.add
-# - submodule_a.add_1
-# - submodule_b.add
-# - submodule_b.add_1
-```
+See [Inspecting Model Structure](../utils/model_inspection.md) for full usage, examples, and a comparison of graph and eager mode op naming.

docs/src/utils/model_inspection.md

@@ -2,9 +2,16 @@
 
 `coreai-opt` configs reference module names, module types, op names, and op types to target specific parts of a model. Before writing a config, you need to know exactly which strings your model exposes. {class}`~coreai_opt.inspection.ModelInspector` discovers these automatically and provides query methods corresponding to each config key type (`op_type_config`, `op_name_config`, `module_name_configs`, `module_type_configs`).
 
-:::{note}
-`ModelInspector` currently supports **graph execution mode only**. Eager mode support is planned. For eager mode op naming, see [How to get names + types](../quantization/config.md#how-to-get-names--types-for-modules-and-ops).
-:::
+## Execution Modes
+
+`ModelInspector` supports two execution modes, selected via the `execution_mode` argument:
+
+- **Graph mode** (`execution_mode="graph"`): Exports the model with `torch.export` and walks the resulting FX graph. Op names are global identifiers assigned during export (for example, `"linear"`, `"linear_1"`). The compressor must be `Quantizer` or `None`.
+- **Eager mode** (`execution_mode="eager"`): Intercepts operations during a live forward pass. Op names are module-qualified identifiers that reflect the module hierarchy (for example, `"linear1.linear"`, `"linear2.linear"`). This mode supports both `Quantizer` and `KMeansPalettizer` as the compressor.
+
+If you plan to compress the model using one of `coreai-opt`'s compression techniques, choose the `execution_mode` you plan to use when compressing for inspection in order to identify the correct op and module names to use in the compression config.
+
+For more information on `graph` mode vs. `eager` mode, see [here](../quantization/overview.md#two-execution-modes-graph-and-eager).
 
 ## Basic Usage
 
@@ -38,97 +45,125 @@ inspector = ModelInspector(
     compressor=Quantizer,
 )
 
-# Print a module-hierarchy tree showing ops, source locations, and connectivity
+# Print a module-hierarchy tree showing ops, connectivity, and source locations
 print(inspector.format_summary())
 ```
 
-Note the use of `compressor=Quantizer` in the list of arguments to `ModelInspector`. This filters the list of ops captured and displayed by `ModelInspector` to only those operations which are registered for compressibility by `Quantizer`. Omitting this argument allows for all ops to be captured and displayed.
+Pass `colorize=False` to suppress ANSI color codes, for example when writing to a file.
+
+:::{note}
+Note the use of `compressor=Quantizer`. This filters the captured and displayed ops to those registered as compressible by `Quantizer`. Omit this argument to capture and display all ops.
+:::
 
 The above code produces output like the following (colors omitted for brevity):
 
 ```text
-Legend:  ■ module_name (module_type)  ◆ op_name [op_type]
-
-(MyModel)
-    module inputs:  linear
-    module outputs: linear_1
+Legend:
+  ■ module_name (module_type)  ◆ op_name [op_type]
+
+  op inputs:  {I: producer[N]}  —  I = op_input_spec index; N = output slot of the producing op
+  op states:     param_name    —  model parameter or buffer
+  op outputs: {N: [consumers]} —  N = output slot index; consumers = ops receiving that output
+  untracked_N                  —  input tensor whose producer was not intercepted (e.g. raw attribute or global tensor); still quantizable via op_input_spec
+  module inputs:  {I: [op[N], ...]}  —  I = module_input_spec index; op[N] = op and its input slot receiving data from outside; absent keys = non-quantizable
+  module outputs: {I: op[N]}         —  I = module_output_spec index; op[N] = op and its output slot leaving the module; absent keys = non-quantizable
+
+(__main__.MyModel)
+    module inputs:  {0: [linear[0]]}
+    module outputs: {0: linear_1[0]}
 ├── ■ linear1 (torch.nn.modules.linear.Linear)
-│       module inputs:  linear
-│       module outputs: linear
+│       module inputs:  {0: [linear[0]]}
+│       module outputs: {0: linear[0]}
 │   └── ◆ linear [linear]
-│         op inputs:  x, linear1_weight, linear1_bias
-│         op outputs: relu
-│         filepath:  my_model.py:16
-│         code:      x = self.linear1(x)
+│         op inputs:  {0: x[0]}
+│         op states:  weight, bias
+│         op outputs: {0: [relu]}
 ├── ■ relu (torch.nn.modules.activation.ReLU)
-│       module inputs:  relu
-│       module outputs: relu
+│       module inputs:  {0: [relu[0]]}
+│       module outputs: {0: relu[0]}
 └── ■ linear2 (torch.nn.modules.linear.Linear)
-        module inputs:  linear_1
-        module outputs: linear_1
+        module inputs:  {0: [linear_1[0]]}
+        module outputs: {0: linear_1[0]}
     └── ◆ linear_1 [linear]
-          op inputs:  relu, linear2_weight, linear2_bias
-          op outputs: output
-          filepath:  my_model.py:18
-          code:      x = self.linear2(x)
+          op inputs:  {0: relu[0]}
+          op states:  weight, bias
+          op outputs: {0: [output]}
 ```
 
-The output shows the model's module hierarchy and the ops within each module. Note in particular that since `relu` is not a registered compressible op by `Quantizer`, it does not show up as an operation within the `ReLU` module.
+Note that `relu` does not appear as an operation (`◆`) within the `relu` module, because `ReLU` is not a compressible op in `Quantizer`. It still appears as a module node (`■`) and in connectivity lines such as `op outputs: {0: [relu]}` and `op inputs: {0: relu[0]}`, because the relu tensor passes through and connects the two linear ops.
 
-Reading the tree:
+## Reading the Tree
 
-- **Module name** and **module type** appear on module lines: `■ module_name (module_type)`. For example, `■ linear1 (torch.nn.modules.linear.Linear)` — `"linear1"` is the module name (usable in `module_name_configs`) and `"torch.nn.modules.linear.Linear"` is the module type (usable in `module_type_configs`).
-- **Op name** and **op type** appear on operation lines: `◆ op_name [op_type]`. For example, `◆ linear_1 [linear]` — `"linear_1"` is the op name (usable in `op_name_config`) and `"linear"` is the op type (usable in `op_type_config`).
-- **Op inputs/outputs** show connectivity between operations.
-- **filepath** and **code** (shown for user-defined modules) show where in your source code the operation originates.
+### Module lines
 
-Using these strings directly in a config:
+Module lines use the form `■ module_name (module_type)`. For example, `■ linear1 (torch.nn.modules.linear.Linear)`:
 
-```python
-config = QuantizerConfig(
-    # Target a specific module by name
-    module_name_configs={
-        "linear1": ModuleQuantizerConfig(...),
-    },
-    # Target all modules of a given type
-    module_type_configs={
-        "torch.nn.modules.linear.Linear": ModuleQuantizerConfig(...),
-    },
-)
+- `"linear1"` is the module name, usable in `module_name_configs`.
+- `"torch.nn.modules.linear.Linear"` is the module type, usable in `module_type_configs`.
 
-# Op-level targeting within a ModuleQuantizerConfig
-config = QuantizerConfig(
-    global_config=ModuleQuantizerConfig(
-        # Target a specific op by name
-        op_name_config={
-            "linear_1": OpQuantizerConfig(...),
-        },
-        # Target all ops of a given type
-        op_type_config={
-            "linear": OpQuantizerConfig(...),
-        },
-    ),
-)
-```
+**Module boundaries** appear indented under the module header:
+
+- `module inputs: {I: [op[N], ...]}` — The activations entering this module from outside. `I` is the position in the module's input spec (matching `module_input_spec` in a config), `op` is the name of the first compressible op inside the module that receives data at that position, and `N` is the input slot on that op. A single external input can fan out to multiple ops. Keys absent from this dict correspond to non-quantizable positions (for example, state tensors or unused arguments).
+- `module outputs: {I: op[N]}` — The activations leaving this module. `I` is the position in the module's output spec, `op` is the compressible op producing that output, and `N` is the op's output slot. Absent keys correspond to non-quantizable positions.
+
+### Op lines
+
+Op lines use the form `◆ op_name [op_type]`. For example, `◆ linear_1 [linear]`:
+
+- `"linear_1"` is the op name, usable in `op_name_config`.
+- `"linear"` is the op type, usable in `op_type_config`.
+
+**Op connectivity** appears indented under the op header:
+
+- `op inputs: {I: producer[N]}` — Activation inputs only (parameters and buffers are on a separate line). `I` is the argument position (matching `op_input_spec` in a config), `producer` is the name of the op that produced this tensor, and `N` is the output slot of that producer. For example, `{0: relu[0]}` means argument 0 comes from output slot 0 of the `relu` op.
+- `op states: param_name, ...` — Model parameters and buffers consumed by this op. This line is omitted if the op takes no states.
+- `op outputs: {N: [consumer1, consumer2, ...]}` — `N` is the output slot index, and the list contains the names of all ops consuming that output.
+- `untracked_N` — Appears in place of a producer name when the input tensor's origin was not intercepted (for example, a raw module attribute or global tensor). These tensors are still quantizable via `op_input_spec`.
+- `filepath` and `code` — Source file and line of the call that produced the op, shown as dim text.
+
+## Eager Mode
 
-Pass `colorize=False` to suppress ANSI color codes (e.g., when writing to a file).
+To inspect using eager mode, pass `execution_mode="eager"`. The same `MyModel` example above yields:
+
+```text
+(__main__.MyModel)
+    module inputs:  {0: [linear1.linear[0]]}
+    module outputs: {0: linear2.linear[0]}
+├── ■ linear1 (torch.nn.modules.linear.Linear)
+│       module inputs:  {0: [linear1.linear[0]]}
+│       module outputs: {0: linear1.linear[0]}
+│   └── ◆ linear1.linear [linear]
+│         op inputs:  {0: input_0}
+│         op states:  weight, bias
+│         op outputs: {0: [relu.relu]}
+│         filepath:  my_model.py:16
+├── ■ relu (torch.nn.modules.activation.ReLU)
+│       module inputs:  {0: [relu.relu[0]]}
+│       module outputs: {0: relu.relu[0]}
+└── ■ linear2 (torch.nn.modules.linear.Linear)
+        module inputs:  {0: [linear2.linear[0]]}
+        module outputs: {0: linear2.linear[0]}
+    └── ◆ linear2.linear [linear]
+          op inputs:  {0: relu.relu[0]}
+          op states:  weight, bias
+          op outputs: {0: [output_0]}
+          filepath:  my_model.py:18
+```
 
 ## Querying Operations by Config Key
 
-Once you have reviewed the full summary to see what names and types are present, you can use query methods to check which operations would be matched by a specific name or type pattern. This is useful for verifying your config will target the intended ops before applying compression.
+Once you have reviewed the full summary to see what names and types are present, use the query methods to check which operations a specific pattern matches. This is useful for verifying that a config targets the intended ops before applying compression.
 
 Each query method returns a tuple of {class}`~coreai_opt.inspection.OpInfo` objects matching the filter. The method names correspond directly to the config keys they help populate.
 
-From the Basic Usage summary, this model exposes:
+From the graph mode summary above, this model exposes:
 
 - **Op types**: `linear`
 - **Op names**: `linear`, `linear_1`
 - **Module types**: `torch.nn.modules.linear.Linear`, `torch.nn.modules.activation.ReLU`
 - **Module names**: `linear1`, `relu`, `linear2`
 
-Op names and module names can be passed as a literal name or as a regex following [Python re syntax](https://docs.python.org/3/library/re.html) for wildcard matching; the pattern is matched against the entire string. The matching methodology is identical to how compression config entries match modules and ops in a model, allowing the user to see exactly which modules or ops would be matched given a particular string.
-
-Each query method returns a tuple of {class}`~coreai_opt.inspection.OpInfo` objects matching the filter:
+Op names and module names can be passed as a literal string or as a regex following [Python re syntax](https://docs.python.org/3/library/re.html) for wildcard matching. The pattern is matched against the full string. The matching behavior is identical to how compression config entries match modules and ops, so you can see exactly which ops a given pattern would select.
 
 **By op type** — exact-string match against `op_type_config` keys:
 
@@ -151,10 +186,10 @@ inspector.get_matched_ops_for_module_name(
 )  # matches the op in module "linear1"
 inspector.get_matched_ops_for_module_name(
     "linear[12]"
-)  # matches ops in modules "linear1" and "linear2"
+)  # matches ops in "linear1" and "linear2"
 ```
 
-Each returned {class}`~coreai_opt.inspection.OpInfo` provides `op_name`, `op_type`, and `module_stack` (the nesting of modules containing the op):
+Each returned {class}`~coreai_opt.inspection.OpInfo` provides `op_name`, `op_type`, `module_stack`, `inputs`, `outputs`, and `is_state`. The `module_stack` is a tuple of {class}`~coreai_opt.inspection.ModuleContext` entries from outermost to innermost module:
 
 ```python
 >>> for op in inspector.get_matched_ops_for_op_type("linear"):
@@ -166,6 +201,37 @@ Each returned {class}`~coreai_opt.inspection.OpInfo` provides `op_name`, `op_typ
   module: linear2 (torch.nn.modules.linear.Linear)
 ```
 
+`OpInfo.inputs` is a tuple of {class}`~coreai_opt.inspection.InputEdge` objects, one per input argument position. Each `InputEdge` carries the producing `OpInfo` and the output slot index (`output_idx`) of that producer. State inputs (parameters, buffers) are interleaved in the tuple at their actual argument positions, and their corresponding `InputEdge` objects have `is_state=True`.
+
+Using these strings directly in a config:
+
+```python
+config = QuantizerConfig(
+    # Target a specific module by name
+    module_name_configs={
+        "linear1": ModuleQuantizerConfig(...),
+    },
+    # Target all modules of a given type
+    module_type_configs={
+        "torch.nn.modules.linear.Linear": ModuleQuantizerConfig(...),
+    },
+)
+
+# Op-level targeting within a ModuleQuantizerConfig
+config = QuantizerConfig(
+    global_config=ModuleQuantizerConfig(
+        # Target a specific op by name
+        op_name_config={
+            "linear_1": OpQuantizerConfig(...),
+        },
+        # Target all ops of a given type
+        op_type_config={
+            "linear": OpQuantizerConfig(...),
+        },
+    ),
+)
+```
+
 ## Navigating the Module Hierarchy
 
 For programmatic access to the inspector's data structures, the {class}`~coreai_opt.inspection.ModelSummary` exposes a {class}`~coreai_opt.inspection.ModuleInfo` tree that mirrors the `nn.Module` hierarchy. These types are publicly exported from `coreai_opt.inspection` for use in custom analysis or tooling.
@@ -180,11 +246,24 @@ linear2: torch.nn.modules.linear.Linear, 1 direct ops
 ```
 
 ```python
-# look up a specific submodule
+# Look up a specific submodule
 linear2_module = root.get_submodule("linear2")
 
-# get all ops under this subtree (depth-first)
+# Get all ops under this subtree (depth-first)
 linear2_ops = linear2_module.all_ops()
 ```
 
 `ModuleInfo` supports the same iteration patterns as `nn.Module`: `children()`, `named_children()`, `modules()`, `named_modules()`, and `get_submodule()`.
+
+`ModuleInfo` also exposes the module boundary connectivity described in the tree:
+
+- `input_ops` — dict mapping module input spec index to a list of {class}`~coreai_opt.inspection.BoundaryEdge` objects, each holding the op and input slot receiving data from outside the module.
+- `output_ops` — dict mapping module output spec index to a single {class}`~coreai_opt.inspection.BoundaryEdge`, holding the op and output slot whose tensor leaves the module.
+
+```python
+# Inspect boundary connectivity for a submodule
+linear1_module = root.get_submodule("linear1")
+for idx, edges in linear1_module.input_ops.items():
+    for edge in edges:
+        print(f"  module input {idx} -> {edge.op.op_name}[{edge.index}]")
+```