Prevent MLX integration install on Intel Macs (#4039)

* Prevent MLX integration install on Intel Macs

MLX wheels are only available for Apple Silicon (ARM64) Macs, not
Intel (x86_64). Added architecture detection using platform.machine()
to check_installation() and get_requirements() to prevent installation
attempts on unsupported Intel Macs.

The integration now properly detects:
- Apple Silicon Macs (arm64/aarch64): Supported
- Intel Macs (x86_64): Not supported
- Linux (any arch): Supported

* Remove unused SUPPORTED_PLATFORMS constant from MLX integration

The SUPPORTED_PLATFORMS constant is no longer needed since we now use
the _is_supported_platform() function which provides more granular
architecture checking (distinguishing Apple Silicon from Intel Macs).

Updated the test to import and use _is_supported_platform() directly
instead of the old SUPPORTED_PLATFORMS constant. This ensures the test
correctly skips on Intel Macs while running on Apple Silicon and Linux.

* Add mlx to mypy overrides

* Move imports for windows test env

Author: strickvl

pyproject.toml

@@ -370,6 +370,8 @@ module = [
     "numba.*",
     "uvloop.*",
     "litellm",
+    "mlx",
+    "mlx.*",
     "jsonref",
 ]
 ignore_missing_imports = true

src/zenml/integrations/mlx/__init__.py

@@ -13,15 +13,30 @@
 #  permissions and limitations under the License.
 """Initialization of the MLX integration."""
 
+import platform
 import sys
 from typing import List, Optional
 
 from zenml.integrations.constants import MLX
 from zenml.integrations.integration import Integration
 
-# MLX can run on Linux, but only with CPU or Cuda devices,
-# see https://ml-explore.github.io/mlx/build/html/install.html#cuda ff.
-SUPPORTED_PLATFORMS = ("darwin", "linux")
+def _is_supported_platform() -> bool:
+    """Check if the current platform supports MLX.
+    
+    MLX requires:
+    - macOS with ARM64 (Apple Silicon)
+    - Linux (any architecture)
+    
+    Returns:
+        True if platform is supported, False otherwise.
+    """
+    if sys.platform == "linux":
+        return True
+    elif sys.platform == "darwin":
+        # MLX only supports Apple Silicon Macs, not Intel
+        machine = platform.machine().lower()
+        return machine in ("arm64", "aarch64")
+    return False
 
 
 class MLXIntegration(Integration):
@@ -31,9 +46,7 @@ class MLXIntegration(Integration):
 
     @classmethod
     def check_installation(cls) -> bool:
-        if sys.platform not in SUPPORTED_PLATFORMS:
-            return False
-        return super().check_installation()
+        return False if not _is_supported_platform() else super().check_installation()
 
     @classmethod
     def get_requirements(
@@ -45,7 +58,9 @@ def get_requirements(
         # similarly on Linux.
         target_os = (target_os or sys.platform).lower()
         if target_os == "darwin":
-            return ["mlx"]
+            # Only return requirements if on Apple Silicon
+            machine = platform.machine().lower()
+            return ["mlx"] if machine in ("arm64", "aarch64") else []
         elif target_os == "linux":
             return ["mlx[cpu]"]
         else:

tests/integration/integrations/mlx/test_mlx_array_materializer.py

@@ -12,22 +12,22 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 
-import sys
-
-import mlx.core as mx
 import pytest
 
 from tests.unit.test_general import _test_materializer
-from zenml.integrations.mlx import SUPPORTED_PLATFORMS
-from zenml.integrations.mlx.materializer import MLXArrayMaterializer
+from zenml.integrations.mlx import _is_supported_platform
 
 
 @pytest.mark.skipif(
-    sys.platform not in SUPPORTED_PLATFORMS,
-    reason="MLX only runs on Apple and Linux",
+    not _is_supported_platform(),
+    reason="MLX only runs on Apple Silicon and Linux (not Intel Macs)",
 )
 def test_mlx_array_materializer():
     """Test the MLX array materializer."""
+    import mlx.core as mx
+
+    from zenml.integrations.mlx.materializer import MLXArrayMaterializer
+
     arr = mx.ones(5)
 
     result = _test_materializer(