diff --git a/changelog.d/63.added b/changelog.d/63.added new file mode 100644 index 0000000..1c0b677 --- /dev/null +++ b/changelog.d/63.added @@ -0,0 +1 @@ +Added support for torchao 0.18.0. diff --git a/pyproject.toml b/pyproject.toml index c9eea89..2a172f9 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/src/coreai_opt/__init__.py b/src/coreai_opt/__init__.py index 34c1d4c..ffcbe9e 100644 --- a/src/coreai_opt/__init__.py +++ b/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", diff --git a/src/coreai_opt/_utils/version_utils.py b/src/coreai_opt/_utils/version_utils.py index 14459de..f145d43 100644 --- a/src/coreai_opt/_utils/version_utils.py +++ b/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}" + ) diff --git a/tests/test_utils/test_version_utils.py b/tests/test_utils/test_version_utils.py new file mode 100644 index 0000000..6880972 --- /dev/null +++ b/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