feat(integrations): Add MLX array materializer (#4027)

* feat(integrations): Add MLX array materializer

This enables saving and loading MLX arrays when using ZenML on Apple platforms.

Structurally, this is very similar to the recently added JAX integration, it's
just the array type that is different, plus we're using MLX's own array IO, so
there is no hard dependency on NumPy.

* tests: Add MLX array materializer test

Modeled after the existing tests, we test that the save-load round trip preserves
information and array values.

* typing: Force mypy to recognize the mx.load() output as single array

* fix: Restrict materializer to Apple + Linux, implement requirement hooks

Since Linux does not support Apple GPU computing, we need to install mlx with
the "cpu" extra on Linux platforms.

Author: nicholasjng

src/zenml/integrations/constants.py

@@ -43,6 +43,7 @@
 LIGHTGBM = "lightgbm"
 # LLAMA_INDEX = "llama_index"
 MLFLOW = "mlflow"
+MLX = "mlx"
 MODAL = "modal"
 NEPTUNE = "neptune"
 NEURAL_PROPHET = "neural_prophet"

src/zenml/integrations/mlx/__init__.py

@@ -0,0 +1,57 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Initialization of the MLX integration."""
+
+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")
+
+
+class MLXIntegration(Integration):
+    """Definition of MLX array integration for ZenML."""
+
+    NAME = MLX
+
+    @classmethod
+    def check_installation(cls) -> bool:
+        if sys.platform not in SUPPORTED_PLATFORMS:
+            return False
+        return super().check_installation()
+
+    @classmethod
+    def get_requirements(
+        cls,
+        target_os: Optional[str] = None,
+        python_version: Optional[str] = None,
+    ) -> List[str]:
+        # sys.platform is "darwin", while platform.system() is "Darwin",
+        # similarly on Linux.
+        target_os = (target_os or sys.platform).lower()
+        if target_os == "darwin":
+            return ["mlx"]
+        elif target_os == "linux":
+            return ["mlx[cpu]"]
+        else:
+            return []
+
+    @classmethod
+    def activate(cls) -> None:
+        """Activates the integration."""
+        from zenml.integrations.mlx import materializer  # noqa

src/zenml/integrations/mlx/materializer.py

@@ -0,0 +1,64 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Implementation of the ZenML MLX materializer."""
+
+import os
+from typing import (
+    Any,
+    ClassVar,
+    Tuple,
+    Type,
+)
+
+import mlx.core as mx
+
+from zenml.enums import ArtifactType
+from zenml.materializers.base_materializer import BaseMaterializer
+
+NUMPY_FILENAME = "data.npy"
+
+
+class MLXArrayMaterializer(BaseMaterializer):
+    """A materializer for MLX arrays."""
+
+    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (mx.array,)
+    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
+
+    def load(self, data_type: Type[Any]) -> mx.array:
+        """Reads data from a `.npy` file, and returns an MLX array.
+
+        Args:
+            data_type: The type of the data to read.
+
+        Returns:
+            The MLX array.
+        """
+        numpy_file = os.path.join(self.uri, NUMPY_FILENAME)
+
+        with self.artifact_store.open(numpy_file, "rb") as f:
+            # loading an .npy file always results in a single array.
+            arr = mx.load(f)
+            assert isinstance(arr, mx.array)
+            return arr
+
+    def save(self, data: mx.array) -> None:
+        """Writes an MLX array to the artifact store as a `.npy` file.
+
+        Args:
+            data: The MLX array to write.
+        """
+        with self.artifact_store.open(
+            os.path.join(self.uri, NUMPY_FILENAME), "wb"
+        ) as f:
+            mx.save(f, data)

tests/integration/integrations/mlx/__init__.py

@@ -0,0 +1,13 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.

tests/integration/integrations/mlx/test_mlx_array_materializer.py

@@ -0,0 +1,38 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  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
+
+
+@pytest.mark.skipif(
+    sys.platform not in SUPPORTED_PLATFORMS,
+    reason="MLX only runs on Apple and Linux",
+)
+def test_mlx_array_materializer():
+    """Test the MLX array materializer."""
+    arr = mx.ones(5)
+
+    result = _test_materializer(
+        step_output_type=mx.array,
+        materializer_class=MLXArrayMaterializer,
+        step_output=arr,
+    )
+    assert mx.allclose(arr, result).item()