mlx poc

Author: williambdean

pytensor/compile/mode.py

@@ -27,6 +27,7 @@
 from pytensor.link.basic import Linker, PerformLinker
 from pytensor.link.c.basic import CLinker, OpWiseCLinker
 from pytensor.link.jax.linker import JAXLinker
+from pytensor.link.mlx.linker import MLXLinker
 from pytensor.link.numba.linker import NumbaLinker
 from pytensor.link.pytorch.linker import PytorchLinker
 from pytensor.link.vm import VMLinker
@@ -50,6 +51,7 @@
     "jax": JAXLinker(),
     "pytorch": PytorchLinker(),
     "numba": NumbaLinker(),
+    "mlx": MLXLinker(),
 }
 
 
@@ -494,13 +496,28 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     ),
 )
 
+MLX = Mode(
+    MLXLinker(),
+    RewriteDatabaseQuery(
+        include=["fast_run"],
+        exclude=[
+            "cxx_only",
+            "BlasOpt",
+            "fusion",
+            "inplace",
+            "scan_save_mem_prealloc",
+        ],
+    ),
+)
+
 
 predefined_modes = {
     "FAST_COMPILE": FAST_COMPILE,
     "FAST_RUN": FAST_RUN,
     "JAX": JAX,
     "NUMBA": NUMBA,
     "PYTORCH": PYTORCH,
+    "MLX": MLX,
 }
 
 _CACHED_RUNTIME_MODES: dict[str, Mode] = {}
