Skip to content

Commit 99d596f

Browse files
committed
feat: warn on import for unsupported torchao/torch pairs
Read torchao's version via importlib.metadata rather than torchao.__version__: on torch < 2.10, importing torchao >= 0.18.0 raises ImportError, so the warning must be emitted before the imports below reach torchao.
1 parent 664c1ea commit 99d596f

1 file changed

Lines changed: 18 additions & 3 deletions

File tree

src/coreai_opt/__init__.py

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,24 @@
88
For deployment via Core AI on Apple Silicon.
99
"""
1010

11-
from . import palettization, pruning, quantization
12-
from ._about import __version__
13-
from .common import CoreMLExportError, ExportBackend
11+
import importlib.metadata
12+
import warnings
13+
14+
import torch
15+
16+
from coreai_opt._utils.version_utils import (
17+
torchao_torch_incompatibility as _torchao_torch_incompatibility,
18+
)
19+
20+
_incompatibility = _torchao_torch_incompatibility(
21+
importlib.metadata.version("torchao"), torch.__version__
22+
)
23+
if _incompatibility:
24+
warnings.warn(_incompatibility, UserWarning, stacklevel=2)
25+
26+
from . import palettization, pruning, quantization # noqa: E402
27+
from ._about import __version__ # noqa: E402
28+
from .common import CoreMLExportError, ExportBackend # noqa: E402
1429

1530
__all__ = [
1631
"CoreMLExportError",

0 commit comments

Comments
 (0)