build: update to latest torchao version

changelog.d/63.added

@@ -0,0 +1 @@
+Added support for torchao 0.18.0.

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     # kmeans1d C++ core (a C++ toolchain must also be present on the host).
     "ninja>=1.11",
     "numpy>=2",
+    "packaging>=23.0",
     "pydantic>=2.0.0",
     "pyyaml>=6.0",
     "rich>=13.0.0",
@@ -48,7 +49,7 @@ dependencies = [
     # version >= 0.15.0 for torch version >= 2.9.0. Opting option 1.
     # These torch versions must be in bounds of torch_2_8, torch_2_9, torch_2_10, and torch_2_11
     "torch>=2.8.0,<=2.11.0",
-    "torchao>=0.15.0,<=0.17.0",
+    "torchao>=0.15.0,<=0.18.0",
     "tqdm>=4.65",
 ]
 [[project.authors]]
@@ -138,7 +139,7 @@ torch_2_10 = [
 ]
 torch_2_11 = [
     "torch==2.11.0",
-    "torchao==0.17.0",
+    "torchao==0.18.0",
     "torchvision==0.26.0",
 ]
 torch_2_8 = [

src/coreai_opt/__init__.py

@@ -8,9 +8,24 @@
 For deployment via Core AI on Apple Silicon.
 """
 
-from . import palettization, pruning, quantization
-from ._about import __version__
-from .common import CoreMLExportError, ExportBackend
+import importlib.metadata
+import warnings
+
+import torch
+
+from coreai_opt._utils.version_utils import (
+    torchao_torch_incompatibility as _torchao_torch_incompatibility,
+)
+
+_incompatibility = _torchao_torch_incompatibility(
+    importlib.metadata.version("torchao"), torch.__version__
+)
+if _incompatibility:
+    warnings.warn(_incompatibility, UserWarning, stacklevel=2)
+
+from . import palettization, pruning, quantization  # noqa: E402
+from ._about import __version__  # noqa: E402
+from .common import CoreMLExportError, ExportBackend  # noqa: E402
 
 __all__ = [
     "CoreMLExportError",

src/coreai_opt/_utils/version_utils.py

@@ -3,8 +3,37 @@
 # 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
 
+from types import ModuleType
+
 from packaging import version
 
 
-def version_ge(module, target_version):
+def version_ge(module: ModuleType, target_version: str) -> bool:
     return version.parse(module.__version__) >= version.parse(target_version)
+
+
+_MIN_TORCHAO_REQUIRING_TORCH_2_11 = "0.18.0"
+_MIN_TORCH_FOR_NEW_TORCHAO = "2.11.0.dev0"
+_TORCHAO_RELEASE_NOTES_URL = "https://github.com/pytorch/ao/releases/tag/v0.18.0"
+
+
+def torchao_torch_incompatibility(torchao_version: str, torch_version: str) -> str | None:
+    """Describe why the installed torchao and torch versions are incompatible.
+
+    Args:
+        torchao_version: The installed torchao version.
+        torch_version: The installed torch version.
+
+    Returns:
+        A message explaining the incompatibility, or ``None`` if the pair is supported.
+    """
+    if version.parse(torchao_version) < version.parse(_MIN_TORCHAO_REQUIRING_TORCH_2_11):
+        return None
+    if version.parse(torch_version) >= version.parse(_MIN_TORCH_FOR_NEW_TORCHAO):
+        return None
+    return (
+        f"torchao {torchao_version} does not support torch<2.11 "
+        f"(found torch {torch_version}). See the torchao "
+        f"{_MIN_TORCHAO_REQUIRING_TORCH_2_11} release notes for more information: "
+        f"{_TORCHAO_RELEASE_NOTES_URL}"
+    )

tests/test_utils/test_version_utils.py

@@ -0,0 +1,46 @@
+# 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
+
+import pytest
+
+from coreai_opt._utils.version_utils import torchao_torch_incompatibility
+
+INCOMPATIBLE = [
+    ("0.18.0", "2.8.0"),
+    ("0.18.0", "2.9.1"),
+    ("0.18.0", "2.10.0"),
+    ("0.18.0", "2.10.0+cu128"),
+    ("0.19.0", "2.10.0"),
+    # A source-built torchao reports a local version, which sorts above the base.
+    ("0.18.0+gitabc1234", "2.10.0"),
+]
+
+COMPATIBLE = [
+    # torch is new enough.
+    ("0.18.0", "2.11.0"),
+    ("0.18.0", "2.11.0+cu128"),
+    ("0.18.0", "2.12.0.dev20260805+cu128"),
+    # A 2.11 pre-release counts as 2.11.
+    ("0.18.0", "2.11.0rc1"),
+    # torchao still supports older torch.
+    ("0.17.0", "2.8.0"),
+    ("0.16.0", "2.10.0"),
+    ("0.15.0", "2.8.0"),
+]
+
+
+@pytest.mark.parametrize(("torchao_version", "torch_version"), INCOMPATIBLE)
+def test_returns_message_for_incompatible_pair(torchao_version, torch_version):
+    message = torchao_torch_incompatibility(torchao_version, torch_version)
+
+    assert message is not None
+    assert torchao_version in message
+    assert torch_version in message
+    assert "https://github.com/pytorch/ao/releases/tag/v0.18.0" in message
+
+
+@pytest.mark.parametrize(("torchao_version", "torch_version"), COMPATIBLE)
+def test_returns_none_for_compatible_pair(torchao_version, torch_version):
+    assert torchao_torch_incompatibility(torchao_version, torch_version) is None