@@ -585,6 +602,8 @@ def get_target_language(mode=None) -> tuple[Literal["py", "c", "numba", "jax"],
         return ("py",)
     if isinstance(linker, CLinker):
         return ("c",)
+    if isinstance(linker, MLXLinker):
+        return ("py",)
 
     if isinstance(linker, VMLinker | OpWiseCLinker):
         return ("c", "py") if config.cxx else ("py",)

pytensor/link/mlx/dispatch/__init__.py

@@ -0,0 +1,5 @@
+# isort: off
+from pytensor.link.mlx.dispatch.basic import mlx_funcify, mlx_typify
+
+import pytensor.link.mlx.dispatch.math
+# isort: on

pytensor/link/mlx/dispatch/basic.py

@@ -0,0 +1,61 @@
+from functools import singledispatch
+from types import NoneType
+
+import mlx.core as mx
+import numpy as np
+
+from pytensor.compile.ops import DeepCopyOp
+from pytensor.graph.fg import FunctionGraph
+from pytensor.link.utils import fgraph_to_python
+
+
+@singledispatch
+def mlx_typify(data, **kwargs):
+    raise NotImplementedError(f"mlx_typify is not implemented for {type(data)}")
+
+
+@mlx_typify.register(np.ndarray)
+@mlx_typify.register(mx.array)
+def mlx_typify_tensor(data, dtype=None, **kwargs):
+    return mx.array(data, dtype=dtype)
+
+
+@mlx_typify.register(slice)
+@mlx_typify.register(NoneType)
+@mlx_typify.register(np.number)
+def mlx_typify_no_conversion_needed(data, **kwargs):
+    return data
+
+
+@singledispatch
+def mlx_funcify(op, node=None, storage_map=None, **kwargs):
+    """Create a MLX compatible function from an PyTensor `Op`."""
+    raise NotImplementedError(
+        f"No MLX conversion for the given `Op`: {op}.\nCheck out `https://github.com/pymc-devs/pytensor/issues/1350` for progress or to request we prioritize this operation"
+    )
+
+
+@mlx_funcify.register(FunctionGraph)
+def mlx_funcify_FunctionGraph(
+    fgraph,
+    node=None,
+    fgraph_name="mlx_funcified_fgraph",
+    conversion_func=mlx_funcify,
+    **kwargs,
+):
+    built_kwargs = {"conversion_func": conversion_func, **kwargs}
+    return fgraph_to_python(
+        fgraph,
+        conversion_func,
+        type_conversion_fn=mlx_typify,
+        fgraph_name=fgraph_name,
+        **built_kwargs,
+    )
+
+
+@mlx_funcify.register(DeepCopyOp)
+def mlx_funcify_DeepCopyOp(op, **kwargs):
+    def deepcopyop(x):
+        return x.copy()
+
+    return deepcopyop

pytensor/link/mlx/dispatch/math.py

@@ -0,0 +1,12 @@
+import mlx.core as mx
+
+from pytensor.link.mlx.dispatch import mlx_funcify
+from pytensor.tensor.math import Dot
+
+
+@mlx_funcify.register(Dot)
+def mlx_funcify_Dot(op, **kwargs):
+    def dot(x, y):
+        return mx.matmul(x, y)
+
+    return dot

pytensor/link/mlx/linker.py

@@ -0,0 +1,113 @@
+from pytensor.link.basic import JITLinker
+from pytensor.link.utils import unique_name_generator
+
+
+class MLXLinker(JITLinker):
+    """A `Linker` that JIT-compiles NumPy-based operations using Apple's MLX."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gen_functors = []
+
+    def fgraph_convert(
+        self,
+        fgraph,
+        order,
+        input_storage,
+        output_storage,
+        storage_map,
+        **kwargs,
+    ):
+        """Convert a PyTensor FunctionGraph to an MLX-compatible function.
+
+        Parameters
+        ----------
+        fgraph : FunctionGraph
+            The function graph to convert
+        order : list
+            The order in which to compute the nodes
+        input_storage : list
+            Storage for the input variables
+        output_storage : list
+            Storage for the output variables
+        storage_map : dict
+            Map from variables to their storage
+
+        Returns
+        -------
+        callable
+            An MLX-compatible function
+        """
+        from pytensor.link.mlx.dispatch import mlx_funcify
+
+        # We want to have globally unique names
+        # across the entire pytensor graph, not
+        # just the subgraph
+        generator = unique_name_generator(["mlx_linker"])
+
+        # Ensure that torch is aware of the generated
+        # code so we can compile without graph breaks
+        def conversion_func_register(*args, **kwargs):
+            functor = mlx_funcify(*args, **kwargs)
+            name = kwargs["unique_name"](functor)
+            self.gen_functors.append((f"_{name}", functor))
+            return functor
+
+        built_kwargs = {
+            "unique_name": generator,
+            "conversion_func": conversion_func_register,
+            **kwargs,
+        }
+        return mlx_funcify(
+            fgraph,
+            input_storage=input_storage,
+            storage_map=storage_map,
+            **built_kwargs,
+        )
+
+    def jit_compile(self, fn):
+        """JIT compile an MLX function.
+
+        Parameters
+        ----------
+        fn : callable
+            The function to compile
+
+        Returns
+        -------
+        callable
+            The compiled function
+        """
+        import mlx.core as mx
+
+        return mx.compile(fn)
+
+    def create_thunk_inputs(self, storage_map):
+        """Create inputs for the MLX thunk.
+
+        Parameters
+        ----------
+        storage_map : dict
+            Map from variables to their storage
+
+        Returns
+        -------
+        list
+            The inputs for the thunk
+        """
+        from numpy.random import Generator, RandomState
+
+        from pytensor.link.mlx.dispatch import mlx_typify
+
+        thunk_inputs = []
+        for n in self.fgraph.inputs:
+            sinput = storage_map[n]
+            # Handle random number generators specially
+            if isinstance(sinput[0], RandomState | Generator):
+                new_value = mlx_typify(
+                    sinput[0], dtype=getattr(sinput[0], "dtype", None)
+                )
+                sinput[0] = new_value
+            thunk_inputs.append(sinput)
+
+        return thunk_inputs

pytensor/link/pytorch/linker.py

@@ -31,16 +31,16 @@ def conversion_func_register(*args, **kwargs):
             **kwargs,
         }
         return pytorch_funcify(
-            fgraph, input_storage=input_storage, storage_map=storage_map, **built_kwargs
+            fgraph,
+            input_storage=input_storage,
+            storage_map=storage_map,
+            **built_kwargs,
         )
 
     def jit_compile(self, fn):
-        import torch
+        import mlx.core as mx
 
-        # flag that tend to help our graphs
-        torch._dynamo.config.capture_dynamic_output_shape_ops = True
-
-        from pytensor.link.pytorch.dispatch import pytorch_typify
+        from pytensor.link.mlx.dispatch import mlx_typify
 
         class wrapper:
             """
@@ -54,7 +54,7 @@ class wrapper:
             """
 
             def __init__(self, fn, gen_functors):
-                self.fn = torch.compile(fn)
+                self.fn = mx.compile(fn)
                 self.gen_functors = gen_functors.copy()
 
             def __call__(self, *inputs, **kwargs):
@@ -65,7 +65,7 @@ def __call__(self, *inputs, **kwargs):
                     setattr(pytensor.link.utils, n[1:], fn)
 
                 # Torch does not accept numpy inputs and may return GPU objects
-                outs = self.fn(*(pytorch_typify(inp) for inp in inputs), **kwargs)
+                outs = self.fn(*(mlx_typify(inp) for inp in inputs), **kwargs)
 
                 # unset attrs
                 for n, _ in self.gen_functors: