diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 39d2410..0f51d09 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -305,6 +305,20 @@ repos:
         types: [python]
         exclude: ^tests/
 
+  - repo: local
+    hooks:
+      - id: check-internal-import-aliases
+        name: Check `_`-prefix aliasing of internal imports
+        description: |
+          Enforce code_style_guide.md §3.3: public modules must alias symbols
+          imported from private modules with a `_` prefix; private modules must
+          not. Fails the commit and reports each violation with a suggested fix;
+          it does not edit files — run the script with `--fix` to apply fixes.
+        entry: python scripts/pre_commit/check_internal_import_aliases.py
+        language: system
+        types: [python]
+        exclude: ^tests/
+
   - repo: local
     hooks:
       - id: towncrier-check
diff --git a/docs/contributing/code_style_guide.md b/docs/contributing/code_style_guide.md
index 2040d0c..f783a3f 100644
--- a/docs/contributing/code_style_guide.md
+++ b/docs/contributing/code_style_guide.md
@@ -21,6 +21,7 @@ This guide documents coding conventions and best practices for the Core AI Optim
     - [Why this matters](#why-this-matters)
     - [Best practice](#best-practice)
     - [When to use this pattern](#when-to-use-this-pattern)
+    - [Enforcement](#enforcement)
   - [3.4 Import Examples](#34-import-examples)
 - [4. Type Annotations](#4-type-annotations)
   - [4.1 When to Annotate](#41-when-to-annotate)
@@ -330,6 +331,17 @@ from proj_pkg.api import private_helper  # ✗ AttributeError - not accessible
 
 - When importing project-internal symbols into public modules to prevent unintended re-export
 
+#### Enforcement
+
+The pre-commit hook `check-internal-import-aliases` (defined in `.pre-commit-config.yaml`, implemented at `scripts/pre_commit/check_internal_import_aliases.py`) enforces this rule. On commit it **fails** when it finds a violation and prints each one with a suggested fix, but it does not edit files — so the commit is blocked until the imports are corrected. Fix them by hand, or run the script with `--fix` to apply the fixable ones automatically:
+
+- Add the missing `_` alias in public modules.
+- Remove the unnecessary `_` alias in private modules.
+- Rename every reference to the bound name in the same file so the rewrite leaves the module compiling.
+- Report (without fixing) when the bound name is shadowed inside a function scope (rare) — resolve those by hand.
+
+Fixing is a separate, opt-in step rather than an automatic commit-time rewrite: auto-editing on commit would interleave this hook's changes with those of the formatting/linting hooks, making them hard to review or revert. Run `python scripts/pre_commit/check_internal_import_aliases.py --fix` (optionally with specific paths) to fix everything in one go, then review and commit the changes.
+
 ### 3.4 Import Examples
 
 **Importing from Public API:**
diff --git a/scripts/pre_commit/check_internal_import_aliases.py b/scripts/pre_commit/check_internal_import_aliases.py
new file mode 100755
index 0000000..4c45207
--- /dev/null
+++ b/scripts/pre_commit/check_internal_import_aliases.py
@@ -0,0 +1,723 @@
+#!/usr/bin/env python3
+
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-Clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+
+"""Enforce `_`-prefix aliasing of project-internal imports per code style guide.
+
+A module is "private" when any segment of its path starts with a single
+underscore (e.g., ``coreai_opt/_utils/helpers.py``,
+``coreai_opt/quantization/_eager/quantizer.py``). When a public module imports
+a symbol from a private module, the bound name must be aliased with a ``_``
+prefix so it cannot be re-imported through the public namespace::
+
+    # public module
+    from coreai_opt._utils import helper as _helper  # OK
+
+When a private module imports from another private module, the alias is
+unnecessary and adds visual noise::
+
+    # private module (path contains `_`)
+    from coreai_opt._utils import helper  # OK, no alias needed
+
+This script flags both kinds of mismatch and, with ``--fix``, rewrites the
+import statement and renames every reference to the bound name within the
+file. See ``docs/contributing/code_style_guide.md`` section 3.3.
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, NamedTuple, Protocol
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+# Keep in sync with the formatter line-length in configs/darker.toml and
+# pyproject.toml ([tool.ruff]); this hook wraps its own rewrites because it
+# runs after the formatter hooks in .pre-commit-config.yaml.
+LINE_LIMIT = 100
+
+
+def _is_private_part(part: str) -> bool:
+    """Return True for single-underscore segments, False for dunder/public."""
+    return part.startswith("_") and not part.startswith("__")
+
+
+def _is_private_module(module: str) -> bool:
+    """Return True if any dotted segment of ``module`` is private."""
+    return any(_is_private_part(p) for p in module.split("."))
+
+
+def _is_private_file(file_path: Path, package_root: Path) -> bool:
+    """Return True if any path segment below ``package_root`` is private."""
+    try:
+        rel = file_path.resolve().relative_to(package_root.resolve())
+    except ValueError:
+        return False
+    parts = list(rel.parts)
+    parts[-1] = Path(parts[-1]).stem
+    return any(_is_private_part(p) for p in parts)
+
+
+def _line_starts(source: str) -> list[int]:
+    """Cumulative UTF-8 byte offset at the start of each (1-indexed) line.
+
+    ``ast`` reports ``col_offset``/``end_col_offset`` as byte offsets into the source,
+    so offsets must be computed and applied in bytes to stay correct on lines that
+    contain multibyte characters.
+    """
+    starts = [0]
+    for line in source.splitlines(keepends=True):
+        starts.append(starts[-1] + len(line.encode("utf-8")))
+    return starts
+
+
+class _Positioned(Protocol):
+    """An AST node carrying source position attributes (any ``ast.expr`` or ``ast.stmt``)."""
+
+    lineno: int
+    col_offset: int
+    end_lineno: int | None
+    end_col_offset: int | None
+
+
+def _node_span(node: _Positioned, line_starts: list[int]) -> tuple[int, int]:
+    """Convert an AST node's (lineno, col_offset, end_lineno, end_col_offset) to byte offsets."""
+    start = line_starts[node.lineno - 1] + node.col_offset
+    end_line = node.end_lineno or node.lineno
+    end_col = node.end_col_offset if node.end_col_offset is not None else node.col_offset
+    end = line_starts[end_line - 1] + end_col
+    return start, end
+
+
+@dataclass(frozen=True)
+class _Edit:
+    """A byte-offset span replacement in the source text."""
+
+    start: int
+    end: int
+    replacement: str
+
+
+@dataclass
+class Violation:
+    """A single import-aliasing rule violation."""
+
+    line: int
+    message: str
+    edits: list[_Edit]
+    fixable: bool
+    skip_reason: str = ""  # why the auto-fix was withheld (empty when fixable)
+
+
+def _is_named_target(target: ast.expr, name: str) -> bool:
+    """Return True if assignment target ``target`` binds ``name`` (handles unpacking)."""
+    if isinstance(target, ast.Name):
+        return target.id == name
+    if isinstance(target, (ast.Tuple, ast.List)):
+        return any(_is_named_target(t, name) for t in target.elts)
+    if isinstance(target, ast.Starred):
+        return _is_named_target(target.value, name)
+    return False
+
+
+def _node_binds_name(node: ast.AST, name: str) -> bool:
+    """Return True if ``node`` binds ``name`` (assignment/loop/with/walrus/comprehension).
+
+    The caller's traversal decides which node kinds reach here: the module-scope pass visits
+    statements only (never walrus/comprehension); the function-scope pass visits every node.
+    """
+    if isinstance(node, ast.Assign):
+        # `name = ...`, chained `a = name = ...`, or unpacking `a, name = ...`
+        return any(_is_named_target(t, name) for t in node.targets)
+    if isinstance(node, (ast.AnnAssign, ast.AugAssign, ast.NamedExpr, ast.For, ast.AsyncFor)):
+        # `name: int = ...`, `name += ...`, walrus `(name := ...)`, or `for name in ...`
+        return _is_named_target(node.target, name)
+    if isinstance(node, (ast.With, ast.AsyncWith)):
+        # `with ... as name:` (optional_vars is None for a bare `with ...:`)
+        return any(
+            item.optional_vars is not None and _is_named_target(item.optional_vars, name)
+            for item in node.items
+        )
+    if isinstance(node, ast.comprehension):
+        # the `for name in ...` clause of a list/set/dict/generator comprehension
+        return _is_named_target(node.target, name)
+    return False
+
+
+def _module_binding_stmt(node: ast.AST, name: str, import_line: int) -> bool:
+    """Return True if statement ``node`` binds ``name`` directly in its enclosing scope."""
+    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+        return node.name == name
+    if isinstance(node, (ast.Import, ast.ImportFrom)):
+        return node.lineno != import_line and any(
+            (alias.asname or alias.name) == name for alias in node.names
+        )
+    if isinstance(node, ast.Try):
+        return any(handler.name == name for handler in node.handlers)
+    return _node_binds_name(node, name)
+
+
+def _module_scope_blocks(node: ast.AST) -> list[list[ast.stmt]]:
+    """Return ``node``'s sub-statement blocks that stay in the enclosing scope.
+
+    Control-flow constructs are transparent; ``def``/``class`` bodies are never
+    returned, so bindings inside nested scopes are not treated as module-level.
+    """
+    if isinstance(node, ast.If):
+        return [node.body, node.orelse]
+    if isinstance(node, (ast.For, ast.AsyncFor, ast.While)):
+        return [node.body, node.orelse]
+    if isinstance(node, (ast.With, ast.AsyncWith)):
+        return [node.body]
+    if isinstance(node, ast.Try):
+        return [node.body, node.orelse, node.finalbody, *(h.body for h in node.handlers)]
+    return []
+
+
+def _has_module_level_binding(tree: ast.Module, name: str, import_line: int) -> bool:
+    """Return True if ``name`` is bound at module scope outside the target import.
+
+    Recurses through control-flow constructs (``if``/``for``/``try``/``with``)
+    but never into ``def``/``class``/``lambda`` bodies — those introduce inner
+    scopes whose bindings don't shadow the module-level import for references in
+    unrelated scopes.
+    """
+
+    def _visit(stmts: list[ast.stmt]) -> bool:
+        for node in stmts:
+            if _module_binding_stmt(node, name, import_line):
+                return True
+            if any(_visit(block) for block in _module_scope_blocks(node)):
+                return True
+        return False
+
+    return _visit(list(tree.body))
+
+
+def _iter_body_nodes(
+    func: ast.FunctionDef | ast.AsyncFunctionDef | ast.Lambda,
+) -> Iterator[ast.AST]:
+    """Yield every descendant node in ``func``'s body (a ``Lambda`` body is a single expr)."""
+    if isinstance(func, ast.Lambda):
+        yield from ast.walk(func.body)
+    else:
+        for stmt in func.body:
+            yield from ast.walk(stmt)
+
+
+def _function_locally_binds(
+    func: ast.FunctionDef | ast.AsyncFunctionDef | ast.Lambda,
+    name: str,
+) -> bool:
+    """Return True if ``func`` binds ``name`` as a parameter or local assignment."""
+    args = func.args
+    params = (*args.posonlyargs, *args.args, *args.kwonlyargs, args.vararg, args.kwarg)
+    if any(param is not None and param.arg == name for param in params):
+        return True
+    return any(_node_binds_name(node, name) for node in _iter_body_nodes(func))
+
+
+def _has_shadowing_function(tree: ast.Module, name: str) -> bool:
+    """Return True if any function locally binds ``name`` *and* references it.
+
+    A function that takes ``helper`` as a parameter, then uses ``helper`` in
+    its body, refers to the parameter — not the module-level import. Renaming
+    those references would be wrong, so we refuse the auto-fix.
+    """
+    for func in ast.walk(tree):
+        if not isinstance(func, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
+            continue
+        if not _function_locally_binds(func, name):
+            continue
+        if any(isinstance(node, ast.Name) and node.id == name for node in _iter_body_nodes(func)):
+            return True
+    return False
+
+
+def _is_safe_to_fix(tree: ast.Module, name: str, import_line: int) -> bool:
+    """Return True if it's safe to rename references to ``name`` across the file."""
+    return not _has_module_level_binding(tree, name, import_line) and not _has_shadowing_function(
+        tree,
+        name,
+    )
+
+
+def _name_reference_edits(
+    tree: ast.Module,
+    old_name: str,
+    new_name: str,
+    import_line: int,
+    line_starts: list[int],
+) -> list[_Edit]:
+    """Return edits that rename every ``Name(id=old_name)`` to ``new_name``.
+
+    Skips the import statement itself (positions inside it are not ``Name`` nodes
+    anyway, but the guard makes intent explicit).
+    """
+    edits: list[_Edit] = []
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Name):
+            continue
+        if node.id != old_name:
+            continue
+        if node.lineno == import_line:
+            continue
+        start, end = _node_span(node, line_starts)
+        edits.append(_Edit(start=start, end=end, replacement=new_name))
+    return edits
+
+
+def _alias_text(name: str, asname: str | None) -> str:
+    """Format an alias as it would appear in source (``name`` or ``name as asname``)."""
+    if asname is None or asname == name:
+        return name
+    return f"{name} as {asname}"
+
+
+def _format_importfrom(
+    module: str,
+    aliases: list[tuple[str, str | None]],
+    indent_col: int,
+) -> str:
+    """Format ``from module import ...`` choosing single-line or paren-wrapped form.
+
+    The first line is assumed to start at column ``indent_col`` (the position of
+    the ``from`` keyword in the original source). Continuation lines include
+    their own leading whitespace so the rewrite preserves the original indent.
+    """
+    parts = [_alias_text(name, asname) for name, asname in aliases]
+    single_line = f"from {module} import {', '.join(parts)}"
+    if indent_col + len(single_line) <= LINE_LIMIT:
+        return single_line
+    body_indent = " " * (indent_col + 4)
+    close_indent = " " * indent_col
+    lines = [f"from {module} import ("]
+    lines.extend(f"{body_indent}{part}," for part in parts)
+    lines.append(f"{close_indent})")
+    return "\n".join(lines)
+
+
+def _format_import(aliases: list[tuple[str, str | None]]) -> str:
+    """Format an ``import a, b as c`` statement from (name, asname) pairs.
+
+    Unlike ``from`` imports, plain ``import`` statements cannot be parenthesized, so this
+    always emits a single line (aliased private-module imports are short in practice).
+    """
+    return "import " + ", ".join(_alias_text(name, asname) for name, asname in aliases)
+
+
+class _AliasChange(NamedTuple):
+    """The rewrite for one imported name whose ``_`` aliasing must change."""
+
+    rename: tuple[str, str]  # (old_bound_name, new_bound_name)
+    message: str
+
+
+def _resolve_alias(
+    alias: ast.alias,
+    module: str,
+    *,
+    file_is_private: bool,
+) -> tuple[tuple[str, str | None], _AliasChange | None]:
+    """Decide the target aliasing for one imported name.
+
+    Returns ``(new_alias, change)`` where ``new_alias`` is the ``(name, asname)``
+    to emit and ``change`` is ``None`` when the import already complies.
+
+    A public module must ``_``-prefix the *bound* name (a custom ``as`` alias included),
+    since any un-prefixed name re-exports the private symbol. A private module only strips
+    the redundant canonical ``_name`` alias; a custom alias there is left as the author chose.
+    """
+    original = alias.name
+    current_alias = alias.asname
+    if original.startswith("_"):
+        # The imported symbol is itself private-named; leave it untouched.
+        return (original, current_alias), None
+
+    bound_name = current_alias or original
+    if file_is_private:
+        # Only the canonical ``name as _name`` alias is redundant noise worth stripping;
+        # renaming a custom alias would change a name the author deliberately picked.
+        if current_alias != "_" + original:
+            return (original, current_alias), None
+        new_asname: str | None = None
+        target_bound = original
+        message = (
+            f"Unnecessary `_` alias for `{original}` imported from "
+            f"private module `{module}` into a private module."
+        )
+    else:
+        # Public module: the bound name must be `_`-prefixed so it cannot be re-imported.
+        if bound_name.startswith("_"):
+            return (original, current_alias), None
+        target_bound = "_" + bound_name
+        new_asname = target_bound
+        message = (
+            f"Missing `_` alias for `{original}` imported from "
+            f"private module `{module}` into a public module. "
+            f"Use `{original} as {target_bound}`."
+        )
+    return (
+        (original, new_asname),
+        _AliasChange(rename=(bound_name, target_bound), message=message),
+    )
+
+
+def _resolve_plain_alias(
+    alias: ast.alias,
+    project_package: str,
+    *,
+    file_is_private: bool,
+) -> _AliasChange | None:
+    """Decide the ``_``-alias change for one ``import module as name`` clause.
+
+    Only an aliased import of an in-project private module into a public module is flagged
+    (``import pkg._priv.mod as name`` -> ``as _name``): a bare ``import pkg._priv.mod`` binds
+    the public top-level package rather than the private module, so it is left alone.
+    """
+    module = alias.name
+    if alias.asname is None or alias.asname.startswith("_") or file_is_private:
+        return None
+    if module != project_package and not module.startswith(project_package + "."):
+        return None
+    if not _is_private_module(module):
+        return None
+    target = "_" + alias.asname
+    message = (
+        f"Missing `_` alias for private module `{module}` bound as `{alias.asname}` "
+        f"in a public module. Use `import {module} as {target}`."
+    )
+    return _AliasChange(rename=(alias.asname, target), message=message)
+
+
+def _statement_edits(
+    node: ast.ImportFrom,
+    module: str,
+    new_aliases: list[tuple[str, str | None]],
+    renames: list[tuple[str, str]],
+    tree: ast.Module,
+    line_starts: list[int],
+) -> list[_Edit]:
+    """Build the import rewrite plus every bound-name rename edit for one statement."""
+    import_start, import_end = _node_span(node, line_starts)
+    replacement = _format_importfrom(module, new_aliases, node.col_offset)
+    edits = [_Edit(start=import_start, end=import_end, replacement=replacement)]
+    for old, new in renames:
+        edits.extend(_name_reference_edits(tree, old, new, node.lineno, line_starts))
+    return edits
+
+
+def _span_has_comment(source_bytes: bytes, start: int, end: int) -> bool:
+    """Return True if the byte span holds a ``#`` — a comment the rewrite would drop.
+
+    Import statements contain no string literals, so any ``#`` within the replaced span
+    is a comment (typically inside a paren-wrapped, multi-line import). A trailing comment
+    on a single-line import falls *after* the span and is preserved regardless.
+    """
+    return b"#" in source_bytes[start:end]
+
+
+def _alias_span_edits(
+    node: ast.ImportFrom,
+    new_aliases: list[tuple[str, str | None]],
+    renames: list[tuple[str, str]],
+    tree: ast.Module,
+    line_starts: list[int],
+) -> list[_Edit]:
+    """Rewrite only the changed ``alias`` nodes in place, preserving surrounding comments.
+
+    Used when the statement span holds a comment (always a paren-wrapped, multi-line
+    import): replacing the whole statement would drop the comment, so each imported name
+    is edited where it sits, leaving the commas, comments, and layout untouched.
+    """
+    edits: list[_Edit] = []
+    for alias, (name, new_asname) in zip(node.names, new_aliases, strict=True):
+        if (name, new_asname) == (alias.name, alias.asname):
+            continue  # this import name already complies
+        start, end = _node_span(alias, line_starts)
+        edits.append(_Edit(start=start, end=end, replacement=_alias_text(name, new_asname)))
+    for old, new in renames:
+        edits.extend(_name_reference_edits(tree, old, new, node.lineno, line_starts))
+    return edits
+
+
+def _build_violation(
+    node: ast.stmt,
+    messages: list[str],
+    edits: list[_Edit],
+    *,
+    all_safe: bool,
+) -> Violation:
+    """Assemble the Violation for one import statement (``edits`` already gated on fixability)."""
+    return Violation(
+        line=node.lineno,
+        message="; ".join(messages),
+        edits=edits,
+        fixable=all_safe,
+        skip_reason="" if all_safe else "bound name is shadowed",
+    )
+
+
+def _check_import_from(
+    node: ast.ImportFrom,
+    project_package: str,
+    tree: ast.Module,
+    line_starts: list[int],
+    source_bytes: bytes,
+    *,
+    file_is_private: bool,
+    fix: bool,
+) -> Violation | None:
+    """Flag a ``from module import ...`` (absolute or relative) that mis-aliases private imports."""
+    # Preserve the written form: relative imports keep their leading dots. They are
+    # always in-package, so they bypass the absolute ``project_package`` prefix check.
+    module = "." * node.level + (node.module or "")
+    is_absolute = node.level == 0
+    if is_absolute and module != project_package and not module.startswith(project_package + "."):
+        return None
+    if not _is_private_module(module):
+        return None
+
+    new_aliases: list[tuple[str, str | None]] = []
+    messages: list[str] = []
+    renames: list[tuple[str, str]] = []  # (old_bound, new_bound) pairs
+    all_safe = True
+    for alias in node.names:
+        new_alias, change = _resolve_alias(alias, module, file_is_private=file_is_private)
+        new_aliases.append(new_alias)
+        if change is None:
+            continue
+        renames.append(change.rename)
+        messages.append(change.message)
+        if not _is_safe_to_fix(tree, change.rename[0], node.lineno):
+            all_safe = False
+    if not renames:
+        return None
+
+    # Only build edits under --fix (report-only mode discards them). The whole-statement
+    # rewrite can re-wrap a long line but drops any comment inside the span; when the
+    # statement carries a comment, edit each name in place instead.
+    edits: list[_Edit] = []
+    if fix and all_safe:
+        import_start, import_end = _node_span(node, line_starts)
+        if _span_has_comment(source_bytes, import_start, import_end):
+            edits = _alias_span_edits(node, new_aliases, renames, tree, line_starts)
+        else:
+            edits = _statement_edits(node, module, new_aliases, renames, tree, line_starts)
+    return _build_violation(node, messages, edits, all_safe=all_safe)
+
+
+def _check_plain_import(
+    node: ast.Import,
+    project_package: str,
+    tree: ast.Module,
+    line_starts: list[int],
+    *,
+    file_is_private: bool,
+    fix: bool,
+) -> Violation | None:
+    """Flag ``import pkg._priv.mod as name`` that binds a private module to a public name."""
+    new_aliases: list[tuple[str, str | None]] = []
+    messages: list[str] = []
+    renames: list[tuple[str, str]] = []  # (old_bound, new_bound) pairs
+    all_safe = True
+    for alias in node.names:
+        change = _resolve_plain_alias(alias, project_package, file_is_private=file_is_private)
+        if change is None:
+            new_aliases.append((alias.name, alias.asname))
+            continue
+        new_aliases.append((alias.name, change.rename[1]))
+        renames.append(change.rename)
+        messages.append(change.message)
+        if not _is_safe_to_fix(tree, change.rename[0], node.lineno):
+            all_safe = False
+    if not renames:
+        return None
+
+    # Only build edits under --fix. A plain ``import`` cannot be parenthesized, so it never
+    # holds an inner comment; a trailing comment sits past the span and survives the rewrite.
+    edits: list[_Edit] = []
+    if fix and all_safe:
+        import_start, import_end = _node_span(node, line_starts)
+        edits.append(
+            _Edit(start=import_start, end=import_end, replacement=_format_import(new_aliases)),
+        )
+        for old, new in renames:
+            edits.extend(_name_reference_edits(tree, old, new, node.lineno, line_starts))
+    return _build_violation(node, messages, edits, all_safe=all_safe)
+
+
+def _check_module(
+    tree: ast.Module,
+    project_package: str,
+    source: str,
+    *,
+    file_is_private: bool,
+    fix: bool,
+) -> list[Violation]:
+    """Build one ``Violation`` per import statement that needs fixing.
+
+    Covers ``from`` imports (absolute and relative) and aliased plain ``import``
+    statements; grouping at the statement level lets us emit one clean replacement each.
+    Rewrite edits are built only when ``fix`` is set — report-only runs skip that work.
+    """
+    line_starts = _line_starts(source)
+    source_bytes = source.encode("utf-8")
+    violations: list[Violation] = []
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ImportFrom):
+            violation = _check_import_from(
+                node,
+                project_package,
+                tree,
+                line_starts,
+                source_bytes,
+                file_is_private=file_is_private,
+                fix=fix,
+            )
+        elif isinstance(node, ast.Import):
+            violation = _check_plain_import(
+                node,
+                project_package,
+                tree,
+                line_starts,
+                file_is_private=file_is_private,
+                fix=fix,
+            )
+        else:
+            continue
+        if violation is not None:
+            violations.append(violation)
+    return violations
+
+
+def _apply_edits(source: str, edits: list[_Edit]) -> str:
+    """Apply edits in reverse byte-offset order so earlier positions stay valid.
+
+    Edits carry UTF-8 byte spans (from ``ast`` positions), so splice on the encoded
+    bytes rather than the ``str`` to stay aligned with multibyte characters.
+    """
+    data = source.encode("utf-8")
+    for edit in sorted(edits, key=lambda e: e.start, reverse=True):
+        data = data[: edit.start] + edit.replacement.encode("utf-8") + data[edit.end :]
+    return data.decode("utf-8")
+
+
+def check_file(
+    file_path: Path,
+    package_root: Path,
+    *,
+    fix: bool,
+) -> tuple[list[Violation], list[Violation]]:
+    """Return (applied, remaining) violations for one file."""
+    try:
+        source = file_path.read_text()
+        tree = ast.parse(source, filename=str(file_path))
+    except (SyntaxError, OSError):
+        return [], []
+
+    file_is_private = _is_private_file(file_path, package_root)
+    violations = _check_module(
+        tree,
+        package_root.name,
+        source,
+        file_is_private=file_is_private,
+        fix=fix,
+    )
+    if not violations or not fix:
+        return [], violations
+
+    fixable = [v for v in violations if v.fixable]
+    unfixable = [v for v in violations if not v.fixable]
+    if not fixable:
+        return [], unfixable
+
+    all_edits = [e for v in fixable for e in v.edits]
+    new_source = _apply_edits(source, all_edits)
+    if new_source != source:
+        file_path.write_text(new_source)
+    return fixable, unfixable
+
+
+def _iter_target_files(paths: list[Path], package_root: Path) -> list[Path]:
+    """Yield Python files inside ``package_root``.
+
+    If ``paths`` is empty, walk ``package_root``. Otherwise, only keep paths
+    that resolve under it (so external files passed by pre-commit are ignored).
+    """
+    if not paths:
+        return sorted(package_root.rglob("*.py"))
+    inside: list[Path] = []
+    for p in paths:
+        if p.suffix != ".py":
+            continue
+        try:
+            p.resolve().relative_to(package_root.resolve())
+        except ValueError:
+            continue
+        inside.append(p)
+    return sorted(set(inside))
+
+
+def main(argv: list[str]) -> int:
+    """Report internal-import aliasing issues; with ``--fix`` apply the suggested fixes.
+
+    Detection (the default, run by the pre-commit hook) never edits files and fails the commit
+    when it finds a violation, so the offender is forced to resolve it — by hand or by running
+    ``--fix`` manually. Keeping the fix opt-in stops its rewrites from getting entangled with
+    the formatting/linting hooks' changes.
+    """
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--fix",
+        action="store_true",
+        help="Apply the suggested fixes in place (default: report violations and fail).",
+    )
+    parser.add_argument(
+        "--root",
+        default="src/coreai_opt",
+        help="Package root used to decide file privacy and bound import paths.",
+    )
+    parser.add_argument("paths", nargs="*", type=Path)
+    args = parser.parse_args(argv)
+
+    package_root = Path(args.root)
+    if not package_root.exists():
+        sys.stderr.write(f"Package root not found: {package_root}\n")
+        return 1
+
+    reported = 0
+    for file_path in _iter_target_files(args.paths, package_root):
+        applied, remaining = check_file(file_path, package_root, fix=args.fix)
+        for v in applied:
+            sys.stdout.write(f"{file_path}:{v.line}: FIXED — {v.message}\n")
+        for v in remaining:
+            hint = "" if v.fixable else f" (needs a manual fix: {v.skip_reason})"
+            sys.stdout.write(f"{file_path}:{v.line}: {v.message}{hint}\n")
+            reported += 1
+
+    if reported:
+        if args.fix:
+            sys.stdout.write(
+                f"\n{reported} issue(s) could not be auto-fixed; resolve them by hand.\n"
+            )
+        else:
+            sys.stdout.write(
+                f"\n{reported} issue(s) found. Fix them before committing — run "
+                f"`python {sys.argv[0]} --fix` to apply the fixable ones automatically.\n"
+            )
+    # Fail the commit whenever anything remains unresolved so it can't slip through.
+    return 1 if reported else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
diff --git a/src/coreai_opt/_utils/insertion/torch_function/modes.py b/src/coreai_opt/_utils/insertion/torch_function/modes.py
index a9ec7a5..a050521 100644
--- a/src/coreai_opt/_utils/insertion/torch_function/modes.py
+++ b/src/coreai_opt/_utils/insertion/torch_function/modes.py
@@ -20,7 +20,7 @@
 from torch.overrides import TorchFunctionMode, _get_current_function_mode
 from torch.utils._pytree import tree_map
 
-from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
+from coreai_opt._utils.spec_utils import PartialConstructor
 from coreai_opt._utils.torch_utils import NamedModule
 from coreai_opt.config import CompressionConfig
 from coreai_opt.config.spec import CompressionSimulatorBase
@@ -65,6 +65,7 @@ class StateParametrizationInfo(NamedTuple):
     Records information for parametrizing a state including module to parametrize, state
     name to parametrize, and the optimizer to use.
     """
+
     module: torch.nn.Module
     state_name: str
     optimizer: torch.nn.Module
@@ -94,17 +95,13 @@ def __init__(
         for name, module in model.named_modules(remove_duplicate=True):
             # If there is no allow list or if the module is in the allow list,
             # register enter/exit hooks
-            if (hooks_allow_list is None or module in hooks_allow_list):
+            if hooks_allow_list is None or module in hooks_allow_list:
                 pre_hook = module.register_forward_pre_hook(self.enter_module(name))
-                hook = module.register_forward_hook(
-                    self.exit_module(name), always_call=True
-                )
+                hook = module.register_forward_hook(self.exit_module(name), always_call=True)
 
             # If the module is not in the allow list, register fallback enter/exit hooks
             else:
-                pre_hook = module.register_forward_pre_hook(
-                    self.fallback_enter_module(name)
-                )
+                pre_hook = module.register_forward_pre_hook(self.fallback_enter_module(name))
                 hook = module.register_forward_hook(
                     self.fallback_exit_module(name), always_call=True
                 )
@@ -122,9 +119,7 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Any:
         # entering/exiting logic should be avoided when calling a module from
         # __torch_function__ itself.
         if not isinstance(_get_current_function_mode(), self.__class__):
-            raise RuntimeError(
-                f"{self.__class__.__name__} is not at the top of the stack."
-            )
+            raise RuntimeError(f"{self.__class__.__name__} is not at the top of the stack.")
         return super().__exit__(exc_type, exc_val, exc_tb)
 
     @abc.abstractmethod
@@ -278,13 +273,12 @@ def f(module: nn.Module, inputs: Any) -> None:
             self.module_boundary_tracker.record_module_boundary_tensors(
                 current_named_module, inputs, "input"
             )
+
         return f
 
     def exit_module(self, name: str) -> Callable:
         def f(module: nn.Module, inputs: Any, outputs: Any) -> None:
-            assert (
-                self.parents[-1].name == name
-            ), f"{self.parents[-1].name} is not {name}."
+            assert self.parents[-1].name == name, f"{self.parents[-1].name} is not {name}."
             if self.current_module not in self.traversed_modules:
                 # No point in keeping track of module boundary tensors if it is an already traversed
                 # module, since repeated traversals will also be skipped during torch function
@@ -317,9 +311,10 @@ def register_all_activations(self) -> None:
         #
         # The inner generator flattens this into (func_name, record) pairs so each
         # individual FunctionPreregistrationRecord is processed independently.
-        for module_name, module_funcs in (
-            self.preregistration_tracker.get_all_pending_registrations().items()
-        ):
+        for (
+            module_name,
+            module_funcs,
+        ) in self.preregistration_tracker.get_all_pending_registrations().items():
             for f_name, pending_record in (
                 (func_name, record)
                 for func_name, pending_records in module_funcs.items()
@@ -327,16 +322,12 @@ def register_all_activations(self) -> None:
             ):
                 # Resolve pending input registrations
                 resolved_inputs = self._resolve_pending_registrations(
-                    pending_record.function,
-                    pending_record.pending_inputs,
-                    boundary_type="input"
+                    pending_record.function, pending_record.pending_inputs, boundary_type="input"
                 )
 
                 # Resolve pending output registrations
                 resolved_outputs = self._resolve_pending_registrations(
-                    pending_record.function,
-                    pending_record.pending_outputs,
-                    boundary_type="output"
+                    pending_record.function, pending_record.pending_outputs, boundary_type="output"
                 )
 
                 self._register_resolved_optimizations(
@@ -348,7 +339,7 @@ def _register_resolved_optimizations(
         input_registrations: list[PendingOptimizerRegistration],
         output_registrations: list[PendingOptimizerRegistration],
         module_name: str,
-        func_name: str
+        func_name: str,
     ):
         """
         Given resolved input and output registrations for a function within a module, register
@@ -414,9 +405,7 @@ def _resolve_pending_registrations(
 
         for pending in pending_registrations:
             pending = self._get_module_component_override_if_applicable(
-                func,
-                pending,
-                boundary_type
+                func, pending, boundary_type
             )
             resolved_registrations.append(pending)
         return resolved_registrations
@@ -425,7 +414,7 @@ def _get_module_component_override_if_applicable(
         self,
         func: Callable,
         pending: PendingOptimizerRegistration,
-        boundary_type: Literal["input", "output"]
+        boundary_type: Literal["input", "output"],
     ) -> PendingOptimizerRegistration:
         """
         For a specific pending optimizer registration, check if the optimizer should be overridden
@@ -449,9 +438,7 @@ def _get_module_component_override_if_applicable(
         # rules, inner module boundary specs will always have higher priority than outer module
         # boundary specs.
         for module_boundaries in module_boundaries_by_module:
-            module_components = self.module_components_dict.get(
-                module_boundaries[0].named_module
-            )
+            module_components = self.module_components_dict.get(module_boundaries[0].named_module)
             if module_components is None:
                 continue
 
@@ -461,8 +448,7 @@ def _get_module_component_override_if_applicable(
                 components_dict = module_components.module_output_components
 
             self._warn_on_multiple_specifications_for_same_tensor(
-                module_boundaries,
-                components_dict
+                module_boundaries, components_dict
             )
 
             # Check for module-level spec override for the corresponding input/output
@@ -470,7 +456,7 @@ def _get_module_component_override_if_applicable(
             override_optimizer, found = get_optimizer_from_components_dict(
                 func,
                 [module_boundary.index for module_boundary in module_boundaries],
-                components_dict
+                components_dict,
             )
 
             if found:
@@ -480,7 +466,7 @@ def _get_module_component_override_if_applicable(
     @staticmethod
     def _warn_on_multiple_specifications_for_same_tensor(
         module_boundaries: list[ModuleBoundaryInfo],
-        components_dict: Mapping[int | str, CompressionSimulatorBase | None]
+        components_dict: Mapping[int | str, CompressionSimulatorBase | None],
     ):
         """
         Given a list of module boundaries for a tensor, check whether multiple indices match
@@ -507,10 +493,9 @@ def _warn_on_multiple_specifications_for_same_tensor(
                     logger.warning(warning_msg)
                     break
 
-
     @staticmethod
     def _split_module_boundary_info_list_by_module(
-        module_boundaries_for_tensor: list[ModuleBoundaryInfo]
+        module_boundaries_for_tensor: list[ModuleBoundaryInfo],
     ) -> list[list[ModuleBoundaryInfo]]:
         """
         Helper function to split the list of module_boundaries_for_tensor into a list of lists where
@@ -538,15 +523,13 @@ def register_all_states(self) -> None:
         for module in self.model.modules():
             for state_name, state in itertools.chain(
                 module.named_parameters(recurse=False, remove_duplicate=False),
-                module.named_buffers(recurse=False, remove_duplicate=False)
+                module.named_buffers(recurse=False, remove_duplicate=False),
             ):
                 optimizer = self.state_resolver.get_optimizer(state)
                 if optimizer:
-                    states_to_parametrize.append(StateParametrizationInfo(
-                        module,
-                        state_name,
-                        optimizer
-                    ))
+                    states_to_parametrize.append(
+                        StateParametrizationInfo(module, state_name, optimizer)
+                    )
 
         # Register all parametrizations
         for state_to_parametrize in states_to_parametrize:
@@ -554,7 +537,7 @@ def register_all_states(self) -> None:
                 state_to_parametrize.module,
                 state_to_parametrize.state_name,
                 state_to_parametrize.optimizer,
-                unsafe=True
+                unsafe=True,
             )
 
     def _preregister_input_and_state_optimization(
@@ -573,8 +556,7 @@ def _preregister_input_and_state_optimization(
         args, kwargs = normalize_args_kwargs(func, args, kwargs)
         pending_registrations: list[PendingOptimizerRegistration] = []
         for tensor_idx, (kwarg_name, tensor) in enumerate(
-            itertools.chain(((None, arg) for arg in args), kwargs.items()
-            )
+            itertools.chain(((None, arg) for arg in args), kwargs.items())
         ):
             components_dict = (
                 op_compression_components.op_state_components
@@ -613,8 +595,8 @@ def _warn_non_quantizable_tensor_setting(
         func_counter: int,
         tensor: torch.nn.Parameter | torch.Tensor,
         idx: int,
-        components_dict: Mapping[int | str, _PartialConstructor | None],
-        is_output: bool
+        components_dict: Mapping[int | str, PartialConstructor | None],
+        is_output: bool,
     ) -> None:
         """
         Check whether module_components is attempting to set a specific non-quantizable
@@ -679,10 +661,7 @@ def _get_op_compression_components(
                 if re.fullmatch(op_name, func_name):
                     return op_name_config
             except re.error as e:
-                error_msg = (
-                    f"Invalid regex pattern '{op_name}' in op_name_config: "
-                    f"{op_name_config}"
-                )
+                error_msg = f"Invalid regex pattern '{op_name}' in op_name_config: {op_name_config}"
                 raise ValueError(error_msg) from e
 
         # Check for op_type match
@@ -706,7 +685,7 @@ def _create_pending_activation_registration(
         activation_type: Literal["input", "output"],
         tensor: torch.Tensor,
         tensor_idx: int,
-        components_dict: Mapping[int | str, _PartialConstructor | None],
+        components_dict: Mapping[int | str, PartialConstructor | None],
         kwarg_name: str | None = None,
     ) -> PendingOptimizerRegistration:
         """
@@ -733,14 +712,11 @@ def _create_pending_activation_registration(
             optimizer_name = self._build_optimizer_name(
                 func_name,
                 self.optimization_type_name,
-                tensor_idx if kwarg_name is None else kwarg_name
+                tensor_idx if kwarg_name is None else kwarg_name,
             )
         else:
             optimizer_name = self._build_optimizer_name(
-                func_name,
-                self.optimization_type_name,
-                "output",
-                tensor_idx
+                func_name, self.optimization_type_name, "output", tensor_idx
             )
 
         return PendingOptimizerRegistration(
@@ -775,7 +751,7 @@ def _preregister_output_optimization(
                     output,
                     idx,
                     op_compression_components.op_output_components,
-                    True
+                    True,
                 )
                 continue
 
@@ -791,11 +767,7 @@ def _preregister_output_optimization(
         return pending_registrations
 
     def __torch_function__(
-        self,
-        func: Callable,
-        types: list,
-        args: tuple = (),
-        kwargs: dict | None = None
+        self, func: Callable, types: list, args: tuple = (), kwargs: dict | None = None
     ) -> Any:
         if kwargs is None:
             kwargs = {}
@@ -880,8 +852,7 @@ def __init__(
         reference_tracker: RegisteredOptimizersTracker,
     ) -> None:
         hooks_allow_list = [
-            module for name, module in model.named_modules()
-            if reference_tracker.has_module(name)
+            module for name, module in model.named_modules() if reference_tracker.has_module(name)
         ]
 
         super().__init__(
@@ -1024,11 +995,7 @@ def f(module: nn.Module, inputs: Any, outputs: Any) -> None:
         return f
 
     def _apply_activation_optimizer_if_exists(
-        self,
-        func: Callable,
-        func_counter: int,
-        value: Any,
-        *name_parts: str | int
+        self, func: Callable, func_counter: int, value: Any, *name_parts: str | int
     ) -> tuple[torch.Tensor, str | None]:
         """Apply activation optimizer to a single tensor.
 
@@ -1043,9 +1010,7 @@ def _apply_activation_optimizer_if_exists(
             else (value, None)
         """
         optimizer_name = self._build_optimizer_name(
-            get_func_name(func, func_counter),
-            self.optimization_type_name,
-            *name_parts
+            get_func_name(func, func_counter), self.optimization_type_name, *name_parts
         )
         if hasattr(self.current_module, optimizer_name):
             optimizer = self.current_module.get_submodule(optimizer_name)
@@ -1055,17 +1020,15 @@ def _apply_activation_optimizer_if_exists(
     def _input_optimization(
         self, func: Callable, func_counter, args: tuple, kwargs: dict
     ) -> tuple[tuple, dict, list[str]]:
-        """ Handle optimizing input activations of func. """
+        """Handle optimizing input activations of func."""
         inputs_quantized: list[str] = []
         normalized_args, normalized_kwargs = normalize_args_kwargs(func, args, kwargs)
         mutable_args = list(normalized_args)
 
         # Process positional arguments
         for idx in range(len(mutable_args)):
-            mutable_args[idx], optimizer_name = (
-                self._apply_activation_optimizer_if_exists(
-                    func, func_counter, mutable_args[idx], idx
-                )
+            mutable_args[idx], optimizer_name = self._apply_activation_optimizer_if_exists(
+                func, func_counter, mutable_args[idx], idx
             )
             if optimizer_name is not None:
                 inputs_quantized.append(optimizer_name)
@@ -1081,9 +1044,7 @@ def _input_optimization(
                 inputs_quantized.append(optimizer_name)
 
         args_to_return, kwargs_to_return = self._get_args_kwargs_to_return(
-            args,
-            mutable_args,
-            normalized_kwargs
+            args, mutable_args, normalized_kwargs
         )
         return args_to_return, kwargs_to_return, inputs_quantized
 
@@ -1101,8 +1062,7 @@ def _get_args_kwargs_to_return(
         num_args_in_normalized_kwargs = len(args) - len(mutable_args)
         normalized_kwargs_list = list(normalized_kwargs.items())
         return_args = (
-            mutable_args +
-            [v for (_, v) in normalized_kwargs_list][:num_args_in_normalized_kwargs]
+            mutable_args + [v for (_, v) in normalized_kwargs_list][:num_args_in_normalized_kwargs]
         )
         return_kwargs = dict(normalized_kwargs_list[num_args_in_normalized_kwargs:])
         return tuple(return_args), return_kwargs
@@ -1110,7 +1070,7 @@ def _get_args_kwargs_to_return(
     def _output_optimization(
         self, func: Callable, func_counter: int, outputs: torch.Tensor | tuple
     ) -> tuple[torch.Tensor | tuple, list[str]]:
-        """ Handle optimizing output activations of func. """
+        """Handle optimizing output activations of func."""
         # Handle both single tensor and tuple inputs
         outputs_quantized: list[str] = []
         was_single_tensor = not isinstance(outputs, tuple)
@@ -1121,10 +1081,8 @@ def _output_optimization(
 
         # Process all outputs
         for idx in range(len(mutable_outputs)):
-            mutable_outputs[idx], optimizer_name = (
-                self._apply_activation_optimizer_if_exists(
-                    func, func_counter, mutable_outputs[idx], "output", idx
-                )
+            mutable_outputs[idx], optimizer_name = self._apply_activation_optimizer_if_exists(
+                func, func_counter, mutable_outputs[idx], "output", idx
             )
             if optimizer_name is not None:
                 outputs_quantized.append(optimizer_name)
@@ -1132,26 +1090,17 @@ def _output_optimization(
         outputs = tuple(mutable_outputs)
         # Return in the same format as input
         return (
-            (outputs[0], outputs_quantized)
-            if was_single_tensor
-            else (outputs, outputs_quantized)
+            (outputs[0], outputs_quantized) if was_single_tensor else (outputs, outputs_quantized)
         )
 
     def __torch_function__(
-        self,
-        func: Callable,
-        types: list,
-        args: tuple = (),
-        kwargs: dict | None = None
+        self, func: Callable, types: list, args: tuple = (), kwargs: dict | None = None
     ) -> Any:
         if kwargs is None:
             kwargs = {}
 
         func_base_name = get_func_base_name(func)
-        if not self.reference_tracker.has_function(
-            self.current_module_name,
-            func_base_name
-        ):
+        if not self.reference_tracker.has_function(self.current_module_name, func_base_name):
             # Functions like detach() can appear here when not seen earlier, if the
             # model was run with no_grad() for instance. If a function was not seen
             # earlier, simply shortcut and return without quantizing.
@@ -1170,13 +1119,8 @@ def __torch_function__(
             )
             out = func(*args, **kwargs)
             # Apply output optimization
-            out, outputs_quantized = self._output_optimization(
-                func, func_counter, out
-            )
+            out, outputs_quantized = self._output_optimization(func, func_counter, out)
             self.registered_optimizers_tracker.record_registration(
-                self.current_module_name,
-                func_base_name,
-                inputs_quantized,
-                outputs_quantized
+                self.current_module_name, func_base_name, inputs_quantized, outputs_quantized
             )
             return out
diff --git a/src/coreai_opt/_utils/insertion/torch_function/state_spec_resolver.py b/src/coreai_opt/_utils/insertion/torch_function/state_spec_resolver.py
index c145642..1d2a9d1 100644
--- a/src/coreai_opt/_utils/insertion/torch_function/state_spec_resolver.py
+++ b/src/coreai_opt/_utils/insertion/torch_function/state_spec_resolver.py
@@ -13,7 +13,7 @@
 import torch
 import torch.nn as nn
 
-from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
+from coreai_opt._utils.spec_utils import PartialConstructor
 from coreai_opt._utils.torch_utils import NamedModule
 from coreai_opt.config.spec import CompressionSimulatorBase
 
@@ -129,7 +129,7 @@ def resolve(
         func: Callable,
         state_tensor: torch.Tensor,
         current_module: NamedModule,
-        components_dict: Mapping[int | str, _PartialConstructor | None],
+        components_dict: Mapping[int | str, PartialConstructor | None],
     ) -> None:
         """Resolve and cache the optimizer for ``state_tensor`` at the current call site.
 
@@ -197,7 +197,7 @@ def _resolve_op_state(
         self,
         func: Callable,
         state_tensor: torch.Tensor,
-        components_dict: Mapping[int | str, _PartialConstructor | None],
+        components_dict: Mapping[int | str, PartialConstructor | None],
     ) -> CompressionSimulatorBase | None:
         """Look up the ``op_state_spec`` optimizer using all local names for ``state_tensor``."""
         local_state_names = self.get_all_local_names(state_tensor)
diff --git a/src/coreai_opt/_utils/insertion/torch_function/types.py b/src/coreai_opt/_utils/insertion/torch_function/types.py
index bac722c..585bce5 100644
--- a/src/coreai_opt/_utils/insertion/torch_function/types.py
+++ b/src/coreai_opt/_utils/insertion/torch_function/types.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
+from coreai_opt._utils.spec_utils import PartialConstructor
 from coreai_opt.config.spec import CompressionSimulatorBase
 
 
@@ -27,15 +27,15 @@ class OpCompressionComponents:
     - None if no compression is applied to that component
     """
 
-    op_input_components: Mapping[
-        int | str, _PartialConstructor[CompressionSimulatorBase] | None
-    ] = field(default_factory=dict)
+    op_input_components: Mapping[int | str, PartialConstructor[CompressionSimulatorBase] | None] = (
+        field(default_factory=dict)
+    )
 
     op_output_components: Mapping[
-        int | str, _PartialConstructor[CompressionSimulatorBase] | None
+        int | str, PartialConstructor[CompressionSimulatorBase] | None
     ] = field(default_factory=dict)
 
-    op_state_components: Mapping[str, _PartialConstructor[CompressionSimulatorBase] | None] = field(
+    op_state_components: Mapping[str, PartialConstructor[CompressionSimulatorBase] | None] = field(
         default_factory=dict
     )
 
@@ -74,15 +74,15 @@ class ModuleCompressionComponents:
     or name to an OpCompressionComponents class.
     """
 
-    weight: Mapping[str, _PartialConstructor[CompressionSimulatorBase] | None] = field(
+    weight: Mapping[str, PartialConstructor[CompressionSimulatorBase] | None] = field(
         default_factory=dict
     )
 
-    input_activation: Mapping[int | str, _PartialConstructor[CompressionSimulatorBase] | None] = (
+    input_activation: Mapping[int | str, PartialConstructor[CompressionSimulatorBase] | None] = (
         field(default_factory=dict)
     )
 
-    output_activation: Mapping[int | str, _PartialConstructor[CompressionSimulatorBase] | None] = (
+    output_activation: Mapping[int | str, PartialConstructor[CompressionSimulatorBase] | None] = (
         field(default_factory=dict)
     )
 
@@ -91,14 +91,14 @@ class ModuleCompressionComponents:
     op_name_components: Mapping[str, OpCompressionComponents] = field(default_factory=dict)
 
     module_input_components: Mapping[
-        int | str, _PartialConstructor[CompressionSimulatorBase] | None
+        int | str, PartialConstructor[CompressionSimulatorBase] | None
     ] = field(default_factory=dict)
 
     module_output_components: Mapping[
-        int | str, _PartialConstructor[CompressionSimulatorBase] | None
+        int | str, PartialConstructor[CompressionSimulatorBase] | None
     ] = field(default_factory=dict)
 
-    module_state_components: Mapping[str, _PartialConstructor[CompressionSimulatorBase] | None] = (
+    module_state_components: Mapping[str, PartialConstructor[CompressionSimulatorBase] | None] = (
         field(default_factory=dict)
     )
 
diff --git a/src/coreai_opt/_utils/insertion/torch_function/utils.py b/src/coreai_opt/_utils/insertion/torch_function/utils.py
index 3111822..31af66d 100644
--- a/src/coreai_opt/_utils/insertion/torch_function/utils.py
+++ b/src/coreai_opt/_utils/insertion/torch_function/utils.py
@@ -16,7 +16,7 @@
 from torch.fx.operator_schemas import create_type_hint, normalize_function
 
 from coreai_opt._utils.config_utils import get_last_matching_spec
-from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
+from coreai_opt._utils.spec_utils import PartialConstructor
 from coreai_opt.config.spec import CompressionSimulatorBase
 
 logger = logging.getLogger(__name__)
@@ -130,7 +130,7 @@ def any_tensor_optimizable(args: list[Any], kwargs: dict[str, Any]) -> bool:
 def get_optimizer_from_components_dict(
     func: Callable,
     tensor_identifiers: int | str | list[int | str],
-    components_dict: Mapping[int | str, _PartialConstructor | None],
+    components_dict: Mapping[int | str, PartialConstructor | None],
 ) -> tuple[CompressionSimulatorBase | None, bool]:
     """Return the appropriate optimizer from ``components_dict``.
 
diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py
index 695b6af..35a4d7e 100644
--- a/src/coreai_opt/_utils/torch_utils.py
+++ b/src/coreai_opt/_utils/torch_utils.py
@@ -17,7 +17,7 @@
 import torch
 import torch.nn.utils.parametrize as P
 
-from coreai_opt._utils.version_utils import version_ge as _version_ge
+from coreai_opt._utils.version_utils import version_ge
 
 logger = logging.getLogger(__name__)
 
@@ -424,7 +424,7 @@ def export_model(
             exported_program = torch.export.export(
                 model, example_inputs, dynamic_shapes=dynamic_shapes
             )
-            if _version_ge(torch, "2.9"):
+            if version_ge(torch, "2.9"):
                 exported_model = exported_program.module(check_guards=False)
             else:
                 exported_model = exported_program.module()
diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index b25b33c..d3324e2 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -10,10 +10,10 @@
 import torch
 import torch.nn as nn
 
-from coreai_opt._utils.registry_utils import ClassRegistryMixin
+from coreai_opt._utils.registry_utils import ClassRegistryMixin as _ClassRegistryMixin
 
 
-class CompressionSimulatorBase(ClassRegistryMixin, nn.Module):
+class CompressionSimulatorBase(_ClassRegistryMixin, nn.Module):
     """
     Abstract base class for compression simulators.
 
diff --git a/src/coreai_opt/coreai_utils/passes/weight_palettization.py b/src/coreai_opt/coreai_utils/passes/weight_palettization.py
index d696211..498dece 100644
--- a/src/coreai_opt/coreai_utils/passes/weight_palettization.py
+++ b/src/coreai_opt/coreai_utils/passes/weight_palettization.py
@@ -14,17 +14,17 @@
 import numpy as np
 
 from coreai_opt.coreai_utils._coreai_imports import (
-    AIProgram,
-    DenseResourceElementsAttr,
-    F16Type,
-    F32Type,
-    InsertionPoint,
-    IntegerType,
-    RankedTensorType,
-    WalkResult,
+    AIProgram as _AIProgram,
+    DenseResourceElementsAttr as _DenseResourceElementsAttr,
+    F16Type as _F16Type,
+    F32Type as _F32Type,
+    InsertionPoint as _InsertionPoint,
+    IntegerType as _IntegerType,
+    RankedTensorType as _RankedTensorType,
+    WalkResult as _WalkResult,
     _get_constant_value_as_np_array,
-    compression_types,
-    coreai,
+    compression_types as _compression_types,
+    coreai as _coreai,
 )
 from coreai_opt.coreai_utils._utils.graph_utils import (
     _apply_compression_transform,
@@ -61,7 +61,7 @@
 
 
 def palettize_weights(
-    coreai_program: AIProgram,
+    coreai_program: _AIProgram,
     lut_dtype: DType | None,
     n_bits: int = 4,
     granularity: CompressionGranularity = CompressionGranularity.PER_TENSOR,
@@ -73,7 +73,7 @@ def palettize_weights(
     enable_fast_kmeans_mode: bool = True,
     rounding_precision: int = 4,
     in_place: bool = False,
-) -> AIProgram:
+) -> _AIProgram:
     """Palettize weights in a Core AI AIProgram (MLIR<CoreAI> IR) by using Core AI ops.
 
     Walks through the IR and palettizes each coreai.constant op that needs to be
@@ -178,7 +178,7 @@ def replace_weight_with_compression_op(
             fused zero_point = lut_zero_point * per_channel_scale.
         """
         if not _should_compress_op(op, weight_num_threshold, _OPS_WEIGHT_NEED_COMPRESSION):
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
         const_weight: Any = op
         weight = _get_constant_value_as_np_array(const_weight)
@@ -194,7 +194,7 @@ def replace_weight_with_compression_op(
                 "The `cluster_dim` is invalid for %s. Skipped this op.",
                 const_weight.name,
             )
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
         if enable_per_channel_scale:
             # Normalize by per channel scales before doing palettization.
@@ -222,7 +222,7 @@ def replace_weight_with_compression_op(
                     "Cannot perform palettization on %s. Skipped this op.",
                     const_weight.name,
                 )
-                return WalkResult.ADVANCE
+                return _WalkResult.ADVANCE
         except ImportError:
             raise
         except Exception as e:
@@ -232,12 +232,12 @@ def replace_weight_with_compression_op(
                 const_weight.name,
                 e,
             )
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
-        with const_weight.context, const_weight.location, InsertionPoint(const_weight):
+        with const_weight.context, const_weight.location, _InsertionPoint(const_weight):
             indices = _create_constant_value_from_np_array(
                 lut_params.indices,  # same shape as weight tensor (for many cases)
-                IntegerType.get_unsigned(n_bits),
+                _IntegerType.get_unsigned(n_bits),
             )
 
             vector_axis = _create_constant_value_from_np_array(
@@ -246,11 +246,11 @@ def replace_weight_with_compression_op(
                     if lut_params.vector_axis is not None
                     else np.int16(0)
                 ),
-                IntegerType.get_signed(16),
+                _IntegerType.get_signed(16),
             )
 
             weight_float_mlir_type = (
-                F32Type.get() if original_weight_type == np.float32 else F16Type.get()
+                _F32Type.get() if original_weight_type == np.float32 else _F16Type.get()
             )
 
             if lut_dtype is not None:
@@ -259,12 +259,12 @@ def replace_weight_with_compression_op(
                 weight_element_type = cast("Any", const_weight.result.type).element_type
 
                 if lut_dtype.is_int():
-                    lut_dtype_builtin = compression_types.string_to_builtin(lut_dtype)
+                    lut_dtype_builtin = _compression_types.string_to_builtin(lut_dtype)
                     ref_mlir_type = _get_string_to_mlir_type()[lut_dtype]
                     quantized_mlir_type = (
-                        IntegerType.get_signed(ref_mlir_type.width)
+                        _IntegerType.get_signed(ref_mlir_type.width)
                         if ref_mlir_type.is_signed
-                        else IntegerType.get_unsigned(ref_mlir_type.width)
+                        else _IntegerType.get_unsigned(ref_mlir_type.width)
                     )
 
                     quant_params = _compute_qparams_by_dtype(
@@ -278,7 +278,7 @@ def replace_weight_with_compression_op(
                             "Failed to compute quantization parameters for %s. Skipped this op.",
                             const_weight.name,
                         )
-                        return WalkResult.ADVANCE
+                        return _WalkResult.ADVANCE
 
                     quantized_lut_data, lut_scale, lut_zero_point = quant_params
                     if lut_zero_point is None:
@@ -335,7 +335,7 @@ def replace_weight_with_compression_op(
                             "Failed to compute quantization parameters for %s. Skipped this op.",
                             const_weight.name,
                         )
-                        return WalkResult.ADVANCE
+                        return _WalkResult.ADVANCE
 
                     quantized_lut_data, lut_scale, _ = quant_params
 
@@ -348,15 +348,15 @@ def replace_weight_with_compression_op(
                         dtype=lut_scale.dtype,
                     )
 
-                    tensor_type = RankedTensorType.get(
+                    tensor_type = _RankedTensorType.get(
                         list(quantized_lut_data.shape), fp8_mlir_type
                     )
-                    lut_quantized_attr = DenseResourceElementsAttr.get_from_buffer(
+                    lut_quantized_attr = _DenseResourceElementsAttr.get_from_buffer(
                         quantized_lut_data,
                         "dense_resource",
                         tensor_type,
                     )
-                    lut_quantized = cast("Any", coreai.ConstantOp(value=lut_quantized_attr).result)
+                    lut_quantized = cast("Any", _coreai.ConstantOp(value=lut_quantized_attr).result)
 
                     lut_scale_const = _create_constant_value_from_np_array(
                         lut_scale_reshaped,
@@ -368,23 +368,23 @@ def replace_weight_with_compression_op(
                         weight_element_type,
                     )
 
-                compressed_weight_quantized = coreai.lut_to_dense(
+                compressed_weight_quantized = _coreai.lut_to_dense(
                     indices=indices,
                     lut=lut_quantized,
                     axis=vector_axis,
                 )
                 # Cast lut_to_dense output from quantized type to float so all
                 # blockwise_shift_scale operands share the same element type.
-                cast_type = RankedTensorType.get(
+                cast_type = _RankedTensorType.get(
                     cast("Any", compressed_weight_quantized.type).shape,
                     weight_float_mlir_type,
                 )
                 compressed_weight_float = cast(
                     "Any",
-                    coreai.CastOp(cast_type, compressed_weight_quantized).result,
+                    _coreai.CastOp(cast_type, compressed_weight_quantized).result,
                 )
                 lut_t = cast("Any", lut_scale_const.type)
-                compressed_weight = coreai.blockwise_shift_scale(
+                compressed_weight = _coreai.blockwise_shift_scale(
                     data=compressed_weight_float,
                     scale=lut_scale_const,
                     offset1=lut_zero_point_const,
@@ -397,7 +397,7 @@ def replace_weight_with_compression_op(
                 # (e.g. coreai.transpose) see the same type contract.
                 compressed_weight = cast(
                     "Any",
-                    coreai.CastOp(
+                    _coreai.CastOp(
                         cast("Any", const_weight.result.type),
                         compressed_weight,
                     ).result,
@@ -409,7 +409,7 @@ def replace_weight_with_compression_op(
                     cast("Any", const_weight.result.type).element_type,
                 )
 
-                compressed_weight = coreai.lut_to_dense(
+                compressed_weight = _coreai.lut_to_dense(
                     indices=indices,
                     lut=lut,
                     axis=vector_axis,
@@ -426,15 +426,15 @@ def replace_weight_with_compression_op(
                     )
                     compressed_weight_float = cast(
                         "Any",
-                        coreai.CastOp(
-                            RankedTensorType.get(
+                        _coreai.CastOp(
+                            _RankedTensorType.get(
                                 cast("Any", compressed_weight.type).shape,
                                 weight_float_mlir_type,
                             ),
                             compressed_weight,
                         ).result,
                     )
-                    compressed_weight = coreai.blockwise_shift_scale(
+                    compressed_weight = _coreai.blockwise_shift_scale(
                         data=compressed_weight_float,
                         scale=scale,
                         offset1=zero_point,
@@ -445,7 +445,7 @@ def replace_weight_with_compression_op(
                     )
                     compressed_weight = cast(
                         "Any",
-                        coreai.CastOp(
+                        _coreai.CastOp(
                             cast("Any", const_weight.result.type),
                             compressed_weight,
                         ).result,
@@ -453,7 +453,7 @@ def replace_weight_with_compression_op(
 
         const_weight.result.replace_all_uses_with(compressed_weight)
 
-        return WalkResult.ADVANCE
+        return _WalkResult.ADVANCE
 
     return _apply_compression_transform(
         coreai_program,
diff --git a/src/coreai_opt/coreai_utils/passes/weight_quantization.py b/src/coreai_opt/coreai_utils/passes/weight_quantization.py
index a1726e2..6775be5 100644
--- a/src/coreai_opt/coreai_utils/passes/weight_quantization.py
+++ b/src/coreai_opt/coreai_utils/passes/weight_quantization.py
@@ -14,17 +14,17 @@
 import numpy as np
 
 from coreai_opt.coreai_utils._coreai_imports import (
-    AIProgram,
-    DenseElementsAttr,
-    DenseResourceElementsAttr,
-    FloatAttr,
-    InsertionPoint,
-    IntegerType,
-    RankedTensorType,
-    WalkResult,
+    AIProgram as _AIProgram,
+    DenseElementsAttr as _DenseElementsAttr,
+    DenseResourceElementsAttr as _DenseResourceElementsAttr,
+    FloatAttr as _FloatAttr,
+    InsertionPoint as _InsertionPoint,
+    IntegerType as _IntegerType,
+    RankedTensorType as _RankedTensorType,
+    WalkResult as _WalkResult,
     _get_constant_value_as_np_array,
-    compression_types,
-    coreai,
+    compression_types as _compression_types,
+    coreai as _coreai,
 )
 from coreai_opt.coreai_utils._utils.graph_utils import (
     _apply_compression_transform,
@@ -79,7 +79,7 @@ def _create_int_quantized_weight(
     quantized_data_val = _create_constant_value_from_np_array(quantized_data, quantized_mlir_type)
     scale_val = _create_constant_value_from_np_array(scale, weight_element_type)
     zero_point_val = _create_constant_value_from_np_array(zero_point, quantized_mlir_type)
-    return coreai.blockwise_shift_scale(
+    return _coreai.blockwise_shift_scale(
         data=quantized_data_val,
         scale=scale_val,
         offset1=zero_point_val,
@@ -104,22 +104,22 @@ def _create_fp_quantized_weight(
     the scale constant is created in that dtype using DenseResourceElementsAttr.
     Otherwise the scale uses the uncompressed weight element type (default behavior).
     """
-    tensor_type = RankedTensorType.get(list(quantized_data.shape), fp_mlir_type)
-    data_attr = DenseResourceElementsAttr.get_from_buffer(
+    tensor_type = _RankedTensorType.get(list(quantized_data.shape), fp_mlir_type)
+    data_attr = _DenseResourceElementsAttr.get_from_buffer(
         quantized_data,
         "dense_resource",
         tensor_type,
     )
-    quantized_data_val = cast("Any", coreai.ConstantOp(value=data_attr).result)
+    quantized_data_val = cast("Any", _coreai.ConstantOp(value=data_attr).result)
 
     if scale_mlir_type is not None:
         scale_cast = np.ascontiguousarray(scale.astype(scale_np_dtype))
         scale_shape = list(scale_cast.shape)
-        scale_tensor_type = RankedTensorType.get(scale_shape, scale_mlir_type)
-        scale_attr = DenseResourceElementsAttr.get_from_buffer(
+        scale_tensor_type = _RankedTensorType.get(scale_shape, scale_mlir_type)
+        scale_attr = _DenseResourceElementsAttr.get_from_buffer(
             scale_cast, "dense_resource", scale_tensor_type
         )
-        scale_val = cast("Any", coreai.ConstantOp(value=scale_attr).result)
+        scale_val = cast("Any", _coreai.ConstantOp(value=scale_attr).result)
 
         # offset2 must be in the output (weight) dtype so blockwise_shift_scale
         # returns a tensor in the original weight precision, not f8E8M0FNU.
@@ -134,13 +134,13 @@ def _create_fp_quantized_weight(
             weight_element_type,
         )
 
-    zero_point_tensor_type = RankedTensorType.get(list(scale.shape), fp_mlir_type)
-    zero_point_attr = DenseElementsAttr.get_splat(
-        zero_point_tensor_type, FloatAttr.get(fp_mlir_type, 0.0)
+    zero_point_tensor_type = _RankedTensorType.get(list(scale.shape), fp_mlir_type)
+    zero_point_attr = _DenseElementsAttr.get_splat(
+        zero_point_tensor_type, _FloatAttr.get(fp_mlir_type, 0.0)
     )
-    zero_point_val = cast("Any", coreai.ConstantOp(value=zero_point_attr).result)
+    zero_point_val = cast("Any", _coreai.ConstantOp(value=zero_point_attr).result)
 
-    return coreai.blockwise_shift_scale(
+    return _coreai.blockwise_shift_scale(
         data=quantized_data_val,
         scale=scale_val,
         offset1=zero_point_val,
@@ -149,7 +149,7 @@ def _create_fp_quantized_weight(
 
 
 def quantize_weights(
-    coreai_program: AIProgram,
+    coreai_program: _AIProgram,
     dtype: DType,
     qscheme: QScheme = QScheme.SYMMETRIC,
     granularity: CompressionGranularity = CompressionGranularity.PER_CHANNEL,
@@ -157,7 +157,7 @@ def quantize_weights(
     weight_num_threshold: int = 1024,
     scale_dtype: DType | None = None,
     in_place: bool = False,
-) -> AIProgram:
+) -> _AIProgram:
     """Quantize weights in a Core AI AIProgram (MLIR<CoreAI> IR) by using Core AI ops.
 
     Walks through the IR and quantizes each coreai.constant op that needs to be
@@ -280,7 +280,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
             coreai.constant (zero_point) ──────┘
         """
         if not _should_compress_op(op, weight_num_threshold, _OPS_WEIGHT_NEED_COMPRESSION):
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
         const_weight: Any = op
         weight = _get_constant_value_as_np_array(const_weight)
@@ -293,7 +293,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
 
         try:
             if dtype.is_int():
-                dtype_builtin = compression_types.string_to_builtin(dtype)
+                dtype_builtin = _compression_types.string_to_builtin(dtype)
                 quant_params = _compute_qparams_by_dtype(
                     weight,
                     dtype_builtin,
@@ -312,7 +312,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
                     "Failed to quantize op %s. Skipped this op.",
                     const_weight.name,
                 )
-                return WalkResult.ADVANCE
+                return _WalkResult.ADVANCE
         except ImportError:
             raise
         except Exception as e:
@@ -321,19 +321,19 @@ def replace_weight_with_compression_op(op: Any) -> Any:
                 const_weight.name,
                 e,
             )
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
         quantized_data, scale, zero_point = quant_params
 
-        with const_weight.context, const_weight.location, InsertionPoint(const_weight):
+        with const_weight.context, const_weight.location, _InsertionPoint(const_weight):
             weight_element_type = cast("Any", const_weight.result.type).element_type
 
             if dtype.is_int():
                 ref_mlir_type = _get_string_to_mlir_type()[dtype]
                 quantized_mlir_type = (
-                    IntegerType.get_signed(ref_mlir_type.width)
+                    _IntegerType.get_signed(ref_mlir_type.width)
                     if ref_mlir_type.is_signed
-                    else IntegerType.get_unsigned(ref_mlir_type.width)
+                    else _IntegerType.get_unsigned(ref_mlir_type.width)
                 )
                 if zero_point is None:
                     zero_point = np.zeros_like(scale, dtype=quantized_data.dtype)
@@ -368,7 +368,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
 
         const_weight.result.replace_all_uses_with(quantized_weight)
 
-        return WalkResult.ADVANCE
+        return _WalkResult.ADVANCE
 
     return _apply_compression_transform(
         coreai_program,
diff --git a/src/coreai_opt/coreai_utils/passes/weight_sparsification.py b/src/coreai_opt/coreai_utils/passes/weight_sparsification.py
index 81dc888..3a0a8d4 100644
--- a/src/coreai_opt/coreai_utils/passes/weight_sparsification.py
+++ b/src/coreai_opt/coreai_utils/passes/weight_sparsification.py
@@ -13,17 +13,17 @@
 import numpy as np
 
 from coreai_opt.coreai_utils._coreai_imports import (
-    AIProgram,
-    DenseElementsAttr,
-    DenseResourceElementsAttr,
-    FloatAttr,
-    InsertionPoint,
-    IntegerType,
-    RankedTensorType,
-    WalkResult,
+    AIProgram as _AIProgram,
+    DenseElementsAttr as _DenseElementsAttr,
+    DenseResourceElementsAttr as _DenseResourceElementsAttr,
+    FloatAttr as _FloatAttr,
+    InsertionPoint as _InsertionPoint,
+    IntegerType as _IntegerType,
+    RankedTensorType as _RankedTensorType,
+    WalkResult as _WalkResult,
     _get_constant_value_as_np_array,
-    compression_types,
-    coreai,
+    compression_types as _compression_types,
+    coreai as _coreai,
 )
 from coreai_opt.coreai_utils._utils.graph_utils import (
     _apply_compression_transform,
@@ -53,7 +53,7 @@
 
 
 def sparsify_weights(
-    coreai_program: AIProgram,
+    coreai_program: _AIProgram,
     target_sparsity: float | None = 0.5,
     block_size: int | None = None,
     n_m_ratio: tuple[int, int] | None = None,
@@ -61,7 +61,7 @@ def sparsify_weights(
     palettize_nbits: int | None = None,
     weight_num_threshold: int = 1024,
     in_place: bool = False,
-) -> AIProgram:
+) -> _AIProgram:
     """Sparsify weights in a Core AI AIProgram (MLIR<CoreAI> IR) by using Core AI ops.
 
     Walks through the IR and sparsifies each coreai.constant op that needs to be
@@ -172,7 +172,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
              ──> coreai.build_sparse_with_bitmask -> coreai.sparse_with_bitmask_to_dense -> output
         """
         if not _should_compress_op(op, weight_num_threshold, _OPS_WEIGHT_NEED_COMPRESSION):
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
         const_weight: Any = op
         weight = _get_constant_value_as_np_array(const_weight)
@@ -209,7 +209,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
                     "Cannot perform sparsification on %s. Skipped this op.",
                     const_weight.name,
                 )
-                return WalkResult.ADVANCE
+                return _WalkResult.ADVANCE
         except ImportError:
             raise
         except Exception as e:
@@ -219,9 +219,9 @@ def replace_weight_with_compression_op(op: Any) -> Any:
                 const_weight.name,
                 e,
             )
-            return WalkResult.ADVANCE
+            return _WalkResult.ADVANCE
 
-        with const_weight.context, const_weight.location, InsertionPoint(const_weight):
+        with const_weight.context, const_weight.location, _InsertionPoint(const_weight):
             require_dequantization = False
 
             if palettize_nbits is not None:
@@ -254,7 +254,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
                     indices_reshaped = lut_params.indices.reshape(original_shape)
                     indices = _create_constant_value_from_np_array(
                         indices_reshaped,
-                        IntegerType.get_unsigned(palettize_nbits),
+                        _IntegerType.get_unsigned(palettize_nbits),
                     )
 
                     lut_data = lut_params.lut.astype(sparse_params.nonzero_data.dtype)
@@ -265,10 +265,10 @@ def replace_weight_with_compression_op(op: Any) -> Any:
 
                     vector_axis = _create_constant_value_from_np_array(
                         np.int16(0),
-                        IntegerType.get_signed(16),
+                        _IntegerType.get_signed(16),
                     )
 
-                    nonzero_data_const = coreai.lut_to_dense(
+                    nonzero_data_const = _coreai.lut_to_dense(
                         indices=indices,
                         lut=lut,
                         axis=vector_axis,
@@ -290,12 +290,14 @@ def replace_weight_with_compression_op(op: Any) -> Any:
 
                     quant_block_sizes = [0] * len(sparse_params.nonzero_data.shape)
                     if quantize_dtype.is_int():
-                        quantize_dtype_builtin = compression_types.string_to_builtin(quantize_dtype)
+                        quantize_dtype_builtin = _compression_types.string_to_builtin(
+                            quantize_dtype
+                        )
                         ref_mlir_type = _get_string_to_mlir_type()[quantize_dtype]
                         quantized_mlir_type = (
-                            IntegerType.get_signed(ref_mlir_type.width)
+                            _IntegerType.get_signed(ref_mlir_type.width)
                             if ref_mlir_type.is_signed
-                            else IntegerType.get_unsigned(ref_mlir_type.width)
+                            else _IntegerType.get_unsigned(ref_mlir_type.width)
                         )
                         quant_params = _compute_qparams_by_dtype(
                             sparse_params.nonzero_data,
@@ -317,7 +319,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
                             "Failed to compute quantization parameters for %s. Skipped this op.",
                             const_weight.name,
                         )
-                        return WalkResult.ADVANCE
+                        return _WalkResult.ADVANCE
 
                     quantized_data, scale, zero_point = quant_params
                     if zero_point is None:
@@ -343,13 +345,13 @@ def replace_weight_with_compression_op(op: Any) -> Any:
                             quantized_mlir_type,
                         )
                     else:
-                        quantized_tensor_type = RankedTensorType.get(
+                        quantized_tensor_type = _RankedTensorType.get(
                             list(quantized_data.shape), fp8_mlir_type
                         )
                         nonzero_data_const = cast(
                             "Any",
-                            coreai.ConstantOp(
-                                value=DenseResourceElementsAttr.get_from_buffer(
+                            _coreai.ConstantOp(
+                                value=_DenseResourceElementsAttr.get_from_buffer(
                                     quantized_data, "dense_resource", quantized_tensor_type
                                 )
                             ).result,
@@ -360,10 +362,10 @@ def replace_weight_with_compression_op(op: Any) -> Any:
                         )
                         zero_point_const = cast(
                             "Any",
-                            coreai.ConstantOp(
-                                value=DenseElementsAttr.get_splat(
-                                    RankedTensorType.get(target_shape, fp8_mlir_type),
-                                    FloatAttr.get(fp8_mlir_type, 0.0),
+                            _coreai.ConstantOp(
+                                value=_DenseElementsAttr.get_splat(
+                                    _RankedTensorType.get(target_shape, fp8_mlir_type),
+                                    _FloatAttr.get(fp8_mlir_type, 0.0),
                                 )
                             ).result,
                         )
@@ -376,16 +378,16 @@ def replace_weight_with_compression_op(op: Any) -> Any:
 
             mask = _create_constant_value_from_np_array(
                 sparse_params.mask,
-                IntegerType.get_unsigned(1),
+                _IntegerType.get_unsigned(1),
             )
-            sparse_tensor = coreai.build_sparse_with_bitmask(
+            sparse_tensor = _coreai.build_sparse_with_bitmask(
                 values=nonzero_data_const,
                 bitmask=mask,
             )
-            compressed_weight = coreai.sparse_with_bitmask_to_dense(sparse_tensor)
+            compressed_weight = _coreai.sparse_with_bitmask_to_dense(sparse_tensor)
 
             if require_dequantization:
-                compressed_weight = coreai.blockwise_shift_scale(
+                compressed_weight = _coreai.blockwise_shift_scale(
                     data=compressed_weight,
                     scale=scale_const,
                     offset1=zero_point_const,
@@ -397,7 +399,7 @@ def replace_weight_with_compression_op(op: Any) -> Any:
 
         const_weight.result.replace_all_uses_with(compressed_weight)
 
-        return WalkResult.ADVANCE
+        return _WalkResult.ADVANCE
 
     return _apply_compression_transform(
         coreai_program,
diff --git a/src/coreai_opt/inspection/_eager_mode.py b/src/coreai_opt/inspection/_eager_mode.py
index ace497f..006ae37 100644
--- a/src/coreai_opt/inspection/_eager_mode.py
+++ b/src/coreai_opt/inspection/_eager_mode.py
@@ -30,7 +30,7 @@
     get_func_base_name,
     get_func_name,
 )
-from coreai_opt._utils.python_utils import fqn as _fqn
+from coreai_opt._utils.python_utils import fqn
 from coreai_opt._utils.torch_utils import NamedModule, flatten_tensors_to_list
 from coreai_opt.base_model_compressor import _BaseModelCompressor
 from coreai_opt.palettization import KMeansPalettizer
@@ -287,7 +287,7 @@ def _capture_output_tensors(self, module: nn.Module, inputs: Any, outputs: Any)
 
     def _get_module_stack(self) -> tuple[ModuleContext, ...]:
         return tuple(
-            ModuleContext(module_name=named_mod.name, module_type=_fqn(type(named_mod.module)))
+            ModuleContext(module_name=named_mod.name, module_type=fqn(type(named_mod.module)))
             for named_mod in self.parents
         )
 
@@ -574,7 +574,7 @@ def parse_ops_for_eager(
     finally:
         mode.remove_hooks()
 
-    root = build_module_tree(_fqn(type(model)), mode.all_ops)
+    root = build_module_tree(fqn(type(model)), mode.all_ops)
 
     # Pre-partition mode.all_ops by module subtree in a single O(N × D) pass, where N is
     # the total op count and D is the average module stack depth. Each op is appended to
diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py
index e75b6ee..8c0c06e 100644
--- a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py
+++ b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py
@@ -12,15 +12,13 @@
 import torch.nn.utils.parametrize as P
 
 from coreai_opt._utils.export_utils import (
-    clear_parametrization_original as _clear_parametrization_original,
-    prepare_mmap_dir as _prepare_mmap_dir,
+    clear_parametrization_original,
+    prepare_mmap_dir,
     validate_coreml_palettization_compatibility,
 )
 from coreai_opt._utils.import_utils import lazy_import_coreai_torch
 from coreai_opt._utils.metadata_utils import CompressionType, MILCompressionMetadata
-from coreai_opt._utils.torch_utils import (
-    mmap_module_state_dict as _mmap_module_state_dict,
-)
+from coreai_opt._utils.torch_utils import mmap_module_state_dict
 from coreai_opt.common import ExportBackend
 from coreai_opt.palettization.spec.fake_palettize import (
     _FakePalettizeImplBase,
@@ -282,11 +280,11 @@ def _import_coreai_torch_modules():
         # drop the leading dot in that case so we don't write ".weight.safetensors".
         stem = f"{module_name}.{param_name}" if module_name else param_name
         path = Path(mmap_dir) / f"{stem}.safetensors"
-        _mmap_module_state_dict(mlir_palett_mod, path)
+        mmap_module_state_dict(mlir_palett_mod, path)
 
     # Replace _FakePalettizeImplBase
     module.parametrizations[param_name][fake_palett_idx] = mlir_palett_mod
-    _clear_parametrization_original(module, param_name)
+    clear_parametrization_original(module, param_name)
 
 
 def _find_fake_palett_parametrization(
@@ -361,7 +359,7 @@ def _process_weight_palettization(
             it via mmap, so large-model finalization does not hold full
             palettized weights in RAM. Only honored for the CoreAI backend.
     """
-    _prepare_mmap_dir(mmap_dir)
+    prepare_mmap_dir(mmap_dir)
 
     for module_name, module in model.named_modules():
         if not P.is_parametrized(module):
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index dee63ab..c42161f 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -45,7 +45,10 @@
     _enable_observer,
 )
 
-from ._prepare_for_export import prepare_for_mil_export, prepare_for_mlir_export
+from ._prepare_for_export import (
+    prepare_for_mil_export as _prepare_for_mil_export,
+    prepare_for_mlir_export as _prepare_for_mlir_export,
+)
 from .kmeans_fake_palettize import _KMeansFakePalettize
 from .supported_ops_registry import _KMeansPalettizerSupportedOpsRegistry
 
@@ -405,10 +408,10 @@ def finalize(
                 pass
 
             case ExportBackend.CoreAI:
-                finalized_model = prepare_for_mlir_export(finalized_model, mmap_dir=mmap_dir)
+                finalized_model = _prepare_for_mlir_export(finalized_model, mmap_dir=mmap_dir)
 
             case ExportBackend.CoreML:
-                finalized_model = prepare_for_mil_export(finalized_model)
+                finalized_model = _prepare_for_mil_export(finalized_model)
 
             case _:
                 msg = f"Unsupported backend: {backend}"
diff --git a/src/coreai_opt/palettization/kmeans/supported_ops_registry.py b/src/coreai_opt/palettization/kmeans/supported_ops_registry.py
index 96a4dd0..0962778 100644
--- a/src/coreai_opt/palettization/kmeans/supported_ops_registry.py
+++ b/src/coreai_opt/palettization/kmeans/supported_ops_registry.py
@@ -11,7 +11,7 @@
 import torch.nn.functional as F
 
 from coreai_opt._utils.insertion.torch_function import (
-    BaseSupportedOpsRegistry,
+    BaseSupportedOpsRegistry as _BaseSupportedOpsRegistry,
 )
 from coreai_opt.palettization.kmeans.kmeans_support_mixins import (
     _ConvPalettizationMixin,
@@ -24,7 +24,7 @@
 _T = TypeVar("_T")
 
 
-class _KMeansPalettizerSupportedOpsRegistry(BaseSupportedOpsRegistry):
+class _KMeansPalettizerSupportedOpsRegistry(_BaseSupportedOpsRegistry):
     """
     Registry for KMeans palettization operations.
 
diff --git a/src/coreai_opt/palettization/spec/granularity.py b/src/coreai_opt/palettization/spec/granularity.py
index 1e166c9..71155dd 100644
--- a/src/coreai_opt/palettization/spec/granularity.py
+++ b/src/coreai_opt/palettization/spec/granularity.py
@@ -9,12 +9,12 @@
 import torch
 from pydantic import BaseModel, ConfigDict, Field, model_serializer
 
-from coreai_opt._utils.registry_utils import ConfigRegistryMixin
+from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin
 
 from .errors import _IncompatibleGranularityError
 
 
-class PalettizationGranularity(BaseModel, ConfigRegistryMixin):
+class PalettizationGranularity(BaseModel, _ConfigRegistryMixin):
     """
     Base class for palettization granularity specifications.
     """
diff --git a/src/coreai_opt/pruning/magnitude_pruner.py b/src/coreai_opt/pruning/magnitude_pruner.py
index c0b3309..7ed79f2 100644
--- a/src/coreai_opt/pruning/magnitude_pruner.py
+++ b/src/coreai_opt/pruning/magnitude_pruner.py
@@ -28,7 +28,10 @@
 from coreai_opt.config.spec.base import CompressionSpec
 from coreai_opt.pruning.spec import PruneImplBase
 
-from ._prepare_for_export import prepare_for_mil_export, prepare_for_mlir_export
+from ._prepare_for_export import (
+    prepare_for_mil_export as _prepare_for_mil_export,
+    prepare_for_mlir_export as _prepare_for_mlir_export,
+)
 from .base_pruner import _BasePruner
 from .config import MagnitudePrunerConfig
 from .supported_ops_registry import _PrunerSupportedOpsRegistry
@@ -162,10 +165,10 @@ def finalize(
                 pass
 
             case ExportBackend.CoreAI:
-                model = prepare_for_mlir_export(model)
+                model = _prepare_for_mlir_export(model)
 
             case ExportBackend.CoreML:
-                model = prepare_for_mil_export(model)
+                model = _prepare_for_mil_export(model)
 
             case _:
                 raise ValueError(f"Unsupported backend: {backend}")
diff --git a/src/coreai_opt/pruning/spec/scheme.py b/src/coreai_opt/pruning/spec/scheme.py
index 7697619..4be410a 100644
--- a/src/coreai_opt/pruning/spec/scheme.py
+++ b/src/coreai_opt/pruning/spec/scheme.py
@@ -11,10 +11,10 @@
 
 from pydantic import BaseModel, ConfigDict, Field, model_serializer
 
-from coreai_opt._utils.registry_utils import ConfigRegistryMixin
+from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin
 
 
-class PruningScheme(BaseModel, ConfigRegistryMixin):
+class PruningScheme(BaseModel, _ConfigRegistryMixin):
     """Base class for pruning scheme specifications.
 
     A pruning scheme defines the structural pattern of sparsity applied
diff --git a/src/coreai_opt/pruning/supported_ops_registry.py b/src/coreai_opt/pruning/supported_ops_registry.py
index 850ad0c..5b5469e 100644
--- a/src/coreai_opt/pruning/supported_ops_registry.py
+++ b/src/coreai_opt/pruning/supported_ops_registry.py
@@ -7,10 +7,12 @@
 
 import torch.nn.functional as F
 
-from coreai_opt._utils.insertion.torch_function import BaseSupportedOpsRegistry
+from coreai_opt._utils.insertion.torch_function import (
+    BaseSupportedOpsRegistry as _BaseSupportedOpsRegistry,
+)
 
 
-class _PrunerSupportedOpsRegistry(BaseSupportedOpsRegistry):
+class _PrunerSupportedOpsRegistry(_BaseSupportedOpsRegistry):
     """Registry of operations that support pruning."""
 
 
diff --git a/src/coreai_opt/quantization/_axis_defaults.py b/src/coreai_opt/quantization/_axis_defaults.py
index cabc0e0..eee3a3a 100644
--- a/src/coreai_opt/quantization/_axis_defaults.py
+++ b/src/coreai_opt/quantization/_axis_defaults.py
@@ -20,9 +20,7 @@
 import torch.nn.utils.parametrize as P
 from torch.fx import GraphModule
 
-from coreai_opt._utils.torch_utils import (
-    ATEN_OP_TO_MODULE_TYPE as _ATEN_OP_TO_MODULE_TYPE,
-)
+from coreai_opt._utils.torch_utils import ATEN_OP_TO_MODULE_TYPE
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
 from coreai_opt.quantization.spec.granularity import (
@@ -247,7 +245,7 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None:
         for module_type_or_op, name in consumers:
             # graph-mode consumers are aten ops, so map them to module types first
             if isinstance(module_type_or_op, torch._ops.OpOverload):
-                module_type = _ATEN_OP_TO_MODULE_TYPE.get(module_type_or_op)
+                module_type = ATEN_OP_TO_MODULE_TYPE.get(module_type_or_op)
             else:
                 module_type = module_type_or_op
 
diff --git a/src/coreai_opt/quantization/_eager/_prepare_for_export.py b/src/coreai_opt/quantization/_eager/_prepare_for_export.py
index 707de99..0b394ae 100644
--- a/src/coreai_opt/quantization/_eager/_prepare_for_export.py
+++ b/src/coreai_opt/quantization/_eager/_prepare_for_export.py
@@ -12,23 +12,20 @@
 import torch.nn as nn
 import torch.nn.utils.parametrize as P
 
-from coreai_opt._utils.export_utils import (
-    clear_parametrization_original as _clear_parametrization_original,
-    prepare_mmap_dir as _prepare_mmap_dir,
-)
+from coreai_opt._utils.export_utils import clear_parametrization_original, prepare_mmap_dir
 from coreai_opt._utils.import_utils import lazy_import_coreai_torch
 from coreai_opt._utils.torch_utils import (
-    get_parent_module_and_attr_name as _get_parent_module_and_attr_name,
-    is_float4_dtype as _is_float4_dtype,
-    mmap_module_state_dict as _mmap_module_state_dict,
+    get_parent_module_and_attr_name,
+    is_float4_dtype,
+    mmap_module_state_dict,
 )
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization._export_utils import (
-    canonicalize_qparam_shape as _canonicalize_qparam_shape,
-    extract_quantization_params as _extract_quantization_params,
-    pack_fp4_to_float4tensor as _pack_fp4_to_float4tensor,
-    select_export_qparams_by_formulation as _select_export_qparams_by_formulation,
-    validate_fp4_export as _validate_fp4_export,
+    canonicalize_qparam_shape,
+    extract_quantization_params,
+    pack_fp4_to_float4tensor,
+    select_export_qparams_by_formulation,
+    validate_fp4_export,
 )
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
 
@@ -49,7 +46,7 @@ def _add_weight_dequantization_parametrization(
     """
 
     # Extract and prepare quantization parameters
-    scale, zero_point, minval = _extract_quantization_params(fake_quant_mod)
+    scale, zero_point, minval = extract_quantization_params(fake_quant_mod)
 
     # Cast scale and minval to appropriate dtype for MLIR backend inference
     _compute_dtype_for_export = fake_quant_mod.qparams_calculator._compute_dtype_for_export
@@ -62,11 +59,11 @@ def _add_weight_dequantization_parametrization(
 
     # Drop one of the offsets so that the export
     # module / runtime selects the right dequant path.
-    zero_point, minval = _select_export_qparams_by_formulation(fake_quant_mod, zero_point, minval)
+    zero_point, minval = select_export_qparams_by_formulation(fake_quant_mod, zero_point, minval)
 
-    if _is_float4_dtype(fake_quant_mod.dtype):
-        _validate_fp4_export(fake_quant_mod, quantized_data)
-        quantized_data = _pack_fp4_to_float4tensor(quantized_data)
+    if is_float4_dtype(fake_quant_mod.dtype):
+        validate_fp4_export(fake_quant_mod, quantized_data)
+        quantized_data = pack_fp4_to_float4tensor(quantized_data)
     if fake_quant_mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu:
         output_dtype = _compute_dtype_for_export
         scale = scale.to(torch.float8_e8m0fnu)
@@ -89,10 +86,10 @@ def _add_weight_dequantization_parametrization(
     if mmap_dir is not None:
         stem = f"{module_name}.{param_name}" if module_name else param_name
         path = Path(mmap_dir) / f"{stem}.safetensors"
-        _mmap_module_state_dict(weight_dequant_mod, path)
+        mmap_module_state_dict(weight_dequant_mod, path)
 
     module.parametrizations[param_name][fake_quant_idx] = weight_dequant_mod
-    _clear_parametrization_original(module, param_name)
+    clear_parametrization_original(module, param_name)
     return weight_dequant_mod
 
 
@@ -113,7 +110,7 @@ def _process_weight_quantization(model: nn.Module, mmap_dir: str | PathLike[str]
             directory and reload it via mmap, so large-model finalization does
             not hold full quantized weights in RAM.
     """
-    _prepare_mmap_dir(mmap_dir)
+    prepare_mmap_dir(mmap_dir)
 
     # Lazy import: coreai_torch is required for MLIR export
     def _import_coreai_torch_modules():
@@ -156,7 +153,7 @@ def _import_coreai_torch_modules():
                 # Shared FakeQuantize across modules (e.g. weight tying): reuse
                 # the same dequant module so post-finalize sharing is preserved.
                 module.parametrizations[param_name][fake_quant_idx] = cached
-                _clear_parametrization_original(module, param_name)
+                clear_parametrization_original(module, param_name)
                 continue
 
             weight_dequant_mod = _add_weight_dequantization_parametrization(
@@ -200,18 +197,18 @@ def _import_coreai_torch_modules():
         if isinstance(module, FakeQuantizeImplBase) and module.quantization_target in (
             CompressionTargetTensor.ACTIVATION,
         ):
-            if _is_float4_dtype(module.dtype):
+            if is_float4_dtype(module.dtype):
                 raise ValueError("FP4 activation quantization is not supported for MLIR export.")
             modules_to_replace.append((name, module))
 
     # Replace each FakeQuantizeImplBase module
     for name, fake_quant_module in modules_to_replace:
         # Extract quantization parameters
-        scale, zero_point, minval = _extract_quantization_params(fake_quant_module)
+        scale, zero_point, minval = extract_quantization_params(fake_quant_module)
 
         # Drop one of the offsets so that the export
         # module / runtime selects the right dequant path.
-        zero_point, minval = _select_export_qparams_by_formulation(
+        zero_point, minval = select_export_qparams_by_formulation(
             fake_quant_module, zero_point, minval
         )
 
@@ -226,11 +223,11 @@ def _import_coreai_torch_modules():
 
         # Canonicalize scale/zero_point/minval to 0-D (per-tensor) or 1-D (per-channel)
         granularity = fake_quant_module.granularity
-        scale = _canonicalize_qparam_shape(scale, granularity)
+        scale = canonicalize_qparam_shape(scale, granularity)
         if zero_point is not None:
-            zero_point = _canonicalize_qparam_shape(zero_point, granularity)
+            zero_point = canonicalize_qparam_shape(zero_point, granularity)
         if minval is not None:
-            minval = _canonicalize_qparam_shape(minval, granularity)
+            minval = canonicalize_qparam_shape(minval, granularity)
 
         axis = fake_quant_module.qparams_calculator._resolved_axis
         axis = axis if axis is not None else 0
@@ -277,7 +274,7 @@ def _import_coreai_torch_modules():
         )
 
         # Replace the module in the model
-        parent_module, attr_name = _get_parent_module_and_attr_name(model, name)
+        parent_module, attr_name = get_parent_module_and_attr_name(model, name)
         setattr(parent_module, attr_name, replacement_module)
 
 
diff --git a/src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py b/src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py
index f0679a6..b3555ed 100644
--- a/src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py
+++ b/src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py
@@ -19,16 +19,14 @@
 
 from coreai_opt._utils.export_utils import validate_coreml_compatibility
 from coreai_opt._utils.metadata_utils import CompressionType, MILCompressionMetadata
-from coreai_opt._utils.torch_utils import (
-    get_parent_module_and_attr_name as _get_parent_module_and_attr_name,
-)
+from coreai_opt._utils.torch_utils import get_parent_module_and_attr_name
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization._export_utils import (
-    convert_dtype_for_torch_quantize as _convert_dtype_for_torch_quantize,
-    create_mil_act_quant_seq as _create_mil_act_quant_seq,
-    extract_quantization_params as _extract_quantization_params,
-    is_module_fake_quant_target as _is_module_fake_quant_target,
-    validate_qformulation_for_mil_export as _validate_qformulation_for_mil_export,
+    convert_dtype_for_torch_quantize,
+    create_mil_act_quant_seq,
+    extract_quantization_params,
+    is_module_fake_quant_target,
+    validate_qformulation_for_mil_export,
 )
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
 
@@ -46,8 +44,8 @@ def _process_weight_quantization(
         fake_quant_param: The fake quantization parametrization to process
 
     """
-    _validate_qformulation_for_mil_export(fake_quant_param)
-    scale, zero_point, _ = _extract_quantization_params(fake_quant_param)
+    validate_qformulation_for_mil_export(fake_quant_param)
+    scale, zero_point, _ = extract_quantization_params(fake_quant_param)
 
     # Remove parametrization but keep fake-quantized weights
     P.remove_parametrizations(
@@ -100,10 +98,10 @@ def _process_activation_quantization(
         fake_quant_mod: The fake quantization module to replace
 
     """
-    _validate_qformulation_for_mil_export(fake_quant_mod)
+    validate_qformulation_for_mil_export(fake_quant_mod)
 
-    scale, zero_point, _ = _extract_quantization_params(fake_quant_mod)
-    converted_dtype, converted_zero_point = _convert_dtype_for_torch_quantize(
+    scale, zero_point, _ = extract_quantization_params(fake_quant_mod)
+    converted_dtype, converted_zero_point = convert_dtype_for_torch_quantize(
         fake_quant_mod.dtype,
         zero_point,
     )
@@ -111,7 +109,7 @@ def _process_activation_quantization(
     # Use non-negative axis for export (None for per-tensor)
     axis = fake_quant_mod.qparams_calculator._resolved_axis
 
-    replacement_module = _create_mil_act_quant_seq(
+    replacement_module = create_mil_act_quant_seq(
         scale=scale,
         zero_point=converted_zero_point,
         dtype=converted_dtype,
@@ -145,13 +143,13 @@ def prepare_for_mil_export(model: nn.Module) -> nn.Module:
         if P.is_parametrized(module):
             for param_name, parametrizations in module.parametrizations.items():
                 for p in parametrizations:
-                    if _is_module_fake_quant_target(p, CompressionTargetTensor.WEIGHT):
+                    if is_module_fake_quant_target(p, CompressionTargetTensor.WEIGHT):
                         validate_coreml_compatibility(
                             CompressionTargetTensor.WEIGHT,
                             p.dtype,
                             f"weight '{param_name}' of module '{module_name}'",
                         )
-        if _is_module_fake_quant_target(module, CompressionTargetTensor.ACTIVATION):
+        if is_module_fake_quant_target(module, CompressionTargetTensor.ACTIVATION):
             validate_coreml_compatibility(
                 CompressionTargetTensor.ACTIVATION,
                 module.dtype,
@@ -164,7 +162,7 @@ def prepare_for_mil_export(model: nn.Module) -> nn.Module:
         if P.is_parametrized(module):
             for param_name, parametrizations in list(module.parametrizations.items()):
                 for p in parametrizations:
-                    if not _is_module_fake_quant_target(p, CompressionTargetTensor.WEIGHT):
+                    if not is_module_fake_quant_target(p, CompressionTargetTensor.WEIGHT):
                         continue
 
                     if _mark_if_not_already_processed(p, processed_fq_ids):
@@ -173,11 +171,11 @@ def prepare_for_mil_export(model: nn.Module) -> nn.Module:
                     _process_weight_quantization(module, param_name, p)
 
         # Handle activation quantization modules
-        if _is_module_fake_quant_target(module, CompressionTargetTensor.ACTIVATION):
+        if is_module_fake_quant_target(module, CompressionTargetTensor.ACTIVATION):
             if _mark_if_not_already_processed(module, processed_fq_ids):
                 continue
 
-            parent_module, attr_name = _get_parent_module_and_attr_name(model, name)
+            parent_module, attr_name = get_parent_module_and_attr_name(model, name)
             _process_activation_quantization(parent_module, attr_name, module)
 
     # Register metadata version
diff --git a/src/coreai_opt/quantization/_eager/_utils.py b/src/coreai_opt/quantization/_eager/_utils.py
index 2e77fad..96e4d94 100644
--- a/src/coreai_opt/quantization/_eager/_utils.py
+++ b/src/coreai_opt/quantization/_eager/_utils.py
@@ -11,8 +11,8 @@
     RegisteredOptimizersTracker,
 )
 from coreai_opt._utils.torch_utils import (
-    get_parent_module_and_attr_name as _get_parent_module_and_attr_name,
-    remove_compression_parametrizations as _remove_compression_parametrizations,
+    get_parent_module_and_attr_name,
+    remove_compression_parametrizations,
 )
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
@@ -71,7 +71,7 @@ def remove_fake_quant_modules(
         int: Number of modules removed.
 
     """
-    num_removed = _remove_compression_parametrizations(model, modules_to_remove)
+    num_removed = remove_compression_parametrizations(model, modules_to_remove)
 
     # Remove activation FQ standalone modules
     ids_to_remove = {id(m) for m in modules_to_remove}
@@ -81,7 +81,7 @@ def remove_fake_quant_modules(
             and module.quantization_target == CompressionTargetTensor.ACTIVATION
             and id(module) in ids_to_remove
         ):
-            parent_module, attr_name = _get_parent_module_and_attr_name(model, name)
+            parent_module, attr_name = get_parent_module_and_attr_name(model, name)
             delattr(parent_module, attr_name)
             num_removed += 1
 
diff --git a/src/coreai_opt/quantization/_eager/quantizer.py b/src/coreai_opt/quantization/_eager/quantizer.py
index 4b8179d..7f00a7f 100644
--- a/src/coreai_opt/quantization/_eager/quantizer.py
+++ b/src/coreai_opt/quantization/_eager/quantizer.py
@@ -20,16 +20,12 @@
     enable_observer,
 )
 
-from coreai_opt._utils.config_utils import (
-    ALL_TENSORS as _ALL_TENSORS,
-)
-from coreai_opt._utils.eager_utils import (
-    EagerCompressionComponentBuilderMixin as _EagerCompressionComponentBuilderMixin,
-)
+from coreai_opt._utils.config_utils import ALL_TENSORS
+from coreai_opt._utils.eager_utils import EagerCompressionComponentBuilderMixin
 from coreai_opt._utils.insertion.torch_function import (
     TorchFunctionEagerHandler,
 )
-from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
+from coreai_opt._utils.spec_utils import PartialConstructor
 from coreai_opt._utils.torch_utils import move_model_to_eval, move_model_to_train
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleCompressionConfig
@@ -37,11 +33,11 @@
 from coreai_opt.config.spec.base import CompressionSpec
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_eager as _apply_weight_axis_defaults,
-    validate_activation_axes as _validate_activation_axes,
+    validate_activation_axes,
 )
 from coreai_opt.quantization._fake_quant_utils import (
-    disable_activation_fake_quant as _disable_activation_fake_quant,
-    enable_weight_fake_quant as _enable_weight_fake_quant,
+    disable_activation_fake_quant,
+    enable_weight_fake_quant,
 )
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config import (
@@ -54,19 +50,15 @@
     FakeQuantizeImplBase,
 )
 
-from ._prepare_for_export import (
-    prepare_for_mlir_export as _prepare_for_mlir_export,
-)
-from ._prepare_for_mil_export import (
-    prepare_for_mil_export as _prepare_for_mil_export,
-)
+from ._prepare_for_export import prepare_for_mlir_export
+from ._prepare_for_mil_export import prepare_for_mil_export
 from ._utils import remove_act_fq_from_reference_tracker, remove_fake_quant_modules
 from .supported_ops_registry import (
     EagerQuantizerSupportedOpsRegistry,
 )
 
 
-class EagerQuantizer(_BaseQuantizer, _EagerCompressionComponentBuilderMixin):
+class EagerQuantizer(_BaseQuantizer, EagerCompressionComponentBuilderMixin):
     """Eager mode quantization-aware training (QAT) quantizer.
 
     Uses `__torch_function__` to trace model execution and insert fake quantizers
@@ -163,7 +155,7 @@ def _validate_config(config: QuantizerConfig) -> None:
         # identification.
         for spec in op_input_specs_to_check + op_output_specs_to_check:
             for key in spec:
-                if isinstance(key, str) and key != _ALL_TENSORS:
+                if isinstance(key, str) and key != ALL_TENSORS:
                     error_msg = (
                         "Only integer indices or '*' are supported for op_input_spec "
                         f"and op_output_spec currently. Got {spec}"
@@ -232,7 +224,7 @@ def _postprocess_prepared_model(model: nn.Module) -> None:
             model (nn.Module): The model after eager prepare().
         """
         _apply_weight_axis_defaults(model)
-        _validate_activation_axes(model)
+        validate_activation_axes(model)
 
     def finalize(
         self,
@@ -282,9 +274,9 @@ def finalize(
             case ExportBackend._TORCH:
                 finalized_model = model
             case ExportBackend.CoreAI:
-                finalized_model = _prepare_for_mlir_export(model, mmap_dir=mmap_dir)
+                finalized_model = prepare_for_mlir_export(model, mmap_dir=mmap_dir)
             case ExportBackend.CoreML:
-                finalized_model = _prepare_for_mil_export(model)
+                finalized_model = prepare_for_mil_export(model)
             case _:
                 msg = f"Unsupported backend {backend}"
                 raise ValueError(msg)
@@ -315,8 +307,8 @@ def calibration_mode(self, model: nn.Module | None = None) -> Generator:
             )
         with move_model_to_eval(self._model):
             self._model.apply(enable_observer)
-            self._model.apply(_enable_weight_fake_quant)
-            self._model.apply(_disable_activation_fake_quant)
+            self._model.apply(enable_weight_fake_quant)
+            self._model.apply(disable_activation_fake_quant)
             try:
                 yield
             finally:
@@ -405,5 +397,5 @@ def _spec_to_partial(
         spec: CompressionSpec | None,
         target: CompressionTargetTensor,
         module_config: ModuleCompressionConfig,
-    ) -> _PartialConstructor | None:
+    ) -> PartialConstructor | None:
         return QuantizationComponentFactory.construct_partial(spec=spec, target=target)
diff --git a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py
index e6a1338..ea9d4f1 100644
--- a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py
+++ b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py
@@ -16,13 +16,8 @@
 from coreai_opt._utils.registry_utils import ClassRegistryMixin
 
 from . import _annotation_utils
-from ._annotation_config import (
-    AnnotationConfig as _AnnotationConfig,
-    AnnotationContext as _AnnotationContext,
-)
-from ._annotation_utils import (
-    OpsListPattern as _OpsListPattern,
-)
+from ._annotation_config import AnnotationConfig, AnnotationContext
+from ._annotation_utils import OpsListPattern
 
 # Generic type variable for match results
 MatchType = TypeVar("MatchType")
@@ -38,7 +33,7 @@
 # 3. Annotation pass context. Holds pass-invariant inputs the annotator may
 #    need (the model's module-name-to-state-names map and the set of shared
 #    observer nodes computed at the start of this annotation pass).
-AnnotatorFunc: TypeAlias = Callable[[MatchType, _AnnotationConfig, _AnnotationContext], Any]
+AnnotatorFunc: TypeAlias = Callable[[MatchType, AnnotationConfig, AnnotationContext], Any]
 
 
 @dataclass(frozen=True)
@@ -56,7 +51,7 @@ class AnnotatorMatchInfo(Generic[MatchType]):
     pattern_length: int
 
 
-def _get_all_patterns_with_activations_appended(starting_pattern: _OpsListPattern):
+def _get_all_patterns_with_activations_appended(starting_pattern: OpsListPattern):
     """
     Given a starting ops list pattern, return a new list of patterns containing the
     combination of the starting pattern with all activation types appended to the
@@ -67,19 +62,19 @@ def _get_all_patterns_with_activations_appended(starting_pattern: _OpsListPatter
         _annotation_utils._supported_activations
         + _annotation_utils._supported_activations_no_inplace
     ):
-        patterns.append(_OpsListPattern(starting_pattern.pattern + [act_fn.__name__]))
+        patterns.append(OpsListPattern(starting_pattern.pattern + [act_fn.__name__]))
     return patterns
 
 
 def _get_all_patterns_from_base_ops(
     base_ops: set[str], use_act: bool = False
-) -> list[_OpsListPattern]:
+) -> list[OpsListPattern]:
     patterns = []
     for base_op in base_ops:
         if use_act:
-            patterns.extend(_get_all_patterns_with_activations_appended(_OpsListPattern([base_op])))
+            patterns.extend(_get_all_patterns_with_activations_appended(OpsListPattern([base_op])))
         else:
-            patterns.append(_OpsListPattern([base_op]))
+            patterns.append(OpsListPattern([base_op]))
     return patterns
 
 
@@ -96,7 +91,7 @@ class BaseAnnotationPattern(Generic[MatchType], ABC):
     """
 
     # Class-level cache for patterns
-    _patterns: list[torch.fx.GraphModule | _OpsListPattern] | None = None
+    _patterns: list[torch.fx.GraphModule | OpsListPattern] | None = None
 
     @classmethod
     @abstractmethod
@@ -111,7 +106,7 @@ def get_annotator_func(cls) -> AnnotatorFunc[MatchType]:
         raise NotImplementedError
 
     @classmethod
-    def get_patterns(cls) -> list[torch.fx.GraphModule | _OpsListPattern]:
+    def get_patterns(cls) -> list[torch.fx.GraphModule | OpsListPattern]:
         """
         Returns a cached list of graph modules representing the patterns
         this annotator can match. Patterns are generated once and cached.
@@ -146,7 +141,7 @@ def _get_single_pattern_length(pattern):
         pattern_len = 0
         # Account for AnnotationPattern classes using get_sequential_partition with a
         # list of types.
-        if isinstance(pattern, _OpsListPattern):
+        if isinstance(pattern, OpsListPattern):
             return len(pattern.pattern)
         for node in pattern.graph.nodes:
             # Sanity check to make sure we account for all node.op types
@@ -170,7 +165,7 @@ def get_pattern_length(cls):
 
     @classmethod
     @abstractmethod
-    def generate_patterns(cls) -> list[torch.fx.GraphModule | _OpsListPattern]:
+    def generate_patterns(cls) -> list[torch.fx.GraphModule | OpsListPattern]:
         """
         Generate the patterns for this annotator. Called once and cached.
 
@@ -183,7 +178,7 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule | _OpsListPattern]:
     @classmethod
     @abstractmethod
     def match_single_pattern(
-        cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | _OpsListPattern
+        cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | OpsListPattern
     ) -> dict[torch.fx.Node, MatchType]:
         """
         Match nodes in model to the provided pattern and return a dictionary mapping
@@ -274,7 +269,7 @@ def get_annotator_func(cls) -> AnnotatorFunc[InternalMatch]:
 
     @classmethod
     def match_single_pattern(
-        cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | _OpsListPattern
+        cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | OpsListPattern
     ) -> dict[torch.fx.Node, InternalMatch]:
         """
         Match nodes in model to the provided pattern and return a dictionary mapping
@@ -320,13 +315,13 @@ def get_annotator_func(cls) -> AnnotatorFunc[tuple[SourcePartition]]:
 
     @classmethod
     def match_single_pattern(
-        cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | _OpsListPattern
+        cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | OpsListPattern
     ) -> dict[torch.fx.Node, tuple[SourcePartition]]:
         """
         Match nodes in model to the provided pattern and return a dictionary mapping
         nodes to matches.
         """
-        if not isinstance(pattern, _OpsListPattern):
+        if not isinstance(pattern, OpsListPattern):
             error_msg = (
                 "NAryActPattern expects patterns of type _OpsListPattern, but "
                 f"got type {type(pattern)}.)"
@@ -372,13 +367,13 @@ def _validate_patterns_length(cls):
 
     @classmethod
     def match_single_pattern(
-        cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | _OpsListPattern
+        cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | OpsListPattern
     ) -> dict[torch.fx.Node, tuple[SourcePartition]]:
         """
         Match nodes in model to the provided pattern and return a dictionary mapping
         nodes to matches.
         """
-        if not isinstance(pattern, _OpsListPattern):
+        if not isinstance(pattern, OpsListPattern):
             error_msg = (
                 "SharedObserverModulePattern expects patterns of type "
                 f"_OpsListPattern, but got type {type(pattern)}.)"
@@ -680,7 +675,7 @@ class MatMulPattern(NAryActPattern):
     """
 
     @classmethod
-    def generate_patterns(cls) -> list[_OpsListPattern]:
+    def generate_patterns(cls) -> list[OpsListPattern]:
         """Returns matmul pattern."""
         return _get_all_patterns_from_base_ops({"matmul"})
 
@@ -692,7 +687,7 @@ class MatMulActPattern(NAryActPattern):
     """
 
     @classmethod
-    def generate_patterns(cls) -> list[_OpsListPattern]:
+    def generate_patterns(cls) -> list[OpsListPattern]:
         """Returns matmul act pattern."""
         return _get_all_patterns_from_base_ops({"matmul"}, use_act=True)
 
@@ -704,7 +699,7 @@ class AddPattern(NAryActPattern):
     """
 
     @classmethod
-    def generate_patterns(cls) -> list[_OpsListPattern]:
+    def generate_patterns(cls) -> list[OpsListPattern]:
         """Returns add pattern."""
         return _get_all_patterns_from_base_ops({"add", "add_"})
 
@@ -727,7 +722,7 @@ class MulPattern(NAryActPattern):
     """
 
     @classmethod
-    def generate_patterns(cls) -> list[_OpsListPattern]:
+    def generate_patterns(cls) -> list[OpsListPattern]:
         """Returns mul pattern."""
         return _get_all_patterns_from_base_ops({"mul", "mul_"})
 
@@ -750,7 +745,7 @@ class SubPattern(NAryActPattern):
     """
 
     @classmethod
-    def generate_patterns(cls) -> list[_OpsListPattern]:
+    def generate_patterns(cls) -> list[OpsListPattern]:
         """Returns sub pattern."""
         # Multiple types of sub need to be listed since torchao doesn't include sub in
         # pt2e.graph_utils._EQUIVALENT_TYPES dict
@@ -769,7 +764,7 @@ class KVCacheUpdatePattern(NAryActPattern):
         """input -> kv-cache-update op -> output"""
 
         @classmethod
-        def generate_patterns(cls) -> list[_OpsListPattern]:
+        def generate_patterns(cls) -> list[OpsListPattern]:
             return _get_all_patterns_from_base_ops({op_name}, use_act=False)
 
     return KVCacheUpdatePattern
diff --git a/src/coreai_opt/quantization/_graph/_annotation_utils.py b/src/coreai_opt/quantization/_graph/_annotation_utils.py
index 27431a3..01bbb61 100644
--- a/src/coreai_opt/quantization/_graph/_annotation_utils.py
+++ b/src/coreai_opt/quantization/_graph/_annotation_utils.py
@@ -26,18 +26,14 @@
 )
 from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
 
-from coreai_opt._utils.config_utils import (
-    ALL_TENSORS as _ALL_TENSORS,
-    ConfigLevel as _ConfigLevel,
-    get_last_matching_spec,
-)
+from coreai_opt._utils.config_utils import ALL_TENSORS, ConfigLevel, get_last_matching_spec
 from coreai_opt._utils.fx_utils import (
     get_local_state_name,
     get_module_boundary_nodes,
     is_coreai_compressed_state_node,
 )
 from coreai_opt._utils.python_utils import get_fn_arg_names
-from coreai_opt._utils.version_utils import version_ge as _version_ge
+from coreai_opt._utils.version_utils import version_ge
 from coreai_opt.config.compression_config import ModuleConfigDict
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization._graph._utils import get_source_module_name
@@ -109,7 +105,7 @@ def _get_aten_graph_module_for_pattern(
         )
 
     exported_program = torch.export.export(pattern, example_inputs, kwargs, strict=False)
-    if _version_ge(torch, "2.9"):
+    if version_ge(torch, "2.9"):
         aten_pattern = exported_program.module(check_guards=False)
     else:
         aten_pattern = exported_program.module()
@@ -1237,7 +1233,7 @@ def annotate_module_level_specs(
             by the full state name.
         model: Model to annotate.
     """
-    for config_level in [_ConfigLevel.MODULE_TYPE, _ConfigLevel.MODULE_NAME]:
+    for config_level in [ConfigLevel.MODULE_TYPE, ConfigLevel.MODULE_NAME]:
         for module_name, module_config in module_configs[config_level].items():
             if _module_config_has_module_level_input_output_spec(module_config):
                 _annotate_nodes_for_module_level_input_output_spec(
@@ -1315,7 +1311,7 @@ def _match_and_annotate_state_node(
     Given a state node, check if any of the module_configs have applicable
     module_state_specs and apply if so.
     """
-    for level in _ConfigLevel.priority_order():
+    for level in ConfigLevel.priority_order():
         # Reversed is needed because two different modules may have shared params where
         # both module configs are setting module_state_spec for the param using their
         # respective local names for the same parameter.
@@ -1388,8 +1384,8 @@ def _get_spec_from_spec_dict(
     for i in identifier:
         if i in spec_dict:
             return (True, spec_dict[i])
-    if _ALL_TENSORS in spec_dict:
-        return (True, spec_dict[_ALL_TENSORS])
+    if ALL_TENSORS in spec_dict:
+        return (True, spec_dict[ALL_TENSORS])
     return (False, None)
 
 
diff --git a/src/coreai_opt/quantization/_graph/_prepare_for_export.py b/src/coreai_opt/quantization/_graph/_prepare_for_export.py
index d7ef4ba..5f83309 100644
--- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py
+++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py
@@ -18,23 +18,20 @@
 from torch.fx import Node
 
 from coreai_opt._utils.export_utils import validate_coreml_compatibility
-from coreai_opt._utils.fx_utils import get_node_type as _get_node_type
+from coreai_opt._utils.fx_utils import get_node_type
 from coreai_opt._utils.import_utils import lazy_import_coreai_torch
 from coreai_opt._utils.metadata_utils import CompressionType, MILCompressionMetadata
-from coreai_opt._utils.torch_utils import (
-    is_float4_dtype as _is_float4_dtype,
-    sanitize_module_name as _sanitize_module_name,
-)
+from coreai_opt._utils.torch_utils import is_float4_dtype, sanitize_module_name
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization._export_utils import (
-    canonicalize_qparam_shape as _canonicalize_qparam_shape,
-    convert_dtype_for_torch_quantize as _convert_dtype_for_torch_quantize,
-    create_mil_act_quant_seq as _create_mil_act_quant_seq,
-    extract_quantization_params as _extract_quantization_params,
-    pack_fp4_to_float4tensor as _pack_fp4_to_float4tensor,
-    select_export_qparams_by_formulation as _select_export_qparams_by_formulation,
-    validate_fp4_export as _validate_fp4_export,
-    validate_qformulation_for_mil_export as _validate_qformulation_for_mil_export,
+    canonicalize_qparam_shape,
+    convert_dtype_for_torch_quantize,
+    create_mil_act_quant_seq,
+    extract_quantization_params,
+    pack_fp4_to_float4tensor,
+    select_export_qparams_by_formulation,
+    validate_fp4_export,
+    validate_qformulation_for_mil_export,
 )
 from coreai_opt.quantization._graph._utils import (
     remove_fake_quant_module,
@@ -266,7 +263,7 @@ def _import_coreai_custom_ops():
         raise ValueError(f"Expected get_attr node for weight quantization, got {input_node.op}")
 
     # Extract and prepare quantization parameters
-    scale, zero_point, minval = _extract_quantization_params(fake_quant_mod)
+    scale, zero_point, minval = extract_quantization_params(fake_quant_mod)
     # Cast scale and minval to appropriate dtype for MLIR backend inference
     _compute_dtype_for_export = fake_quant_mod.qparams_calculator._compute_dtype_for_export
     scale = scale.to(dtype=_compute_dtype_for_export)
@@ -279,12 +276,12 @@ def _import_coreai_custom_ops():
 
     # Drop one of the offsets so that the export
     # module / runtime selects the right dequant path.
-    zero_point, minval = _select_export_qparams_by_formulation(fake_quant_mod, zero_point, minval)
+    zero_point, minval = select_export_qparams_by_formulation(fake_quant_mod, zero_point, minval)
 
     # Convert to export dtypes for MLIR ops
-    if _is_float4_dtype(fake_quant_mod.dtype):
-        _validate_fp4_export(fake_quant_mod, quantized_data)
-        quantized_data = _pack_fp4_to_float4tensor(quantized_data)
+    if is_float4_dtype(fake_quant_mod.dtype):
+        validate_fp4_export(fake_quant_mod, quantized_data)
+        quantized_data = pack_fp4_to_float4tensor(quantized_data)
     if fake_quant_mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu:
         scale = scale.to(torch.float8_e8m0fnu)
 
@@ -338,7 +335,7 @@ def _process_mlir_activation_quantization(
         node: The fake quantization node to replace
         fake_quant_mod: The fake quantization module
     """
-    if _is_float4_dtype(fake_quant_mod.dtype):
+    if is_float4_dtype(fake_quant_mod.dtype):
         raise ValueError("FP4 activation quantization is not supported for MLIR export.")
 
     def _import_coreai_custom_ops():
@@ -353,11 +350,11 @@ def _import_coreai_custom_ops():
         raise ValueError(f"Node {node} has no input arguments")
 
     # Extract and prepare quantization parameters
-    scale, zero_point, minval = _extract_quantization_params(fake_quant_mod)
+    scale, zero_point, minval = extract_quantization_params(fake_quant_mod)
 
     # Drop the offset the active formulation doesn't consume so the runtime op
     # selects the right dequant path.
-    zero_point, minval = _select_export_qparams_by_formulation(fake_quant_mod, zero_point, minval)
+    zero_point, minval = select_export_qparams_by_formulation(fake_quant_mod, zero_point, minval)
 
     # Cast scale and minval to appropriate dtype for MLIR backend inference
     _compute_dtype_for_export = fake_quant_mod.qparams_calculator._compute_dtype_for_export
@@ -370,11 +367,11 @@ def _import_coreai_custom_ops():
 
     # Canonicalize scale/zero_point/minval to 0-D (per-tensor) or 1-D (per-channel)
     granularity = fake_quant_mod.granularity
-    scale = _canonicalize_qparam_shape(scale, granularity)
+    scale = canonicalize_qparam_shape(scale, granularity)
     if zero_point is not None:
-        zero_point = _canonicalize_qparam_shape(zero_point, granularity)
+        zero_point = canonicalize_qparam_shape(zero_point, granularity)
     if minval is not None:
-        minval = _canonicalize_qparam_shape(minval, granularity)
+        minval = canonicalize_qparam_shape(minval, granularity)
 
     # Register buffers and get buffer names
     base_name = node.name.replace(".", "_")
@@ -456,9 +453,9 @@ def _process_mil_weight_quantization(
         modules: Mapping of module names to modules
 
     """
-    _validate_qformulation_for_mil_export(fake_quant_mod)
+    validate_qformulation_for_mil_export(fake_quant_mod)
     # Extract quantization parameters
-    scale, zero_point, _ = _extract_quantization_params(fake_quant_mod)
+    scale, zero_point, _ = extract_quantization_params(fake_quant_mod)
     module_name, param_name = _get_weight_input_names(fake_quant_node, fake_quant_mod)
 
     # Get the module that owns the weight parameter
@@ -508,10 +505,10 @@ def _process_mil_activation_quantization(
         fake_quant_mod: The fake quantization module
 
     """
-    _validate_qformulation_for_mil_export(fake_quant_mod)
+    validate_qformulation_for_mil_export(fake_quant_mod)
 
-    scale, zero_point, _ = _extract_quantization_params(fake_quant_mod)
-    converted_dtype, converted_zero_point = _convert_dtype_for_torch_quantize(
+    scale, zero_point, _ = extract_quantization_params(fake_quant_mod)
+    converted_dtype, converted_zero_point = convert_dtype_for_torch_quantize(
         fake_quant_mod.dtype,
         zero_point,
     )
@@ -519,7 +516,7 @@ def _process_mil_activation_quantization(
     # Use non-negative axis for export (None for per-tensor)
     axis = fake_quant_mod.qparams_calculator._resolved_axis
 
-    sequential_module = _create_mil_act_quant_seq(
+    sequential_module = create_mil_act_quant_seq(
         scale=scale,
         zero_point=converted_zero_point,
         dtype=converted_dtype,
@@ -527,7 +524,7 @@ def _process_mil_activation_quantization(
     )
 
     # Add Sequential module to the model
-    module_name = f"{MIL_ACT_QUANT_MODULE_PREFIX}{_sanitize_module_name(fake_quant_node.name)}"
+    module_name = f"{MIL_ACT_QUANT_MODULE_PREFIX}{sanitize_module_name(fake_quant_node.name)}"
     model.add_module(module_name, sequential_module)
 
     # Replace fake quant node with call_module to the Sequential
@@ -732,7 +729,7 @@ def _import_coreai_custom_ops():
     for op_node in model.graph.nodes:
         if op_node.op != "call_function":
             continue
-        if _get_node_type(op_node, warn_on_failure=False) != op_type:
+        if get_node_type(op_node, warn_on_failure=False) != op_type:
             continue
 
         if len(op_node.all_input_nodes) <= quant_input_idx:
diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py
index 4c066e5..09e0a48 100644
--- a/src/coreai_opt/quantization/_graph/quantizer.py
+++ b/src/coreai_opt/quantization/_graph/quantizer.py
@@ -40,29 +40,19 @@
 )
 from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
 
-from coreai_opt._utils.config_utils import (
-    ALL_TENSORS as _ALL_TENSORS,
-    ConfigLevel as _ConfigLevel,
-)
-from coreai_opt._utils.fx_utils import (
-    get_node_type as _get_node_type,
-    normalize_module_fqn,
-)
-from coreai_opt._utils.torch_utils import (
-    export_model as _export_model,
-    move_model_to_eval,
-    move_model_to_train,
-)
-from coreai_opt._utils.version_utils import version_ge as _version_ge
+from coreai_opt._utils.config_utils import ALL_TENSORS, ConfigLevel
+from coreai_opt._utils.fx_utils import get_node_type, normalize_module_fqn
+from coreai_opt._utils.torch_utils import export_model, move_model_to_eval, move_model_to_train
+from coreai_opt._utils.version_utils import version_ge
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleConfigDict, _build_module_alias_map
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_graph as _apply_weight_axis_defaults,
-    validate_activation_axes as _validate_activation_axes,
+    validate_activation_axes,
 )
 from coreai_opt.quantization._fake_quant_utils import (
-    disable_activation_fake_quant as _disable_activation_fake_quant,
-    enable_weight_fake_quant as _enable_weight_fake_quant,
+    disable_activation_fake_quant,
+    enable_weight_fake_quant,
 )
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config import (
@@ -78,32 +68,29 @@
 
 from ._annotation_config import AnnotationConfig, AnnotationContext
 from ._annotation_pattern_registry import (
-    AnnotatorMatchInfo as _AnnotatorMatchInfo,
-    SharedObserverModulePattern as _SharedObserverModulePattern,
+    AnnotatorMatchInfo,
+    SharedObserverModulePattern,
     _AnnotationPatternRegistry,
     _make_kv_cache_update_pattern,
 )
 from ._annotation_utils import (
     _get_input_qspec_map,
     adjust_output_qspec_for_qscheme_and_propagate,
-    annotate_module_level_specs as _annotate_module_level_specs,
+    annotate_module_level_specs,
     is_node_annotated,
 )
-from ._conv_bn_utils import (
-    fold_conv_bn_weights as _fold_conv_bn_weights,
-    remove_conv_bn_zeros_like_dtype as _remove_conv_bn_zeros_like_dtype,
-)
+from ._conv_bn_utils import fold_conv_bn_weights, remove_conv_bn_zeros_like_dtype
 from ._prepare_for_export import (
     _move_cache_dequant_to_output,
     prepare_for_mil_export,
     prepare_for_mlir_export,
 )
 from ._utils import (
-    force_per_tensor_for_channel_altering_ops as _force_per_tensor_for_channel_altering_ops,
-    get_source_module_name as _get_source_module_name,
-    remove_fake_quant_nodes as _remove_fake_quant_nodes,
-    restore_kwargs as _restore_kwargs,
-    strip_non_aten_metadata_kwargs as _strip_non_aten_metadata_kwargs,
+    force_per_tensor_for_channel_altering_ops,
+    get_source_module_name,
+    remove_fake_quant_nodes,
+    restore_kwargs,
+    strip_non_aten_metadata_kwargs,
 )
 
 logger = logging.getLogger(__name__)
@@ -160,7 +147,7 @@ class _RankedAnnotation(NamedTuple):
 
     node: torch.fx.Node
     config: OpQuantizerConfig
-    match: _AnnotatorMatchInfo
+    match: AnnotatorMatchInfo
     priority: int
 
 
@@ -232,7 +219,7 @@ def _get_shared_observer_nodes(self, model: torch.fx.GraphModule) -> set[torch.f
         shared_observer_annotators = [
             a_class
             for a_class in self._all_patterns()
-            if issubclass(a_class, _SharedObserverModulePattern)
+            if issubclass(a_class, SharedObserverModulePattern)
         ]
         shared_observer_nodes = set()
         for annotator in shared_observer_annotators:
@@ -285,7 +272,7 @@ def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
                 context,
             )
 
-        _annotate_module_level_specs(
+        annotate_module_level_specs(
             self._module_configs, self._module_name_to_state_names_map, model
         )
 
@@ -316,7 +303,7 @@ def _override_cache_op_annotations(
         for node in model.graph.nodes:
             if node.op != "call_function":
                 continue
-            op_type = _get_node_type(node, warn_on_failure=False)
+            op_type = get_node_type(node, warn_on_failure=False)
             kc = self._kv_cache_quant_configs.get(op_type)
             if kc is None:
                 continue
@@ -331,13 +318,13 @@ def _override_cache_op_annotations(
     def _match_all_annotators(
         self,
         model: torch.fx.GraphModule,
-    ) -> dict[torch.fx.Node, list[_AnnotatorMatchInfo]]:
+    ) -> dict[torch.fx.Node, list[AnnotatorMatchInfo]]:
         """
         Given an exported model, use all registered annotators to match nodes in the
         model. Build a dictionary mapping nodes to _AnnotatorMatchInfo objects
         containing information about matched patterns.
         """
-        node_to_annotator_match_info_dict: dict[torch.fx.Node, list[_AnnotatorMatchInfo]] = {}
+        node_to_annotator_match_info_dict: dict[torch.fx.Node, list[AnnotatorMatchInfo]] = {}
         for annotator_class in self._all_patterns():
             # Match patterns for this annotator across entire model
             annotator_node_match_dict = annotator_class._match_all_patterns(model)
@@ -356,8 +343,8 @@ def _match_all_annotators(
     def _sort_nodes_in_annotation_order(
         self,
         model: torch.fx.GraphModule,
-        node_to_annotator_match_info_dict: dict[torch.fx.Node, list[_AnnotatorMatchInfo]],
-    ) -> list[tuple[torch.fx.Node, OpQuantizerConfig, _AnnotatorMatchInfo]]:
+        node_to_annotator_match_info_dict: dict[torch.fx.Node, list[AnnotatorMatchInfo]],
+    ) -> list[tuple[torch.fx.Node, OpQuantizerConfig, AnnotatorMatchInfo]]:
         """
         Produce a list of tuples of (node, config to apply, annotator match info) which
         will later be iterated to annotate matches.
@@ -376,7 +363,7 @@ def _sort_nodes_in_annotation_order(
         )
 
         nodes_with_annotation_match_info: list[
-            tuple[torch.fx.Node, OpQuantizerConfig, _AnnotatorMatchInfo]
+            tuple[torch.fx.Node, OpQuantizerConfig, AnnotatorMatchInfo]
         ] = []
         for config_level_node_dict in config_level_node_dicts:
             for op_config_level in _OpConfigLevel.priority_order():
@@ -391,7 +378,7 @@ def _sort_nodes_in_annotation_order(
     def _get_config_level_node_dicts(
         self,
         model: torch.fx.GraphModule,
-        node_to_annotator_match_info_dict: dict[torch.fx.Node, list[_AnnotatorMatchInfo]],
+        node_to_annotator_match_info_dict: dict[torch.fx.Node, list[AnnotatorMatchInfo]],
     ) -> tuple[NodeConfigDict, NodeConfigDict, NodeConfigDict]:
         """
         Create and return three dicts corresponding to config levels for mapping nodes
@@ -411,9 +398,9 @@ def _get_config_level_node_dicts(
 
         # Precompute {canonical_key: insertion_index} once per config level so that
         # _set_config_to_use_for_node can look up a key's priority in O(1)
-        config_key_index: dict[_ConfigLevel, dict[object, int]] = {
+        config_key_index: dict[ConfigLevel, dict[object, int]] = {
             level: {key: idx for idx, key in enumerate(self._module_configs[level].keys())}
-            for level in (_ConfigLevel.MODULE_NAME, _ConfigLevel.MODULE_TYPE)
+            for level in (ConfigLevel.MODULE_NAME, ConfigLevel.MODULE_TYPE)
         }
 
         # Iterating through the nodes in topological ordering guarantees that when
@@ -426,8 +413,8 @@ def _get_config_level_node_dicts(
                 if self._set_config_to_use_for_node(
                     node,
                     module_name_node_config_dict,
-                    _ConfigLevel.MODULE_NAME,
-                    config_key_index[_ConfigLevel.MODULE_NAME],
+                    ConfigLevel.MODULE_NAME,
+                    config_key_index[ConfigLevel.MODULE_NAME],
                 ):
                     continue
 
@@ -435,14 +422,14 @@ def _get_config_level_node_dicts(
                 if self._set_config_to_use_for_node(
                     node,
                     module_type_node_config_dict,
-                    _ConfigLevel.MODULE_TYPE,
-                    config_key_index[_ConfigLevel.MODULE_TYPE],
+                    ConfigLevel.MODULE_TYPE,
+                    config_key_index[ConfigLevel.MODULE_TYPE],
                 ):
                     continue
 
                 # Take a shortcut for global since the config will be the same for all
                 # other modules (no need to name match).
-                global_config = list(self._module_configs[_ConfigLevel.GLOBAL].values())[0]
+                global_config = list(self._module_configs[ConfigLevel.GLOBAL].values())[0]
 
                 config, op_config_level = self._get_config_for_node(node, global_config)
                 # GLOBAL level has a single config, so priority is trivially 0.
@@ -456,7 +443,7 @@ def _set_config_to_use_for_node(
         self,
         node: torch.fx.Node,
         node_config_dict: NodeConfigDict,
-        config_level: _ConfigLevel,
+        config_level: ConfigLevel,
         config_key_index: dict[object, int],
     ) -> bool:
         """
@@ -468,7 +455,7 @@ def _set_config_to_use_for_node(
         within-level priority during the sort phase: lower index = higher
         precedence.
         """
-        qualified_name = _get_source_module_name(node)
+        qualified_name = get_source_module_name(node)
         if qualified_name is None:
             return False
 
@@ -505,7 +492,7 @@ def _get_config_for_node(
 
         # Check for op_type match
         for op_type, op_type_config in reversed(module_level_config.op_type_config.items()):
-            node_type = _get_node_type(node)
+            node_type = get_node_type(node)
             if node_type is not None and op_type == node_type:
                 return op_type_config, _OpConfigLevel.OP_TYPE
 
@@ -521,8 +508,8 @@ def _get_config_for_node(
     def _expand_and_sort_nodes_for_pattern_length(
         self,
         node_to_config_dict: dict[torch.fx.Node, _NodePriorityConfig],
-        node_to_annotator_match_info_dict: dict[torch.fx.Node, list[_AnnotatorMatchInfo]],
-    ) -> list[tuple[torch.fx.Node, OpQuantizerConfig, _AnnotatorMatchInfo]]:
+        node_to_annotator_match_info_dict: dict[torch.fx.Node, list[AnnotatorMatchInfo]],
+    ) -> list[tuple[torch.fx.Node, OpQuantizerConfig, AnnotatorMatchInfo]]:
         """
         Return a list of lists consisting of (node, config, annotator match info).
 
@@ -725,7 +712,7 @@ def _validate_config(config: QuantizerConfig) -> None:
         # identification.
         for spec in specs_to_check_for_index_or_wildcard:
             for key in spec:
-                if isinstance(key, str) and key != _ALL_TENSORS:
+                if isinstance(key, str) and key != ALL_TENSORS:
                     error_msg = (
                         "Only integer indices or '*' are supported for op and module "
                         f"input and output specs currently. Got {spec}"
@@ -736,7 +723,7 @@ def _validate_config(config: QuantizerConfig) -> None:
         # valid)
         for spec in specs_to_check_for_zero_or_wildcard:
             for key in spec:
-                if key not in [_ALL_TENSORS, 0]:
+                if key not in [ALL_TENSORS, 0]:
                     error_msg = (
                         "op_output_qspec currently supports setting for '*' or 0 "
                         f"tensor only. Got {spec}"
@@ -855,7 +842,7 @@ def _validate_kv_cache_quant_ops(
             matched = [
                 n
                 for n in exported_model.graph.nodes
-                if n.op == "call_function" and _get_node_type(n, warn_on_failure=False) == op
+                if n.op == "call_function" and get_node_type(n, warn_on_failure=False) == op
             ]
             if not matched:
                 raise ValueError(
@@ -975,7 +962,7 @@ def prepare(
                 preserved_attrs[attr] = getattr(self._model, attr)
 
         # Export the model to FX GraphModule
-        exported_model = _export_model(
+        exported_model = export_model(
             self._model, example_inputs, dynamic_shapes, export_with_no_grad
         )
 
@@ -986,15 +973,15 @@ def prepare(
         # torchao < 0.16.0 asserts annotated nodes have empty kwargs,
         # so we strip metadata kwargs from non-aten nodes before and restore after.
 
-        torchao_requires_empty_kwargs = not _version_ge(torchao, "0.16.0")
+        torchao_requires_empty_kwargs = not version_ge(torchao, "0.16.0")
         if torchao_requires_empty_kwargs:
-            saved_kwargs = _strip_non_aten_metadata_kwargs(exported_model.graph)
+            saved_kwargs = strip_non_aten_metadata_kwargs(exported_model.graph)
         try:
             prepared_model = prepare_qat_pt2e(exported_model, quantizer)
         except Exception as e:
             raise type(e)(f"prepare_qat_pt2e call failed, with error: {e}") from e
         if torchao_requires_empty_kwargs:
-            _restore_kwargs(prepared_model.graph, saved_kwargs)
+            restore_kwargs(prepared_model.graph, saved_kwargs)
 
         # Apply post-processing fixes to the prepared model
         self._postprocess_prepared_model(prepared_model)
@@ -1173,8 +1160,8 @@ def calibration_mode(self, model: torch.fx.GraphModule | None = None):
         # Enable observers; keep weight FQ on, disable activation FQ so observers
         # see the effect of quantized weights on activation ranges.
         self._model.apply(enable_observer)
-        self._model.apply(_enable_weight_fake_quant)
-        self._model.apply(_disable_activation_fake_quant)
+        self._model.apply(enable_weight_fake_quant)
+        self._model.apply(disable_activation_fake_quant)
 
         # Move model to eval mode and save original state
         with move_model_to_eval(self._model):
@@ -1284,7 +1271,7 @@ def _get_fake_quantize_modules(
             # (e.g. when the FQ feeds directly into the graph output node).
             source = None
             for neighbor in list(node.users) + list(node.args):
-                source = _get_source_module_name(neighbor)
+                source = get_source_module_name(neighbor)
                 if source is not None:
                     break
             if source is not None:
@@ -1304,17 +1291,17 @@ def _postprocess_prepared_model(model: torch.fx.GraphModule) -> None:
                 prepare_qat_pt2e().
         """
         # Fix dtype in Conv+BN decomposition for non-float32 models
-        _remove_conv_bn_zeros_like_dtype(model)
+        remove_conv_bn_zeros_like_dtype(model)
 
         # Channel-altering ops (flatten, reshape, transpose, etc.) share observers
         # with their inputs but invalidate axis semantics, so per-channel/per-block
         # granularity is mathematically incorrect. Force per-tensor.
-        _force_per_tensor_for_channel_altering_ops(model)
+        force_per_tensor_for_channel_altering_ops(model)
 
         # Apply weight axis defaults for per channel and per block quantization
         _apply_weight_axis_defaults(model)
 
-        _validate_activation_axes(model)
+        validate_activation_axes(model)
 
     @staticmethod
     def _remove_disabled_fake_quant_nodes(model: torch.fx.GraphModule) -> None:
@@ -1327,7 +1314,7 @@ def _remove_disabled_fake_quant_nodes(model: torch.fx.GraphModule) -> None:
                 if isinstance(mod, FakeQuantizeImplBase) and mod.is_disabled():
                     disabled_fq_nodes.add(node)
         if disabled_fq_nodes:
-            _remove_fake_quant_nodes(model, disabled_fq_nodes)
+            remove_fake_quant_nodes(model, disabled_fq_nodes)
 
     @staticmethod
     def _attach_preserved_attrs_to_model(
@@ -1359,4 +1346,4 @@ def _post_conversion_process(model: torch.fx.GraphModule) -> torch.fx.GraphModul
 
         """
         # Apply conv+bn folding
-        return _fold_conv_bn_weights(model)
+        return fold_conv_bn_weights(model)
diff --git a/src/coreai_opt/quantization/quantizer.py b/src/coreai_opt/quantization/quantizer.py
index 6d0de32..627cf14 100644
--- a/src/coreai_opt/quantization/quantizer.py
+++ b/src/coreai_opt/quantization/quantizer.py
@@ -23,7 +23,7 @@
 from coreai_opt._utils.export_utils import (
     validate_mmap_backend_and_device as _validate_mmap_backend_and_device,
 )
-from coreai_opt._utils.torch_utils import get_module_name
+from coreai_opt._utils.torch_utils import get_module_name as _get_module_name
 from coreai_opt.common import ExportBackend
 from coreai_opt.quantization._eager import EagerQuantizer as _EagerQuantizer
 from coreai_opt.quantization._graph import GraphQuantizer as _GraphQuantizer
@@ -287,7 +287,7 @@ def _maybe_apply_fn_to_fqs(self, fn: callable, module: nn.Module | None = None)
             self._quantizer._model.apply(fn)
             return
 
-        prefix = get_module_name(self._quantizer._model, module)
+        prefix = _get_module_name(self._quantizer._model, module)
         if prefix is None:
             raise ValueError(f"Module {module} is not a submodule of the prepared model.")
 
diff --git a/src/coreai_opt/quantization/spec/granularity.py b/src/coreai_opt/quantization/spec/granularity.py
index af0657a..3ee1dde 100644
--- a/src/coreai_opt/quantization/spec/granularity.py
+++ b/src/coreai_opt/quantization/spec/granularity.py
@@ -11,11 +11,11 @@
 import torch
 from pydantic import BaseModel, ConfigDict, Field, model_serializer
 
-from coreai_opt._utils.registry_utils import ConfigRegistryMixin
+from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin
 from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError
 
 
-class QuantizationGranularity(BaseModel, ConfigRegistryMixin):
+class QuantizationGranularity(BaseModel, _ConfigRegistryMixin):
     """
     Base class for quantization granularity specifications.
     """
diff --git a/src/coreai_opt/quantization/spec/range_calculator.py b/src/coreai_opt/quantization/spec/range_calculator.py
index 66a5adc..d021c99 100644
--- a/src/coreai_opt/quantization/spec/range_calculator.py
+++ b/src/coreai_opt/quantization/spec/range_calculator.py
@@ -9,12 +9,12 @@
 import torch.nn as nn
 from torchao.quantization.quant_primitives import _get_reduction_params
 
-from coreai_opt._utils.registry_utils import ClassRegistryMixin
+from coreai_opt._utils.registry_utils import ClassRegistryMixin as _ClassRegistryMixin
 
 from .granularity import QuantizationGranularity
 
 
-class RangeCalculatorBase(ClassRegistryMixin, nn.Module):
+class RangeCalculatorBase(_ClassRegistryMixin, nn.Module):
     """
     Base class and registry for classes used to compute the range
     of a given tensor.
diff --git a/src/coreai_opt/quantization/spec/spec.py b/src/coreai_opt/quantization/spec/spec.py
index 52a8372..95c28b9 100644
--- a/src/coreai_opt/quantization/spec/spec.py
+++ b/src/coreai_opt/quantization/spec/spec.py
@@ -19,7 +19,7 @@
 )
 
 from coreai_opt._utils.torch_utils import (
-    get_n_bits_from_dtype,
+    get_n_bits_from_dtype as _get_n_bits_from_dtype,
     is_float4_dtype as _is_float4_dtype,
 )
 from coreai_opt.config.spec import CompressionSpec, CompressionType
@@ -601,7 +601,7 @@ def get_n_bits_from_dtype(cls, dtype: torch.dtype) -> int:
         Raises:
             RuntimeError: If unable to extract bits from the dtype
         """
-        return get_n_bits_from_dtype(dtype)
+        return _get_n_bits_from_dtype(dtype)
 
     @classmethod
     def get_target_dtype(cls, dtype: torch.dtype) -> torch.dtype:
diff --git a/tests/devtools/test_check_internal_import_aliases.py b/tests/devtools/test_check_internal_import_aliases.py
new file mode 100644
index 0000000..90c6f55
--- /dev/null
+++ b/tests/devtools/test_check_internal_import_aliases.py
@@ -0,0 +1,236 @@
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-Clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+"""Golden (before, after) cases for the internal-import-alias pre-commit hook.
+
+Each case maps a short label to a ``(before, after)`` pair: the module source before the hook
+runs, and the source after ``--fix`` is applied. ``after == before`` means the hook makes no
+change — either the import already complies, it is out of scope, or the violation is real but
+cannot be auto-fixed (reported for a manual fix instead). The label doubles as the pytest id,
+so the two dicts read as a quick catalogue of everything the hook covers.
+
+The rule is direction-sensitive, so the cases are split by the *importing* module:
+- public modules (no ``_`` path segment) must alias private-module imports with a ``_`` prefix;
+- private modules (a ``_`` path segment) must not carry that redundant alias.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+from _test_helpers import load_script
+
+_SCRIPT = (
+    Path(__file__).parent.parent.parent
+    / "scripts"
+    / "pre_commit"
+    / "check_internal_import_aliases.py"
+)
+_hook = load_script(_SCRIPT)
+
+
+def _apply_fix(tmp_path: Path, before: str, *, private: bool) -> str:
+    """Write ``before`` into a public/private ``coreai_opt`` module; return it after ``--fix``."""
+    package_root = tmp_path / "coreai_opt"
+    module = (package_root / "_internal" / "mod.py") if private else (package_root / "mod.py")
+    module.parent.mkdir(parents=True, exist_ok=True)
+    module.write_text(before)
+    _hook.check_file(module, package_root, fix=True)
+    return module.read_text()
+
+
+# label -> (before, after) for a PUBLIC module importing from a private module.
+PUBLIC_MODULE_CASES: dict[str, tuple[str, str]] = {
+    "absolute-missing-alias": (
+        "from coreai_opt._utils import helper\n\nhelper()\n",
+        "from coreai_opt._utils import helper as _helper\n\n_helper()\n",
+    ),
+    "absolute-already-aliased-noop": (
+        "from coreai_opt._utils import helper as _helper\n\n_helper()\n",
+        "from coreai_opt._utils import helper as _helper\n\n_helper()\n",
+    ),
+    # A custom alias is prefixed, not replaced (the author's chosen name is kept).
+    "custom-alias-gets-prefixed": (
+        "from coreai_opt._utils import helper as h\n\nh()\n",
+        "from coreai_opt._utils import helper as _h\n\n_h()\n",
+    ),
+    "custom-alias-already-prefixed-noop": (
+        "from coreai_opt._utils import helper as _h\n\n_h()\n",
+        "from coreai_opt._utils import helper as _h\n\n_h()\n",
+    ),
+    "two-names-single-line": (
+        "from coreai_opt._utils import helper, other\n\nhelper()\nother()\n",
+        "from coreai_opt._utils import helper as _helper, other as _other\n\n_helper()\n_other()\n",
+    ),
+    # Adding the aliases pushes past the line limit, so the statement is paren-wrapped.
+    "long-import-wraps": (
+        "from coreai_opt._utils import clear_parametrization_original, prepare_mmap_dir\n\n"
+        "clear_parametrization_original()\nprepare_mmap_dir()\n",
+        "from coreai_opt._utils import (\n"
+        "    clear_parametrization_original as _clear_parametrization_original,\n"
+        "    prepare_mmap_dir as _prepare_mmap_dir,\n"
+        ")\n\n"
+        "_clear_parametrization_original()\n_prepare_mmap_dir()\n",
+    ),
+    # Relative import of a private module -> aliased, leading dots preserved.
+    "relative-import": (
+        "from ._utils import helper\n\nhelper()\n",
+        "from ._utils import helper as _helper\n\n_helper()\n",
+    ),
+    "relative-deeper-private": (
+        "from .sub._priv import helper\n\nhelper()\n",
+        "from .sub._priv import helper as _helper\n\n_helper()\n",
+    ),
+    # Relative import of a public sibling -> unchanged (target is not private).
+    "relative-public-sibling-noop": (
+        "from .utils import helper\n\nhelper()\n",
+        "from .utils import helper\n\nhelper()\n",
+    ),
+    # Plain `import ... as name` binding a private module to a public name -> prefixed.
+    "plain-import-aliased": (
+        "import coreai_opt._utils.helpers as helpers\n\nhelpers.run()\n",
+        "import coreai_opt._utils.helpers as _helpers\n\n_helpers.run()\n",
+    ),
+    # Bare `import ...` binds the public top-level package `coreai_opt` -> unchanged.
+    "plain-import-bare-noop": (
+        "import coreai_opt._utils.helpers\n\ncoreai_opt._utils.helpers.run()\n",
+        "import coreai_opt._utils.helpers\n\ncoreai_opt._utils.helpers.run()\n",
+    ),
+    "public-module-import-noop": (
+        "from coreai_opt.utils import helper\n\nhelper()\n",
+        "from coreai_opt.utils import helper\n\nhelper()\n",
+    ),
+    # A different top-level package (prefix look-alike) is not our package -> unchanged.
+    "sibling-top-level-package-noop": (
+        "from coreai_opt_benchmarking._bench import helper\n\nhelper()\n",
+        "from coreai_opt_benchmarking._bench import helper\n\nhelper()\n",
+    ),
+    "imported-symbol-already-private-noop": (
+        "from coreai_opt._utils import _internal_helper\n\n_internal_helper()\n",
+        "from coreai_opt._utils import _internal_helper\n\n_internal_helper()\n",
+    ),
+    # Trailing comment sits outside the statement span, so it survives the rewrite.
+    "single-line-trailing-comment": (
+        "from coreai_opt._utils import helper  # noqa: F401\n\nhelper()\n",
+        "from coreai_opt._utils import helper as _helper  # noqa: F401\n\n_helper()\n",
+    ),
+    # Comment inside a multi-line import -> aliased in place so the comment is preserved.
+    "multiline-inner-comment": (
+        "from coreai_opt._utils import (\n"
+        "    helper,  # the sole public entry point\n"
+        "    other,\n"
+        ")\n\n"
+        "helper()\nother()\n",
+        "from coreai_opt._utils import (\n"
+        "    helper as _helper,  # the sole public entry point\n"
+        "    other as _other,\n"
+        ")\n\n"
+        "_helper()\n_other()\n",
+    ),
+    # Bound name shadowed by a function parameter -> reported but NOT auto-fixed.
+    "shadowed-param-not-fixed": (
+        "from coreai_opt._utils import helper\n\n\ndef run(helper):\n    return helper\n",
+        "from coreai_opt._utils import helper\n\n\ndef run(helper):\n    return helper\n",
+    ),
+    # Bound name rebound at module scope -> unsafe to rename, so reported but NOT auto-fixed.
+    "module-level-rebind-not-fixed": (
+        "from coreai_opt._utils import helper\n\nhelper = helper or None\n",
+        "from coreai_opt._utils import helper\n\nhelper = helper or None\n",
+    ),
+    # Multibyte character earlier on the line -> rename lands at the correct byte offset.
+    "multibyte-offset": (
+        "from coreai_opt._utils import helper\n\n\n"
+        "def run():\n    café = 1\n    return café + helper()\n",
+        "from coreai_opt._utils import helper as _helper\n\n\n"
+        "def run():\n    café = 1\n    return café + _helper()\n",
+    ),
+}
+
+# label -> (before, after) for a PRIVATE module importing from a private module.
+PRIVATE_MODULE_CASES: dict[str, tuple[str, str]] = {
+    # Redundant canonical `_` alias -> stripped, references renamed back.
+    "strip-redundant-alias": (
+        "from coreai_opt._utils import helper as _helper\n\n_helper()\n",
+        "from coreai_opt._utils import helper\n\nhelper()\n",
+    ),
+    "no-alias-noop": (
+        "from coreai_opt._utils import helper\n\nhelper()\n",
+        "from coreai_opt._utils import helper\n\nhelper()\n",
+    ),
+    # A deliberate custom name is left alone in a private module.
+    "custom-alias-noop": (
+        "from coreai_opt._utils import helper as h\n\nh()\n",
+        "from coreai_opt._utils import helper as h\n\nh()\n",
+    ),
+    "relative-strip-alias": (
+        "from ._sibling import helper as _helper\n\n_helper()\n",
+        "from ._sibling import helper\n\nhelper()\n",
+    ),
+    # Alias stripped in place inside a multi-line import so the comment is preserved.
+    "multiline-inner-comment-strip": (
+        "from coreai_opt._utils import (\n"
+        "    helper as _helper,  # note\n"
+        "    other as _other,\n"
+        ")\n\n"
+        "_helper()\n_other()\n",
+        "from coreai_opt._utils import (\n"
+        "    helper,  # note\n"
+        "    other,\n"
+        ")\n\n"
+        "helper()\nother()\n",
+    ),
+}
+
+
+@pytest.mark.parametrize(
+    ("before", "after"),
+    PUBLIC_MODULE_CASES.values(),
+    ids=PUBLIC_MODULE_CASES.keys(),
+)
+def test_public_module_cases(tmp_path: Path, before: str, after: str) -> None:
+    assert _apply_fix(tmp_path, before, private=False) == after
+
+
+@pytest.mark.parametrize(
+    ("before", "after"),
+    PRIVATE_MODULE_CASES.values(),
+    ids=PRIVATE_MODULE_CASES.keys(),
+)
+def test_private_module_cases(tmp_path: Path, before: str, after: str) -> None:
+    assert _apply_fix(tmp_path, before, private=True) == after
+
+
+def _write_module(tmp_path: Path, source: str) -> Path:
+    """Write ``source`` into a public ``coreai_opt`` module under ``tmp_path``."""
+    package_root = tmp_path / "coreai_opt"
+    package_root.mkdir(parents=True, exist_ok=True)
+    (package_root / "mod.py").write_text(source)
+    return package_root
+
+
+def test_detection_fails_and_reports(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None:
+    """Detection (no --fix) exits non-zero on a violation so the commit is blocked."""
+    root = _write_module(tmp_path, "from coreai_opt._utils import helper\n\nhelper()\n")
+    exit_code = _hook.main(["--root", str(root)])
+    assert exit_code == 1
+    assert "Missing `_` alias" in capsys.readouterr().out
+
+
+def test_detection_passes_when_clean(tmp_path: Path) -> None:
+    """Detection exits 0 when there is nothing to flag."""
+    root = _write_module(tmp_path, "value = 1\n")
+    assert _hook.main(["--root", str(root)]) == 0
+
+
+def test_fix_resolves_and_then_passes(tmp_path: Path) -> None:
+    """--fix applies the fix (exit 0); a following detection run is then clean."""
+    root = _write_module(tmp_path, "from coreai_opt._utils import helper\n\nhelper()\n")
+    assert _hook.main(["--root", str(root)]) == 1  # blocks before fixing
+    assert _hook.main(["--fix", "--root", str(root)]) == 0  # applies the fix
+    assert _hook.main(["--root", str(root)]) == 0  # clean afterward
+    assert (root / "mod.py").read_text() == (
+        "from coreai_opt._utils import helper as _helper\n\n_helper()\n"
+    )