Add MLX dispatch for gamma (#2333) #7849
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: Tests | |
| on: | |
| push: | |
| branches: | |
| - main | |
| - v2 | |
| pull_request: | |
| branches: | |
| - main | |
| - v2 | |
| # Cancels all previous workflow runs for pull requests that have not completed. | |
| concurrency: | |
| # The concurrency group contains the workflow name and the branch name for pull requests | |
| # or the commit hash for any other events. | |
| group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }} | |
| cancel-in-progress: true | |
| # Deny all permissions by default; jobs opt into the minimum they need. | |
| permissions: {} | |
| jobs: | |
| changes: | |
| name: "Check for changes" | |
| runs-on: ubuntu-latest | |
| permissions: | |
| contents: read | |
| outputs: | |
| changes: ${{ steps.changes.outputs.src }} | |
| steps: | |
| - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 | |
| with: | |
| fetch-depth: 0 | |
| persist-credentials: false | |
| - uses: dorny/paths-filter@7b450fff21473bca461d4b92ce414b9d0420d706 # v4.0.2 | |
| id: changes | |
| with: | |
| filters: | | |
| python: &python | |
| - 'pytensor/**/*.py' | |
| - 'tests/**/*.py' | |
| - 'pytensor/**/*.pyx' | |
| - 'tests/**/*.pyx' | |
| - '*.py' | |
| src: | |
| - *python | |
| - 'pytensor/**/*.c' | |
| - 'tests/**/*.c' | |
| - 'pytensor/**/*.h' | |
| - 'tests/**/*.h' | |
| - '.github/workflows/*.yml' | |
| - 'setup.cfg' | |
| - 'requirements.txt' | |
| - '.pre-commit-config.yaml' | |
| style: | |
| name: Check code style | |
| needs: changes | |
| runs-on: ubuntu-latest | |
| permissions: | |
| contents: read | |
| if: ${{ needs.changes.outputs.changes == 'true' }} | |
| strategy: | |
| matrix: | |
| python-version: ["3.12", "3.14"] | |
| steps: | |
| - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 | |
| with: | |
| persist-credentials: false | |
| - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 | |
| test: | |
| name: "mode ${{ matrix.default-mode }} : py${{ matrix.python-version }} : ${{ matrix.os }} : ${{ matrix.part[0] }}" | |
| needs: | |
| - changes | |
| - style | |
| runs-on: ${{ matrix.os }} | |
| permissions: | |
| contents: read | |
| if: ${{ needs.changes.outputs.changes == 'true' && needs.style.result == 'success' }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| # default-mode: "NUMBA" is actually a no-op, to make sure we're testing default config settings | |
| default-mode: ["CVM", "NUMBA", "FAST_COMPILE"] | |
| python-version: ["3.12", "3.14"] | |
| os: ["ubuntu-latest"] | |
| install-jax: [0] | |
| install-torch: [0] | |
| install-mlx: [0] | |
| install-xarray: [0] | |
| free-threading: [0] | |
| part: | |
| - [ "*rest", "tests --ignore=tests/scan --ignore=tests/tensor --ignore=tests/xtensor --ignore=tests/link/numba" ] | |
| - [ "scan", "tests/scan" ] | |
| - [ "tensor *rest", "tests/tensor --ignore=tests/tensor/test_basic.py --ignore=tests/tensor/test_elemwise.py --ignore=tests/tensor/test_math.py --ignore=tests/tensor/test_math_scipy.py --ignore=tests/tensor/test_blas.py --ignore=tests/tensor/signal --ignore=tests/tensor/rewriting --ignore=tests/tensor/linalg --ignore=tests/tensor/test_pad.py --ignore tests/tensor/test_blockwise.py --ignore tests/tensor/test_subtensor.py" ] | |
| - [ "tensor basic+elemwise+blockwise", "tests/tensor/test_basic.py tests/tensor/test_elemwise.py tests/tensor/test_blockwise.py " ] | |
| - [ "tensor math+subtensor", "tests/tensor/test_math.py tests/tensor/test_subtensor.py" ] | |
| - [ "tensor scipy+blas+pad", "tests/tensor/test_math_scipy.py tests/tensor/test_blas.py tests/tensor/signal tests/tensor/test_pad.py" ] | |
| - [ "tensor rewriting", "tests/tensor/rewriting" ] | |
| - [ "tensor linalg", "tests/tensor/linalg" ] | |
| exclude: | |
| - python-version: "3.12" | |
| default-mode: "FAST_COMPILE" | |
| include: | |
| - part: ["doctests", "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link --ignore=pytensor/ipython.py"] | |
| default-mode: "CVM" | |
| python-version: "3.12" | |
| os: "ubuntu-latest" | |
| - part: ["numba link", "tests/link/numba --ignore=tests/link/numba/linalg"] | |
| default-mode: "CVM" | |
| python-version: "3.12" | |
| os: "ubuntu-latest" | |
| - part: ["numba link linalg", "tests/link/numba/linalg"] | |
| default-mode: "CVM" | |
| python-version: "3.13" | |
| os: "ubuntu-latest" | |
| - part: ["jax link", "tests/link/jax"] | |
| install-jax: 1 | |
| default-mode: "CVM" | |
| python-version: "3.14" | |
| os: "ubuntu-latest" | |
| - part: ["pytorch link", "tests/link/pytorch"] | |
| install-torch: 1 | |
| default-mode: "CVM" | |
| python-version: "3.12" | |
| os: "ubuntu-latest" | |
| - part: ["xtensor", "tests/xtensor"] | |
| install-xarray: 1 | |
| default-mode: "CVM" | |
| python-version: "3.14" | |
| os: "ubuntu-latest" | |
| - part: ["mlx link", "tests/link/mlx"] | |
| install-mlx: 1 | |
| default-mode: "CVM" | |
| python-version: "3.12" | |
| os: "macos-15" | |
| - part: ["macos smoke test", "tests/tensor/test_elemwise.py tests/tensor/test_math_scipy.py tests/tensor/test_blas.py"] | |
| default-mode: "CVM" | |
| python-version: "3.14" | |
| os: "macos-15" | |
| # conftest's pytest_sessionfinish asserts the GIL was never re-enabled. | |
| - part: ["free-threading smoke test", "tests/tensor/test_math_scipy.py"] | |
| default-mode: "FAST_RUN" | |
| python-version: "3.14" | |
| os: "ubuntu-latest" | |
| free-threading: 1 | |
| steps: | |
| - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 | |
| with: | |
| fetch-depth: 0 | |
| persist-credentials: false | |
| - name: Set up Python ${{ matrix.python-version }} | |
| uses: mamba-org/setup-micromamba@d7c9bd84e824b79d2af72a2d4196c7f4300d3476 # v3.0.0 | |
| with: | |
| environment-name: pytensor-test | |
| micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved | |
| init-shell: bash | |
| post-cleanup: "all" | |
| # Create the env with python-freethreading up front so it resolves to the | |
| # cp*t build from the start; adding it later conflicts with an already-solved | |
| # regular CPython. | |
| create-args: >- | |
| python=${{ matrix.python-version }} | |
| ${{ matrix.free-threading == 1 && 'python-freethreading' || '' }} | |
| - name: Create matrix id | |
| id: matrix-id | |
| env: | |
| MATRIX_CONTEXT: ${{ toJson(matrix) }} | |
| run: | | |
| echo $MATRIX_CONTEXT | |
| export MATRIX_ID=`echo $MATRIX_CONTEXT | sha256sum | cut -c 1-32` | |
| echo $MATRIX_ID | |
| echo "id=$MATRIX_ID" >> $GITHUB_OUTPUT | |
| - name: Install dependencies | |
| shell: micromamba-shell {0} | |
| run: | | |
| # Install the exact numba cap from pyproject.toml so CI validates the | |
| # version the bump-numba workflow pinned (fails loudly if conda-forge | |
| # doesn't have it yet, rather than silently testing an older release). | |
| NUMBA_VERSION=$(sed -n 's/.*"numba>=[0-9.]*,<=\([0-9.]*\)".*/\1/p' pyproject.toml | head -1) | |
| echo "Installing numba==${NUMBA_VERSION}" | |
| # python-freethreading constrains the solve to the free-threaded CPython build | |
| if [[ $FREE_THREADING == "1" ]]; then FT_PKG="python-freethreading"; else FT_PKG=""; fi | |
| if [[ $OS == "macos-15" ]]; then | |
| micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba==${NUMBA_VERSION}" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx libblas=*=*accelerate rich; | |
| elif [[ $FREE_THREADING == "1" ]]; then | |
| # mkl/mkl-service have no free-threaded (cp*t) builds, so rely on the default (openblas) BLAS. | |
| micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba==${NUMBA_VERSION}" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx rich; | |
| else | |
| micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba==${NUMBA_VERSION}" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx mkl mkl-service rich; | |
| fi | |
| if [[ $FREE_THREADING == "1" ]]; then python -c 'import sysconfig; assert sysconfig.get_config_var("Py_GIL_DISABLED") == 1, "Expected a free-threaded interpreter"'; fi | |
| if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" && pip install "jax>=0.8,<0.9.1" jaxlib numpyro equinox 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>=0.30,<0.32"; fi | |
| if [[ $INSTALL_XARRAY == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" xarray xarray-einstats; fi | |
| pip install -e ./ | |
| micromamba list && pip freeze | |
| python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))' | |
| if [[ $OS == "macos-15" ]]; then | |
| python -c 'import pytensor; assert pytensor.config.blas__ldflags.startswith("-framework Accelerate"), "Blas flags are not set to MacOS Accelerate"'; | |
| elif [[ $FREE_THREADING != "1" ]]; then | |
| # Blas flags are only used by the C backend; the free-threading job runs the numba backend. | |
| python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'; | |
| fi | |
| env: | |
| PYTHON_VERSION: ${{ matrix.python-version }} | |
| INSTALL_JAX: ${{ matrix.install-jax }} | |
| INSTALL_TORCH: ${{ matrix.install-torch }} | |
| INSTALL_XARRAY: ${{ matrix.install-xarray }} | |
| INSTALL_MLX: ${{ matrix.install-mlx }} | |
| FREE_THREADING: ${{ matrix.free-threading }} | |
| OS: ${{ matrix.os}} | |
| - name: Run tests | |
| shell: micromamba-shell {0} | |
| run: | | |
| if [[ $DEFAULT_MODE == "FAST_COMPILE" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,mode=FAST_COMPILE; fi | |
| if [[ $DEFAULT_MODE == "CVM" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,linker=cvm; fi | |
| # Disable the C backend on the free-threaded build so no test imports a C | |
| # extension, which would re-enable the GIL (checked in conftest). | |
| if [[ $FREE_THREADING == "1" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,cxx=; fi | |
| export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe | |
| python -m pytest -r A --verbose --runslow --durations=50 --cov=pytensor/ --cov-report=xml:coverage/coverage-${MATRIX_ID}.xml --no-cov-on-fail $PART --benchmark-skip | |
| env: | |
| MATRIX_ID: ${{ steps.matrix-id.outputs.id }} | |
| MKL_THREADING_LAYER: GNU | |
| MKL_NUM_THREADS: 1 | |
| OMP_NUM_THREADS: 1 | |
| PART: ${{ matrix.part[1] }} | |
| DEFAULT_MODE: ${{ matrix.default-mode }} | |
| FREE_THREADING: ${{ matrix.free-threading }} | |
| # The macOS runners are virtualized and abort when a kernel is executed on the | |
| # MLX GPU stream, so the tests that need one are skipped here and run locally. | |
| PYTENSOR_MLX_SKIP_GPU: 1 | |
| - name: Upload coverage file | |
| uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 | |
| with: | |
| name: coverage-${{ steps.matrix-id.outputs.id }} | |
| path: coverage/coverage-${{ steps.matrix-id.outputs.id }}.xml | |
| benchmarks: | |
| name: "Benchmarks" | |
| needs: | |
| - changes | |
| - style | |
| runs-on: ubuntu-latest | |
| permissions: | |
| contents: read | |
| if: ${{ needs.changes.outputs.changes == 'true' && needs.style.result == 'success' }} | |
| strategy: | |
| fail-fast: false | |
| steps: | |
| - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 | |
| with: | |
| fetch-depth: 0 | |
| persist-credentials: false | |
| - name: Set up Python 3.12 | |
| uses: mamba-org/setup-micromamba@d7c9bd84e824b79d2af72a2d4196c7f4300d3476 # v3.0.0 | |
| with: | |
| environment-name: pytensor-test | |
| micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved | |
| init-shell: bash | |
| post-cleanup: "all" | |
| - name: Install dependencies | |
| shell: micromamba-shell {0} | |
| run: | | |
| NUMBA_VERSION=$(sed -n 's/.*"numba>=[0-9.]*,<=\([0-9.]*\)".*/\1/p' pyproject.toml | head -1) | |
| echo "Installing numba==${NUMBA_VERSION}" | |
| micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" mkl numpy scipy pip mkl-service cython pytest "numba==${NUMBA_VERSION}" jax jaxlib pytest-benchmark | |
| pip install -e ./ | |
| micromamba list && pip freeze | |
| python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))' | |
| python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"' | |
| env: | |
| PYTHON_VERSION: 3.12 | |
| - name: Run benchmarks | |
| shell: micromamba-shell {0} | |
| run: | | |
| export PYTENSOR_FLAGS=on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe | |
| python -m pytest --runslow --benchmark-only | |
| all-checks: | |
| if: ${{ always() }} | |
| runs-on: ubuntu-latest | |
| name: "All tests" | |
| needs: [changes, style, test] | |
| steps: | |
| - name: Check build matrix status | |
| if: ${{ needs.changes.outputs.changes == 'true' && (needs.style.result != 'success' || needs.test.result != 'success') }} | |
| run: exit 1 | |
| upload-coverage: | |
| runs-on: ubuntu-latest | |
| name: "Upload coverage" | |
| permissions: | |
| contents: read | |
| needs: [changes, all-checks] | |
| if: ${{ needs.changes.outputs.changes == 'true' && needs.all-checks.result == 'success' }} | |
| steps: | |
| - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 | |
| with: | |
| persist-credentials: false | |
| - name: Set up Python | |
| uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 | |
| with: | |
| python-version: "3.14" | |
| - name: Install dependencies | |
| run: | | |
| python -m pip install -U coverage>=5.1 coveralls | |
| - name: Download coverage file | |
| uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 | |
| with: | |
| pattern: coverage-* | |
| path: coverage | |
| merge-multiple: true | |
| - name: Upload coverage to Codecov | |
| uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 | |
| with: | |
| directory: ./coverage/ | |
| fail_ci_if_error: true | |
| token: ${{ secrets.CODECOV_TOKEN }} |