Add MLX backend support for Nutpie compilation #908
Workflow file for this run
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: CI | |
| on: | |
| push: | |
| branches: | |
| - main | |
| - master | |
| tags: | |
| - "*" | |
| pull_request: | |
| workflow_dispatch: | |
| permissions: | |
| contents: read | |
| jobs: | |
| build_linux: | |
| name: build linux (${{ matrix.platform.target }}) | |
| runs-on: ${{ matrix.platform.runner }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| platform: | |
| - runner: ubuntu-22.04 | |
| target: x86_64 | |
| - runner: ubuntu-22.04-arm | |
| target: aarch64 | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 | |
| with: | |
| enable-cache: false | |
| - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 | |
| with: | |
| python-version: | | |
| 3.12 | |
| 3.13 | |
| 3.14 | |
| - name: Build wheels | |
| uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 | |
| with: | |
| target: ${{ matrix.platform.target }} | |
| args: --release --out dist --interpreter 3.12 3.13 3.14 --zig | |
| sccache: false | |
| manylinux: auto | |
| before-script-linux: | | |
| dnf install -y clang-libs clang || sudo apt install llvm-dev libclang-dev clang | |
| - name: Upload wheels | |
| uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 | |
| with: | |
| name: wheels-linux-${{ matrix.platform.target }} | |
| path: dist | |
| test_linux: | |
| name: test linux (${{ matrix.target }}, ${{ matrix.python-version }}, ${{ matrix.suite }}) | |
| runs-on: ${{ matrix.runner }} | |
| needs: build_linux | |
| # stan failures block release; pymc/flow are expected red against | |
| # released pymc, so allow them to fail without gating. | |
| continue-on-error: ${{ matrix.suite != 'stan' }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| target: [x86_64, aarch64] | |
| suite: [stan, pymc, flow] | |
| python-version: ["3.12", "3.13", "3.14"] | |
| include: | |
| - target: x86_64 | |
| runner: ubuntu-22.04 | |
| - target: aarch64 | |
| runner: ubuntu-22.04-arm | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 | |
| with: | |
| enable-cache: false | |
| - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 | |
| with: | |
| name: wheels-linux-${{ matrix.target }} | |
| path: dist | |
| - name: pytest | |
| shell: bash | |
| run: | | |
| set -e | |
| python3 -m venv .venv | |
| source .venv/bin/activate | |
| uv pip install pytest pytest-timeout pytest-arraydiff | |
| uv pip install nutpie --no-deps --find-links dist --no-index | |
| case "${{ matrix.suite }}" in | |
| stan) | |
| uv pip install 'nutpie[stan]' | |
| pytest -m "stan and not flow" --arraydiff | |
| ;; | |
| pymc) | |
| uv pip install 'nutpie[pymc]' jax --constraint .github/uv-constraints-main.txt | |
| pytest -m "pymc and not flow" --arraydiff | |
| ;; | |
| flow) | |
| uv pip install 'nutpie[all]' --constraint .github/uv-constraints-main.txt | |
| pytest -m flow --arraydiff | |
| ;; | |
| esac | |
| # pyarrow doesn't currently seem to work on musllinux | |
| #musllinux: | |
| # runs-on: ${{ matrix.platform.runner }} | |
| # strategy: | |
| # fail-fast: false | |
| # matrix: | |
| # platform: | |
| # - runner: ubuntu-22.04 | |
| # target: x86_64 | |
| # - runner: ubuntu-22.04 | |
| # target: aarch64 | |
| # steps: | |
| # - uses: actions/checkout@v6 | |
| # - uses: actions/setup-python@v6 | |
| # with: | |
| # python-version: "3.12" | |
| # - name: Install uv | |
| # uses: astral-sh/setup-uv@v7 | |
| # - name: Build wheels | |
| # uses: PyO3/maturin-action@v1 | |
| # with: | |
| # target: ${{ matrix.platform.target }} | |
| # args: --release --out dist --find-interpreter | |
| # sccache: ${{ !startsWith(github.ref, 'refs/tags/') }} | |
| # manylinux: musllinux_1_2 | |
| # before-script-linux: | | |
| # dnf install -y clang-libs clang || apt install llvm-dev libclang-dev clang | |
| # - name: Upload wheels | |
| # uses: actions/upload-artifact@v7 | |
| # with: | |
| # name: wheels-musllinux-${{ matrix.platform.target }} | |
| # path: dist | |
| # - name: pytest | |
| # if: ${{ startsWith(matrix.platform.target, 'x86_64') }} | |
| # uses: addnab/docker-run-action@v3 | |
| # with: | |
| # image: alpine:latest | |
| # options: -v ${{ github.workspace }}:/io -w /io | |
| # run: | | |
| # set -e | |
| # apk add py3-pip py3-virtualenv curl make clang | |
| # curl -LsSf https://astral.sh/uv/install.sh | sh | |
| # source $HOME/.local/bin/env | |
| # python3 -m virtualenv .venv | |
| # source .venv/bin/activate | |
| # # No numba packages for alpine | |
| # uv pip install 'nutpie[stan]' --find-links dist --force-reinstall | |
| # uv pip install pytest | |
| # pytest | |
| # - name: pytest | |
| # if: ${{ !startsWith(matrix.platform.target, 'x86') }} | |
| # uses: uraimo/run-on-arch-action@v2 | |
| # with: | |
| # arch: ${{ matrix.platform.target }} | |
| # distro: alpine_latest | |
| # githubToken: ${{ github.token }} | |
| # install: | | |
| # apk add py3-virtualenv curl make clang | |
| # curl -LsSf https://astral.sh/uv/install.sh | sh | |
| # source $HOME/.local/bin/env | |
| # run: | | |
| # set -e | |
| # python3 -m virtualenv .venv | |
| # source $HOME/.local/bin/env | |
| # source .venv/bin/activate | |
| # uv pip install pytest | |
| # # No numba packages for alpine | |
| # uv pip install 'nutpie[stan]' --find-links dist --force-reinstall | |
| # pytest | |
| build_windows: | |
| name: build windows | |
| runs-on: windows-latest | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 | |
| with: | |
| python-version: | | |
| 3.12 | |
| 3.13 | |
| 3.14 | |
| architecture: x64 | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 | |
| with: | |
| enable-cache: false | |
| - name: Install LLVM and Clang | |
| uses: KyleMayes/install-llvm-action@ebc0426251bc40c7cd31162802432c68818ab8f0 # v2.0.9 | |
| with: | |
| version: "15.0" | |
| directory: ${{ runner.temp }}/llvm | |
| - name: Set up TBB | |
| run: | | |
| Add-Content $env:GITHUB_PATH "$(pwd)/stan/lib/stan_math/lib/tbb" | |
| - name: Build wheels | |
| uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 | |
| env: | |
| LIBCLANG_PATH: ${{ runner.temp }}/llvm/lib | |
| with: | |
| target: x64 | |
| args: --release --out dist --interpreter 3.12 3.13 3.14 | |
| sccache: false | |
| - name: Upload wheels | |
| uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 | |
| with: | |
| name: wheels-windows-x64 | |
| path: dist | |
| test_windows: | |
| name: test windows (${{ matrix.python-version }}, ${{ matrix.suite }}) | |
| runs-on: windows-latest | |
| needs: build_windows | |
| continue-on-error: ${{ matrix.suite != 'stan' }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| suite: [stan, pymc, flow] | |
| python-version: ["3.12", "3.13", "3.14"] | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 | |
| with: | |
| enable-cache: false | |
| - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 | |
| with: | |
| name: wheels-windows-x64 | |
| path: dist | |
| - name: pytest | |
| shell: bash | |
| run: | | |
| set -e | |
| python3 -m venv .venv | |
| source .venv/Scripts/activate | |
| uv pip install pytest pytest-timeout pytest-arraydiff | |
| uv pip install nutpie --no-deps --find-links dist --no-index | |
| case "${{ matrix.suite }}" in | |
| stan) | |
| uv pip install "nutpie[stan]" | |
| pytest -m "stan and not flow" --arraydiff | |
| ;; | |
| pymc) | |
| uv pip install "nutpie[pymc]" jax --constraint .github/uv-constraints-main.txt | |
| pytest -m "pymc and not flow" --arraydiff | |
| ;; | |
| flow) | |
| uv pip install "nutpie[all]" --constraint .github/uv-constraints-main.txt | |
| pytest -m flow --arraydiff | |
| ;; | |
| esac | |
| build_macos: | |
| name: build macos (${{ matrix.platform.target }}) | |
| runs-on: ${{ matrix.platform.runner }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| platform: | |
| - runner: macos-15-intel | |
| target: x86_64 | |
| - runner: macos-14 | |
| target: aarch64 | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 | |
| with: | |
| python-version: | | |
| 3.12 | |
| 3.13 | |
| 3.14 | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 | |
| with: | |
| enable-cache: false | |
| - uses: maxim-lobanov/setup-xcode@ed7a3b1fda3918c0306d1b724322adc0b8cc0a90 # v1.7.0 | |
| with: | |
| xcode-version: latest-stable | |
| - name: Build wheels | |
| uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 | |
| with: | |
| target: ${{ matrix.platform.target }} | |
| args: --release --out dist --interpreter 3.12 3.13 3.14 --zig | |
| sccache: false | |
| - name: Upload wheels | |
| uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 | |
| with: | |
| name: wheels-macos-${{ matrix.platform.target }} | |
| path: dist | |
| test_macos: | |
| name: test macos (${{ matrix.target }}, ${{ matrix.python-version }}, ${{ matrix.suite }}) | |
| runs-on: ${{ matrix.runner }} | |
| needs: build_macos | |
| continue-on-error: ${{ matrix.suite != 'stan' }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| target: [x86_64, aarch64] | |
| suite: [stan, pymc, flow] | |
| python-version: ["3.12", "3.13", "3.14"] | |
| include: | |
| - target: x86_64 | |
| runner: macos-15-intel | |
| - target: aarch64 | |
| runner: macos-14 | |
| exclude: | |
| # numba isn't available on macOS x86_64 | |
| - target: x86_64 | |
| suite: pymc | |
| - target: x86_64 | |
| suite: flow | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 | |
| with: | |
| enable-cache: false | |
| - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 | |
| with: | |
| name: wheels-macos-${{ matrix.target }} | |
| path: dist | |
| - name: pytest | |
| run: | | |
| set -e | |
| python3 -m venv .venv | |
| source .venv/bin/activate | |
| uv pip install pytest pytest-timeout pytest-arraydiff | |
| uv pip install nutpie --no-deps --find-links dist --no-index | |
| case "${{ matrix.suite }}" in | |
| stan) | |
| uv pip install 'nutpie[stan]' | |
| pytest -m "stan and not flow" --arraydiff | |
| ;; | |
| 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) | |
| uv pip install 'nutpie[all]' --constraint .github/uv-constraints-main.txt | |
| pytest -m flow --arraydiff | |
| ;; | |
| esac | |
| test_pymc_dev: | |
| name: pymc (dev) | |
| runs-on: ubuntu-22.04 | |
| needs: build_linux | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 | |
| with: | |
| enable-cache: false | |
| - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 | |
| with: | |
| python-version: "3.12" | |
| - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 | |
| with: | |
| name: wheels-linux-x86_64 | |
| path: dist | |
| - name: pytest | |
| run: | | |
| python3 -m venv .venv | |
| source .venv/bin/activate | |
| uv pip install nutpie --no-deps --find-links dist --no-index | |
| # --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: | |
| runs-on: ubuntu-22.04 | |
| permissions: | |
| contents: read | |
| pages: write | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 | |
| with: | |
| python-version: "3.12" | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 | |
| with: | |
| enable-cache: false | |
| - name: Install system dependencies | |
| run: sudo apt-get update && sudo apt-get install -y clang libclang-dev | |
| - name: Install Quarto | |
| uses: quarto-dev/quarto-actions/setup@8a96df13519ee81fd526f2dfca5962811136661b # v2.2.0 | |
| - name: Build and install nutpie | |
| run: | | |
| python3 -m venv .venv | |
| source .venv/bin/activate | |
| uv pip install maturin | |
| maturin build --release --out dist | |
| - name: Install doc dependencies | |
| run: | | |
| source .venv/bin/activate | |
| uv pip install nutpie --no-deps --find-links dist --no-index | |
| # --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 | |
| run: | | |
| source .venv/bin/activate | |
| cd docs | |
| quarto render | |
| - name: Upload docs artifact | |
| if: ${{ github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/') }} | |
| uses: actions/upload-pages-artifact@fc324d3547104276b827a68afc52ff2a11cc49c9 # v5.0.0 | |
| with: | |
| path: docs/_site | |
| deploy-docs: | |
| name: Deploy Docs | |
| runs-on: ubuntu-latest | |
| if: ${{ github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/') }} | |
| needs: [docs] | |
| permissions: | |
| pages: write | |
| id-token: write | |
| environment: | |
| name: github-pages | |
| url: ${{ steps.deployment.outputs.page_url }} | |
| steps: | |
| - name: Deploy to GitHub Pages | |
| id: deployment | |
| uses: actions/deploy-pages@cd2ce8fcbc39b97be8ca5fce6e763baed58fa128 # v5.0.0 | |
| sdist: | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 | |
| with: | |
| persist-credentials: false | |
| - name: Build sdist | |
| uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 | |
| with: | |
| command: sdist | |
| args: --out dist | |
| - name: Upload sdist | |
| uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 | |
| with: | |
| name: wheels-sdist | |
| path: dist | |
| release: | |
| name: Release | |
| runs-on: ubuntu-latest | |
| if: ${{ startsWith(github.ref, 'refs/tags/') || github.event_name == 'workflow_dispatch' }} | |
| needs: | |
| [ | |
| build_linux, | |
| build_windows, | |
| build_macos, | |
| sdist, | |
| test_linux, | |
| test_windows, | |
| test_macos, | |
| test_pymc_dev, | |
| docs, | |
| ] | |
| environment: | |
| name: pypi | |
| permissions: | |
| # Use to sign the release artifacts | |
| id-token: write | |
| # Used to upload release artifacts | |
| contents: write | |
| # Used to generate artifact attestation | |
| attestations: write | |
| steps: | |
| - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 | |
| - name: Generate artifact attestation | |
| uses: actions/attest@59d89421af93a897026c735860bf21b6eb4f7b26 # v4.1.0 | |
| with: | |
| subject-path: "wheels-*/*" | |
| - name: Publish to PyPI | |
| if: ${{ startsWith(github.ref, 'refs/tags/') }} | |
| uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 | |
| with: | |
| command: upload | |
| args: --non-interactive --skip-existing wheels-*/* | |
| - name: Upload to GitHub Release | |
| uses: softprops/action-gh-release@b4309332981a82ec1c5618f44dd2e27cc8bfbfda # v3.0.0 | |
| with: | |
| files: | | |
| wasm-wheels/*.whl | |
| prerelease: ${{ contains(github.ref, 'alpha') || contains(github.ref, 'beta') }} |