Merge branch 'main' into u/usimha/add-mnist-magnitude-pruning-tutorial

Author: u-simha

.gitignore

@@ -47,7 +47,6 @@ uv.lock
 *.nc
 .python-version
 /.tox
-deps/
 scratch/
 /local_stash
 /model_weights

.pre-commit-config.yaml

@@ -23,6 +23,7 @@ repos:
         language: system
         entry: scripts/pre_commit/add_license_header.py --license-file configs/BSD-3-LICENSE-HEADER-TEMPLATE --start-year 2026
         files: '\.(py|sh|js|css|html)$|(^|/)(GNUmakefile|[Mm]akefile)$'
+        exclude: ^src/coreai_opt/deps/
 
   # ----------------------------------------------------------------------------
   # 0.2 Update LICENSE year

NOTICE.txt

@@ -0,0 +1,25 @@
+Copyright 2026 Apple Inc.
+
+This project contains content adapted from kmeans1d (https://github.com/dstein64/kmeans1d), the license for which follows:
+
+MIT License
+
+Copyright (c) 2019 Daniel Steinberg
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

changelog.d/180525445.fixed

@@ -0,0 +1 @@
+Fix setting of qscheme and float_range for fixed output range ops

changelog.d/31.changed

@@ -0,0 +1 @@
+Replace the coremltools-based 1D k-means used by palettization with a vendored C++ core that is JIT-compiled at runtime via `torch.utils.cpp_extension`. `coremltools` is no longer a runtime dependency (it is now an optional dependency, installable via the `coreml` extra). This requires a C++ compiler to be available on the host at runtime.

docs/src/examples/resnet50.md

@@ -53,9 +53,9 @@ prepared_model = quantizer.prepare(example_inputs)
 
 Calibration is necessary when activation quantization is enabled. In order to determine proper quantization parameters for activation quantizers, representative data must be passed through the prepared model in the `calibration_mode()` context.
 
-Inside `calibration_mode()`, fake quantization is disabled and observers track tensor ranges seen at each activation quantizer.
-Each forward pass updates the activation scales without introducing quantization noise into the output.
-On exit, observers are disabled and fake quantization is re-enabled.
+Inside `calibration_mode()`, activation fake quantization is disabled while weight fake quantization stays on, and observers track tensor ranges seen at each activation quantizer.
+Each forward pass updates the activation scales using activations produced with quantized weights upstream, without injecting activation quantization noise into the observed values.
+On exit, observers are disabled and activation fake quantization is re-enabled.
 
 ```python
 with quantizer.calibration_mode():

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/advanced.md

@@ -166,17 +166,28 @@ The `qscheme` controls how these bins are distributed around zero, by determinin
 
 ## Quantization Defaults for Known-Range Activations
 
-In graph mode, certain ops have known output ranges. For these ops, the user's `qscheme` setting is not respected — the activation is always treated as asymmetric or symmetric depending on the op, regardless of what the user configured. The treatment also differs between `relu` and the `sigmoid` / `tanh` family. For `relu`, only `qscheme` is overridden; `dtype`, scale, and zero point are still derived from the user's spec and calibration data. For `sigmoid` and `tanh`, scale, zero point, **and** `dtype` are pinned to fixed values (always `torch.uint8`, ignoring whatever `dtype` the user configured).
+In graph mode, certain activation ops have analytically known output ranges. For these ops, the quantizer overrides the `qscheme` and `float_range` of the qparams calculator at prepare time, regardless of what the user configured. The user's `dtype` is always preserved — these adjustments do not change the number of bits or the signed/unsigned choice.
 
-| Op        | Output range | Always treated as | Scale     | Zero point |
-| --------- | ------------ | ----------------- | --------- | ---------- |
-| `relu`    | \[0, ∞)      | asymmetric        | dynamic   | dynamic    |
-| `sigmoid` | [0, 1]       | asymmetric        | `1 / 256` | `0`        |
-| `tanh`    | [-1, 1]      | symmetric         | `2 / 256` | `128`      |
+The scale and zero point values in the table below assume the default `int8` dtype. For other dtypes, the same formulas apply with the appropriate `quant_min` / `quant_max`.
 
-**Relu**: Treated as asymmetric. The user's `qscheme` is ignored, but `dtype`, scale, and zero point are still derived from the user's spec and calibration data. The zero point follows `zero_point = quant_min - round(min_val / scale)`. Since `relu`'s observed min is always `0`, the zero point very commonly ends up near `quant_min` (e.g., `-128` for `int8`).
+| Op            | Output range        | `qscheme`  | `float_range` | Scale (int8) | Zero point (int8) |
+| ------------- | ------------------- | ---------- | ------------- | ------------ | ----------------- |
+| `hardsigmoid` | [0, 1]              | asymmetric | (0, 1)        | 1 / 255      | −128              |
+| `hardtanh`    | Depends (see below) | Depends    | Depends       | Depends      | Depends           |
+| `relu`        | \[0, ∞)             | asymmetric | (0, None)     | dynamic      | −128              |
+| `relu6`       | [0, 6]              | asymmetric | (0, 6)        | 6 / 255      | −128              |
+| `sigmoid`     | [0, 1]              | asymmetric | (0, 1)        | 1 / 255      | −128              |
+| `tanh`        | [−1, 1]             | symmetric  | (−1, 1)       | 2 / 255      | 0                 |
 
-> **Motivation for asymmetric `relu` and `sigmoid`**: Both ops produce non-negative outputs. With symmetric quantization, the zero point sits at the center of the quantized range, placing half the bins in negative territory that these ops never produce. Those bins are effectively wasted — no floating-point value will ever map to them, reducing quantization resolution by half. Asymmetric treatment shifts the zero point toward the edge of the range so all bins cover values the op actually produces.
+**Relu**: The lower bound of `float_range` is pinned to 0 and `qscheme` is set to asymmetric. Because the observed minimum is always 0, the zero point is fixed at `quant_min` (−128 for int8) and stays there regardless of calibration data. The upper bound remains `None` (data-driven), so the scale continues to update during calibration.
+
+**Sigmoid and hardsigmoid**: Both `qscheme` and `float_range` are fully pinned. Scale and zero point are entirely determined by the dtype and the fixed output range — calibration data has no effect on them.
+
+**Tanh**: `qscheme` (symmetric) and `float_range` (−1, 1) are fully pinned. Scale and zero point are entirely determined by the dtype and range.
+
+**Hardtanh**: Bounds are read from the op's node arguments at prepare time, so the effective range and qscheme depend on how the op was configured. If `min_val == −max_val` the range is symmetric around zero and `qscheme` is set to symmetric; otherwise `qscheme` is set to asymmetric. Both ends of `float_range` are pinned to the configured bounds. `relu6` is a special case of `hardtanh(0, 6)` and is handled identically.
+
+> **Motivation for asymmetric treatment**: Symmetric quantization places the zero point at the center of the quantized range. For `relu`, `sigmoid`, and `hardsigmoid`, whose outputs are always non-negative, symmetric quantization places half the bins in negative territory that the op never produces — wasting half the available resolution. Asymmetric quantization shifts the zero point to the edge of the range so that all bins cover values the op actually generates. For `tanh` and symmetric `hardtanh`, the output is centered at zero so both halves of the range are used equally, and symmetric quantization is appropriate.
 
 Eager mode does not perform these adjustments — all activations are quantized uniformly using the user-configured spec.
 

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/quantization/overview.md

@@ -111,9 +111,10 @@ Right after `prepare()`, the activation scales come only from `example_inputs`,
 This is what the context manager handles:
 
 - Inside the context:
-  - fake-quantization is turned **off**: forward pass gives the same output as the unquantized model. Hence without distorting the outputs, the quantization params can be computed.
-  - range observers are turned **on**: this means that each forward pass updates the observed activation ranges, and hence the activation quantization scales.
-- After exiting the context manager, observers are turned back off and fake-quantization back on, leaving the model ready for evaluation.
+  - activation fake-quantization is turned **off**: activation observers see undistorted activation values, so the observed ranges (and resulting scales) reflect the true distribution rather than already-quantized values.
+  - weight fake-quantization stays **on**: activations flowing into each observer are produced with quantized weights upstream, matching what the deployed model will actually see.
+  - range observers are turned **on**: each forward pass updates the observed activation ranges, and hence the activation quantization scales.
+- After exiting the context manager, observers are turned back off and activation fake-quantization back on, leaving the model ready for evaluation.
 
 A small amount of representative data is typically enough.