Add MLX backend

.github/workflows/test.yml

@@ -81,6 +81,7 @@ jobs:
  install-numba: [0]
  install-jax: [0]
  install-torch: [0]
+ install-mlx: [0]
  install-xarray: [0]
  part:
  - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/xtensor"
@@ -106,6 +107,7 @@ jobs:
  install-numba: 0
  install-jax: 0
  install-torch: 0
+ install-mlx: 0
  install-xarray: 0
  - install-numba: 1
  os: "ubuntu-latest"
@@ -149,7 +151,16 @@ jobs:
  fast-compile: 0
  float32: 0
  part: "tests/xtensor"
- - os: macos-15
+ - os: "macos-15"
+ python-version: "3.11"
+ fast-compile: 0
+ float32: 0
+ install-mlx: 1
+ install-numba: 0
+ install-jax: 0
+ install-torch: 0
+ part: "tests/link/mlx"
+ - os: "macos-15"
  python-version: "3.13"
  fast-compile: 0
  float32: 0
@@ -194,6 +205,7 @@ jobs:
  if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" "numba>=0.57"; fi
  if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" jax jaxlib numpyro equinox && pip install tfp-nightly; fi
  if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" pytorch pytorch-cuda=12.1 "mkl<=2024.0" -c pytorch -c nvidia; fi
+ if [[ $INSTALL_MLX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" mlx; fi
  if [[ $INSTALL_XARRAY == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" xarray xarray-einstats; fi
 
  pip install -e ./
@@ -210,6 +222,7 @@ jobs:
  INSTALL_JAX: ${{ matrix.install-jax }}
  INSTALL_TORCH: ${{ matrix.install-torch}}
  INSTALL_XARRAY: ${{ matrix.install-xarray }}
+ INSTALL_MLX: ${{ matrix.install-mlx }}
  OS: ${{ matrix.os}}
 
  - name: Run tests

.gitignore

@@ -27,7 +27,6 @@ __pycache__
 \#*\#
 build
 compiled/*.cpp
-core.*
 cutils_ext.cpp
 dist
 doc/.build/

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(),
 }
 
 
@@ -504,13 +506,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] = {}

pytensor/link/mlx/__init__.py

@@ -0,0 +1 @@
+from pytensor.link.mlx.linker import MLXLinker

pytensor/link/mlx/dispatch/__init__.py

@@ -0,0 +1,13 @@
+# isort: off
+from pytensor.link.mlx.dispatch.basic import mlx_funcify, mlx_typify
+
+import pytensor.link.mlx.dispatch.math
+import pytensor.link.mlx.dispatch.basic
+import pytensor.link.mlx.dispatch.elemwise
+import pytensor.link.mlx.dispatch.shape
+import pytensor.link.mlx.dispatch.subtensor
+import pytensor.link.mlx.dispatch.core
+import pytensor.link.mlx.dispatch.signal
+import pytensor.link.mlx.dispatch.signal.conv
+import pytensor.link.mlx.dispatch.blockwise
+# isort: on

pytensor/link/mlx/dispatch/basic.py

@@ -0,0 +1,101 @@
+import warnings
+from copy import deepcopy
+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 import Constant
+from pytensor.graph.fg import FunctionGraph
+from pytensor.link.utils import fgraph_to_python
+from pytensor.raise_op import Assert, CheckAndRaise
+
+
+@singledispatch
+def mlx_typify(data, **kwargs):
+ raise NotImplementedError(f"mlx_typify is not implemented for {type(data)}")
+
+
+@mlx_typify.register(np.ndarray)
+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(mx.array)
+def mlx_typify_no_conversion_needed(data, **kwargs):
+ return data
+
+
+@mlx_typify.register(int)
+@mlx_typify.register(float)
+def mlx_typify_python_scalar(data, **kwargs):
+ return mx.array(data)
+
+
+@mlx_typify.register(bool)
+@mlx_typify.register(np.bool_)
+def mlx_typify_bool(data, **kwargs):
+ return bool(data)
+
+
+@mlx_typify.register(np.integer)
+@mlx_typify.register(np.floating)
+@mlx_typify.register(np.complexfloating)
+def mlx_typify_numpy_scalar(data, **kwargs):
+ return mx.array(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 deepcopy(x)
+
+ return deepcopyop
+
+
+@mlx_funcify.register(Assert)
+@mlx_funcify.register(CheckAndRaise)
+def mlx_funcify_CheckAndRaise(op, node, **kwargs):
+ conds = node.inputs[1:]
+ if any(isinstance(cond, Constant) and not bool(cond.data) for cond in conds):
+ raise op.exc_type(op.msg)
+
+ warnings.warn(
+ f"""Skipping `{type(op).__name__}` Op (assertion: {op.msg}) as MLX tracing would remove it.""",
+ stacklevel=2,
+ )
+
+ def assert_fn(x, *inputs):
+ return x
+
+ return assert_fn

pytensor/link/mlx/dispatch/blockwise.py

@@ -0,0 +1,35 @@
+import mlx.core as mx
+
+from pytensor.link.mlx.dispatch import mlx_funcify
+from pytensor.tensor.blockwise import Blockwise
+
+
+@mlx_funcify.register(Blockwise)
+def funcify_Blockwise(op: Blockwise, node, **kwargs):
+ # 2) Otherwise, get the core python function for this Blockwise
+ core_node = op._create_dummy_core_node(node.inputs)
+ core_f = mlx_funcify(op.core_op, core_node)
+
+ # 3) Determine how many inputs correspond to batch dimensions
+ n_batch = op.batch_ndim(node)
+
+ # 4) Handle case where no vectorization is needed
+ if n_batch == 0:
+ return core_f
+
+ # 5) Vectorize using mx.vmap over any batched inputs
+ in_axes: list[int | None] = []
+ for inp, sig in zip(node.inputs, op.inputs_sig):
+ batch_ndim = inp.type.ndim - len(sig)
+ if batch_ndim == 0:
+ in_axes.append(None)
+ continue
+
+ batch_bcast = inp.type.broadcastable[:batch_ndim]
+ # If all batch dims are broadcastable (size 1), treat input as static
+ in_axes.append(0 if not all(batch_bcast) else None)
+
+ if not any(axis == 0 for axis in in_axes):
+ return core_f
+
+ return mx.vmap(core_f, in_axes=tuple(in_axes))