Add MLX backend support for Nutpie compilation

.github/uv-constraints-dev.txt

@@ -1,2 +1,13 @@
-pytensor @ git+https://github.com/pymc-devs/pytensor.git@main
+# Test against pymc's development branch. We deliberately do NOT also pin
+# pytensor@main: pymc main caps the pytensor it supports (currently
+# pytensor>=3.0.7,<3.1), and whenever pytensor main has moved ahead of that cap
+# (e.g. pytensor main is >=3.1) the two git-main pins become mutually
+# unsatisfiable. uv then silently backtracks nutpie to an old PyPI release that
+# lacks the dev/docs extra, so pytest/jupyter never get installed. Let pymc@main
+# pull the (released) pytensor it is actually compatible with instead.
 pymc @ git+https://github.com/pymc-devs/pymc.git@main
+# pandas 3.0 wheels are built against the numpy 2.5 C ABI but only declare
+# numpy>=1.26, while numba caps numpy<2.5. The resolver then pairs pandas 3.0
+# with numpy 2.4, and the ABI mismatch segfaults in native code (e.g.
+# pd.date_range in test_zarr_store). Pin until numba supports numpy>=2.5.
+pandas < 3.0.0

.github/uv-constraints-main.txt

@@ -1 +1,6 @@
 arviz < 1.0.0
+# pandas 3.0 wheels are built against the numpy 2.5 C ABI but only declare
+# numpy>=1.26, while numba caps numpy<2.5. The resolver then pairs pandas 3.0
+# with numpy 2.4, and the ABI mismatch segfaults in native code (e.g.
+# pd.date_range in test_zarr_store). Pin until numba supports numpy>=2.5.
+pandas < 3.0.0

.github/workflows/ci.yml

@@ -359,6 +359,10 @@ jobs:
               ;;
             pymc)
               uv pip install 'nutpie[pymc]' jax --constraint .github/uv-constraints-main.txt
+              # mlx<0.31: https://github.com/ml-explore/mlx/issues/3329
+              if [ "${{ matrix.target }}" = "aarch64" ]; then
+                uv pip install 'mlx>=0.29.0,<0.31'
+              fi
               pytest -m "pymc and not flow" --arraydiff
               ;;
             flow)
@@ -391,7 +395,10 @@ jobs:
           python3 -m venv .venv
           source .venv/bin/activate
           uv pip install nutpie --no-deps --find-links dist --no-index
-          uv pip install 'nutpie[dev]' --constraint .github/uv-constraints-dev.txt
+          # --find-links dist keeps the locally-built wheel as the preferred
+          # candidate; without it uv re-resolves nutpie from PyPI and downgrades
+          # to a release that lacks the 'dev' extra (so pytest isn't installed).
+          uv pip install 'nutpie[dev]' --find-links dist --constraint .github/uv-constraints-dev.txt
           pytest -m "pymc" --arraydiff
 
   docs:
@@ -424,7 +431,11 @@ jobs:
         run: |
           source .venv/bin/activate
           uv pip install nutpie --no-deps --find-links dist --no-index
-          uv pip install --constraint .github/uv-constraints-dev.txt "nutpie[docs]"
+          # --find-links dist keeps the locally-built wheel as the preferred
+          # candidate; without it uv re-resolves nutpie from PyPI and downgrades
+          # to a release that lacks the 'docs' extra (so jupyter/nbformat aren't
+          # installed and `quarto render` fails to start the kernel).
+          uv pip install --find-links dist --constraint .github/uv-constraints-dev.txt "nutpie[docs]"
       - name: Render docs
         env:
           TBB_CXX_TYPE: clang

.gitignore

@@ -33,3 +33,8 @@ reports
 benchmarks*
 reports*
 results*
+
+# ai folders
+.ai/
+.cursor/
+.claude/

pyproject.toml

@@ -35,13 +35,15 @@ Repository = "https://github.com/pymc-devs/nutpie"
 stan = ["bridgestan >= 2.7.0", "stanio >= 0.5.1"]
 pymc = ["pymc >= 5.20.1", "numba >= 0.60.0"]
 pymc-jax = ["pymc >= 5.20.1", "jax >= 0.4.27"]
