diff --git a/.github/uv-constraints-dev.txt b/.github/uv-constraints-dev.txt index d9d99a3..0901999 100644 --- a/.github/uv-constraints-dev.txt +++ b/.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 diff --git a/.github/uv-constraints-main.txt b/.github/uv-constraints-main.txt index 7b50d73..10d3ffe 100644 --- a/.github/uv-constraints-main.txt +++ b/.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 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 98377cc..43c0e48 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/.gitignore b/.gitignore index 82b98be..3bd64ed 100644 --- a/.gitignore +++ b/.gitignore @@ -33,3 +33,8 @@ reports benchmarks* reports* results* + +# ai folders +.ai/ +.cursor/ +.claude/ diff --git a/pyproject.toml b/pyproject.toml index 5203291..d34093b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ 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", @@ -42,6 +43,7 @@ dev = [ "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", ] diff --git a/python/nutpie/compile_pymc.py b/python/nutpie/compile_pymc.py index 13245f7..e917447 100644 --- a/python/nutpie/compile_pymc.py +++ b/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 diff --git a/python/nutpie/compiled_pyfunc.py b/python/nutpie/compiled_pyfunc.py index 07602ee..5737c33 100644 --- a/python/nutpie/compiled_pyfunc.py +++ b/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, ) diff --git a/tests/test_pymc.py b/tests/test_pymc.py index e901e70..14e8c08 100644 --- a/tests/test_pymc.py +++ b/tests/test_pymc.py @@ -15,9 +15,28 @@ import nutpie import nutpie.compile_pymc +# Check if MLX is available (macOS only, optional dependency) +MLX_AVAILABLE = find_spec("mlx") is not None + +# Build backend list dynamically based on availability +backend_params = [ + ("numba", None), + ("jax", "pytensor"), + ("jax", "jax"), +] + +# Only add MLX backends if MLX is available +if MLX_AVAILABLE: + backend_params.extend( + [ + ("mlx", "pytensor"), + ("mlx", "mlx"), + ] + ) + parameterize_backends = pytest.mark.parametrize( "backend, gradient_backend", - [("numba", None), ("jax", "pytensor"), ("jax", "jax")], + backend_params, ) @@ -303,6 +322,20 @@ def test_pymc_model_with_coordinate(backend, gradient_backend): @pytest.mark.pymc @parameterize_backends def test_pymc_model_store_extra(backend, gradient_backend): + if backend == "mlx" and gradient_backend == "pytensor": + # PyTensor's MLX linker mis-handles uneven Split ops, which appear in + # the gradient graph of the ZeroSumNormal/Dirichlet transforms below. + # Its Split dispatch passes a numpy array of split indices to + # ``mx.split``, which MLX interprets as an equal-split *count* and + # raises "Array split does not result in sub arrays with equal size" + # (a Python list of indices would work). The mlx/mlx combination is + # unaffected because it differentiates with MLX instead of PyTensor. + pytest.xfail( + "PyTensor's MLX linker Split dispatch is broken for uneven splits " + "(pytensor/link/mlx/dispatch/core.py passes a numpy array to " + "mx.split); only affects gradient_backend='pytensor'." + ) + with pm.Model() as model: model.add_coord("foo", length=5) model.add_coord("bar", length=4) @@ -552,6 +585,43 @@ def test_deterministic_sampling_jax(): return trace.posterior.a.values.ravel() +# NOTE: Unlike the numba/jax variants, this is intentionally *not* an +# array_compare test. MLX evaluates the logp and its gradient on the Metal GPU +# in float32 (float64 is unsupported on the GPU), and NUTS is a chaotic +# integrator: the float32 rounding -- which varies with the GPU model, the +# Metal/MLX version and kernel fusion -- accumulates across the leapfrog steps +# and makes the draws differ between machines. So instead of comparing against +# bit-exact reference values we validate statistical correctness against the +# analytic posterior, which is machine-independent. +@pytest.mark.pymc +def test_deterministic_sampling_mlx(): + if not MLX_AVAILABLE: + pytest.skip("MLX not installed") + + # No observed data, so the posterior of ``a`` is exactly its + # HalfNormal(sigma=1) prior. + with pm.Model() as model: + pm.HalfNormal("a") + + compiled = nutpie.compile_pymc_model(model, backend="mlx", gradient_backend="mlx") + trace = nutpie.sample( + compiled, chains=2, seed=123, draws=1000, tune=1000, progress_bar=False + ) + a = trace.posterior.a.values + + assert a.shape == (2, 1000) + assert np.all(np.isfinite(a)) + assert (a >= 0).all() # HalfNormal support + + # Analytic moments of HalfNormal(sigma=1). Tolerances are generous (~4x the + # Monte Carlo error) so the test is robust to the cross-machine float32 + # differences described above while still catching a genuinely broken logp. + expected_mean = np.sqrt(2.0 / np.pi) # ~0.7979 + expected_std = np.sqrt(1.0 - 2.0 / np.pi) # ~0.6028 + assert a.mean() == pytest.approx(expected_mean, abs=0.05) + assert a.std() == pytest.approx(expected_std, abs=0.05) + + @pytest.mark.pymc def test_zarr_store(tmp_path): coords = { @@ -619,6 +689,12 @@ def tmp_path(): def test_dims_model(backend, gradient_backend): import pymc.dims as pmd + if backend == "mlx": + pytest.xfail( + "PyTensor's MLX linker does not yet implement XTensorFromTensor; " + "see https://github.com/pymc-devs/pytensor/issues/1350" + ) + coords = {"a": range(3), "b": range(5)} with pm.Model(coords=coords) as model: print(model.dim_lengths)