+pymc-mlx = ["pymc >= 5.20.1", "mlx >= 0.29.0,<0.31"]
 nnflow = ["flowjax >= 17.1.0", "equinox >= 0.11.12"]
 dev = [
   "bridgestan >= 2.7.0",
   "stanio >= 0.5.1",
   "pymc >= 5.20.1",
   "numba >= 0.60.0",
   "jax >= 0.4.27",
+  "mlx >= 0.29.0,<0.31",
   "flowjax >= 17.0.2",
   "pytest",
   "pytest-timeout",
@@ -53,6 +55,7 @@ all = [
   "pymc >= 5.20.1",
   "numba >= 0.60.0",
   "jax >= 0.4.27",
+  "mlx >= 0.29.0,<0.31",
   "flowjax >= 17.1.0",
   "equinox >= 0.11.12",
 ]

python/nutpie/compile_pymc.py

@@ -520,11 +520,151 @@ def expand(_x, **shared):
     )
 
 
+def _compile_pymc_model_mlx(
+    model,
+    *,
+    gradient_backend=None,
+    pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
+    var_names: Iterable[str] | None = None,
+    **kwargs,
+):
+    if find_spec("mlx") is None:
+        raise ImportError(
+            "MLX is not installed in the current environment. "
+            "Please install it with something like "
+            "'pip install mlx' "
+            "and restart your kernel in case you are in an interactive session."
+        )
+    import mlx.core as mx
+
+    # mlx>=0.31 segfaults inside Compiled::eval_gpu on the sampler worker
+    # thread; see https://github.com/ml-explore/mlx/issues/3329.
+    from importlib.metadata import PackageNotFoundError, version as _pkg_version
+
+    try:
+        _mlx_version_str = _pkg_version("mlx")
+    except PackageNotFoundError:
+        _mlx_version_str = None
+    if _mlx_version_str is not None:
+        _mlx_version = tuple(
+            int(p) for p in _mlx_version_str.split(".")[:2] if p.isdigit()
+        )
+        if _mlx_version >= (0, 31):
+            raise RuntimeError(
+                f"MLX {_mlx_version_str} is not supported by nutpie's MLX "
+                "backend due to a known SIGSEGV in compiled Metal kernels "
+                "(see https://github.com/ml-explore/mlx/issues/3329). "
+                "Please install mlx>=0.29,<0.31."
+            )
+
+    if gradient_backend is None:
+        gradient_backend = "pytensor"
+    elif gradient_backend not in ["mlx", "pytensor"]:
+        raise ValueError(f"Unknown gradient backend: {gradient_backend}")
+
+    (
+        n_dim,
+        _,
+        logp_fn_pt,
+        expand_fn_pt,
+        initial_point_fn,
+        shape_info,
+        reparameterized_names,
+    ) = _make_functions(
+        model,
+        mode="MLX",
+        compute_grad=gradient_backend == "pytensor",
+        join_expanded=False,
+        pymc_initial_point_fn=pymc_initial_point_fn,
+        var_names=var_names,
+    )
+
+    logp_fn = logp_fn_pt.vm.jit_fn
+    expand_fn = expand_fn_pt.vm.jit_fn
+
+    logp_shared_names = [var.name for var in logp_fn_pt.get_shared()]
+    expand_shared_names = [var.name for var in expand_fn_pt.get_shared()]
+
+    if gradient_backend == "mlx":
+        inner_logp_fn = logp_fn
+
+        def logp_fn_mlx_grad(x, *shared):
+            return mx.value_and_grad(lambda x: inner_logp_fn(x, *shared)[0])(x)
+
+        logp_fn = mx.compile(logp_fn_mlx_grad)
+
+    shared_data = {}
+    shared_vars = {}
+    seen = set()
+    for val in [*logp_fn_pt.get_shared(), *expand_fn_pt.get_shared()]:
+        if val.name in shared_data and val not in seen:
+            raise ValueError(f"Shared variables must have unique names: {val.name}")
+        shared_data[val.name] = mx.array(val.get_value())
+        shared_vars[val.name] = val
+        seen.add(val)
+
+    def make_logp_func():
+        def logp(_x, **shared):
+            # nutpie operates in float64, but MLX evaluates on the Metal GPU
+            # where float64 is unsupported, so the logp and its gradient are
+            # necessarily computed in float32. We cast explicitly rather than
+            # relying on ``mx.array``'s implicit float64->float32 downcast, so
+            # the precision contract cannot silently change with the MLX
+            # version or a global default-dtype override. This single-precision
+            # evaluation is the root cause of MLX sampling not being
+            # bit-reproducible across machines; see test_deterministic_sampling_mlx.
+            _x_mlx = mx.array(_x, dtype=mx.float32)
+            logp, grad = logp_fn(_x_mlx, *[shared[name] for name in logp_shared_names])
+            return float(logp), np.asarray(grad, dtype="float64", order="C")
+
+        return logp
+
+    names, slices, shapes = shape_info
+    # TODO do not cast to float64
+    dtypes = [np.dtype("float64")] * len(names)
+
+    def make_expand_func(seed1, seed2, chain):
+        # TODO handle seeds
+        def expand(_x, **shared):
+            # Match the logp closure: cast explicitly to float32 (MLX has no
+            # float64 on the GPU) instead of relying on the implicit downcast.
+            _x_mlx = mx.array(_x, dtype=mx.float32)
+            values = expand_fn(_x_mlx, *[shared[name] for name in expand_shared_names])
+            return {
+                name: np.asarray(val, order="C", dtype=dtype).reshape(shape)
+                for name, val, dtype, shape in zip(
+                    names, values, dtypes, shapes, strict=True
+                )
+            }
+
+        return expand
+
+    dims, coords = _prepare_dims_and_coords(model, shape_info, reparameterized_names)
+
+    return from_pyfunc(
+        ndim=n_dim,
+        make_logp_fn=make_logp_func,
+        make_expand_fn=make_expand_func,
+        make_initial_point_fn=initial_point_fn,
+        expanded_dtypes=dtypes,
+        expanded_shapes=shapes,
+        expanded_names=names,
+        shared_data=shared_data,
+        dims=dims,
+        coords=coords,
+        raw_logp_fn=None,
+        reparameterized_names=reparameterized_names,
+        # MLX is not thread-safe; see https://github.com/ml-explore/mlx/issues/2133.
+        force_single_core=True,
+        shared_data_converter=mx.array,
+    )
+
+
 def compile_pymc_model(
     model: "pm.Model",
     *,
-    backend: Literal["numba", "jax"] = "numba",
-    gradient_backend: Literal["pytensor", "jax"] = "pytensor",
+    backend: Literal["numba", "jax", "mlx"] = "numba",
+    gradient_backend: Literal["pytensor", "jax", "mlx"] = "pytensor",
     initial_points: dict[Union["Variable", str], np.ndarray | float | int]
     | None = None,
     jitter_rvs: set["TensorVariable"] | None = None,
@@ -541,11 +681,11 @@ def compile_pymc_model(
     ----------
     model : pymc.Model
         The model to compile.
-    backend : ["jax", "numba"]
+    backend : ["jax", "numba", "mlx"]
         The pytensor backend that is used to compile the logp function.
-    gradient_backend: ["pytensor", "jax"]
+    gradient_backend: ["pytensor", "jax", "mlx"]
         Which library is used to compute the gradients. This can only be changed
-        to "jax" if the jax backend is used.
+        to "jax" if the jax backend is used, or "mlx" if the mlx backend is used.
     jitter_rvs : set
         The set (or list or tuple) of random variables for which a U(-1, +1)
         jitter should be added to the initial value. Only available for
@@ -585,7 +725,7 @@ def compile_pymc_model(
     from pymc.model.transform.optimization import freeze_dims_and_data
 
     if freeze_model is None:
-        freeze_model = backend == "jax"
+        freeze_model = backend in ["jax", "mlx"]
 
     if freeze_model:
         model = freeze_dims_and_data(model)
@@ -604,8 +744,10 @@ def compile_pymc_model(
     initial_point_fn = _wrap_with_lock(initial_point_fn)
 
     if backend.lower() == "numba":
-        if gradient_backend == "jax":
-            raise ValueError("Gradient backend cannot be jax when using numba backend")
+        if gradient_backend in ["jax", "mlx"]:
+            raise ValueError(
+                f"Gradient backend cannot be {gradient_backend} when using numba backend"
+            )
         return _compile_pymc_model_numba(
             model=model,
             pymc_initial_point_fn=initial_point_fn,
@@ -620,8 +762,16 @@ def compile_pymc_model(
             var_names=var_names,
             **kwargs,
         )
+    elif backend.lower() == "mlx":
+        return _compile_pymc_model_mlx(
+            model=model,
+            gradient_backend=gradient_backend,
+            pymc_initial_point_fn=initial_point_fn,
+            var_names=var_names,
+            **kwargs,
+        )
     else:
-        raise ValueError(f"Backend must be one of numba and jax. Got {backend}")
+        raise ValueError(f"Backend must be one of numba, jax, and mlx. Got {backend}")
 
 
 def _wrap_with_lock(func: Callable) -> Callable:
@@ -668,7 +818,7 @@ def _compute_shapes(model) -> dict[str, tuple[int, ...]]:
 def _make_functions(
     model: "pm.Model",
     *,
-    mode: Literal["JAX", "NUMBA"],
+    mode: Literal["JAX", "NUMBA", "MLX"],
     compute_grad: bool,
     join_expanded: bool,
     pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
@@ -690,7 +840,7 @@ def _make_functions(
     model: pymc.Model
         The model to compile
     mode: str
-        Pytensor compile mode. One of "NUMBA" or "JAX"
+        Pytensor compile mode. One of "NUMBA", "JAX", or "MLX"
     compute_grad: bool
         Whether to compute gradients using pytensor. Must be True if mode is
         "NUMBA", otherwise False implies Jax will be used to compute gradients

python/nutpie/compiled_pyfunc.py

@@ -23,6 +23,8 @@ class PyFuncModel(CompiledModel):
     _coords: dict[str, Any]
     _raw_logp_fn: Callable | None
     _transform_adapt_args: dict | None = None
+    _force_single_core: bool = False
+    _shared_data_converter: Callable[[Any], Any] | None = None
 
     @property
     def shapes(self) -> dict[str, tuple[int, ...]]:
@@ -42,7 +44,12 @@ def with_data(self, **updates):
                 raise ValueError(f"Unknown data variable: {name}")
 
         updated = self._shared_data.copy()
-        updated.update(**updates)
+        if self._shared_data_converter is not None:
+            for name, value in updates.items():
+                updated[name] = self._shared_data_converter(value)
+        else:
+            updated.update(**updates)
+
         return dataclasses.replace(self, _shared_data=updated)
 
     def with_transform_adapt(self, **kwargs):
@@ -58,6 +65,8 @@ def _make_sampler(
         extra_callback_rate,
         store,
     ):
+        if self._force_single_core:
+            cores = 1
         model = self._make_model(init_mean)
         return _lib.PySampler.from_pyfunc(
             settings,
@@ -120,6 +129,8 @@ def from_pyfunc(
     make_transform_adapter=None,
     raw_logp_fn=None,
     reparameterized_names=None,
+    force_single_core: bool = False,
+    shared_data_converter: Callable[[Any], Any] | None = None,
 ):
     if coords is None:
         coords = {}
@@ -152,4 +163,6 @@ def from_pyfunc(
         _shared_data=shared_data,
         _raw_logp_fn=raw_logp_fn,
         reparameterized_names=reparameterized_names,
+        _force_single_core=force_single_core,
+        _shared_data_converter=shared_data_converter,
     )