diff --git a/.flake8 b/.flake8 index 0dede3f1d..959e9a22c 100644 --- a/.flake8 +++ b/.flake8 @@ -11,6 +11,6 @@ extend-ignore = # E203 whitespace before ':' (to be compatible with black) per-file-ignores = scripts/create_pickle.py:F403,F405, - graphblas/tests/*.py:T201, + graphblas/tests/*.py:T201,B043, graphblas/core/ss/matrix.py:SIM113, graphblas/**/__init__.py:F401, diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 5ace4600a..94e8d4b16 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -4,3 +4,12 @@ updates: directory: "/" schedule: interval: "weekly" + # Wait a week after a release before proposing it, so a compromised or + # broken release has time to be noticed and yanked. + cooldown: + default-days: 7 + # One PR per week for all action bumps instead of one PR per action. + groups: + github-actions: + patterns: + - "*" diff --git a/.github/workflows/audit_versions.yml b/.github/workflows/audit_versions.yml new file mode 100644 index 000000000..162713824 --- /dev/null +++ b/.github/workflows/audit_versions.yml @@ -0,0 +1,80 @@ +# Checks the CI dependency pools in scripts/ci_pick_versions.py against the world. +# +# `ci_pick_versions.py --validate` runs in pre-commit and proves the picker agrees with +# the constraints it documents. It cannot notice when a constraint has drifted from +# reality, because that happens when a package releases, not when we commit. So this +# workflow runs weekly, and on changes to the picker or the audits themselves. +# +# A failure here does not mean CI is broken right now. It means a pool or a table no +# longer matches what PyPI or conda-forge publishes, and the next random pick that +# lands on the stale combination will fail for reasons that look unrelated. + +name: Audit dependency versions + +on: + schedule: + # Mondays at 06:00 UTC, before most of the week's PRs + - cron: "0 6 * * 1" + workflow_dispatch: + pull_request: + paths: + - scripts/ci_pick_versions.py + - scripts/audit_pypi_versions.py + - scripts/audit_conda_versions.py + - .github/workflows/audit_versions.yml + push: + branches: + - main + paths: + - scripts/ci_pick_versions.py + - scripts/audit_pypi_versions.py + - scripts/audit_conda_versions.py + - .github/workflows/audit_versions.yml + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +jobs: + pypi: + name: Audit against PyPI + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.12" + - run: python -m pip install --upgrade pip packaging + - name: Stress-test the picker + run: python scripts/ci_pick_versions.py --validate + - name: Audit pools against PyPI + run: python scripts/audit_pypi_versions.py --samples 200 + + conda: + name: Audit against conda-forge + runs-on: ubuntu-latest + # conda-forge is what CI installs from, so this is the one that binds. It needs a + # conda to search with, which is why it is a separate, slower job. + defaults: + run: + shell: bash -el {0} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: conda-incubator/setup-miniconda@be893c923ea9cf1cf7cd510fbdde27c7e18cbdcb # v4.1.0 + with: + auto-update-conda: true + python-version: "3.12" + channels: conda-forge + conda-remove-defaults: true + channel-priority: strict + activate-environment: audit + - run: python -m pip install --upgrade pip packaging + - name: Audit pools against conda-forge + run: python scripts/audit_conda_versions.py --samples 100 diff --git a/.github/workflows/debug.yml b/.github/workflows/debug.yml index 6c2b202b1..bb8503570 100644 --- a/.github/workflows/debug.yml +++ b/.github/workflows/debug.yml @@ -9,13 +9,16 @@ on: required: false default: false +permissions: + contents: read + jobs: build: runs-on: ubuntu-latest strategy: fail-fast: false matrix: - pyver: [3.10] + pyver: [3.11] testopts: - "--blocking" # - "--non-blocking --record --runslow" @@ -26,7 +29,7 @@ jobs: # - "conda-forge" steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false @@ -71,5 +74,5 @@ jobs: ( for attr in Matrix Scalar Vector Recorder _agg agg base binary descriptor dtypes exceptions expr ffi formatting infix init io lib mask matrix monoid op operator scalar semiring tests unary vector recorder _ss ss ; do echo python -c \"from graphblas import $attr\" ; if ! python -c "from graphblas import $attr" ; then exit 1 ; fi ; done ) - name: Setup tmate session - uses: mxschmitt/action-tmate@v3 + uses: mxschmitt/action-tmate@35b54afac29c97fb54faba5b513f8fbd1882f113 # v3.24 if: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.debug_enabled }} diff --git a/.github/workflows/imports.yml b/.github/workflows/imports.yml index e24d0d4db..9adfb6b0e 100644 --- a/.github/workflows/imports.yml +++ b/.github/workflows/imports.yml @@ -6,6 +6,14 @@ on: branches: - main +permissions: + contents: read + +# Cancels a run's identical twin only; see the note in test_and_build.yml. +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.run_id }}-${{ github.event.pull_request.head.sha || github.sha }} + cancel-in-progress: true + jobs: rngs: runs-on: ubuntu-latest @@ -13,46 +21,36 @@ jobs: os: ${{ steps.os.outputs.selected }} pyver: ${{ steps.pyver.outputs.selected }} steps: + # Plain bash instead of a marketplace action; see the note on the rngs job in + # test_and_build.yml. To weight a choice, repeat it; to skip one, remove it. - name: RNG for os - uses: ddradar/choose-random-action@v3.0.0 id: os - with: - contents: | - ubuntu-latest - macos-latest - windows-latest - weights: | - 1 - 1 - 1 + run: | + choices=(ubuntu-latest macos-latest windows-latest) + selected=${choices[$((RANDOM % ${#choices[@]}))]} + echo "selected: $selected (from: ${choices[*]})" + echo "selected=$selected" >> "$GITHUB_OUTPUT" - name: RNG for Python version - uses: ddradar/choose-random-action@v3.0.0 id: pyver - with: - contents: | - 3.10 - 3.11 - 3.12 - 3.13 - weights: | - 1 - 1 - 1 - 1 + run: | + choices=(3.11 3.12 3.13 3.14 3.15) + selected=${choices[$((RANDOM % ${#choices[@]}))]} + echo "selected: $selected (from: ${choices[*]})" + echo "selected=$selected" >> "$GITHUB_OUTPUT" test_imports: needs: rngs runs-on: ${{ needs.rngs.outputs.os }} # runs-on: ${{ matrix.os }} # strategy: # matrix: - # python-version: ["3.10", "3.11", "3.12", "3.13"] + # python-version: ["3.11", "3.12", "3.13", "3.14", "3.15"] # os: ["ubuntu-latest", "macos-latest", "windows-latest"] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: ${{ needs.rngs.outputs.pyver }} # python-version: ${{ matrix.python-version }} diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 655a576e5..567dceb83 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -16,11 +16,11 @@ jobs: name: pre-commit-hooks runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: - python-version: "3.10" - - uses: pre-commit/action@v3.0.1 + python-version: "3.11" + - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 diff --git a/.github/workflows/publish_pypi.yml b/.github/workflows/publish_pypi.yml index 32926c5c8..fd2f3af10 100644 --- a/.github/workflows/publish_pypi.yml +++ b/.github/workflows/publish_pypi.yml @@ -5,38 +5,53 @@ on: tags: - "20*" +permissions: + contents: read + jobs: build_and_deploy: runs-on: ubuntu-latest - if: github.repository == 'python-graphblas/python-graphblas' defaults: run: shell: bash -l {0} steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: - python-version: "3.10" + python-version: "3.11" - name: Install build dependencies run: | python -m pip install --upgrade pip python -m pip install build twine - name: Build wheel and sdist run: python -m build --sdist --wheel - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 with: name: releases path: dist if-no-files-found: error - name: Check with twine run: python -m twine check --strict dist/* - - name: Publish to PyPI - uses: pypa/gh-action-pypi-publish@v1.12.4 + + publish: + needs: build_and_deploy + if: startsWith(github.ref, 'refs/tags/') && github.repository == 'python-graphblas/python-graphblas' + runs-on: ubuntu-latest + # The trusted publisher registered on PyPI names this environment, which also asks a + # maintainer to approve the upload. + environment: pypi + permissions: + # Only this job can mint the OIDC token that PyPI exchanges for upload rights. + id-token: write + steps: + - uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2 with: - user: __token__ - password: ${{ secrets.PYPI_TOKEN }} + name: releases + path: dist + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3 diff --git a/.github/workflows/security.yml b/.github/workflows/security.yml new file mode 100644 index 000000000..e7840ac2c --- /dev/null +++ b/.github/workflows/security.yml @@ -0,0 +1,58 @@ +name: Security + +# Run the security hooks of .pre-commit-config.yaml, with its versions and arguments, where +# pre-commit.ci cannot: with network access, and weekly, for advisories published since. +on: + pull_request: + paths: + - pyproject.toml + - .pre-commit-config.yaml + - .github/** + schedule: + # Mondays at 06:41 UTC + - cron: "41 6 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + pysentry: + name: pysentry + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.14" + # Advisories against the floors in pyproject.toml + - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 + with: + extra_args: pysentry --all-files + + zizmor: + name: zizmor + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.14" + # With a token, zizmor also runs its online audits, such as advisories against the actions + # these workflows use (known-vulnerable-actions) and pins to commits outside an action's + # repository (impostor-commit). + - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 + env: + GH_TOKEN: ${{ github.token }} + with: + extra_args: zizmor --all-files diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index bfc17834b..ecd82e2ab 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -17,9 +17,17 @@ on: branches: - main -# concurrency: -# group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} -# cancel-in-progress: true +permissions: + contents: read + +# GitHub can deliver a PR's pull_request event twice for one push when the PR's base +# branch moved in the same push, as it does for every PR above the bottom of a restacked +# stack. Keyed on the PR and its head commit, the later run cancels its identical twin, +# and a run for an earlier commit is left to finish, as before (a group keyed on the +# branch would cancel it). Pushes to main key on the run id, so none is ever cancelled. +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.run_id }}-${{ github.event.pull_request.head.sha || github.sha }} + cancel-in-progress: true jobs: rngs: @@ -49,34 +57,25 @@ jobs: mapnumpy: ${{ steps.mapnumpy.outputs.selected }} backend: ${{ steps.backend.outputs.selected }} steps: + # Plain bash instead of a marketplace action: resolving a third-party action is a + # network dependency that has failed with repeated 5xx errors ("Internal Server Error + # occurred while resolving ddradar/choose-random-action"), and a failure here kills + # the whole matrix. To weight a choice, repeat it (e.g. `A A B` picks A twice as + # often); to skip a choice, remove it from the list. - name: RNG for mapnumpy - uses: ddradar/choose-random-action@v3.0.0 id: mapnumpy - with: - contents: | - A - B - C - D - weights: | - 1 - 1 - 1 - 1 + run: | + choices=(A B C D) + selected=${choices[$((RANDOM % ${#choices[@]}))]} + echo "selected: $selected (from: ${choices[*]})" + echo "selected=$selected" >> "$GITHUB_OUTPUT" - name: RNG for backend - uses: ddradar/choose-random-action@v3.0.0 id: backend - with: - contents: | - E - F - G - H - weights: | - 1 - 1 - 1 - 1 + run: | + choices=(E F G H) + selected=${choices[$((RANDOM % ${#choices[@]}))]} + echo "selected: $selected (from: ${choices[*]})" + echo "selected=$selected" >> "$GITHUB_OUTPUT" build_and_test: needs: rngs runs-on: ${{ matrix.os }} @@ -93,51 +92,55 @@ jobs: os: ["ubuntu-latest", "macos-latest", "windows-latest"] slowtask: ["pytest_normal", "pytest_bizarro", "notebooks"] # repeat: [1, 2, 3] # For stress testing + include: + # Additional architectures: run pytest_normal only to limit CI cost + - os: macos-15-intel # macOS x86_64 (Intel) + slowtask: pytest_normal + - os: ubuntu-24.04-arm # Linux aarch64 + slowtask: pytest_normal + # Uncomment these temporarily for stress testing + # - os: macos-15-intel # macOS x86_64 (Intel) + # slowtask: pytest_bizarro + # - os: ubuntu-24.04-arm # Linux aarch64 + # slowtask: pytest_bizarro + # - os: macos-15-intel # macOS x86_64 (Intel) + # slowtask: notebooks + # - os: ubuntu-24.04-arm # Linux aarch64 + # slowtask: notebooks env: - # Wheels on OS X come with an OpenMP that conflicts with OpenMP from conda-forge. - # Setting this is a workaround. + # Wheels on OS X come with an OpenMP that conflicts with OpenMP from conda-forge: + # libomp aborts when a second copy starts in the process, unless this is set. Two + # copies can still deadlock or crash in an OpenMP barrier, so "Update env" keeps + # conda-forge's libomp out of macOS wheel draws where it can. KMP_DUPLICATE_LIB_OK: ${{ contains(matrix.os, 'macos') && 'TRUE' || 'FALSE' }} steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false + # Plain bash instead of a marketplace action; see the note on the rngs job above. + # To weight a choice, repeat it; to skip a choice, remove it from the list. - name: RNG for Python version - uses: ddradar/choose-random-action@v3.0.0 id: pyver - with: - # We should support major Python versions for at least 36 months as per SPEC 0 - # We may be able to support pypy if anybody asks for it - # 3.9.16 0_73_pypy - contents: | - 3.10 - 3.11 - 3.12 - 3.13 - weights: | - 1 - 1 - 1 - 1 + # We should support major Python versions for at least 36 months as per SPEC 0 + # We may be able to support pypy if anybody asks for it + # 3.9.16 0_73_pypy + run: | + choices=(3.11 3.12 3.13 3.14 3.15) + selected=${choices[$((RANDOM % ${#choices[@]}))]} + echo "selected: $selected (from: ${choices[*]})" + echo "selected=$selected" >> "$GITHUB_OUTPUT" - name: RNG for source of python-suitesparse-graphblas - uses: ddradar/choose-random-action@v3.0.0 id: sourcetype - with: - # Weights must be natural numbers, so set weights to very large to skip one - # (such as if 'upstream' is known to not work). - contents: | - conda-forge - wheel - source - upstream - weights: | - 1 - 1 - 1 - 1 + # Remove 'upstream' from the list if it is known to not work. + run: | + choices=(conda-forge wheel source upstream) + selected=${choices[$((RANDOM % ${#choices[@]}))]} + echo "selected: $selected (from: ${choices[*]})" + echo "selected=$selected" >> "$GITHUB_OUTPUT" - name: Setup conda - uses: conda-incubator/setup-miniconda@v3 + uses: conda-incubator/setup-miniconda@be893c923ea9cf1cf7cd510fbdde27c7e18cbdcb # v4.1.0 id: setup_conda with: auto-update-conda: true @@ -146,252 +149,180 @@ jobs: conda-remove-defaults: ${{ contains(steps.pyver.outputs.selected, 'pypy') && 'false' || 'true' }} channel-priority: ${{ contains(steps.pyver.outputs.selected, 'pypy') && 'flexible' || 'strict' }} activate-environment: graphblas - auto-activate-base: false + # setup-miniconda deprecated ``auto-activate-base`` for this name, and warns + # on every run that still uses the old one. + auto-activate: false - name: Update env + id: update_env + env: + # Step outputs reach the script through the environment, not ``${{ }}``, so + # bash never parses them as code (zizmor's template-injection audit). + PYVER: ${{ steps.pyver.outputs.selected }} + SOURCETYPE: ${{ steps.sourcetype.outputs.selected }} run: | - # Install dependencies based on the needs of the job. - # Don't panic! This may look scary at a glance, but each line makes sense. - # - # First let's randomly get versions of dependencies to install. - # Consider removing old versions when they become problematic or very old (>=2 years). - - # Randomly choosing versions of dependencies based on Python version works surprisingly well... - if [[ ${{ startsWith(steps.pyver.outputs.selected, '3.10') }} == true ]]; then - nxver=$(python -c 'import random ; print(random.choice(["=2.8", "=3.0", "=3.1", "=3.2", "=3.3", "=3.4", ""]))') - npver=$(python -c 'import random ; print(random.choice(["=1.24", "=1.25", "=1.26", "=2.0", "=2.1", "=2.2", ""]))') - spver=$(python -c 'import random ; print(random.choice(["=1.9", "=1.10", "=1.11", "=1.12", "=1.13", "=1.14", "=1.15", ""]))') - pdver=$(python -c 'import random ; print(random.choice(["=1.5", "=2.0", "=2.1", "=2.2", ""]))') - akver=$(python -c 'import random ; print(random.choice(["=1.10", "=2.0", "=2.1", "=2.2", "=2.3", "=2.4", "=2.5", "=2.6", "=2.7", ""]))') - fmmver=$(python -c 'import random ; print(random.choice(["=1.4", "=1.5", "=1.6", "=1.7", ""]))') - yamlver=$(python -c 'import random ; print(random.choice(["=5.4", "=6.0", ""]))') - sparsever=$(python -c 'import random ; print(random.choice(["=0.14", "=0.15", ""]))') - elif [[ ${{ startsWith(steps.pyver.outputs.selected, '3.11') }} == true ]]; then - nxver=$(python -c 'import random ; print(random.choice(["=2.8", "=3.0", "=3.1", "=3.2", "=3.3", "=3.4", ""]))') - npver=$(python -c 'import random ; print(random.choice(["=1.24", "=1.25", "=1.26", "=2.0", "=2.1", "=2.2", ""]))') - spver=$(python -c 'import random ; print(random.choice(["=1.9", "=1.10", "=1.11", "=1.12", "=1.13", "=1.14", "=1.15", ""]))') - pdver=$(python -c 'import random ; print(random.choice(["=1.5", "=2.0", "=2.1", "=2.2", ""]))') - akver=$(python -c 'import random ; print(random.choice(["=1.10", "=2.0", "=2.1", "=2.2", "=2.3", "=2.4", "=2.5", "=2.6", "=2.7", ""]))') - fmmver=$(python -c 'import random ; print(random.choice(["=1.4", "=1.5", "=1.6", "=1.7", ""]))') - yamlver=$(python -c 'import random ; print(random.choice(["=5.4", "=6.0", ""]))') - sparsever=$(python -c 'import random ; print(random.choice(["=0.14", "=0.15", ""]))') - elif [[ ${{ startsWith(steps.pyver.outputs.selected, '3.12') }} == true ]]; then - nxver=$(python -c 'import random ; print(random.choice(["=3.2", "=3.3", "=3.4", ""]))') - npver=$(python -c 'import random ; print(random.choice(["=1.26", "=2.0", "=2.1", "=2.2", ""]))') - spver=$(python -c 'import random ; print(random.choice(["=1.11", "=1.12", "=1.13", "=1.14", "=1.15", ""]))') - pdver=$(python -c 'import random ; print(random.choice(["=2.1", "=2.2", ""]))') - akver=$(python -c 'import random ; print(random.choice(["=2.4", "=2.5", "=2.6", "=2.7", ""]))') - fmmver=$(python -c 'import random ; print(random.choice(["=1.7", ""]))') - yamlver=$(python -c 'import random ; print(random.choice(["=6.0", ""]))') - sparsever=$(python -c 'import random ; print(random.choice(["=0.14", "=0.15", ""]))') - else # Python 3.13 - nxver=$(python -c 'import random ; print(random.choice(["=3.4", ""]))') - npver=$(python -c 'import random ; print(random.choice(["=2.1", "=2.2", ""]))') - spver=$(python -c 'import random ; print(random.choice(["=1.14", "=1.15", ""]))') - pdver=$(python -c 'import random ; print(random.choice(["=2.2", ""]))') - akver=$(python -c 'import random ; print(random.choice(["=2.7", ""]))') - fmmver=NA # Not yet supported - yamlver=$(python -c 'import random ; print(random.choice(["=6.0", ""]))') - sparsever=NA # Not yet supported - fi - - # But there may be edge cases of incompatibility we need to handle (more handled below) - if [[ ${{ steps.sourcetype.outputs.selected }} == "source" ]]; then - # TODO: there are currently issues with some numpy versions when - # installing python-suitesparse-grphblas from source. - npver="" - spver="" - pdver="" + # Pick random compatible dependency versions using the version picker script. + # See scripts/ci_pick_versions.py for version pools and constraint logic. + pyver=$(python -c "import os; v=os.environ['PYVER'].split('.');print(f'{v[0]}.{v[1]}')") + # Assign first, then eval. `eval "$(cmd)"` reports eval's status, not cmd's, and + # this step's shell has no `-e`, so a picker that exits non-zero (a pool table it + # rejects, say) would leave every *ver variable unset and quietly install the + # latest of everything: a green run that tested nothing it claims to have tested. + if ! picks=$(python scripts/ci_pick_versions.py --python "$pyver" --source "$SOURCETYPE") ; then + echo "::error::ci_pick_versions.py failed; refusing to install unpinned dependencies" + exit 1 fi + eval "$picks" - # We can have a tight coupling with python-suitesparse-graphblas. - # That is, we don't need to support versions of it that are two years old. - # But, it's still useful for us to test with different versions! + # Set up psg conda package variable (conda-forge installs via conda, others via pip). + # For conda-forge, psg must always be installed; an empty psgver just means "latest". psg="" - if [[ ${{ steps.sourcetype.outputs.selected}} == "upstream" ]] ; then - # Upstream needs to build with numpy 2 - psgver="" - if [[ ${{ startsWith(steps.pyver.outputs.selected, '3.13') }} == true ]]; then - npver=$(python -c 'import random ; print(random.choice(["=2.1", "=2.2", ""]))') - else - npver=$(python -c 'import random ; print(random.choice(["=2.0", "=2.1", "=2.2", ""]))') - fi - elif [[ ${{ startsWith(steps.pyver.outputs.selected, '3.13') }} == true ]] ; then - if [[ ${{ steps.sourcetype.outputs.selected}} == "conda-forge" ]] ; then - psgver=$(python -c 'import random ; print(random.choice(["=9.3.1.0", ""]))') - psg=python-suitesparse-graphblas${psgver} - else - psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", ""]))') - fi - elif [[ ${{ startsWith(steps.pyver.outputs.selected, '3.12') }} == true ]] ; then - if [[ ${{ steps.sourcetype.outputs.selected}} == "conda-forge" ]] ; then - if [[ $npver == =1.* ]] ; then - psgver=$(python -c 'import random ; print(random.choice(["=8.2.0.1", "=8.2.1.0"]))') - else - psgver=$(python -c 'import random ; print(random.choice(["=9.3.1.0", ""]))') - fi - psg=python-suitesparse-graphblas${psgver} - else - if [[ $npver == =1.* ]] ; then - psgver=$(python -c 'import random ; print(random.choice(["==8.2.0.1", "==8.2.1.0"]))') - else - psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", ""]))') - fi - fi - # python-suitsparse-graphblas support is the same for Python 3.10 and 3.11 - elif [[ ${{ steps.sourcetype.outputs.selected}} == "conda-forge" ]] ; then - if [[ $npver == =1.* ]] ; then - psgver=$(python -c 'import random ; print(random.choice(["=7.4.0", "=7.4.1", "=7.4.2", "=7.4.3.0", "=7.4.3.1", "=7.4.3.2", "=8.0.2.1", "=8.2.0.1", "=8.2.1.0"]))') - else - psgver=$(python -c 'import random ; print(random.choice(["=9.3.1.0", ""]))') - fi + if [[ $SOURCETYPE == "conda-forge" ]] ; then psg=python-suitesparse-graphblas${psgver} - elif [[ ${{ steps.sourcetype.outputs.selected}} == "wheel" ]] ; then - if [[ $npver == =1.* ]] ; then - psgver=$(python -c 'import random ; print(random.choice(["==7.4.3.2", "==8.0.2.1", "==8.2.0.1", "==8.2.1.0"]))') - else - psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", ""]))') - fi - elif [[ ${{ steps.sourcetype.outputs.selected}} == "source" ]] ; then - # These should be exact versions - if [[ $npver == =1.* ]] ; then - psgver=$(python -c 'import random ; print(random.choice(["==7.4.0.0", "==7.4.1.0", "==7.4.2.0", "==7.4.3.0", "==7.4.3.1", "==7.4.3.2", "==8.0.2.1", "==8.2.0.1", "==8.2.1.0"]))') - else - psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", ""]))') - fi fi + # Wheel and source draws install psg in "Build extension module", where this + # step's shell variables no longer exist, so psgver goes there as an output. + echo "psgver=$psgver" >> "$GITHUB_OUTPUT" - # Numba is tightly coupled to numpy versions - if [[ ${npver} == "=1.26" ]] ; then - numbaver=$(python -c 'import random ; print(random.choice(["=0.58", "=0.59", "=0.60", "=0.61", ""]))') - if [[ ${spver} == "=1.9" ]] ; then - spver=$(python -c 'import random ; print(random.choice(["=1.10", "=1.11", ""]))') - fi - elif [[ ${npver} == "=1.25" ]] ; then - numbaver=$(python -c 'import random ; print(random.choice(["=0.58", "=0.59", "=0.60", "=0.61", ""]))') - elif [[ ${npver} == "=1.24" || ${{ startsWith(steps.pyver.outputs.selected, '3.11') }} == true ]] ; then - numbaver=$(python -c 'import random ; print(random.choice(["=0.57", "=0.58", "=0.59", "=0.60", "=0.61", ""]))') - else - numbaver="" - fi - # Only numba >=0.59 support Python 3.12 - if [[ ${{ startsWith(steps.pyver.outputs.selected, '3.12') }} == true ]] ; then - numbaver=$(python -c 'import random ; print(random.choice(["=0.59", "=0.60", "=0.61", ""]))') + # Wheels bundle their own GraphBLAS. Otherwise give a pinned psg the GraphBLAS it was + # built for: python-suitesparse-graphblas X.Y.Z.W wraps SuiteSparse:GraphBLAS X.Y.Z, + # and psg's test_version_tracks_libgraphblas checks that they match. conda-forge's + # psg X.Y.Z.W accepts any graphblas >=X.Y.Z,=2.2.2 supports NumPy 2 - pdver=$(python -c 'import random ; print(random.choice(["=2.2", ""]))') - - # Only awkward >=2.6.3 supports NumPy 2 - if [[ ${{ startsWith(steps.pyver.outputs.selected, '3.13') }} == true ]] ; then - akver=$(python -c 'import random ; print(random.choice(["=2.7", ""]))') - else - akver=$(python -c 'import random ; print(random.choice(["=2.6", "=2.7", ""]))') - fi - - # Only scipy >=1.13 supports NumPy 2 - if [[ $spver == "=1.9" || $spver == "=1.10" || $spver == "=1.11" || $spver == "=1.12" ]] ; then - spver="=1.13" - fi + # A psg wheel on macOS bundles its own libomp, and with conda-forge's libomp in the + # same process too, a GraphBLAS parallel region can hang or crash (see + # KMP_DUPLICATE_LIB_OK). conda-forge's libomp comes in with OpenBLAS's openmp build + # and with Numba's default "omp" threading layer, so take OpenBLAS's pthreads build + # and Numba's "workqueue" layer instead. conda-forge has pthreads builds for macOS + # from libopenblas 0.3.30 on only, which scipy 1.9-1.11 do not accept, so those + # draws keep the openmp build. + blas="" + if [[ $SOURCETYPE == "wheel" && ${{ contains(matrix.os, 'macos') }} == true ]] ; then + case $spver in + =1.9 | =1.10 | =1.11) ;; + *) blas="libopenblas=*=*pthreads*" ;; + esac + echo "NUMBA_THREADING_LAYER=workqueue" >> "$GITHUB_ENV" fi - fmm=fast_matrix_market${fmmver} + # Decide whether to install numba and sparse. + # Random skip: pypy always, 20% on Windows, 40% on other OSes (unless notebooks). awkward=awkward${akver} - - # Don't install numba and sparse for some versions - if [[ ${{ contains(steps.pyver.outputs.selected, 'pypy') || - startsWith(steps.pyver.outputs.selected, '3.14') }} == true || + if [[ ${{ contains(steps.pyver.outputs.selected, 'pypy') }} == true || ( ${{ matrix.slowtask != 'notebooks'}} == true && ( ( ${{ matrix.os == 'windows-latest' }} == true && $(python -c 'import random ; print(random.random() < .2)') == True ) || ( ${{ matrix.os == 'windows-latest' }} == false && $(python -c 'import random ; print(random.random() < .4)') == True ))) ]] then - # Some packages aren't available for pypy or Python 3.13; randomly otherwise (if not running notebooks) echo "skipping numba" numba="" numbaver=NA sparse="" sparsever=NA - if [[ ${{ contains(steps.pyver.outputs.selected, 'pypy') }} ]]; then + # Compare explicitly: `[[ false ]]` tests a non-empty string, not a + # boolean, so testing the substituted value on its own is always + # true. Matches the `== true` form used for the numba skip above. + if [[ ${{ contains(steps.pyver.outputs.selected, 'pypy') }} == true ]]; then awkward="" akver=NA - fmm="" - fmmver=NA - # Be more flexible until we determine what versions are supported by pypy npver="" spver="" pdver="" yamlver="" fi - elif [[ ${npver} == =2.* ]] ; then - # Don't install numba for unsupported versions of numpy + elif [[ ${numbaver} == "NA" ]] ; then + # Version picker already determined numba is incompatible numba="" - numbaver=NA sparse="" sparsever=NA else numba=numba${numbaver} - sparse=sparse${sparsever} - fi - - # sparse does not yet support Python 3.13 - if [[ ${{ startsWith(steps.pyver.outputs.selected, '3.13') }} == true ]] ; then - sparse="" - sparsever=NA - fi - # fast_matrix_market does not yet support Python 3.13 or osx-arm64 - if [[ ${{ startsWith(steps.pyver.outputs.selected, '3.13') }} == true || - ${{ matrix.os == 'macos-latest' }} == true ]] - then - fmm="" - fmmver=NA + if [[ $sparsever == "NA" ]] ; then + sparse="" + else + sparse=sparse${sparsever} + fi fi echo "versions: np${npver} sp${spver} pd${pdver} ak${akver} nx${nxver} numba${numbaver} yaml${yamlver} sparse${sparsever} psg${psgver}" set -x # echo on - $(command -v mamba || command -v conda) install -c nodefaults \ - packaging pytest coverage pytest-randomly cffi donfig tomli c-compiler make \ + # conda-forge ships two ABIs for CPython 3.13+, and this step did not + # restate ``python``, so the solver was free to replace the interpreter + # setup-miniconda installed: a 3.14 draw on ubuntu-24.04-arm became + # 3.14.7-*_cp314t, where pandas 2.3 raises a GIL RuntimeWarning that + # filterwarnings=error turns into a collection error. Restating the + # version pins the interpreter and ``python-gil`` pins python_abi to + # the GIL build; setup-miniconda's own ``python-version`` above is + # equally ABI-free, so this covers that step too. + # python-suitesparse-graphblas pins the same pair in + # .github/workflows/test.yml, and asks for python-freethreading on the + # matrix entries that want the other ABI. + if [[ $PYVER == *pypy* ]] ; then + interpreter=pypy + else + interpreter="python=$PYVER python-gil" + fi + # No -c: the channels are the ones setup-miniconda configured. On the command + # line, nodefaults is not the environment-file keyword but a channel on + # anaconda.org, which -c would put first. + $(command -v mamba || command -v conda) install \ + packaging pip pytest coverage pytest-randomly cffi donfig tomli c-compiler make \ pyyaml${yamlver} ${sparse} pandas${pdver} scipy${spver} numpy${npver} ${awkward} \ - networkx${nxver} ${numba} ${fmm} ${psg} \ + networkx${nxver} ${numba} ${psg} \ ${{ matrix.slowtask == 'pytest_bizarro' && 'black' || '' }} \ - ${{ matrix.slowtask == 'notebooks' && 'matplotlib nbconvert jupyter "ipython>=7" drawsvg' || '' }} \ + ${{ matrix.slowtask == 'notebooks' && '"matplotlib${mplver}" nbconvert jupyter "ipython>=7" drawsvg' || '' }} \ ${{ steps.sourcetype.outputs.selected == 'upstream' && 'cython' || '' }} \ - ${{ steps.sourcetype.outputs.selected != 'wheel' && '"graphblas>=7.4,<9.4"' || '' }} \ - ${{ contains(steps.pyver.outputs.selected, 'pypy') && 'pypy' || '' }} \ + ${graphblas} ${blas:+"$blas"} \ + ${interpreter} \ ${{ matrix.os == 'windows-latest' && 'cmake' || 'm4' }} \ # ${{ matrix.os != 'windows-latest' && 'pytest-forked' || '' }} # to investigate crashes - name: Build extension module + env: + SOURCETYPE: ${{ steps.sourcetype.outputs.selected }} + PSGVER: ${{ steps.update_env.outputs.psgver }} run: | - if [[ ${{ steps.sourcetype.outputs.selected }} == "wheel" ]]; then + # ``python -m pip`` ensures the conda env's interpreter drives the + # install, not a system ``pip`` on PATH. + if [[ $SOURCETYPE == "wheel" ]]; then # Add --pre if installing a pre-release - pip install --no-deps --only-binary ":all:" suitesparse-graphblas${psgver} + python -m pip install --no-deps --only-binary ":all:" "suitesparse-graphblas${PSGVER}" # Add the below line to the conda install command above if installing from test.pypi.org # ${{ steps.sourcetype.outputs.selected == 'wheel' && 'setuptools setuptools-git-versioning wheel cython' || '' }} \ - # pip install --no-deps --only-binary ":all:" --index-url https://test.pypi.org/simple/ "suitesparse-graphblas>=7.4.3" - elif [[ ${{ steps.sourcetype.outputs.selected }} == "source" ]]; then + # python -m pip install --no-deps --only-binary ":all:" --index-url https://test.pypi.org/simple/ "suitesparse-graphblas>=7.4.3" + elif [[ $SOURCETYPE == "source" ]]; then # Add --pre if installing a pre-release - pip install --no-deps --no-binary suitesparse-graphblas suitesparse-graphblas${psgver} + python -m pip install --no-deps --no-binary suitesparse-graphblas "suitesparse-graphblas${PSGVER}" # Add the below line to the conda install command above if installing from test.pypi.org # ${{ steps.sourcetype.outputs.selected == 'source' && 'setuptools setuptools-git-versioning wheel cython' || '' }} \ - # pip install --no-deps --no-build-isolation --no-binary suitesparse-graphblas --index-url https://test.pypi.org/simple/ suitesparse-graphblas==7.4.3.3 - elif [[ ${{ steps.sourcetype.outputs.selected }} == "upstream" ]]; then - pip install --no-deps git+https://github.com/GraphBLAS/python-suitesparse-graphblas.git@main#egg=suitesparse-graphblas + # python -m pip install --no-deps --no-build-isolation --no-binary suitesparse-graphblas --index-url https://test.pypi.org/simple/ suitesparse-graphblas==7.4.3.3 + elif [[ $SOURCETYPE == "upstream" ]]; then + python -m pip install --no-deps git+https://github.com/GraphBLAS/python-suitesparse-graphblas.git@main#egg=suitesparse-graphblas fi - pip install --no-deps -e . + python -m pip install --no-deps -e . - name: python-suitesparse-graphblas tests run: | # Don't use our conftest.py ; allow `test_print_jit_config` to fail if it doesn't exist (cd .. pytest --pyargs suitesparse_graphblas -s -k test_print_jit_config || true pytest -v --pyargs suitesparse_graphblas || true) - - name: Print platform and sysconfig variables - run: | - python -c "import platform ; print(platform.uname())" - python -c "import pprint, sysconfig ; pprint.pprint(sysconfig.get_config_vars())" + - name: JIT diagnostics + run: python scripts/jit_diagnostics.py || true - name: Unit tests + # A hang (e.g. two libomp copies deadlocking in an OpenMP barrier on macOS) would + # otherwise hold the runner for GitHub's 6-hour limit. + timeout-minutes: 120 run: | A=${{ needs.rngs.outputs.mapnumpy == 'A' || '' }} ; B=${{ needs.rngs.outputs.mapnumpy == 'B' || '' }} C=${{ needs.rngs.outputs.mapnumpy == 'C' || '' }} ; D=${{ needs.rngs.outputs.mapnumpy == 'D' || '' }} @@ -424,6 +355,7 @@ jobs: coverage run -m pytest --color=yes --randomly -v ${args} \ ${{ matrix.slowtask == 'pytest_normal' && '--runslow' || '' }} - name: Unit tests (bizarro scalars) + timeout-minutes: 120 run: | # Run tests again with Scalars being C scalars by default find graphblas -type f -name "*.py" -print0 | xargs -0 sed -i -s \ @@ -464,6 +396,7 @@ jobs: git checkout . # Undo changes to scalar default - name: Miscellaneous tests if: matrix.slowtask == 'pytest_normal' + timeout-minutes: 30 run: | # Test (and cover) automatic initialization coverage run -a graphblas/tests/test_auto_init.py @@ -487,11 +420,24 @@ jobs: # coverage run -a -m pytest --color=yes -x --no-mapnumpy -k test_npmonoid graphblas/tests/test_numpyops.py --runslow - name: More tests for coverage if: matrix.slowtask == 'notebooks' && matrix.os == 'windows-latest' + timeout-minutes: 30 run: | # We use 'notebooks' slow task b/c it should have numba installed coverage run -a -m pytest --color=yes --runslow --no-mapnumpy -p no:randomly -v -k 'test_commutes or test_bool_doesnt_get_too_large or test_npbinary or test_npmonoid or test_npsemiring' coverage run -a -m pytest --color=yes --runslow --mapnumpy -p no:randomly -k 'test_bool_doesnt_get_too_large or test_npunary or test_binaryop_monoid_numpy' coverage run -a -m pytest --color=yes -x --no-mapnumpy --runslow -k test_binaryop_attributes_numpy graphblas/tests/test_op.py + - name: JIT tests for coverage + if: matrix.slowtask == 'notebooks' && matrix.os != 'windows-latest' + timeout-minutes: 30 + run: | + # Most of test_ssjit.py needs the suitesparse backend, --runslow, Numba and a + # GraphBLAS that can JIT-compile, all in one pytest run. In one run in four or five, + # no job draws all four, so its coverage came and went. 'notebooks' jobs install + # Numba unless the picker rules it out. Windows is left out: conda-forge's GraphBLAS + # there names the build machine's compiler, so it never compiles. The tests skip + # what a draw cannot do: the JIT tests on psg's wheels (built without the JIT) and + # on GraphBLAS 7, the UDT tests without Numba. + coverage run -a -m pytest --color=yes --runslow --backend=suitesparse -p no:randomly -v graphblas/tests/test_ssjit.py - name: Auto-generated code check if: matrix.slowtask == 'pytest_bizarro' run: | @@ -504,9 +450,15 @@ jobs: coverage xml coverage report --show-missing - name: codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@303a32d7a59b442fa8d48b6a1cc6825c09c847a5 # v7.1.1 + with: + # Codecov takes tokenless uploads only from forks. A PR from a branch of this + # repository, or a push to main, is refused with "Token required because + # branch is protected". + token: ${{ secrets.CODECOV_TOKEN }} - name: Notebooks Execution check if: matrix.slowtask == 'notebooks' + timeout-minutes: 30 run: | # Run notebooks only if numba is installed if python -c 'import numba' 2> /dev/null ; then diff --git a/.github/zizmor.yml b/.github/zizmor.yml deleted file mode 100644 index 61f32c2e0..000000000 --- a/.github/zizmor.yml +++ /dev/null @@ -1,16 +0,0 @@ -rules: - use-trusted-publishing: - # TODO: we should update to use trusted publishing - ignore: - - publish_pypi.yml - excessive-permissions: - # It is probably good practice to use narrow permissions - ignore: - - debug.yml - - imports.yml - - publish_pypi.yml - - test_and_build.yml - template-injection: - # We use templates pretty heavily - ignore: - - test_and_build.yml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 43e28b8fe..8515f7f41 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,26 +10,32 @@ ci: autoupdate_schedule: quarterly autoupdate_commit_msg: "chore: update pre-commit hooks" autofix_commit_msg: "style: pre-commit fixes" - skip: [pylint, no-commit-to-branch] + # pysentry downloads advisory databases at run time, and pre-commit.ci + # runs hooks without network access, so it runs locally and in security.yml. + # zizmor does run on pre-commit.ci, offline, which skips its online audits; + # security.yml runs those. + skip: [pylint, no-commit-to-branch, pysentry] fail_fast: false default_language_version: python: python3 repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + rev: v6.0.0 hooks: - id: check-added-large-files - id: check-case-conflict + - id: check-executables-have-shebangs + - id: check-illegal-windows-names - id: check-merge-conflict - # - id: check-symlinks - id: check-ast + - id: check-json # no JSON files yet; enabled for the future - id: check-toml - id: check-yaml - - id: check-executables-have-shebangs - id: check-vcs-permalinks + - id: debug-statements - id: destroyed-symlinks - id: detect-private-key - - id: debug-statements + - id: forbid-submodules - id: end-of-file-fixer exclude_types: [svg] - id: mixed-line-ending @@ -37,114 +43,189 @@ repos: - id: name-tests-test args: ["--pytest-test-first"] - repo: https://github.com/abravalheri/validate-pyproject - rev: v0.23 + rev: v0.27 hooks: - id: validate-pyproject name: Validate pyproject.toml # I don't yet trust ruff to do what autoflake does - repo: https://github.com/PyCQA/autoflake - rev: v2.3.1 + rev: v2.4.0 hooks: - id: autoflake args: [--in-place] # We can probably remove `isort` if we come to trust `ruff --fix`, # but we'll need to figure out the configuration to do this in `ruff` - repo: https://github.com/pycqa/isort - rev: 6.0.0 + rev: 9.0.2 hooks: - id: isort # Let's keep `pyupgrade` even though `ruff --fix` probably does most of it - repo: https://github.com/asottile/pyupgrade - rev: v3.19.1 + rev: v3.22.0 hooks: - id: pyupgrade - args: [--py310-plus] + args: [--py311-plus] - repo: https://github.com/MarcoGorelli/auto-walrus - rev: 0.3.4 + rev: 0.4.1 hooks: - id: auto-walrus args: [--line-length, "100"] - - repo: https://github.com/psf/black - rev: 25.1.0 + - repo: https://github.com/psf/black-pre-commit-mirror + rev: 26.10.1 hooks: - id: black - id: black-jupyter - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.9.6 + rev: v0.17.3 hooks: - - id: ruff + - id: ruff-check args: [--fix-only, --show-fixes] # Let's keep `flake8` even though `ruff` does much of the same. # `flake8-bugbear` and `flake8-simplify` have caught things missed by `ruff`. - repo: https://github.com/PyCQA/flake8 - rev: 7.1.2 + rev: 7.5.0 hooks: - id: flake8 args: ["--config=.flake8"] additional_dependencies: &flake8_dependencies # These versions need updated manually - - flake8==7.1.2 - - flake8-bugbear==24.12.12 - - flake8-simplify==0.21.0 + - flake8==7.5.0 + - flake8-bugbear==26.10.27 + - flake8-simplify==0.31.1 - repo: https://github.com/codespell-project/codespell - rev: v2.4.1 + rev: v2.4.3 hooks: - id: codespell - types_or: [python, rst, markdown] - additional_dependencies: [tomli] + types_or: [python, markdown, rst, toml, yaml] + additional_dependencies: + - tomli; python_version<'3.11' files: ^(graphblas|docs)/ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.9.6 + rev: v0.17.3 hooks: - - id: ruff + - id: ruff-check + # - id: ruff-format # we prefer black; uncomment to compare with ruff's formatter - repo: https://github.com/sphinx-contrib/sphinx-lint - rev: v1.0.0 + rev: v1.0.2 hooks: - id: sphinx-lint args: [--enable, all, "--disable=line-too-long,leaked-markup"] # `pyroma` may help keep our package standards up to date if best practices change. # This is probably a "low value" check though and safe to remove if we want faster pre-commit. - repo: https://github.com/regebro/pyroma - rev: "4.2" + # 5.1 is still in beta (5.1b2) and a pre-release pin defeats pre-commit.ci's + # autoupdate "avoid prereleases" logic; stay on 5.0.1 until 5.1 GAs. + rev: "5.0.1" hooks: - id: pyroma args: [-n, "10", .] - repo: https://github.com/shellcheck-py/shellcheck-py - rev: "v0.10.0.1" + rev: "v0.11.0.1-1" hooks: - id: shellcheck - repo: https://github.com/rbubley/mirrors-prettier - rev: v3.5.1 + rev: v3.9.10 hooks: - id: prettier + args: [--prose-wrap=preserve] + - repo: https://github.com/pre-commit/pygrep-hooks + rev: v1.10.0 + hooks: + - id: rst-backticks + - id: rst-directive-colons + - id: rst-inline-touching-normal + - id: python-check-blanket-noqa + - id: python-check-blanket-type-ignore + - id: python-no-eval + - id: python-no-log-warn + - id: text-unicode-replacement-char - repo: https://github.com/ComPWA/taplo-pre-commit rev: v0.9.3 hooks: - id: taplo-format + args: ["--option", "column_width=100"] - repo: https://github.com/rhysd/actionlint - rev: v1.7.7 + rev: v1.7.12 hooks: - id: actionlint + - repo: https://github.com/zizmorcore/zizmor-pre-commit + rev: v1.32.0 + hooks: + - id: zizmor - repo: https://github.com/python-jsonschema/check-jsonschema - rev: 0.31.1 + rev: 0.39.0 hooks: - id: check-dependabot - id: check-github-workflows - id: check-readthedocs - repo: https://github.com/adrienverge/yamllint - rev: v1.35.1 + rev: v1.38.0 hooks: - id: yamllint - - repo: https://github.com/woodruffw/zizmor-pre-commit - rev: v1.3.1 + - repo: https://github.com/pysentry/pysentry-pre-commit + rev: v0.5.0 hooks: - - id: zizmor + - id: pysentry + # Run when pyproject.toml or this file changes rather than at every commit, where + # it fails offline. security.yml also runs it weekly. + always_run: false + files: ^(pyproject\.toml|\.pre-commit-config\.yaml)$ + # Audit the floors in pyproject.toml. Without `--direct-only`, pysentry resolves the + # newest versions with uv (which the hook installs), and no longer sees a floor that + # allows a vulnerable release. + args: + - --format=human + - --fail-on=low + - --sources=pypa,osv,pypi + - --direct-only - repo: meta hooks: - id: check-hooks-apply - id: check-useless-excludes - repo: local hooks: + - id: disallow-caps + name: Disallow improper capitalization + language: pygrep + entry: PyBind|Numpy\>|Cmake|CCache|Github|PyTest|\|PyLint + exclude: (.pre-commit-config.yaml|docs/pages/guides/style\.md)$ + - id: disallow-words + name: Disallow certain words + language: pygrep + entry: "[Ff]alsey" + exclude: .pre-commit-config.yaml$ + - id: disallow-bad-permalinks + name: Disallow _ in permalinks + language: pygrep + entry: "^permalink:.*_.*" + - id: disallow-ai-dashes + # AI-ism: em-dash, en-dash, " -- " separator, or "word--jam". + # See ``scripts/find_aiisms.sh`` for the manual / branch-diff variant + # and ``ai/CLAUDE.md`` "Writing Style" for the rationale. + name: Disallow em/en-dash and ' -- ' as a separator + language: pygrep + entry: "—|–| -- |[A-Za-z0-9]--" + types_or: [python, rst, markdown, toml, yaml, shell] + # README intentionally uses en-dash. docs/conf.py + jit_diagnostics.py + # use ``# --- Section ---`` banner comments / print output, which the + # regex also matches. find_aiisms.sh contains the matching regex. + exclude: | + (?x)^( + README\.md + | docs/conf\.py + | scripts/jit_diagnostics\.py + | scripts/find_aiisms\.sh + | \.pre-commit-config\.yaml + )$ + # Stress-tests every combination the CI dependency picker can produce against + # the constraints it documents. ~3s for 160k combinations, and it only runs + # when the picker itself changes. + - id: ci-pick-versions + name: Validate CI dependency version picker + entry: python scripts/ci_pick_versions.py --validate + language: system + files: ^scripts/ci_pick_versions\.py$ + pass_filenames: false # Add `--hook-stage manual` to pre-commit command to run (very slow) # It's probably better (and faster!) to simply run `pylint graphblas/some/file.py` - id: pylint @@ -156,7 +237,7 @@ repos: args: [graphblas/] pass_filenames: false - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + rev: v6.0.0 hooks: - id: no-commit-to-branch # no commit directly to main # diff --git a/.readthedocs.yml b/.readthedocs.yml index b35d401db..1bcd4baef 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -6,12 +6,13 @@ sphinx: configuration: docs/conf.py fail_on_warning: True -# Try to make build faster with mamba. See: -# https://docs.readthedocs.io/en/stable/guides/conda.html#making-builds-faster-with-mamba +# Read the Docs dropped the ubuntu-20.04 image, and mambaforge is superseded by +# miniforge3 (whose conda solves with libmamba anyway). Valid values are listed at: +# https://docs.readthedocs.io/en/stable/config-file/v2.html#build-tools-python build: - os: "ubuntu-20.04" + os: "ubuntu-24.04" tools: - python: "mambaforge-4.10" + python: "miniforge3-latest" conda: environment: docs/env.yml diff --git a/README.md b/README.md index 1080314c7..3b84afef2 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ For algorithms, see - **SuiteSparse:GraphBLAS User Guide:** [https://github.com/DrTimothyAldenDavis/GraphBLAS/raw/stable/Doc/GraphBLAS_UserGuide.pdf](https://github.com/DrTimothyAldenDavis/GraphBLAS/raw/stable/Doc/GraphBLAS_UserGuide.pdf) - **Source:** [https://github.com/python-graphblas/python-graphblas](https://github.com/python-graphblas/python-graphblas) - **Bug reports:** [https://github.com/python-graphblas/python-graphblas/issues](https://github.com/python-graphblas/python-graphblas/issues) -- **Github discussions:** [https://github.com/python-graphblas/python-graphblas/discussions](https://github.com/python-graphblas/python-graphblas/discussions) +- **GitHub discussions:** [https://github.com/python-graphblas/python-graphblas/discussions](https://github.com/python-graphblas/python-graphblas/discussions) - **Weekly community call:** [python-graphblas#247](https://github.com/python-graphblas/python-graphblas/issues/247) or [https://scientific-python.org/calendars/](https://scientific-python.org/calendars/) - **Chat via Discord:** [https://discord.com/invite/vur45CbwMz](https://discord.com/invite/vur45CbwMz) in the [#graphblas channel](https://discord.com/channels/786703927705862175/1024732940233605190) @@ -58,8 +58,7 @@ The following are not required by python-graphblas, but may be needed for certai - `pandas` – required for nicer `__repr__`; - `matplotlib` – required for basic plotting of graphs; - `scipy` – used in `io` module to read/write `scipy.sparse` format; -- `networkx` – used in `io` module to interface with `networkx` graphs; -- `fast-matrix-market` - for faster read/write of Matrix Market files with `gb.io.mmread` and `gb.io.mmwrite`. +- `networkx` – used in `io` module to interface with `networkx` graphs. ## Description @@ -179,7 +178,7 @@ s(accum) << v.reduce(op) ## Creating new Vectors / Matrices ```python -A = Matrix.new(dtype, num_rows, num_cols) # new_type +A = Matrix(dtype, num_rows, num_cols) # new_type B = A.dup() # dup A = Matrix.from_coo([row_indices], [col_indices], [values]) # build ``` @@ -226,7 +225,7 @@ Python-graphblas requires `numba` which enables compiling user-defined Python fu Example customized UnaryOp: ```python -from graphblas import unary +from graphblas import unary, Vector def force_odd_func(x): if x % 2 == 0: @@ -268,7 +267,7 @@ import graphblas as gb # scipy.sparse matrices A = gb.io.from_scipy_sparse(m) -m = gb.io.to_scipy_sparse(m, format="csr") +m = gb.io.to_scipy_sparse(A, format="csr") # networkx graphs A = gb.io.from_networkx(g) diff --git a/dev-requirements.txt b/dev-requirements.txt index a281672ec..a986ece58 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -6,7 +6,7 @@ pyyaml pandas # For I/O awkward -fast_matrix_market +# fast_matrix_market is deprecated (no longer maintained; last supported Python is 3.12) networkx scipy sparse @@ -29,5 +29,6 @@ mypy # For building docs nbsphinx numpydoc -pydata-sphinx-theme -sphinx-panels +pydata-sphinx-theme>=0.15 +sphinx>=7 +sphinx-design diff --git a/docs/_static/custom.css b/docs/_static/custom.css index f7dd59b74..98331d655 100644 --- a/docs/_static/custom.css +++ b/docs/_static/custom.css @@ -5,8 +5,12 @@ margin-bottom: 30px; } -.intro-card:hover { - box-shadow: 0.2rem 0.5rem 1rem var(--pst-color-link) !important; +html[data-theme="light"] .intro-card:hover { + box-shadow: 0 0.5rem 1rem rgba(0, 0, 0, 0.18) !important; +} + +html[data-theme="dark"] .intro-card:hover { + box-shadow: 0 0.5rem 1rem var(--pst-color-link) !important; } .intro-card .card-header { diff --git a/docs/_static/matrix.css b/docs/_static/matrix.css index 1937178e5..b773003cd 100644 --- a/docs/_static/matrix.css +++ b/docs/_static/matrix.css @@ -5,7 +5,7 @@ table.matrix { border: 0px; } -/* Disable a horizintal line from the default stylesheet */ +/* Disable a horizontal line from the default stylesheet */ .table.matrix > :not(caption) > * > * { border-bottom-width: 0px; } @@ -48,7 +48,7 @@ table.matrix > tbody tr td:last-of-type { /* prevents empty cells from collapsing, especially empty rows */ table.matrix > tbody tr td:empty::before { - /* basicaly fills empty cells with   */ + /* basically fills empty cells with   */ content: "\00a0\00a0\00a0"; visibility: hidden; } diff --git a/docs/api_reference/index.rst b/docs/api_reference/index.rst index 84e7d65eb..219e19bae 100644 --- a/docs/api_reference/index.rst +++ b/docs/api_reference/index.rst @@ -8,6 +8,8 @@ API Reference :maxdepth: 2 collections + types operators io + utilities exceptions diff --git a/docs/api_reference/io.rst b/docs/api_reference/io.rst index 1cfc98516..9c8b8e2cb 100644 --- a/docs/api_reference/io.rst +++ b/docs/api_reference/io.rst @@ -73,3 +73,7 @@ Visualization ~~~~~~~~~~~~~ .. autofunction:: graphblas.viz.draw + +.. autofunction:: graphblas.viz.spy + +.. autofunction:: graphblas.viz.datashade diff --git a/docs/api_reference/operators.rst b/docs/api_reference/operators.rst index 8836bb638..6853831cd 100644 --- a/docs/api_reference/operators.rst +++ b/docs/api_reference/operators.rst @@ -31,8 +31,20 @@ IndexUnaryOp .. autoclass:: graphblas.core.operator.IndexUnaryOp() :members: +IndexBinaryOp +~~~~~~~~~~~~~ + +.. autoclass:: graphblas.core.operator.IndexBinaryOp() + :members: + SelectOp ~~~~~~~~ .. autoclass:: graphblas.core.operator.SelectOp() :members: + +Aggregator +~~~~~~~~~~ + +.. autoclass:: graphblas.core.operator.Aggregator() + :members: diff --git a/docs/api_reference/types.rst b/docs/api_reference/types.rst new file mode 100644 index 000000000..e8187caa5 --- /dev/null +++ b/docs/api_reference/types.rst @@ -0,0 +1,22 @@ +Types +----- + +DataType +~~~~~~~~ + +The object returned by :func:`~graphblas.dtypes.register_new` and +:func:`~graphblas.dtypes.register_anonymous`, and carried by every collection's +``.dtype``. The JIT introspection properties below report what SuiteSparse +actually registered for a user-defined type; see :doc:`../user_guide/udt` for +the operator-side counterparts on a typed operator (``op.jit_c_name``, +``op.jit_c_source``). + +.. autoclass:: graphblas.dtypes.DataType() + :members: jit_c_name, jit_c_definition + +Registering user-defined types +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autofunction:: graphblas.dtypes.register_new + +.. autofunction:: graphblas.dtypes.register_anonymous diff --git a/docs/api_reference/utilities.rst b/docs/api_reference/utilities.rst new file mode 100644 index 000000000..8e6e10486 --- /dev/null +++ b/docs/api_reference/utilities.rst @@ -0,0 +1,13 @@ +Utilities +--------- + +Initialization +~~~~~~~~~~~~~~~ + +.. autofunction:: graphblas.init + +Recorder +~~~~~~~~ + +.. autoclass:: graphblas.Recorder + :members: diff --git a/docs/conf.py b/docs/conf.py index 283f6d047..b6b269114 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -19,7 +19,7 @@ # -- Project information ----------------------------------------------------- project = "python-graphblas" -copyright = "2020-2023, Anaconda, Inc. and contributors" +copyright = "2020-2026, Anaconda, Inc. and contributors" author = "Anaconda, Inc" # The full version, including alpha/beta/rc tags @@ -35,7 +35,7 @@ # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. -extensions = ["sphinx.ext.autodoc", "numpydoc", "sphinx_panels", "nbsphinx"] +extensions = ["sphinx.ext.autodoc", "numpydoc", "sphinx_design", "nbsphinx"] html_css_files = ["custom.css", "matrix.css"] html_js_files = ["custom.js"] diff --git a/docs/env.yml b/docs/env.yml index 78a50afbe..e47eea420 100644 --- a/docs/env.yml +++ b/docs/env.yml @@ -3,7 +3,7 @@ channels: - conda-forge - nodefaults dependencies: - - python=3.10 + - python=3.13 - pip # python-graphblas dependencies - donfig @@ -16,8 +16,8 @@ dependencies: - pandas - scipy>=1.7.0 # docs dependencies - - commonmark # For RTD - nbsphinx - numpydoc - - pydata-sphinx-theme=0.13.1 - - sphinx-panels=0.6.0 + - pydata-sphinx-theme>=0.15 + - sphinx>=7 + - sphinx-design diff --git a/docs/getting_started/index.rst b/docs/getting_started/index.rst index 3726131d2..cbd0bb6a6 100644 --- a/docs/getting_started/index.rst +++ b/docs/getting_started/index.rst @@ -30,11 +30,10 @@ Optional Dependencies The following are not required by ``python-graphblas``, but may be needed for certain functionality to work. - - `pandas `__ -- required for nicer ``__repr__`` - - `matplotlib `__ -- required for basic plotting of graphs - - `scipy `__ -- used in ``io`` module to read/write ``scipy.sparse`` format - - `networkx `__ -- used in ``io`` module to interface with networkx graphs - - `fast-matrix-market `__ -- for faster read/write of Matrix Market files with ``gb.io.mmread`` and ``gb.io.mmwrite`` + - `pandas `__: required for nicer ``__repr__`` + - `matplotlib `__: required for basic plotting of graphs + - `scipy `__: used in ``io`` module to read/write ``scipy.sparse`` format + - `networkx `__: used in ``io`` module to interface with networkx graphs GraphBLAS Fundamentals ---------------------- diff --git a/docs/index.rst b/docs/index.rst index ba752d558..b18582087 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -19,60 +19,51 @@ Licensed under the `Apache 2.0 license `_). + With multiprocessing, dask, pytest-xdist and the like, set + ``GRAPHBLAS_JIT_C_CONTROL=run`` in the environment the workers start in, so that none + of them compiles: UDT work raises only the default ``"run"``, not one you set. (At the + default, the first UDT operator that needs a kernel turns the JIT on for the whole + process, so such processes race too.) Setting ``"run"`` costs UDTs their compiled + kernels: UDT operators call their Numba functions instead + (typically 2 to 3 times slower for elementwise operations), and registering types or + operators from C source fails. To keep the kernels, give each process its own cache + instead, by setting the ``GRAPHBLAS_CACHE_PATH`` environment variable before graphblas + initializes. The default will change once SuiteSparse publishes compiled kernels + atomically. +- The first use of each new kernel compiles it, typically in 0.2 to 1 second. Kernels are + cached on disk, in ``gb.ss.config["jit_cache_path"]`` (by default under + ``~/.SuiteSparse`` on Linux and macOS), so later runs load them instead. +- The cache is keyed only on the SuiteSparse:GraphBLAS version, so two builds of one + version (a conda build and a source build, say) can load each other's kernels, and get + wrong results or errors if their internals differ. Give each build its own cache path. + +If you hit a bug, freeze or crash that may involve the JIT, set +``GRAPHBLAS_JIT_C_CONTROL=run`` (which UDT work leaves alone) and delete the cache +directory, then try again. Please +`report it `_ either way. The +``suitesparse-vanilla`` backend leaves the control as the SuiteSparse build sets it. diff --git a/docs/user_guide/io.rst b/docs/user_guide/io.rst index f27b40bd3..91223a085 100644 --- a/docs/user_guide/io.rst +++ b/docs/user_guide/io.rst @@ -53,7 +53,11 @@ A python-graphblas Matrix can be created from a NetworkX graph using ``gb.io.fro The dtype and weight label can be specified. A list of nodes can also be provided to only pull out a subset of the NetworkX graph. -``gb.io.to_networkx()`` goes the other direction. It always returns a DiGraph with weights labelled as "weight". +``gb.io.to_networkx()`` goes the other direction for a square Matrix. It returns a DiGraph with one +node per row, numbered by row index (rows with no entries included), and the values stored under the +edge attribute ``edge_attribute`` ("weight" by default, or ``None`` to leave them out). A Matrix with +huge dimensions, such as a hypersparse one, has too many rows to make each a node: pass +``create_isolates=False`` to make nodes only of the indices that have edges. Scipy.Sparse @@ -138,7 +142,7 @@ Matrix Market files The `Matrix Market file format `_ is a common file format for storing sparse arrays in human-readable ASCII. -Matrix Market files--also called MM files--often use ".mtx" file extension. +Matrix Market files (also called MM files) often use ".mtx" file extension. For example, many datasets in MM format can be found in `the SuiteSparse Matrix Collection `_. Use ``gb.io.mmread()`` to read a Matrix Market file to a python-graphblas Matrix, @@ -146,5 +150,3 @@ and ``gb.io.mmwrite()`` to write a Matrix to a Matrix Market file. These names match the equivalent functions in `scipy.sparse `_. ``scipy`` is required to be installed to read Matrix Market files. -If ``fast_matrix_market`` is installed, it will be used by default for -`much better performance `_. diff --git a/docs/user_guide/operations.rst b/docs/user_guide/operations.rst index 18d0352d7..b41af344c 100644 --- a/docs/user_guide/operations.rst +++ b/docs/user_guide/operations.rst @@ -8,9 +8,9 @@ Matrix Multiply The GraphBLAS spec contains three methods for matrix multiplication, depending on whether the inputs are Matrix or Vector. - - **mxm** -- Matrix-Matrix multiplication - - **mxv** -- Matrix-Vector multiplication - - **vxm** -- Vector-Matrix multiplication + - **mxm**: Matrix-Matrix multiplication + - **mxv**: Matrix-Vector multiplication + - **vxm**: Vector-Matrix multiplication These three methods exist on python-graphblas collections, but the preferred approach is using the ``@`` symbol, which indicates matrix multiplication in Python. diff --git a/docs/user_guide/operators.rst b/docs/user_guide/operators.rst index 8bb5e9fa8..3a3b2845a 100644 --- a/docs/user_guide/operators.rst +++ b/docs/user_guide/operators.rst @@ -19,21 +19,21 @@ Example usage: Common unary operators are: - - **identity** -- returns the input unchanged - - **abs** -- absolute value - - **ainv** -- additive inverse (f(x) = -x) - - **minv** -- multiplicative inverse (f(x) = 1/x) - - **lnot** -- logical not (True -> False) - - **bnot** -- binary not (110010 -> 001101) - - **one** -- result is always 1.0 (or equivalent for the dtype) - - **round** -- floating-point round function - - **floor** -- floating-point floor function - - **ceil** -- floating-point ceiling function - - **sin** -- trigonometric sine function - - **cos** -- trigonometric cosine function - - **tan** -- trigonometric tangent function - - **exp** -- exponential - - **log** -- logarithm + - **identity**: returns the input unchanged + - **abs**: absolute value + - **ainv**: additive inverse (f(x) = -x) + - **minv**: multiplicative inverse (f(x) = 1/x) + - **lnot**: logical not (True -> False) + - **bnot**: binary not (110010 -> 001101) + - **one**: result is always 1.0 (or equivalent for the dtype) + - **round**: floating-point round function + - **floor**: floating-point floor function + - **ceil**: floating-point ceiling function + - **sin**: trigonometric sine function + - **cos**: trigonometric cosine function + - **tan**: trigonometric tangent function + - **exp**: exponential + - **log**: logarithm Unary operators are located in the ``graphblas.unary`` namespace. Additional unary operators registered from numpy are located in ``graphblas.unary.numpy``. @@ -56,32 +56,32 @@ Example usage: Common binary operators are: - - **pair** -- result is always 1.0 (or equivalent for the dtype) - - **first** -- f(a, b) = a - - **second** -- f(a, b) = b - - **min** -- min(a, b) - - **max** -- max(a, b) - - **eq** -- a == b - - **ne** -- a != b - - **gt** -- a > b - - **lt** -- a < b - - **ge** -- a >= b - - **le** -- a <= b - - **plus** -- a + b - - **minus** -- a - b - - **times** -- a * b - - **truediv** -- a / b - - **fmod** -- a % b - - **pow** -- a ** b - - **atan2** -- math.atan2(a, b) - - **lor** -- logical or (a | b) - - **land** -- logical and (a & b) - - **lxor** -- logical xor (a ^ b) - - **lxnor** -- logical xnor ~(a ^ b) - - **bor** -- binary or - - **band** -- binary and - - **bxor** -- binary xor - - **bxnor** -- binary xnor + - **pair**: result is always 1.0 (or equivalent for the dtype) + - **first**: f(a, b) = a + - **second**: f(a, b) = b + - **min**: min(a, b) + - **max**: max(a, b) + - **eq**: a == b + - **ne**: a != b + - **gt**: a > b + - **lt**: a < b + - **ge**: a >= b + - **le**: a <= b + - **plus**: a + b + - **minus**: a - b + - **times**: a * b + - **truediv**: a / b + - **fmod**: a % b + - **pow**: a ** b + - **atan2**: math.atan2(a, b) + - **lor**: logical or (a | b) + - **land**: logical and (a & b) + - **lxor**: logical xor (a ^ b) + - **lxnor**: logical xnor ~(a ^ b) + - **bor**: binary or + - **band**: binary and + - **bxor**: binary xor + - **bxnor**: binary xnor Binary operators are located in the ``graphblas.binary`` namespace. Additional binary operators registered from numpy are located in ``graphblas.binary.numpy``. @@ -107,15 +107,15 @@ Example usage: Common monoids are: - - **any** -- return either input - - **min** -- min(a, b) - - **max** -- max(a, b) - - **plus** -- a + b - - **times** -- a * b - - **land** -- a & b - - **lor** -- a | b - - **lxor** -- a ^ b - - **lxnor** -- ~(a ^ b) + - **any**: return either input + - **min**: min(a, b) + - **max**: max(a, b) + - **plus**: a + b + - **times**: a * b + - **land**: a & b + - **lor**: a | b + - **lxor**: a ^ b + - **lxnor**: ~(a ^ b) Monoids are located in the ``graphblas.monoid`` namespace. Additional monoids registered from numpy are located in ``graphblas.monoid.numpy``. @@ -149,6 +149,8 @@ Common semirings are: - **min_second** - **max_first** - **max_second** + - **plus_first** (sum the left operand's values over the connection pattern) + - **plus_second** (sum the right operand's values over the connection pattern) - **plus_min** - **lor_land** - **land_lor** @@ -156,6 +158,43 @@ Common semirings are: Semirings are located in the ``graphblas.semiring`` namespace. Additional semirings registered from numpy are located in ``graphblas.semiring.numpy``. +The ``first`` and ``second`` semirings +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``first`` binary operator returns its left input and ``second`` returns its right input, +each ignoring the other. Combined with the ``plus`` monoid they give two semirings that read +the values of only one input during a multiply: + + - ``plus_first``: for ``C << A.mxm(B, semiring.plus_first)``, each ``C[i, k]`` is the sum of + ``A[i, j]`` over the ``j`` where both ``A[i, j]`` and ``B[j, k]`` are present. ``B`` + contributes only its structure; its stored values are never read. + - ``plus_second``: the mirror image. Each ``C[i, k]`` sums ``B[j, k]`` over those same ``j``, + and ``A`` contributes only its structure. + +These are the natural choice when one operand is a boolean or *iso* (single-valued) adjacency +matrix, where a ``plus_times`` multiply either multiplies by one (a no-op) or rescales every +term by the same constant. ``first`` and ``second`` skip the product and accumulate one side's +values directly over the connection pattern. On a 0/1 matrix, ``plus_first`` counts the +length-two paths between each pair of nodes. ``plus`` has no BOOL form, so for two boolean +matrices give the semiring an integer type, as in ``A.mxm(A, semiring.plus_first[int])``. + +.. code-block:: python + + from graphblas import Matrix, semiring + + # A: 0->1 (10), 0->2 (20); B: 1->0 (3), 2->0 (7) + # A @ B reaches node 0 from node 0 through j = 1 and j = 2. + A = Matrix.from_coo([0, 0], [1, 2], [10, 20], nrows=3, ncols=3) + B = Matrix.from_coo([1, 2], [0, 0], [3, 7], nrows=3, ncols=3) + + A.mxm(B, semiring.plus_first).new() # C[0, 0] == 30 (10 + 20, taken from A) + A.mxm(B, semiring.plus_second).new() # C[0, 0] == 10 ( 3 + 7, taken from B) + A.mxm(B, semiring.plus_times).new() # C[0, 0] == 170 (10*3 + 20*7) + +The same ``first`` / ``second`` choice pairs with other monoids: ``min_first`` and ``min_second`` +(listed above) carry a value along an edge without combining it, the pattern behind +label-propagation traversals. + IndexUnary Operators -------------------- @@ -187,26 +226,26 @@ Example usage with a thunk parameter: Defined IndexUnary operators are: - - **index** -- return the vector index - - **rowindex** -- return the matrix row index - - **colindex** -- return the matrix column index - - **diagindex** -- return the matrix diagonal index (i.e. column - row) - - **tril** -- lower triangle matrix (True if column >= row) - - **triu** -- upper triangle matrix (True if column <= row) - - **diag** -- matrix diagonal (True if row == column) - - **offdiag** -- matrix off-diagonal (True if row != column) - - **indexle** -- vector index <= thunk - - **indexgt** -- vector index > thunk - - **rowle** -- matrix row index <= thunk - - **rowgt** -- matrix row index > thunk - - **colle** -- matrix column index <= thunk - - **colgt** -- matrix column index > thunk - - **valueeq** -- value == thunk - - **valuene** -- value != thunk - - **valuelt** -- value < thunk - - **valuele** -- value <= thunk - - **valuegt** -- value > thunk - - **valuege** -- value >= thunk + - **index**: return the vector index + - **rowindex**: return the matrix row index + - **colindex**: return the matrix column index + - **diagindex**: return the matrix diagonal index (i.e. column - row) + - **tril**: lower triangle matrix (True if column >= row) + - **triu**: upper triangle matrix (True if column <= row) + - **diag**: matrix diagonal (True if row == column) + - **offdiag**: matrix off-diagonal (True if row != column) + - **indexle**: vector index <= thunk + - **indexgt**: vector index > thunk + - **rowle**: matrix row index <= thunk + - **rowgt**: matrix row index > thunk + - **colle**: matrix column index <= thunk + - **colgt**: matrix column index > thunk + - **valueeq**: value == thunk + - **valuene**: value != thunk + - **valuelt**: value < thunk + - **valuele**: value <= thunk + - **valuegt**: value > thunk + - **valuege**: value >= thunk IndexUnary operators are located in two places. @@ -221,6 +260,39 @@ IndexUnary operators are located in two places. (i.e. all except rowindex, colindex, and diagindex). Calling the operators in the select namespace will perform a ``select`` operation. +IndexBinary Operators +--------------------- + +IndexBinary operators are the two-input analogue of IndexUnary operators. The function +receives both values, the indices of *each* value, and a thunk parameter: +``f(x, ix, jx, y, iy, jy, theta) -> z``, where ``ix, jx`` are the row/column indices of +``x`` and ``iy, jy`` are those of ``y``. + +There are no built-in IndexBinary operators; all are user-defined (see +:doc:`udf`). Binding a ``theta`` value to an IndexBinaryOp produces a +regular BinaryOp usable directly in ``ewise_mult`` / ``ewise_add``, or as +the multiplier of a Semiring for ``mxm`` / ``mxv`` / ``vxm``. Like other +BinaryOps, its type is chosen when it runs, from both inputs and ``theta``. + +Example usage: + +.. code-block:: python + + def discounted_sum(x, ix, jx, y, iy, jy, theta): + return (x + y) * theta + + gb.indexbinary.register_new("discounted_sum", discounted_sum) + + binop = gb.indexbinary.discounted_sum(0.5) # bind theta + C << A.ewise_mult(B, binop) + + # For mxm, wrap in a Semiring + sr = gb.semiring.register_anonymous(gb.monoid.plus, binop) + D << A.mxm(B, sr) + +IndexBinary operators are located in the ``graphblas.indexbinary`` namespace. +They require SuiteSparse:GraphBLAS 9.4 or newer. + Aggregators ----------- @@ -298,7 +370,7 @@ This facilitates writing more succinct code such as: cur_min(accum=op.min) << op.min_plus(A @ B).reduce_rowwise(op.min) -In the case of name conflicts (ex. binary.min and monoid.min), only one will exist in the the +In the case of name conflicts (ex. binary.min and monoid.min), only one will exist in the ``graphblas.op`` namespace. However, almost all functions which require a specific kind of operator have a mechanism to convert from an identically named operator of a different type. diff --git a/docs/user_guide/types.rst b/docs/user_guide/types.rst index aaeba1c6b..2a0649e38 100644 --- a/docs/user_guide/types.rst +++ b/docs/user_guide/types.rst @@ -30,6 +30,42 @@ When a data type is needed in an API call, a string or numpy or numba dtype may instead of the actual data type object. Additionally, the Python builtin ``bool``, ``int``, and ``float`` may be used. ``int`` indicates INT64 and ``float`` indicates FP64. +Python numbers in operations +---------------------------- + +A Python ``bool``, ``int``, ``float`` or ``complex`` used as an operand, such as +the ``1`` in ``v + 1`` or ``v.apply(binary.times, right=0.5)``, is *weak*, as in +numpy 2 (NEP 50). It takes the other operand's data type when that type's kind +can hold it, and its own default type (BOOL, INT64, FP64 or FC64) when not. So +with an INT8 vector ``v``, ``v + 1`` is INT8 and ``v * 0.5`` is FP64, while +``v + 300`` raises ``OverflowError`` rather than wrapping, and ``v * 2`` wraps +in INT8 as it does in numpy. An FP32 vector times ``0.5`` is FP32, plus ``1j`` +is FC32, and plus ``1e300`` is infinity, with a ``RuntimeWarning`` as in numpy. numpy scalars, 0-d arrays and ``Scalar`` objects keep their own +type: ``v + np.int64(1)`` is INT64, and so is ``np.int64(1) + v``, since numpy's +operators leave an operation with a python-graphblas object to python-graphblas. +python-graphblas applies this rule itself, +so it does not depend on the installed version of numpy. User-defined types +follow the same rule, field by field (see :doc:`udt`). + +A comparison has no result type to keep, so an integer outside the operand's +range compares exactly, as in numpy: ``v < 300`` is True for every element, as +is ``w < 2**63`` for an INT64 vector ``w``. +Nor does an operator whose result has another type raise: ``v / 300`` is FP64 +whatever the integer, and ``v.apply(binary.isclose, right=300)`` is False for +every element. A user-defined operator whose result has another type gets the +integer exactly where the operand's type and the integer's share an integer +type (INT16 for ``300`` here), and raises ``OverflowError`` where they would +share only a floating-point type, which can round the integer: beside an INT64 +vector, ``2**63`` would meet the elements in FP64, where ``2**63 - 1`` and +``2**63`` are equal. An +FP32 vector ``== 0.1`` compares in FP32, and ``== 1e300`` compares with +infinity, with a ``RuntimeWarning`` as in numpy. The defaults of ``ewise_union`` +follow the rule, and so does the ``fill_value`` of ``to_dense``, except that a +value that does not fit widens the result instead of raising. A literal given +to a typed operator, such as ``binary.plus["INT64"]``, takes that operator's +type, and a ``select`` or ``IndexUnaryOp`` thunk that is not compared with the +values, such as a row index, keeps its own type. + User-defined Types ------------------ diff --git a/docs/user_guide/udf.rst b/docs/user_guide/udf.rst index e7b984b44..b9ad278b8 100644 --- a/docs/user_guide/udf.rst +++ b/docs/user_guide/udf.rst @@ -1,55 +1,247 @@ -User-defined Functions -====================== +User-defined Functions (UDFs) +============================= -1. Show example of register_new -2. Discuss commonality for other operators -3. Discuss register_anonymous +python-graphblas lets you write custom operators in Python. ``numba`` +JIT-compiles them to native machine code so they run inside +SuiteSparse:GraphBLAS at C speed. This guide covers every operator type +that accepts a user function: ``UnaryOp``, ``BinaryOp``, ``IndexUnaryOp``, +``SelectOp``, and ``IndexBinaryOp``. -Python-graphblas requires ``numba`` which enables compiling user-defined Python functions -to native machine code for use by the GraphBLAS backend. This provides functions which are -very performant. - -Example user-defined UnaryOp: +A first example +--------------- .. code-block:: python - from graphblas import unary + from graphblas import unary, Vector - def force_odd_func(x): + def force_odd(x): if x % 2 == 0: return x + 1 return x - unary.register_new("force_odd", force_odd_func) + unary.register_new("force_odd", force_odd) v = Vector.from_coo([0, 1, 3, 4, 5], [1, 2, 3, 8, 14]) w = v.apply(unary.force_odd).new() + # w = [1, 3, _, 3, 9, 15] + +``register_new`` vs ``register_anonymous`` +------------------------------------------ + +- ``register_new(name, func)`` puts the operator on ``gb.unary.{name}`` (or + ``gb.binary.{name}``, etc.). Use this for operators you'll reference by + name across files. +- ``register_anonymous(func)`` returns the operator without adding it to a + namespace. Use this for one-off operators or operators created inside + another function. -.. csv-table:: w - :class: matrix - :header: 0,1,2,3,4,5 +Lambdas are auto-registered as anonymous wherever a UnaryOp is expected: - 1,3,,3,9,15 +.. code-block:: python + + v.apply(lambda x: x % 5 - 2).new() +Operator signatures +------------------- -Similar methods exist for BinaryOp and IndexUnaryOp. User-defined Monoids and Semirings are -constructed out of previously defined and built-in UnaryOps and BinaryOps. +The function signature depends on the operator type: -Auto-registration of Lambdas ----------------------------- +================== ==================================================== ===================================== +Operator Function signature Returns +================== ==================================================== ===================================== +``UnaryOp`` ``f(x) -> z`` a value +``BinaryOp`` ``f(x, y) -> z`` a value +``IndexUnaryOp`` ``f(x, ix, jx, theta) -> z`` a value (often bool, for select) +``SelectOp`` same as IndexUnaryOp ``bool`` +``IndexBinaryOp`` ``f(x, ix, jx, y, iy, jy, theta) -> z`` a value (requires SS >= 9.4) +================== ==================================================== ===================================== -As a convenience, any lambda expression used in place of a UnaryOp will be automatically -compiled as registered anonymously. +``ix, jx`` are the row and column indices of ``x``; same for ``iy, jy`` and ``y``. +``theta`` is a scalar parameter bound when the operator is used (see +:ref:`udf_indexbinary` below). -Example lambda usage: +Parameterized UDFs +------------------ + +To parameterize a UDF (e.g., a scale factor known only at call time), write a +factory function that returns the actual operator, and register it with +``parameterized=True``: .. code-block:: python - v.apply(lambda x: x % 5 - 2).new() + from graphblas import binary + + def make_scaled_add(scale): + def inner(x, y): + return scale * (x + y) + return inner + + binary.register_new("scaled_add", make_scaled_add, parameterized=True) + + # Bind the parameter, then use the resulting op: + op = binary.scaled_add(2.0) + result = v.ewise_mult(w, op).new() + +User-defined types (UDTs) +------------------------- + +Pass ``is_udt=True`` when your UDF operates on user-defined record or array +types. See :doc:`udt` for the full story. + +.. code-block:: python + + from graphblas import binary, dtypes + import numpy as np + + edge_dtype = dtypes.register_anonymous( + np.dtype([("weight", np.float64), ("hops", np.int32)], align=True), + "Edge", + ) + + def combine_edges(x, y): + return (x["weight"] + y["weight"], x["hops"] + y["hops"]) + + binary.register_new("combine_edges", combine_edges, is_udt=True) + +The function can return a tuple matching the record's fields, an existing +record value, or a numpy array for array UDTs. + +.. _udf_indexbinary: -.. csv-table:: - :class: matrix - :header: 0,1,2,3,4,5 +IndexBinaryOps and the theta parameter +-------------------------------------- + +An IndexBinaryOp receives row and column indices for both operands plus a +scalar ``theta``. Binding ``theta`` produces a regular ``BinaryOp`` that +works in ``ewise_mult``, ``ewise_add``, and the other elementwise paths: + +.. code-block:: python + + from graphblas import indexbinary, Matrix + + def discounted_sum(x, ix, jx, y, iy, jy, theta): + return (x + y) * theta + + indexbinary.register_new("discounted_sum", discounted_sum) + + bound = indexbinary.discounted_sum(0.5) # theta = 0.5 + A = Matrix.from_coo([0, 1], [0, 1], [1, 2]) + B = Matrix.from_coo([0, 1], [0, 1], [3, 4]) + C = A.ewise_mult(B, bound).new() + # C is FP64: C[0,0] = (1+3)*0.5 = 2.0; C[1,1] = (2+4)*0.5 = 3.0 + +Like any other BinaryOp, the bound op takes its type each time it runs: the +type that holds both inputs and ``theta``, as ``apply`` types an IndexUnaryOp +from its input and thunk. One type serves x, y and ``theta``. ``theta`` has +the type of its value: a Python int is INT64 and a float FP64, and a numpy +scalar or a ``Scalar`` keeps its own type. With no ``theta``, ``theta`` is 0 +(``False``), which changes no type. Above, the INT64 matrices and the float +``theta`` run in FP64. To choose the type yourself, name it: +``indexbinary.discounted_sum[float](0.5)``, or +``indexbinary.discounted_sum(0.5, dtype=float)``. The inputs and ``theta`` are +then cast to that type. + +An op registered with ``is_udt=True`` is compiled for the types of x, y and +``theta`` as they are, so none of them is cast: ``op(0.5)`` on INT64 inputs +gets the float 0.5. A numpy array ``theta`` has the array UDT of its dtype +and shape. A tuple has no type of its own, so give a record or array +``theta`` as a ``Scalar`` of its UDT (``Scalar.from_value((1, 2.5), udt)``). +``op[X, Y, T](theta)`` names all three types, and ``op[X, Y](theta)`` gives +``theta`` the type of y. + +This differs from an IndexUnaryOp in ``apply``, which makes a number thunk an +element of the UDT data's type. Beside UDT data, a number ``theta`` stays a +number: with ``x + theta`` on INT32[3] data, ``op(5)`` adds an INT64 to each +element, which gives an int64 array, not an INT32[3] element, so the op +raises ``UdfParseError`` when it is used. To make ``theta`` the data's type, +pass a ``Scalar`` of that type, or a numpy value of the element type +(``op(np.int32(5))``), or name the type: ``op(5, dtype=udt)`` binds +``op[udt](5)``, which makes 5 an element of ``udt``. + +A ``Scalar`` ``theta`` is copied when it is bound, so changing the ``Scalar`` +afterwards does not change the bound op. An ``apply`` expression, by +contrast, reads its ``Scalar`` thunk when it runs. + +To use a bound IndexBinaryOp as the multiplier in ``mxm`` / ``mxv`` / +``vxm``, wrap it in a Semiring. ``Semiring.register_anonymous`` accepts a +bound IBO directly: + +.. code-block:: python + + from graphblas import monoid, semiring + + sr = semiring.register_anonymous(monoid.plus, bound) + C = A.mxm(B, sr).new() + +The Semiring is typed each time it runs, as the bound op is. One built from +an op bound at a type (``discounted_sum[float](0.5)``) is monomorphic in that +op's input/output types instead (SuiteSparse builds exactly one +``GrB_Semiring`` for that type pair, rather than the type-polymorphic table +the standard semirings carry). Per SuiteSparse, monoids themselves cannot be +built from an IndexBinaryOp; only the multiplier slot accepts one. + +IndexBinaryOps require SuiteSparse:GraphBLAS >= 9.4. + +What numba accepts +------------------ + +The function body is compiled by ``numba.njit`` and must be pure +numerical Python, with these constraints: + +- No closures over Python objects (capture scalars, not lists or dicts). +- ``numpy`` array operations and standard library calls like ``math.sin`` + work; complex Python types (sets, dicts) do not. +- Records (UDT fields) are accessed by name: ``x["weight"]``, not ``x.weight``. +- Tuple returns work for UDTs (one tuple element per field). + +When compilation fails, you get a ``UdfParseError`` with the actionable +diagnostic line pulled out of Numba's typing pass, rather than the full +multi-hundred-line traceback. The most common causes: + +- Referencing a field that doesn't exist on the record UDT. +- Returning a tuple whose length doesn't match the record's field count + (the error names the expected arity). +- Calling a function Numba doesn't support in nopython mode. + +See :doc:`udt` for UDT-specific guidance. + +Lazy registration +----------------- + +Pass ``lazy=True`` to defer Numba compilation until the operator is first +used. This is useful for libraries that register many operators at import +time: + +.. code-block:: python - -1,0,,1,1,2 + unary.register_new("heavy_op", heavy_func, lazy=True) + # heavy_func isn't compiled yet; the operator object exists, but no + # numba.njit has run. + +The compile happens on first lookup (``unary.heavy_op[int]``). + +Compilation caching +------------------- + +A handful of built-in operators are themselves UDFs (``binary.floordiv``, +``rfloordiv``, ``absfirst``, ``abssecond``, and ``rpow``). Because these are +plain module-level functions, python-graphblas compiles them with Numba's +``cache=True``, so the Numba half of the compile is paid once per machine +instead of once per Python process. Numba writes the cache (``.nbi`` index and +``.nbc`` data files) into the ``__pycache__`` directory next to the operator +module. A warm cache shaves roughly 20% off a single op's first-touch build and +about 45% when several are used; the rest is Numba/LLVM startup and per-process +object linking, which caching cannot remove. + +If that ``__pycache__`` directory is not writable (a read-only install, for +example), Numba silently falls back to its per-user cache directory +(``~/.cache/numba`` on Linux, ``~/Library/Caches/numba`` on macOS; set +``NUMBA_CACHE_DIR`` to override). No warning is emitted for the fallback. If +neither is writable, these operators are compiled uncached, once per process. + +Operators you create with ``register_new`` or ``register_anonymous`` are never +passed ``cache=True``: Numba can only cache a function with a stable on-disk +source, which a lambda or an interactively defined function does not have. A +registered operator's compile is paid once per process, and ``lazy=True`` is +the tool for deferring it. diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst new file mode 100644 index 000000000..88b77086a --- /dev/null +++ b/docs/user_guide/udt.rst @@ -0,0 +1,518 @@ + +User-defined Types (UDTs) +========================= + +python-graphblas supports user-defined types (record-style structs and +fixed-shape arrays) as the value type of any ``Scalar``, ``Vector``, or +``Matrix``. Built-in arithmetic operators automatically lift to UDTs +field-by-field, and SuiteSparse:GraphBLAS JIT-compiles dedicated C kernels +for them when possible. + +What is a UDT +------------- + +A UDT is any ``numpy.dtype`` you register with python-graphblas. There are three shapes. + +**Record UDTs** have heterogeneous fields, like a C struct: + +.. code-block:: python + + import numpy as np + from graphblas import dtypes + + edge_dtype = dtypes.register_anonymous( + np.dtype([("weight", np.float64), ("hops", np.int32)], align=True), + "Edge", + ) + +**Array UDTs** are fixed-shape, homogeneous values, like an inline C array: + +.. code-block:: python + + point3 = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "Point3") + +Multi-dimensional shapes work too (``np.dtype((np.float64, (2, 4)))``); the +layout is flattened row-major in C. An array of arrays has the same layout, so +it registers as the same UDT: ``np.dtype((point3.np_type, (2,)))`` is +``FP64[2, 3]``. Record fields that nest arrays are flattened the same way. + +**Dataclass UDTs** are record UDTs derived from a ``@dataclass``: + +.. code-block:: python + + from dataclasses import dataclass + + @dataclass + class Edge: + weight: float + hops: int + + edge_dtype = dtypes.register_anonymous(Edge) + +Field annotations may be real types (``int``, ``float``) or string forms +(``"int"``, e.g. under ``from __future__ import annotations``). Compound +annotations like ``Optional[int]`` raise from ``lookup_dtype``. + +Registration +------------ + +Two forms: + +.. code-block:: python + + # register_anonymous returns the DataType but does not add it to gb.dtypes. + udt = dtypes.register_anonymous(numpy_dtype_or_dataclass, "MyUdt") + + # register_new is the same, but also assigns to gb.dtypes.MyUdt. + udt = dtypes.register_new("MyUdt", numpy_dtype_or_dataclass) + +Field type rules: + +- Numeric scalar types (``int``, ``float``, ``bool``, ``complex``, the + corresponding numpy scalar types) are supported. +- Fixed-shape arrays (``("pos", np.float64, (3,))``) and records (a nested + struct) are supported as fields, and built-in operators lift through them. + In the dict and dataclass forms, write an array field as ``"FP64[3]"``; the + dict form also takes an existing UDT, such as ``{"id": int, "pos": point3}``. +- Numba cannot type an array of ``bool``, ``datetime64``, ``timedelta64``, + ``str`` or ``bytes``, so with Numba installed, registering a UDT that holds + one raises ``ValueError``. Use an array of ``uint8`` for ``bool`` and + ``bytes`` (one per byte), of ``uint32`` for ``str`` (one per character), and + of ``int64`` for the others. +- Text (``str`` or ``bytes``, such as ``np.dtype("S5")``), ``datetime64`` and + ``timedelta64`` may be fields, or a whole UDT that is neither record nor + array. Such a UDT can be stored and extracted. With Numba, it can also be + compared (``==``, ``!=`` and ``isequal``), and ``first``, ``second``, ``pair`` + and ``any`` work on it, but arithmetic operators do not lift to it: ``plus`` + and ``min`` raise ``KeyError``. Without Numba, only a ``Scalar``'s ``==``, + ``!=`` and ``isequal``, and ``pair``, work on it. +- Python objects (``np.dtype(object)``) are not supported. +- Dataclass annotation strings (e.g., ``"int"``) resolve through + ``lookup_dtype``. + +Anonymous UDTs share one ``DataType`` per ``numpy.dtype``. Re-registering the +same dtype under a different name updates the Python-side ``name`` but does +*not* change the SuiteSparse-side ``GxB_JIT_C_NAME``, which is frozen at first +registration. See :ref:`udt_jit_introspection` below. + +Working with UDT values +----------------------- + +Construct, get, set, and iterate as usual; values are numpy structured scalars +or arrays, depending on the UDT shape: + +.. code-block:: python + + from graphblas import Vector + + v = Vector(edge_dtype, size=3) + v[0] = (1.5, 4) # tuple matches the record fields + v[1] = (2.5, 7) + print(v[0].new().value) # numpy.void; access fields by name (e.g., ['weight']) + +For array UDTs, pass a sequence or numpy array: + +.. code-block:: python + + p = Vector(point3, size=2) + p[0] = (1.0, 2.0, 3.0) + p[1] = np.array([4.0, 5.0, 6.0]) + +Built-in operators on UDTs +-------------------------- + +The following operators auto-lift to any UDT on first use: + +- BinaryOps: ``plus``, ``minus``, ``times``, ``truediv``, ``floordiv``, + ``min``, ``max``, ``eq``, ``ne``. The positional selectors ``first``, + ``second``, ``any``, and ``pair`` work too, since they don't touch field + values. +- UnaryOps: ``ainv``, ``abs``. +- Monoids: ``plus``, ``times``, ``min``, ``max``, ``any``. +- Semirings combining a UDT-lifting monoid with a UDT-lifting BinaryOp + (e.g., ``plus_times``, ``min_plus``, ``max_times``, ``any_first``). +- Aggregators built on those monoids: ``sum``, ``prod``, ``min``, ``max``, + ``any_value``, ``count``. The positional ``agg.ss.first`` and + ``agg.ss.last`` work on any dtype, including UDTs. + +For example: + +.. code-block:: python + + from graphblas import binary, monoid, semiring, agg + + plus_edge = binary.plus[edge_dtype] # field-wise add + sum_edges = monoid.plus[edge_dtype] # field-wise additive monoid + plus_times = semiring.plus_times[edge_dtype] # field-wise multiply-add semiring + total = v.reduce(agg.sum[edge_dtype]).new() # field-wise reduce + +How results are typed +~~~~~~~~~~~~~~~~~~~~~ + +A lifted op types each field or element as it types a Vector of that built-in +dtype, so it never narrows. Below, ``v`` is a Vector of an ``INT8[3]`` UDT. + +.. list-table:: + :header-rows: 1 + :widths: 25 45 30 + + * - Operands + - Rule + - Example + * - One UDT + - Each field or element as on its built-in dtype. + - ``v / v`` is ``FP64[3]`` + * - Two UDTs + - Records pair when they have the same field names, in the same order, at + every level, and the same shape for each field. Array UDTs pair when + their shapes broadcast as numpy's do. Any other pair raises + ``KeyError``. + - ``INT8[3]`` with ``INT64[2, 3]`` gives ``INT64[2, 3]`` + * - Result type + - An operand's type when it holds the result; else the promoted layout, + which is the UDT registered with that layout if there is one, else an + anonymous UDT. + - ``{x: INT64, y: FP64}`` with ``{x: FP64, y: INT64}`` gives + ``{x: FP64, y: FP64}`` + * - A built-in Vector, ``Scalar``, numpy scalar or numpy array + - Strong: keeps its dtype. A scalar combines with every field or element; + an array pairs as an array UDT of its own shape. A record's array field + keeps its shape, so a literal for it must broadcast to that shape, or + ``ValueError`` names both: ``np.array([[1, 2, 3]])`` does not fit an + ``INT8[3]`` field, though beside an ``INT8[3]`` UDT it gives + ``INT64[1, 3]``. + - ``v + np.int64(1)`` is ``INT64[3]`` + * - A Python number, or a tuple, list or dict of them laid out like an + element + - Weak, field by field: a field's dtype when its kind holds the number, + else INT64, FP64 or FC64 (numpy 2's NEP 50 rule, on numpy 1 too, and + extended to sequences). An int out of an integer field's range raises + ``OverflowError``, except for ``truediv``, whose fields are floats + whatever the int (``v / 300`` is ``FP64[3]``). + - ``v + 1`` is ``INT8[3]``, ``v * 2.5`` and ``v + (0.5, 1, 2)`` are + ``FP64[3]``, ``v + 300`` raises + * - ``eq`` and ``ne`` + - Type a literal as above, as built-in comparisons do, except that a + literal with an int out of a field's range equals no element, rather + than raise. Return ``BOOL``. Array elements broadcast, as numpy's ``==`` + does, and so does ``==`` on a Scalar, which returns a bool. ``isequal`` + compares as ``np.array_equal``: elements must have the same shape, + apart from leading axes of length 1. + - ``v == 300`` is ``False``; an ``FP32[3]`` UDT ``== 0.1`` compares in + FP32; a Scalar ``s`` of ``[1, 1, 1]`` has ``s == 1`` but not + ``s.isequal(1)`` + * - A literal that must become an element: beside a user-defined op, an + ``IndexUnaryOp`` thunk, an ``ewise_union`` default + - Its type by the rules above must fit the UDT: the UDT itself, or a + numpy value or ``Scalar`` that casts to it safely. A weak int does not + fit a bool field: write ``True`` or ``False``. A value that does not fit + but converts exactly, such as ``0`` for a bool field or ``2.0`` for an + int field, still converts with a ``DeprecationWarning`` and will raise + in a future version; any other raises ``ValueError``, and a float too + large for a float field raises ``OverflowError``. + - ``v.ewise_union(w, binary.plus, 0.5, 0)`` raises + +A Monoid used element-wise (``apply``, ``ewise_mult``, ``ewise_union``) uses +its BinaryOp, for literals and for two different UDTs. ``min``, ``max`` and ``floordiv`` reject UDTs +with a complex field; use a custom op for those. + +Each element, including each element of an array field, is computed as the +built-in op computes it, and the Numba cfunc and a C JIT kernel give the same +bits: ``min`` and ``max`` ignore NaN (C's ``fmin`` and ``fmax``) and resolve a +tie between ``-0.0`` and ``0.0`` to ``-0.0`` for ``min`` and ``0.0`` for +``max``, ``floordiv`` is numpy's, with 0 for an integer divided by zero, and +``truediv`` is numpy's, with numpy's infinities for a zero divisor. A complex +quotient takes Smith's method, as numpy's does, but divides where numpy +multiplies by a reciprocal, so it can differ from numpy's in the last bit (and +stays finite where numpy's overflows). A few edge cases can differ from built-in +vectors, where the UDT gives numpy's answer: ``INT64_MIN // -1`` is +``INT64_MIN`` (an ``INT64`` vector gives 0), a complex product or quotient with +an infinite or NaN part, a complex quotient with a ``-0.0`` divisor, and the +signed-zero tie, which SuiteSparse's built-in ``min`` leaves to the C library. + +Storing results +~~~~~~~~~~~~~~~ + +GraphBLAS casts built-in dtypes when it stores a result, but it cannot cast a +UDT, so python-graphblas does it, without ever making a temporary copy of the +data. + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - Store + - What happens + * - Into an object of another UDT type of the same layout: records that + pair as operands do, or arrays of the same shape apart from leading + axes of length 1 + - Each field or element is cast as GraphBLAS casts that pair of built-in + dtypes: integers wrap; a float stored as an integer truncates toward + zero, saturates at the bounds, and is 0 if NaN; a complex number loses + its imaginary part; anything nonzero is ``True``. + * - ...from an existing object: ``w << v``, ``v.dup(dtype=...)``, with a + mask, an accumulator or ``replace`` + - One pass through a cast op. + * - ...from an element-wise lifted op: ``apply``, ``ewise_mult``, + ``ewise_union``, and infix operators that use them + - The op writes the object's type itself, in the same call: + ``int_udt_vec += 0.5`` truncates, as it does for an ``INT8`` vector. + * - ...from anything else: ``mxm``, ``reduce`` into a Vector, ``ewise_add``, + a user-defined op, or a whole Vector or Matrix assigned into part of an + object + - ``DomainMismatch``: this would need a converted copy of the whole + result. Make it explicitly: compute the result, then store it + (``w << expr.new()``); to assign, convert the value with + ``.dup(dtype=...)`` first. + * - ...from a ``Scalar`` or one element + - Computed in its own type, then cast. With an accumulator, the value is + cast before it accumulates, as GraphBLAS converts a value of a built-in + dtype, and so is a literal: an ``INT8[3]`` Scalar's ``s += 0.5`` adds 0, + as an ``INT8`` Scalar's does. + * - Between a UDT and a built-in dtype, or layouts that do not correspond + - ``DomainMismatch``. + +For the same reason, ``ewise_add`` on Vectors and Matrices raises +``DomainMismatch`` when an operand's type differs from the result's, as for two +different UDTs or ``truediv`` on an integer UDT: it would copy a value that has +no partner into the result type. ``ewise_union`` and ``ewise_mult`` work. Infix +``+`` on two such operands takes the union instead, with each operand's zero for +a missing value, as ``-`` takes 0. That zero is ``-0.0`` in a float field, so an +entry of one operand alone keeps its value, as ``ewise_add`` would copy it; +only a ``-0.0`` beside an integer field of the other operand becomes ``0.0``. A +Scalar is one element, so ``ewise_add`` converts its operands to the result type +instead. + +.. code-block:: python + + v = Vector(int8x3, size=2) # an INT8[3] UDT + f = Vector(fp32x3, size=2) # an FP32[3] UDT + (v * 2.5).new() # FP64[3]: the literal is weak, 2.5 needs FP64 + (v * f).new() # FP32[3], as INT8 times FP32 + v += 1 # stays INT8[3] + v << v * 2.5 # computes in FP64, stores as INT8 (truncates) + (v + f).new() # FP32[3], by ewise_union; v.ewise_add(f) raises + +Composite aggregators (``agg.hypot``, ``agg.L1norm``, ``agg.Linfnorm``, +``agg.sum_of_squares``, ``agg.sum_of_inverses``) do *not* auto-lift to UDTs; +they reference scalar-only binary ops that don't generalize trivially. +Use a custom monoid plus reduction if you need this on UDT-valued data. + +Custom UDFs over UDTs +--------------------- + +Write a function with explicit field access, register it with ``is_udt=True``: + +.. code-block:: python + + from graphblas import binary + + def merge_edges(x, y): + # Returns a new Edge with the lower weight and summed hops. + return (min(x["weight"], y["weight"]), x["hops"] + y["hops"]) + + op = binary.register_new("merge_edges", merge_edges, is_udt=True) + + a = Vector(edge_dtype, size=2) + a[0] = (1.5, 2) + a[1] = (3.0, 4) + b = Vector(edge_dtype, size=2) + b[0] = (1.0, 3) + b[1] = (2.5, 1) + + c = a.ewise_mult(b, op[edge_dtype]).new() + # c[0] = (1.0, 5); c[1] = (2.5, 5) + +UDT UDFs can return a tuple matching the field layout, an existing record +value, or a numpy array (for array UDTs). + +For nested record UDTs the tuple is *flat over the leaves*. Given +``[("id", int32), ("pt", [("x", float64), ("y", float64)])]``, the UDF should +return ``(id, x, y)``, not ``(id, (x, y))``. Returning an existing record value +(e.g., one of the inputs) is also fine and preserves the nested shape. + +Every field needs a value on every path through the UDF. A record cannot hold +``None``, so a UDF that can return ``None`` for a field +(``v if cond else None``) is rejected when the op is typed. So is an array, +tuple, or list returned for a scalar field. + +For array UDTs each operand arrives as a numpy view of that element's values, in +the UDT's declared shape: a ``np.dtype((np.float64, (2, 4)))`` UDT hands the UDF +a 2-by-4 array, indexable as ``x[i, j]``. Array expressions work as written, and +the UDF may return one of its operands or build a new array that fills the +element, either at the element's own shape or at one that broadcasts to it (a +``(1,)`` return fills every slot of a ``(6,)`` element): + +.. code-block:: python + + def midpoint(x, y): + return (x + y) / 2 + + op = binary.register_new("midpoint", midpoint, is_udt=True) + + a = Vector(point3, size=1) + a[0] = [0.0, 2.0, 4.0] + b = Vector(point3, size=1) + b[0] = [10.0, 20.0, 30.0] + + c = a.ewise_mult(b, op[point3]).new() + # c[0] = [5.0, 11.0, 17.0] + +A return that cannot fill the element, such as ``x[:2]`` from a 3-element UDT, +is rejected with a ``UdfParseError`` when the op is typed for the UDT +(``op[point3]``, or the first operation that uses it). The same goes for an +array returned into an array-typed field of a record UDT. A tuple or list also +fills such a field, but only a one-dimensional one and only at its exact +length; unlike an array, it does not broadcast. Numba does not know the shape +of an array the UDF builds, or the length of a list, so the check runs the UDF +once on sample values; a UDF that raises on those values is not checked. + +If your UDF references a field that doesn't exist, or returns the wrong arity, +you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced +from Numba's typing pass instead of a 200-line traceback. + +.. _udt_jit_introspection: + +JIT and introspection +--------------------- + +When all field names are valid C identifiers (not C reserved words), and the +field types map to C primitives, SuiteSparse JIT-compiles a dedicated C kernel +for each ``op[udt]`` lookup. A lifted op on a UDT and a built-in type, or on two +UDT types that pair, compiles too, under a name such as +``plus_7_int64_t_gbudt_Edge``. The kernel is cached on disk +and reused across processes (keyed by content hash, so renaming a UDT doesn't +invalidate). JIT lets SuiteSparse inline the kernel into its eWise and reduce +templates, eliminating the per-element function-call overhead the Numba +function-pointer fallback incurs. Elementwise operations on UDTs are +typically **2-3x** faster. + +Inspect what SuiteSparse is JIT-ing from Python: + +.. code-block:: python + + udt.jit_c_name # 'gbudt_Edge', or None if not JIT-able + udt.jit_c_definition # 'typedef struct { ... } gbudt_Edge ;' + + op = binary.plus[udt] + op.jit_c_name # 'plus_gbudt_Edge' + op.jit_c_source # full C source SS will compile + + monoid.plus[udt].jit_c_source # same kernel + semiring.plus_times[udt].jit_c_source # the multiplier + semiring.plus_times[udt].monoid.jit_c_source # the additive monoid + agg.sum[udt].jit_c_source # walks to monoid.plus + agg.count[udt].jit_c_source # None (composite agg) + +JIT is skipped when: + +- A field name is a C reserved word (``class``, ``return``, ...), a macro or + typedef of the C headers the kernels include (``NULL``, ``FILE``, ``M_PI``, + ``complex``, ``INT_MAX``, ``SIZE_MAX``, ``BUFSIZ``, ``SIGINT``, ...), or + starts with ``GB_``, ``GrB_`` or ``GxB_``, which GraphBLAS keeps for its own + names. The op falls through to the Numba cfunc path. +- A field type isn't in the numpy-to-C map (rare; the standard numeric + scalar types all map). +- A record has an array-valued field. Built-in operators still lift to it + through the Numba cfunc. +- The numpy layout doesn't match what a C compiler would produce. The most + common case is a packed record with mixed-width fields (e.g., + ``np.dtype([("a", int32), ("b", float64)])`` without ``align=True``). + Pass ``align=True`` to ``np.dtype``, or use the dict / dataclass form, + which auto-aligns. +- The JIT compiler isn't usable after auto-fix (see below). +- An op on a UDT and another type is an array UDT that broadcasts to another + shape, a record with an array-valued field, or ``eq`` or ``ne``; or an op on + one UDT gives another type (``v / v`` on ``INT8[3]`` is ``FP64[3]``). These + keep the Numba cfunc without a warning. + +SuiteSparse knows a UDT by ``gbudt_`` and its name, which cannot collide with +an identifier in the C code of a kernel, whatever the name. A name starting +with ``_gb``, which python-graphblas keeps for the types it makes itself, is +the exception: it is the C name as it is. A UDT without a name, or whose name +is not a C identifier or would make that C name too long to leave room for its +ops' names (over 118 characters), gets a synthetic ``_gbudt_NNN`` C name and +still takes the JIT path. Pass ``name=`` only for +readable JIT cache filenames and introspection. To name such a UDT in C source +of your own, as for ``gb.binary.ss.register_new``, use its ``jit_c_name``, not +its Python name, which SuiteSparse does not know. A type registered from C +source (``gb.dtypes.ss.register_new``) keeps the name you give it, so give it +118 characters or fewer: the built-in operators lifted to it are named with +that name after them (``floordiv_``, ``plus_``, ...), and a longer one leaves +them names SuiteSparse refuses, so they run without the JIT. + +When the cause is the UDT itself (the first four cases above), the first time +auto-lift produces a non-JIT'd op for a given ``(op, dtype)`` pair in a +process, a ``graphblas.exceptions.NoJITWarning`` (a subclass of +``UserWarning``) names the cause. An unusable compiler does not warn; +``gb.ss.fix_jit_config()`` returns ``False`` when it cannot compile, and +``gb.ss.jit_compiler_is_usable()`` is the cheap check. Silence the warning by +category or by message:: + + import warnings + from graphblas.exceptions import NoJITWarning + warnings.filterwarnings("ignore", category=NoJITWarning) + # or, equivalently, match the message text: + warnings.filterwarnings("ignore", message="UDT operator running without JIT") + +In every skipped case the operation still works correctly via the Numba +function-pointer path; you only lose the JIT speedup. A field named like a +macro that only some platform's C headers define (beyond the C standard's) +is not skipped, and its kernel fails to compile; rename the field. + +JIT compiler auto-fix +~~~~~~~~~~~~~~~~~~~~~ + +A conda-installed ``python-suitesparse-graphblas`` bakes in the build host's +compiler path (e.g., ``/Users/runner/...``), which doesn't exist on user +machines. With the bogus default, SuiteSparse emits the JIT ``.c`` source +but never compiles a ``.dylib`` or ``.so``, and silently falls back to the +cfunc path. The 2-3x JIT speedup is silently lost. + +python-graphblas repairs the compiler path when ``graphblas.ss`` is +imported, which the first access to ``gb.ss`` does (and at init, when the JIT +control is configured to ``'on'``). If ``jit_c_compiler_name`` doesn't exist on +disk, it is replaced with one from ``$CONDA_PREFIX/bin/`` (or from +``sysconfig`` for pure-pip installs). The +import changes nothing else, so reading ``gb.ss.about`` does not change what +later operations compute. + +``jit_c_control`` starts at ``graphblas.config["jit_c_control"]``, ``'run'`` by +default (run cached kernels only; no compile, no load from disk); see +:ref:`suitesparse_jit`. That default ``'run'`` is raised to ``'on'`` (compile, +load, and run) by each of these, since each asks SuiteSparse to compile C source. +It stays ``'on'`` for the rest of the process, for every operation, built-in ones +included: + +- a UDT operator that gets C source, such as the first ``binary.plus[udt]`` or + ``binary.plus[udt, "INT64"]``; +- ``gb.dtypes.ss.register_new``, which registers a type from a C typedef; +- ``gb.binary.ss.register_new`` and its ``unary``, ``indexunary``, ``select``, + and ``indexbinary`` siblings, which register an operator from a C definition. + +A control you set is kept: ``'off'``, ``'pause'``, ``'load'`` (load and run +cached kernels, never compile), and ``'run'`` too when it was set before init +(``GRAPHBLAS_JIT_C_CONTROL``, ``gb.config.set``, or a YAML config file). With +nothing configured, a ``'run'`` set later in ``gb.ss.config`` is the default +again. Under a ``'run'`` you set, UDT operators use their Numba functions, and +the registrations above fail with an error saying how to opt in: +``gb.ss.config["jit_c_control"] = "on"``, or ``None`` for the default. +SuiteSparse itself sets ``'load'`` after a compile fails, so keeping it means a +failing compile is tried once, not again at every later request. + +Call the helper manually to re-fix or verify:: + + gb.ss.fix_jit_config() # repair compiler path, set 'on', probe + gb.ss.jit_compiler_is_usable() # cheap check: True iff path exists + +Pickle and serialize +-------------------- + +UDT-typed Matrices and Vectors pickle round-trip in the same process. For +cross-process work (``multiprocessing.spawn``, ``dask``, etc.) the receiving +process must also have the UDT registered under the same name. UDTs registered +via ``register_new`` re-register automatically on unpickle. + +``ss.serialize`` and ``ss.deserialize`` work too: the dtype name is stored +in the blob and re-resolved on load, falling back to ``GrB_NAME`` for UDTs +whose C identifier differs from the Python name. diff --git a/environment.yml b/environment.yml index 2bae0b76e..b5d999cb8 100644 --- a/environment.yml +++ b/environment.yml @@ -45,8 +45,9 @@ dependencies: # For building docs - nbsphinx - numpydoc - - pydata-sphinx-theme - - sphinx-panels + - pydata-sphinx-theme>=0.15 + - sphinx>=7 + - sphinx-design # For building logo - drawsvg - cairosvg @@ -55,7 +56,6 @@ dependencies: # - black # - black-jupyter # - codespell - # - commonmark # - cython # - cytoolz # - distributed diff --git a/graphblas/__init__.py b/graphblas/__init__.py index 63110eeeb..e15604058 100644 --- a/graphblas/__init__.py +++ b/graphblas/__init__.py @@ -25,7 +25,45 @@ def get_config(): import donfig import yaml - config = donfig.Config("graphblas") + class Config(donfig.Config): + """donfig Config that rejects silent attribute writes. + + Options are set with ``config.set(name=value)`` (optionally as a + ``with config.set(name=value):`` block) and read with ``config[name]``. + Plain attribute assignment (``config.name = value``) would otherwise + create a dead instance attribute and leave the real option unchanged, + so we raise instead. + """ + + def __init__(self, *args, **kwargs): + # Writes are unrestricted until donfig's own __init__ finishes, and + # the attributes it set there become the allowlist. Deriving that + # list beats hardcoding donfig's internals: a donfig release that + # adds or renames a field can't then break attribute access here. + object.__setattr__(self, "_initializing", True) + super().__init__(*args, **kwargs) + object.__setattr__(self, "_donfig_attrs", frozenset(self.__dict__) - {"_initializing"}) + object.__setattr__(self, "_initializing", False) + + def __setattr__(self, key, value): + # An instance built without running __init__ has no allowlist yet, + # so stay permissive rather than raising a confusing AttributeError + # from the guard. + if self.__dict__.get("_initializing", True) or key in self._donfig_attrs: + object.__setattr__(self, key, value) + return + if key in self.config: + raise AttributeError( + f"Cannot set config option {key!r} by attribute assignment. " + f"Use `graphblas.config.set({key}=...)` to change it (optionally " + f"in a `with` block for a scoped change) and " + f"`graphblas.config[{key!r}]` to read it." + ) + raise AttributeError( + f"Unknown config option {key!r}; known options are {sorted(self.config)}." + ) + + config = Config("graphblas") path = Path(__file__).parent / "graphblas.yaml" with path.open() as f: defaults = yaml.safe_load(f) @@ -36,6 +74,10 @@ def get_config(): config = get_config() del get_config +# None until a backend is initialized. Touching a special attribute such as +# gb.Matrix auto-initializes, as does an explicit init(), and sets this to +# "suitesparse" or "suitesparse-vanilla". Reading gb.backend is not itself a +# special-attribute access, so it never triggers initialization. backend = None _init_params = None _SPECIAL_ATTRS = { @@ -49,6 +91,7 @@ def get_config(): "core", "dtypes", "exceptions", + "indexbinary", "indexunary", "io", "monoid", @@ -136,6 +179,9 @@ def _init(backend_arg, blocking, automatic=False): # Already initialized with these parameters; nothing more to do return + if backend_arg == "suitesparse": + # Checked before anything is initialized, so a bad value changes nothing + jit_c_control = _jit_c_control_config() backend = backend_arg if backend in {"suitesparse", "suitesparse-vanilla"}: try: @@ -196,6 +242,35 @@ class Lib: core.ffi = ffi core.lib = lib core.NULL = ffi.NULL + if backend == "suitesparse": + from .core.ss.jit_config import _set_jit_c_control_at_init + + _set_jit_c_control_at_init(jit_c_control) + + +_JIT_C_CONTROLS = ("on", "load", "run", "pause", "off") + + +def _jit_c_control_config(): + """``config["jit_c_control"]``, the SuiteSparse JIT control set at init (see graphblas.yaml). + + None when it is not set, which is the default: ``'run'``, raised by UDT work. + """ + val = config.get("jit_c_control") + if val is None: + return None + # YAML 1.1 reads an unquoted on/off as a bool, so a config file saying + # ``jit_c_control: off`` arrives here as False. + if val is True or val is False: + return "on" if val else "off" + if isinstance(val, str) and val.lower() in _JIT_C_CONTROLS: + return val.lower() + raise ValueError( + "Bad value for graphblas.config['jit_c_control'] (or the GRAPHBLAS_JIT_C_CONTROL " + f"environment variable): {val!r}. Must be one of: " + + ", ".join(repr(x) for x in _JIT_C_CONTROLS) + + ", or unset" + ) # Ideally this is in operator.py, but lives here to avoid circular references diff --git a/graphblas/agg/__init__.py b/graphblas/agg/__init__.py index da7c13591..d25017279 100644 --- a/graphblas/agg/__init__.py +++ b/graphblas/agg/__init__.py @@ -109,7 +109,9 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/binary/__init__.py b/graphblas/binary/__init__.py index 1b8985f73..4c2b7a267 100644 --- a/graphblas/binary/__init__.py +++ b/graphblas/binary/__init__.py @@ -66,7 +66,9 @@ def __getattr__(key): f"module {__name__!r} unable to compile UDF for {key!r}; " "install numba for UDF support" ) - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 24a49ba1a..95d89c670 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -2,13 +2,14 @@ from .. import backend, config from .. import replace as replace_singleton -from ..dtypes import BOOL from ..exceptions import check_status from . import NULL from .descriptor import lookup as descriptor_lookup from .expr import AmbiguousAssignOrExtract, Updater -from .mask import Mask +from .literals import _literal_store_dtype +from .mask import Mask, _check_mask from .operator import UNKNOWN_OPCLASS, binary_from_string, find_opclass, get_typed_op +from .udt import _fused_store_expr, _scalar_accum_operand, _store_cast_op from .utils import _Pointer, libget, output_type _recorder = ContextVar("recorder") @@ -166,21 +167,19 @@ def _expect_op(self, op, values, *, within, **kwargs): AmbiguousAssignOrExtract._expect_type = _expect_type -def _check_mask(mask, output=None): - if not isinstance(mask, Mask): - # Convert bool objects to value masks - if output_type(mask).__name__ in {"Vector", "Matrix"}: - if mask.dtype != BOOL: - raise TypeError( - f"Mask must be boolean objects (got {mask.dtype}) " - "or indicate values (M.V) or structure (M.S)" - ) - mask = mask.V # auto-compute (will raise if disabled) - else: - raise TypeError(f"Invalid mask: {type(mask)}") - if output is not None and output.ndim == 1 and mask.parent.ndim != 1: - raise TypeError(f"Mask object must be type Vector; got {type(mask.parent)}") - return mask +# Curated hints for common attribute-access mistakes on Vector/Matrix/Scalar. +# Only consulted from __getattr__, which fires solely on a genuine attribute +# miss, so the normal (slotted) attribute hot path is untouched. +_INSTANCE_ATTR_HINTS = { + # `.new()` resolves expressions (e.g. `A.mxm(B).new()`); a concrete + # object is copied with `.dup()`. + "new": ( + "`.new()` resolves an expression (e.g. `A.mxm(B).new()`); a concrete " + "object has no `.new()`. Use `.dup()` to copy this object." + ), + # transpose is the `.T` property, not a method. + "transpose": "transpose is the `.T` property, e.g. `A.T` (not a method call).", +} class BaseType: @@ -188,6 +187,24 @@ class BaseType: __slots__ = "gb_obj", "dtype", "name", "__weakref__" # Flag for operations which depend on scalar vs vector/matrix _is_scalar = False + # numpy's operators return NotImplemented for an operand of higher priority than theirs + # (ndarray's is 0, np.matrix's 10), so with a numpy scalar or array on the left Python calls + # our reflected method, and `np.int16(2) - s` is the expression `s - np.int16(2)` would be, + # not numpy's value of `s.__array__()`. Ufuncs such as `np.sin(s)` still use __array__. + # Expressions, TransposedMatrix and the index expressions set it too. + __array_priority__ = 1000 + + def __getattr__(self, name): + # Python calls this only after normal lookup fails (slots and methods resolve + # first), so it is free on the hot path. It fails both for a missing name and + # for a property that raises AttributeError, such as ``v.ss`` on the vanilla + # backend or ``__index__`` on a float Scalar. Add hints for common mistakes, + # and otherwise look the name up again, so a property raises its own message. + if (hint := _INSTANCE_ATTR_HINTS.get(name)) is not None: + raise AttributeError( + f"{type(self).__name__!r} object has no attribute {name!r}; {hint}" + ) + return object.__getattribute__(self, name) def __call__( self, @@ -266,28 +283,40 @@ def __or__(self, other): from .infix import _ewise_infix_expr, _ewise_mult_expr_types if isinstance(other, _ewise_mult_expr_types): - raise TypeError("XXX") + raise TypeError( + "Cannot mix `|` (ewise_add) with `&` (ewise_mult) expressions due to Python" + " operator precedence. Use explicit method calls instead." + ) return _ewise_infix_expr(self, other, method="ewise_add", within="__or__") def __ror__(self, other): from .infix import _ewise_infix_expr, _ewise_mult_expr_types if isinstance(other, _ewise_mult_expr_types): - raise TypeError("XXX") + raise TypeError( + "Cannot mix `|` (ewise_add) with `&` (ewise_mult) expressions due to Python" + " operator precedence. Use explicit method calls instead." + ) return _ewise_infix_expr(other, self, method="ewise_add", within="__ror__") def __and__(self, other): from .infix import _ewise_add_expr_types, _ewise_infix_expr if isinstance(other, _ewise_add_expr_types): - raise TypeError("XXX") + raise TypeError( + "Cannot mix `&` (ewise_mult) with `|` (ewise_add) expressions due to Python" + " operator precedence. Use explicit method calls instead." + ) return _ewise_infix_expr(self, other, method="ewise_mult", within="__and__") def __rand__(self, other): from .infix import _ewise_add_expr_types, _ewise_infix_expr if isinstance(other, _ewise_add_expr_types): - raise TypeError("XXX") + raise TypeError( + "Cannot mix `&` (ewise_mult) with `|` (ewise_add) expressions due to Python" + " operator precedence. Use explicit method calls instead." + ) return _ewise_infix_expr(other, self, method="ewise_mult", within="__rand__") def __matmul__(self, other): @@ -329,6 +358,7 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * if expr._is_scalar and self._is_scalar: # Extract element (s << v[1]) if accum is not None: + expr = _scalar_accum_operand(self, expr) self(**opts) << self.ewise_add(expr, accum) return expr.parent._extract_element( @@ -350,6 +380,10 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * expr = expr._extract_delayed() elif type(expr) is type(self): # Simple assignment (w << v) + if (cast_op := _store_cast_op(expr.dtype, self.dtype, is_value=True)) is not None: + expr = expr.apply(cast_op) + self._update(expr, mask, accum, replace, input_mask, opts=opts) + return if self._is_scalar: if accum is not None: self(**opts) << self.ewise_add(expr, accum) @@ -374,6 +408,7 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * # s << (v @ v) expr = expr._to_expr() elif accum is not None: + expr = _scalar_accum_operand(self, expr) self(**opts) << self.ewise_add(expr, accum) return else: @@ -387,16 +422,22 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * from .matrix import Matrix, MatrixExpression, TransposedMatrix if type(expr) is TransposedMatrix and type(self) is Matrix: - # Transpose (C << A.T) - expr = MatrixExpression( - "transpose", - "GrB_transpose", - [expr], - expr_repr="{0}", - dtype=expr.dtype, - nrows=expr._nrows, - ncols=expr._ncols, - ) + if ( + cast_op := _store_cast_op(expr.dtype, self.dtype, is_value=True) + ) is not None: + # C << A.T of another UDT type: the cast reads A transposed. + expr = expr.apply(cast_op) + else: + # Transpose (C << A.T) + expr = MatrixExpression( + "transpose", + "GrB_transpose", + [expr], + expr_repr="{0}", + dtype=expr.dtype, + nrows=expr._nrows, + ncols=expr._ncols, + ) elif isinstance(expr, InfixExprBase): # w << (v & v) expr = expr._to_expr() @@ -406,7 +447,7 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * if type(expr) is Scalar: scalar = expr else: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, expr) try: scalar = Scalar.from_value(expr, dtype, is_cscalar=None, name="") except TypeError: @@ -434,6 +475,21 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * if input_mask is not None: raise TypeError("`input_mask` argument may only be used for extract") + # With an accumulator too: it is typed on this object's dtype, so the + # result still has to be that type. Before the Aggregator branch, which + # hands GraphBLAS the output as it is: agg.first on another UDT read its + # bytes as this object's type. + if (cast_op := _store_cast_op(expr.dtype, self.dtype)) is not None: + if self._is_scalar: + # One element: compute it in its own type, then cast it as it is + # stored, so the accumulator applies to the cast value. + result = expr.new(name="s_result", **opts) + self._update(result.apply(cast_op), mask, accum, replace, opts=opts) + return + if expr.cfunc_name is not None: + # A recipe (cfunc_name None) stores through ``<<`` itself, and + # each of its stores is checked there. + expr = _fused_store_expr(expr, self.dtype) if expr.op is not None and expr.op.opclass == "Aggregator": updater = self(mask=mask, accum=accum, replace=replace, **opts) expr.op._new(updater, expr) @@ -448,7 +504,11 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * complement = False structure = False else: - mask = _check_mask(mask, self) + # Assignment (`method_name == "__setitem__"`) may target a Matrix + # row/column with a Vector mask, so only enforce the strict + # mask-kind match for full-tensor operations. + strict_kind = expr.method_name != "__setitem__" + mask = _check_mask(mask, self, strict_kind=strict_kind) complement = mask.complement structure = mask.structure @@ -531,10 +591,12 @@ class BaseExpression: "expr_repr", "dtype", "_value", + "_right_literal", "__weakref__", ) output_type = None _is_scalar = False + __array_priority__ = BaseType.__array_priority__ # numpy's operators defer to ours def __init__( self, @@ -566,6 +628,9 @@ def __init__( self.dtype = op.return_type else: self.dtype = dtype + # The right operand as the caller gave it, a literal a comparison typed + # beside the values; the select helpers compare indices with it instead. + self._right_literal = None self._value = None def new(self, dtype=None, *, mask=None, name=None, **opts): @@ -599,7 +664,9 @@ def _new(self, dtype, mask, name, is_cscalar=None, **opts): elif mask is None: output.update(self, **opts) else: - mask = _check_mask(mask, output) + # `.new()` always builds a full output matching this expression, so + # the mask kind must match the output dimensions. + mask = _check_mask(mask, output, strict_kind=True) output(mask=mask, **opts).update(self) return output diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 2d4178b14..97cbf1f62 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -1,21 +1,36 @@ +import dataclasses import warnings from ast import literal_eval import numpy as np -from numpy import promote_types, result_type +from numpy import promote_types from .. import backend, dtypes from ..core import NULL, _has_numba, ffi, lib if _has_numba: import numba + from numba.np.numpy_support import as_dtype as _numba_as_dtype # Default assumption unless FC32/FC64 are found in lib _supports_complex = hasattr(lib, "GrB_FC64") or hasattr(lib, "GxB_FC64") class DataType: - __slots__ = "name", "gb_obj", "gb_name", "c_type", "numba_type", "np_type", "__weakref__" + __slots__ = ( + "name", + "gb_obj", + "gb_name", + "c_type", + "numba_type", + "np_type", + # ``(c_name, c_definition)`` actually registered with SuiteSparse, + # or ``None``. Captured at first registration. ``GxB_JIT_C_DEFINITION`` + # is one-shot in SS, so this is the authoritative record of what SS + # will use for JIT, regardless of any later Python-side rename. + "_jit_c_info", + "__weakref__", + ) def __init__(self, name, gb_obj, gb_name, c_type, numba_type, np_type): self.name = name @@ -24,6 +39,7 @@ def __init__(self, name, gb_obj, gb_name, c_type, numba_type, np_type): self.c_type = c_type self.numba_type = numba_type self.np_type = np.dtype(np_type) if np_type is not None else None + self._jit_c_info = None def __repr__(self): return self.name @@ -65,6 +81,21 @@ def _is_anonymous(self): def _is_udt(self): return self.gb_name is None + @property + def jit_c_name(self): + """The C type name SuiteSparse uses for this UDT, or ``None``. + + Returns the name SuiteSparse actually registered (not the latest + Python-side ``self.name``, which can diverge after a re-register). Is + ``None`` for built-in types and for UDTs not expressible in C. + """ + return self._jit_c_info[0] if self._jit_c_info is not None else None + + @property + def jit_c_definition(self): + """The C struct typedef SuiteSparse uses for this UDT, or ``None``.""" + return self._jit_c_info[1] if self._jit_c_info is not None else None + @staticmethod def _deserialize(name, dtype, is_anonymous): if is_anonymous: @@ -74,8 +105,135 @@ def _deserialize(name, dtype, is_anonymous): return register_new(name, dtype) -def register_new(name, dtype): - if not name.isidentifier(): +def _get_udt_c_type_info(datatype): + """Return ``(c_name, c_typedef)`` for a UDT, or ``None`` if not expressible in C. + + The returned ``c_name`` is ``gbudt_`` and ``datatype.name``, or a fresh + ``_gbudt_NNN`` name when the user didn't supply a name, or supplied one + that isn't a valid C identifier or is too long for SuiteSparse + (``udt_codegen._pick_c_type_name``), so the UDT can still take the JIT + path. Returns ``None`` only for the cases that remain unrepresentable in + C: an unsupported field type (object, datetime, string), or a field name + that collides with a C reserved word or isn't an identifier. Supports + record UDTs (including nested records) and array UDTs (e.g., ``FP64[3]``). + """ + from .operator.udt_codegen import _udt_c_typedef + + if datatype.np_type is None: + return None + return _udt_c_typedef(datatype.name, datatype.np_type) + + +def _set_udt_jit_c_definition(datatype): + """Set JIT C name and typedef on a UDT so SuiteSparse can JIT-compile kernels. + + Without these, SuiteSparse treats UDTs as opaque byte arrays and cannot + JIT-compile optimized kernels for operations on them. We generate a C + struct definition from the numpy dtype and set both ``GxB_JIT_C_NAME`` + and ``GxB_JIT_C_DEFINITION`` on the ``GrB_Type``. Both are needed for + SuiteSparse's eWise and reduce JIT templates to expand correctly. + + Setting ``GxB_JIT_C_NAME`` causes ``GxB_deserialize_type_name`` to + return this short name rather than the numpy repr stored in + ``GrB_NAME``, so the deserialize path in ``Matrix.ss.deserialize`` and + ``Vector.ss.deserialize`` falls back to ``GrB_NAME`` when the short + name doesn't resolve to a known dtype. This preserves round-tripping + for anonymous-by-name UDTs. + + Supports both record UDTs (``{'x': float, 'y': float}``) and array UDTs + (``np.dtype((np.float64, (3,)))``). + """ + from .operator.udt_codegen import ( + _has_jit_set, + _jit_c_type_defs, + _synthetic_udt_counter, + _udt_c_typedef, + ) + + if not _has_jit_set or not datatype._is_udt: + return + info = _get_udt_c_type_info(datatype) + if info is None: + return + c_name, typedef = info + # A name already given to another layout (two UDTs registered under one + # name) gets a fresh one: SuiteSparse would share their JIT kernels. + while _jit_c_type_defs.setdefault(c_name, typedef) != typedef: + c_name = f"_gbudt_{next(_synthetic_udt_counter)}" + info = _udt_c_typedef(c_name, datatype.np_type, is_c_name=True) + c_name, typedef = info + lib.GrB_Type_set_String(datatype._carg, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) + lib.GrB_Type_set_String( + datatype._carg, ffi.new("char[]", typedef.encode()), lib.GxB_JIT_C_DEFINITION + ) + # Remember what SS actually has. ``GxB_JIT_C_DEFINITION`` is one-shot + # (returns ``GrB_ALREADY_SET`` on a second call), so a later rename of + # the Python DataType does not change what SS uses for JIT, and the + # introspection properties must reflect what SS has rather than the + # renamed value. + datatype._jit_c_info = info + + +# Elements Numba cannot type in an array (its NestedArray needs a bitwidth they lack), and +# what to store instead +_NUMBA_UNTYPED_ARRAY_ELEMENTS = { + "b": ("bool", "uint8"), + "M": ("datetime64", "int64"), + "m": ("timedelta64", "int64"), + "U": ("str", "uint32 (one per character)"), + "S": ("bytes", "uint8 (one per byte)"), +} + + +def _numba_untyped_array_element(dtype): + """The kind of the first array element in dtype that Numba cannot type, or None.""" + if dtype.subdtype is not None: + base = dtype.subdtype[0] + if base.kind in _NUMBA_UNTYPED_ARRAY_ELEMENTS: + return base.kind + return _numba_untyped_array_element(base) + for field_dtype, *_ in (dtype.fields or {}).values(): + if (kind := _numba_untyped_array_element(field_dtype)) is not None: + return kind + return None + + +def _numba_type(dtype): + try: + return numba.typeof(dtype).dtype + except AttributeError as exc: + # Numba's own error is "'Boolean' object has no attribute 'bitwidth'" and the like + kind = _numba_untyped_array_element(dtype) + if kind is None: # pragma: no cover (safety) + raise + element, instead = _NUMBA_UNTYPED_ARRAY_ELEMENTS[kind] + raise ValueError( + f"Numba, which python-graphblas uses for UDTs, cannot type an array of {element}, " + f"as in {dtype}; use an array of {instead} instead" + ) from exc + + +def register_new(name, dtype=None): + """Register a user-defined type and put it on the ``gb.dtypes`` namespace. + + Symmetric with :func:`register_anonymous`. Accepts the same dtype forms + (numpy dtype, dict, string, dataclass class/instance). When ``dtype`` is a + dataclass and no explicit ``name`` is given, the dataclass class name is + used; passing ``register_new(MyDataclass)`` is shorthand for + ``register_new(MyDataclass.__name__, MyDataclass)``. + """ + if dtype is None: + # Sole-argument call. Accept a dataclass (use its class name) or a + # ``DataType``; otherwise complain with a useful message. + if dataclasses.is_dataclass(name): + dtype = name + name = name.__name__ if isinstance(name, type) else type(name).__name__ + else: + raise TypeError( + "register_new() requires both `name` and `dtype`, or a single " + "@dataclass class/instance whose class name becomes the dtype name" + ) + if not isinstance(name, str) or not name.isidentifier(): raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}") if name in _registry or hasattr(dtypes, name): raise ValueError(f"{name!r} name for dtype is unavailable") @@ -86,6 +244,60 @@ def register_new(name, dtype): def register_anonymous(dtype, name=None): + """Register a user-defined type without adding it to the ``gb.dtypes`` namespace. + + Returns the existing ``DataType`` if one is already registered for this + ``np.dtype``. The new ``name`` (if given) replaces the old Python-side + name; the SuiteSparse-side C name set at first registration does not + change. See ``DataType.jit_c_name``. + + Nested subarrays are flattened first, so ``np.dtype((np.dtype((np.float64, + (3,))), (2,)))``, two ``FP64[3]`` elements, registers as ``FP64[2, 3]`` and + returns the same ``DataType``. This applies inside records too. + + Parameters + ---------- + dtype : np.dtype | dict | str | @dataclass + The dtype to register. Supported forms: + + - A ``np.dtype`` directly (record or array, including multi-dim). + - A dict like ``{"x": int, "y": float}``: keys become field names, + and values resolve through ``lookup_dtype``. + - A string like ``"INT64[3, 4]"`` for array UDTs. + - A ``@dataclass`` class or instance: fields become record fields, + and the class name is used as the default UDT name. Field type + annotations may be either real types (``int``) or strings + (``"int"``), e.g., with ``from __future__ import annotations`` or + PEP 649. Compound annotations like ``Optional[int]`` raise from + ``lookup_dtype``. + + Field types must resolve through ``lookup_dtype``. Python built-ins + (``int``, ``float``, ``bool``, ``complex``), the corresponding numpy + scalar types, and graphblas dtype names are supported. Object, + string, and nested-UDT fields are not. + name : str, optional + A human-readable name for this UDT. Used in error messages and + ``repr``. Any name takes the SuiteSparse JIT path: its C name is + ``gbudt_`` and the name, or a synthetic one when that is not a C + identifier or is too long. + + Returns + ------- + DataType + The registered UDT. + """ + # Convert a ``@dataclass`` (class or instance) directly to a numpy + # record dtype. ``lookup_dtype`` handles both type-object and + # string-annotation forms of field types. + if dataclasses.is_dataclass(dtype): + fields = dataclasses.fields(dtype) + if not fields: + raise ValueError("dataclass must have at least one field to convert to a UDT") + if name is None: + # Default the UDT name to the dataclass class name for both + # class form (``MyDC``) and instance form (``MyDC(...)``). + name = dtype.__name__ if isinstance(dtype, type) else type(dtype).__name__ + dtype = np.dtype([(f.name, lookup_dtype(f.type).np_type) for f in fields], align=True) try: dtype = np.dtype(dtype) except TypeError: @@ -102,16 +314,31 @@ def register_anonymous(dtype, name=None): dtype = np.dtype((base_dtype.np_type, shape)) else: raise + dtype = _flatten_subarrays(dtype) if dtype in _registry: - # Always use the same object, but use the latest name + # Always use the same object, but use the latest name. The + # Python-side ``rv.name`` updates here, but the SuiteSparse-side + # ``GxB_JIT_C_NAME`` and ``GxB_JIT_C_DEFINITION`` are one-shot + # (return ``GrB_ALREADY_SET``) and keep the first name. + # ``rv._jit_c_info`` was pinned at first registration; that's what + # ``jit_c_name`` and ``jit_c_definition`` return, so introspection + # reflects what SS actually has regardless of any later Python-side + # rename. JIT codegen for ops on this UDT also uses the pinned name + # (see ``_make_jit_c_definition``). rv = _registry[dtype] if name is not None: if rv.gb_name is not None and name != rv.gb_name: raise ValueError("dtype must not be a builtin type") - rv.name = name # Rename an existing object (a little weird, but okay) + rv.name = name return rv if dtype.hasobject: raise ValueError("dtype must not allow Python objects") + if dtype.itemsize == 0: + # SuiteSparse 7 takes a size of 0 as 1 (later versions raise InvalidValue), which + # made ``str`` a UDT that holds nothing, so a str literal was not a TypeError + raise ValueError(f"dtype must not have a size of 0: {dtype}") + # Before GrB_Type_new, so a dtype Numba cannot type does not leave a GrB_Type behind + numba_type = _numba_type(dtype) if _has_numba else None from ..exceptions import check_status_carg @@ -131,16 +358,13 @@ def register_anonymous(dtype, name=None): # We don't yet have C definitions np_repr = _dtype_to_string(dtype).encode() if len(np_repr) > lib.GxB_MAX_NAME_LEN: - msg = ( - f"UDT repr is too large to serialize ({len(repr(dtype).encode())} > " - f"{lib.GxB_MAX_NAME_LEN})." - ) + msg = f"UDT repr is too large to serialize ({len(np_repr)} > {lib.GxB_MAX_NAME_LEN})." if name is not None: np_repr = name.encode()[: lib.GxB_MAX_NAME_LEN] else: np_repr = np_repr[: lib.GxB_MAX_NAME_LEN] warnings.warn( - f"{msg}. It will use the following name, " + f"{msg} It will use the following name, " f"and the dtype may need to be specified when deserializing: {np_repr}", stacklevel=2, ) @@ -152,13 +376,11 @@ def register_anonymous(dtype, name=None): # For now, let's use "opaque" unsigned bytes for the c type. if name is None: name = _default_name(dtype) - numba_type = numba.typeof(dtype).dtype if _has_numba else None rv = DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype) _registry[gb_obj] = rv _registry[dtype] = rv - if _has_numba: - _registry[numba_type] = rv - _registry[numba_type.name] = rv + # Set JIT C type definition so SuiteSparse can JIT-compile kernels for this UDT. + _set_udt_jit_c_definition(rv) return rv @@ -368,6 +590,12 @@ def lookup_dtype(key, value=None): return _registry[key] except (KeyError, TypeError): pass + if _has_numba and isinstance(key, (numba.types.Record, numba.types.NestedArray)): + # A UDT's Numba type is not a registry key. Convert it to its numpy + # dtype, which is already flat, and look that up, so a layout no UDT + # has registered yet (e.g., a record's array field returned by a UDF) + # resolves the same way regardless of what else is registered. + return lookup_dtype(_numba_as_dtype(key)) if value is not None and hasattr(value, "dtype") and value.dtype in _registry: return _registry[value.dtype] # np.dtype(x) accepts some weird values; we may want to guard against some @@ -385,6 +613,28 @@ def lookup_dtype(key, value=None): raise ValueError(f"Unknown dtype: {key} of type {type(key)}") +def _raise_dtype_or_arraylike(cls_name, dtype, exc): + """Turn a failed constructor dtype lookup into a helpful error. + + Call from ``Vector``/``Matrix`` after ``lookup_dtype(dtype)`` has already + raised ``exc``. The constructors take a dtype as the first argument, so a + common mistake is to pass the data instead (``Vector([1, 2, 3])``). When + ``dtype`` is array-like data rather than a dtype spec, point at the ``from_*`` + constructors; otherwise re-raise the original ``Unknown dtype`` error. Valid + list/tuple dtype specs (structured and subarray dtypes) never reach here + because ``lookup_dtype`` accepts them, but a malformed one does, so the + lookup's error stays attached as the cause. + """ + if isinstance(dtype, (list, tuple, np.ndarray)): + raise TypeError( + f"{cls_name}() expects a dtype as the first argument, not " + f"{type(dtype).__name__} data. If you meant to pass values, use a " + f"constructor such as {cls_name}.from_coo(...) or " + f"{cls_name}.from_dense(...)." + ) from exc + raise exc + + def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False): """Returns a type that can hold both type1 and type2. @@ -393,14 +643,15 @@ def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False): unify(INT8, UINT16) -> INT32 unify(BOOL, UINT16) -> UINT16 unify(FP32, INT32) -> FP64 + + ``is_left_scalar`` and ``is_right_scalar`` no longer change the result. A + Python number is typed before it gets here (``literals._literal_operand``), and + a typed scalar promotes as any operand does, as numpy 2 promotes it. They + used to stand the scalar in as a 0-d array of value 0, which numpy 1 typed + by value, so ``int8_vec + 300`` wrapped there. """ if type1 is type2: return type1 - if is_left_scalar: - if not is_right_scalar: - return lookup_dtype(result_type(np.array(0, type1.np_type), type2.np_type)) - elif is_right_scalar: - return lookup_dtype(result_type(type1.np_type, np.array(0, type2.np_type))) return lookup_dtype(promote_types(type1.np_type, type2.np_type)) @@ -421,6 +672,98 @@ def _default_name(dtype): return repr(dtype) +def _flatten_subarrays(np_type): + """Collapse each run of nested subarray dtypes in ``np_type`` into one subarray. + + numpy keeps ``np.dtype((np.dtype((np.float64, (3,))), (2,)))`` layered, as + two ``FP64[3]`` elements, and compares it unequal to the flat + ``np.dtype((np.float64, (2, 3)))``. Nothing else tells them apart: the bytes + are the same, an array of either expands to a ``(..., 2, 3)`` float64 array, + and Numba types both as ``nestedarray(float64, (2, 3))``. Registering the flat + form gives each layout one DataType. Records are rebuilt only when a field + changes, keeping their offsets, titles, itemsize, and alignment flag. + """ + if np_type.subdtype is not None: + base, shape = np_type.subdtype + flat_base = _flatten_subarrays(base) + if flat_base.subdtype is not None: + flat_base, inner_shape = flat_base.subdtype + shape += inner_shape + elif flat_base is base: + return np_type + return np.dtype((flat_base, shape)) + if np_type.names is not None: + fields = [np_type.fields[name] for name in np_type.names] + formats = [_flatten_subarrays(field[0]) for field in fields] + if all(new is field[0] for new, field in zip(formats, fields, strict=True)): + return np_type + spec = { + "names": list(np_type.names), + "formats": formats, + "offsets": [field[1] for field in fields], + "itemsize": np_type.itemsize, + } + if any(len(field) > 2 for field in fields): + spec["titles"] = [field[2] if len(field) > 2 else None for field in fields] + return np.dtype(spec, align=np_type.isalignedstruct) + return np_type + + +def _view_if_same_layout(values, np_type): + """Return ``values`` viewed as ``np_type`` if the two differ only in subarray nesting. + + A record whose field nests subarrays has a different dtype from its + registered, flattened form, and numpy casts between the two wrongly in + either direction: ``astype``, ``np.array(..., dtype=)``, and item assignment + all scramble the field's values. The bytes already match, so a view is exact. + """ + if ( + isinstance(values, (np.ndarray, np.void)) + and values.dtype != np_type + and values.dtype.itemsize == np_type.itemsize + and _flatten_subarrays(values.dtype) == np_type + ): + return values.view(np_type) + return values + + +def _nats_in_leaf_units(values, np_type): + """Return ``values`` with each NaT duration in it in the unit of the leaf it fills. + + numpy before 2.3 converts a NaT ``timedelta64`` scalar, alone or in a tuple + or list, into another unit as the integer it is stored as: NaT ms into + ``timedelta64[s]`` is 9223372036854774 s, and NaT s into ``timedelta64[ms]`` + is 0 ms. Converting an array (``astype``) keeps NaT, as does a date, so each + NaT scalar is converted that way first. ``values`` holds elements of + ``np_type``: a list is a sequence of them, and a tuple is a record's fields + (or, for any other dtype, a sequence too). + """ + if not _has_duration_leaf(np_type): + return values + while np_type.subdtype is not None: + np_type = np_type.subdtype[0] + if isinstance(values, list) or (isinstance(values, tuple) and np_type.names is None): + items = [_nats_in_leaf_units(x, np_type) for x in values] + return items if isinstance(values, list) else tuple(items) + if isinstance(values, tuple) and len(values) == len(np_type.names): + return tuple( + _nats_in_leaf_units(x, np_type.fields[name][0]) + for x, name in zip(values, np_type.names, strict=True) + ) + if np_type.kind == "m" and isinstance(values, np.timedelta64) and np.isnat(values): + return np.asarray(values).astype(np_type)[()] + return values + + +def _has_duration_leaf(np_type): + """Whether ``np_type`` has a timedelta64 leaf: a record field, an array's element, or itself.""" + if np_type.subdtype is not None: + return _has_duration_leaf(np_type.subdtype[0]) + if np_type.names is not None: + return any(_has_duration_leaf(np_type.fields[name][0]) for name in np_type.names) + return np_type.kind == "m" + + def _dtype_to_string(dtype): """Convert a dtype to a string that can be safely evaluated to recreate the dtype. @@ -440,13 +783,53 @@ def _dtype_to_string(dtype): np_type = dtype.np_type s = str(np_type) try: - if np.dtype(literal_eval(s)) == np_type: # pragma: no branch (safety) + if np.dtype(literal_eval(s)) == np_type: return s except Exception: pass - if np.dtype(np_type.str) != np_type: # pragma: no cover (safety) - raise ValueError(f"Unable to reliably convert dtype to string and back: {dtype}") - return repr(np_type.str) + if np.dtype(np_type.str) == np_type: + return repr(np_type.str) + # Some nested dtypes (e.g. ``align=True`` outer + ``packed`` inner) don't + # round-trip via ``str(np_type)`` because numpy's reconstruction enforces + # the outer's alignment on the inner. Fall back to an explicit per-field + # dict using literal offsets/itemsize, which encodes the exact layout + # without invoking ``align``. ``literal_eval`` accepts the result, so + # ``_string_to_dtype`` doesn't need a special case for it. + return repr(_dtype_to_explicit_dict(np_type)) + + +def _dtype_to_explicit_dict(np_type): + """Encode ``np_type`` as a plain dict/tuple/str tree of literals. + + Layout-preserving fallback for :func:`_dtype_to_string`. Records become a + ``{names, formats, offsets, itemsize}`` dict whose nested ``formats`` may + be further nested dicts; array dtypes become a ``(base_str, shape)`` + tuple. The result round-trips through ``np.dtype(literal_eval(s))`` + without ever using ``aligned=True``, so the original byte layout is + preserved verbatim even for the aligned-outer / packed-inner cases. + """ + if np_type.names is not None: + formats = [] + for name in np_type.names: + sub = np_type.fields[name][0] + if sub.names is not None or sub.subdtype is not None: + formats.append(_dtype_to_explicit_dict(sub)) + else: + formats.append(sub.str) + return { + "names": list(np_type.names), + "formats": formats, + "offsets": [np_type.fields[name][1] for name in np_type.names], + "itemsize": np_type.itemsize, + } + if np_type.subdtype is not None: + base, shape = np_type.subdtype + if base.names is not None: + base_repr = _dtype_to_explicit_dict(base) + else: + base_repr = base.str + return (base_repr, shape) + return np_type.str def _string_to_dtype(s): diff --git a/graphblas/core/expr.py b/graphblas/core/expr.py index efec2db5f..1e16c44b4 100644 --- a/graphblas/core/expr.py +++ b/graphblas/core/expr.py @@ -281,6 +281,8 @@ def __lshift__(self, expr): def __eq__(self, other): raise TypeError(f"__eq__ not defined for objects of type {type(self)}.") + __hash__ = None + def __bool__(self): raise TypeError(f"__bool__ not defined for objects of type {type(self)}.") @@ -288,6 +290,7 @@ def __bool__(self): class AmbiguousAssignOrExtract: __slots__ = "parent", "resolved_indexes", "_value", "__weakref__" _is_scalar = False + __array_priority__ = 1000 # numpy's operators defer to ours; see BaseType def __init__(self, parent, resolved_indexes): self.parent = parent @@ -319,7 +322,7 @@ def new(self, dtype=None, *, mask=None, input_mask=None, name=None, **opts): if input_mask is not None: if mask is not None: raise TypeError("mask and input_mask arguments cannot both be given") - from .base import _check_mask + from .mask import _check_mask input_mask = _check_mask(input_mask, self.parent) mask = self._input_mask_to_mask(input_mask, **opts) @@ -473,6 +476,8 @@ def update(self, expr): def __eq__(self, other): raise TypeError(f"__eq__ not defined for objects of type {type(self)}.") + __hash__ = None + def __bool__(self): raise TypeError(f"__bool__ not defined for objects of type {type(self)}.") @@ -480,6 +485,7 @@ def __bool__(self): class InfixExprBase: __slots__ = "left", "right", "_expr", "__weakref__" _is_scalar = False + __array_priority__ = 1000 # numpy's operators defer to ours; see BaseType def __init__(self, left, right): self.left = left @@ -507,6 +513,10 @@ def _to_expr(self): self._expr = getattr(self.left, self.method_name)(self.right) return self._expr + def _call_left(self, method_name, *args, **kwargs): + """Call the left operand's ``method_name`` with the right operand, as ``op(x & y)`` does.""" + return getattr(self.left, method_name)(self.right, *args, **kwargs) + def _get_value(self, attr=None, default=None): expr = self._to_expr() return expr._get_value(attr=attr, default=default) diff --git a/graphblas/core/formatting.py b/graphblas/core/formatting.py index 0b6252101..d4ac8cfe4 100644 --- a/graphblas/core/formatting.py +++ b/graphblas/core/formatting.py @@ -4,6 +4,7 @@ from .. import backend, config, monoid, unary from ..dtypes import BOOL from ..exceptions import OutOfMemory +from .infix import _is_literal from .matrix import Matrix, TransposedMatrix from .vector import Vector @@ -119,7 +120,7 @@ """ -def _update_matrix_dataframe(df, matrix, rows, row_offset, columns, column_offset, *, mask=None): +def _update_matrix_array(arr, matrix, rows, row_offset, columns, column_offset, *, mask=None): if rows is None and columns is None: if mask is None: submatrix = matrix @@ -167,13 +168,17 @@ def _update_matrix_dataframe(df, matrix, rows, row_offset, columns, column_offse np_type = submatrix.dtype.np_type if submatrix.dtype._is_udt and np_type.subdtype is not None: vals = vals.tolist() - df.values[rows, cols] = vals + if isinstance(vals, np.ndarray) and vals.dtype.names is not None: + # Structured array: convert numpy.void elements to tuples for consistent display + arr[rows, cols] = [tuple(v) for v in vals] + else: + arr[rows, cols] = vals if np.issubdtype(np_type, np.inexact): nulls = np.isnan(vals) - df.values[rows[nulls], cols[nulls]] = "nan" + arr[rows[nulls], cols[nulls]] = "nan" -def _update_vector_dataframe(df, vector, columns, column_offset, *, mask=None): +def _update_vector_array(arr, vector, columns, column_offset, *, mask=None): if columns is None: if mask is None: subvector = vector @@ -205,9 +210,13 @@ def _update_vector_dataframe(df, vector, columns, column_offset, *, mask=None): np_type = subvector.dtype.np_type if subvector.dtype._is_udt and np_type.subdtype is not None: vals = vals.tolist() - df.values[0, cols] = vals + if isinstance(vals, np.ndarray) and vals.dtype.names is not None: + # Structured array: convert numpy.void elements to tuples for consistent display + arr[0, cols] = [tuple(v) for v in vals] + else: + arr[0, cols] = vals if np.issubdtype(np_type, np.inexact): - df.values[0, cols[np.isnan(vals)]] = "nan" + arr[0, cols[np.isnan(vals)]] = "nan" def _get_max_columns(): @@ -244,11 +253,11 @@ def _get_matrix_dataframe(matrix, max_rows, min_rows, max_columns, *, mask=None) max_columns = _get_max_columns() rows, row_groups = _get_chunk(matrix._nrows, min_rows, max_rows) columns, column_groups = _get_chunk(matrix._ncols, max_columns, max_columns) - df = pd.DataFrame(columns=columns, index=rows) + arr = np.full((len(rows), len(columns)), np.nan, dtype=object) for row_group, row_offset in row_groups: for column_group, column_offset in column_groups: - _update_matrix_dataframe( - df, + _update_matrix_array( + arr, matrix, row_group, row_offset, @@ -256,6 +265,7 @@ def _get_matrix_dataframe(matrix, max_rows, min_rows, max_columns, *, mask=None) column_offset, mask=mask, ) + df = pd.DataFrame(arr, columns=columns, index=rows) if ( (mask is None or mask.structure) and df.shape != matrix.shape @@ -306,9 +316,10 @@ def _get_vector_dataframe(vector, max_rows, min_rows, max_columns, *, mask=None) if max_columns is None: # pragma: no branch max_columns = _get_max_columns() columns, column_groups = _get_chunk(vector._size, max_columns, max_columns) - df = pd.DataFrame(columns=columns, index=[""]) + arr = np.full((1, len(columns)), np.nan, dtype=object) for column_group, column_offset in column_groups: - _update_vector_dataframe(df, vector, column_group, column_offset, mask=mask) + _update_vector_array(arr, vector, column_group, column_offset, mask=mask) + df = pd.DataFrame(arr, columns=columns, index=[""]) if ( (mask is None or mask.structure) and df.size != vector._size @@ -368,7 +379,7 @@ def matrix_info(matrix, *, mask=None, expr=None, for_html=True): vals = [matrix._nvals, matrix._nrows, matrix._ncols, matrix.dtype.name] if expr is None and backend == "suitesparse": keys.append("format") - if type(matrix) is Matrix: + if not matrix._is_transposed: vals.append(get_format(matrix)) else: # TransposedMatrix vals.append(get_format(matrix._matrix, is_transposed=True)) @@ -707,11 +718,7 @@ def _format_infix_expression(expr, header, expr_name): ) computed = "" - if config.get("autocompute") and ( - expr.method_name not in {"ewise_add", "ewise_mult"} - or expr.left.dtype == BOOL - and expr.right.dtype == BOOL - ): + if config.get("autocompute") and _computes_itself(expr): computed = get_expr_result(expr, html=True) if "__EXPR__" in computed: return computed.replace("__EXPR__", topline) @@ -730,17 +737,39 @@ def _format_infix_expression(expr, header, expr_name): ) +def _operand_dtype(operand): + """The dtype to show for an operand of a Scalar infix expression. + + A literal, as in ``s & 1``, has no dtype until an op types it, so its own + type is shown. + """ + if _is_literal(operand): + return type(operand).__name__ + return operand.dtype + + +def _computes_itself(expr): + """Whether infix ``expr`` computes without an op, as ``x @ y`` and BOOL ``x & y`` do.""" + if expr.method_name not in {"ewise_add", "ewise_mult"}: + return True + # A literal beside a Scalar, as in `s & True`, is BOOL beside BOOL if it is a bool + return all( + type(x) in {bool, np.bool_} if _is_literal(x) else x.dtype == BOOL + for x in (expr.left, expr.right) + ) + + def format_scalar_infix_expression(expr): expr_repr = expr._format_expr() name = f"gb.{type(expr).__name__}" header = create_header( expr_repr, ["left_dtype", "right_dtype"], - [expr.left.dtype, expr.right.dtype], + [_operand_dtype(expr.left), _operand_dtype(expr.right)], name=name, quote=False, ) - arg_string = get_result_string(expr) + arg_string = get_infix_result_string(expr) return ( f"{header}\n\n" f"{arg_string}" @@ -754,17 +783,13 @@ def format_scalar_infix_expression_html(expr): header = create_header_html( expr_html, ["left_dtype", "right_dtype"], - [expr.left.dtype, expr.right.dtype], + [_operand_dtype(expr.left), _operand_dtype(expr.right)], ) return _format_infix_expression(expr, header, expr_html) def get_infix_result_string(expr): - if ( - expr.method_name not in {"ewise_add", "ewise_mult"} - or expr.left.dtype == BOOL - and expr.right.dtype == BOOL - ): + if _computes_itself(expr): arg_string = get_result_string(expr) else: arg_string = "" diff --git a/graphblas/core/infix.py b/graphblas/core/infix.py index 24c109639..2bdbb00c6 100644 --- a/graphblas/core/infix.py +++ b/graphblas/core/infix.py @@ -1,3 +1,5 @@ +import numpy as np + from .. import backend, binary from ..dtypes import BOOL from ..exceptions import DimensionMismatch @@ -6,9 +8,10 @@ from . import automethods, recorder, utils from .base import _expect_op, _expect_type from .expr import InfixExprBase +from .literals import _LITERAL_KINDS, _literal_operand from .mask import Mask from .matrix import Matrix, MatrixExpression, TransposedMatrix -from .scalar import Scalar, ScalarExpression +from .scalar import Scalar, ScalarExpression, _literal_scalar from .utils import output_type, wrapdoc from .vector import Vector, VectorExpression @@ -16,15 +19,26 @@ InfixExprBase._expect_type = _expect_type +def _typed_operands(self, op): + """Return the operands of ``self``, with a literal beside a Scalar typed as ``op``'s operand.""" + left, right = self.left, self.right + if _is_literal(left): + left = _literal_scalar(right, left, op, within=self.method_name, argname="left") + elif _is_literal(right): + right = _literal_scalar(left, right, op, within=self.method_name, argname="right") + return left, right + + def _ewise_add_to_expr(self): if self._expr is not None: return self._expr - if self.left.dtype == BOOL and self.right.dtype == BOOL: - self._expr = self.left.ewise_add(self.right, lor) + left, right = _typed_operands(self, lor) + if left.dtype == BOOL and right.dtype == BOOL: + self._expr = left.ewise_add(right, lor) return self._expr raise TypeError( "Bad dtypes for `x | y`! Automatic computation of `x | y` infix expressions is only valid " - f"for BOOL dtypes. The argument dtypes are {self.left.dtype} and {self.right.dtype}.\n\n" + f"for BOOL dtypes. The argument dtypes are {left.dtype} and {right.dtype}.\n\n" "When auto-computed for boolean dtypes, `x | y` performs ewise_add (union) using LOR.\n\n" "Typical usage is to create an ewise_add expression such as `monoid.plus(x | y)`." ) @@ -33,18 +47,31 @@ def _ewise_add_to_expr(self): def _ewise_mult_to_expr(self): if self._expr is not None: return self._expr - if self.left.dtype == BOOL and self.right.dtype == BOOL: - self._expr = self.left.ewise_mult(self.right, land) + left, right = _typed_operands(self, land) + if left.dtype == BOOL and right.dtype == BOOL: + self._expr = left.ewise_mult(right, land) return self._expr raise TypeError( "Bad dtypes for `x & y`! Automatic computation of `x & y` infix expressions is only valid " - f"for BOOL dtypes. The argument dtypes are {self.left.dtype} and {self.right.dtype}.\n\n" + f"for BOOL dtypes. The argument dtypes are {left.dtype} and {right.dtype}.\n\n" "When auto-computed for boolean dtypes, `x & y` performs ewise_mult (intersection) using " "LAND.\n\n" "Typical usage is to create an ewise_mult expression such as `monoid.times(x & y)`." ) +def _is_literal(operand): + """Whether an operand of an infix expression is a literal, which only a Scalar takes.""" + return output_type(operand) not in {Scalar, Vector, Matrix, TransposedMatrix} + + +def _operand_name(operand, *, html=False): + if _is_literal(operand): + # As in `s & 1`; it shows its value, as an unnamed Scalar does + return repr(operand) + return operand._name_html if html else operand.name + + class ScalarInfixExpr(InfixExprBase): __slots__ = () ndim = 0 @@ -70,6 +97,23 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): def is_grbscalar(self): return not self.is_cscalar + def _format_expr(self): + return f"{_operand_name(self.left)} {self._infix} {_operand_name(self.right)}" + + def _format_expr_html(self): + left = _operand_name(self.left, html=True) + right = _operand_name(self.right, html=True) + return f"{left} {self._infix} {right}" + + def _call_left(self, method_name, *args, **kwargs): + if _is_literal(self.left): + # A literal stays on the left: `binary.minus(1 & s)` is `1 - s`. Scalar's method + # types it beside `s` once the op is known, as it types a literal on the right. + return getattr(Scalar, method_name)(self.left, self.right, *args, **kwargs) + return super()._call_left(method_name, *args, **kwargs) + + __hash__ = None + # Begin auto-generated code: Scalar __and__ = wrapdoc(Scalar.__and__)(property(automethods.__and__)) __array__ = wrapdoc(Scalar.__array__)(property(automethods.__array__)) @@ -134,10 +178,16 @@ class ScalarEwiseAddExpr(ScalarInfixExpr): # Don't allow e.g. `plus(x | y & z)` def __and__(self, other): - raise TypeError("XXX") + raise TypeError( + "Cannot mix `&` (ewise_mult) into an `|` (ewise_add) infix chain due to Python" + " operator precedence. Use explicit method calls instead." + ) def __rand__(self, other): - raise TypeError("XXX") + raise TypeError( + "Cannot mix `&` (ewise_mult) into an `|` (ewise_add) infix chain due to Python" + " operator precedence. Use explicit method calls instead." + ) class ScalarEwiseMultExpr(ScalarInfixExpr): @@ -155,10 +205,16 @@ class ScalarEwiseMultExpr(ScalarInfixExpr): # Don't allow e.g. `plus(x | y & z)` def __or__(self, other): - raise TypeError("XXX") + raise TypeError( + "Cannot mix `|` (ewise_add) into an `&` (ewise_mult) infix chain due to Python" + " operator precedence. Use explicit method calls instead." + ) def __ror__(self, other): - raise TypeError("XXX") + raise TypeError( + "Cannot mix `|` (ewise_add) into an `&` (ewise_mult) infix chain due to Python" + " operator precedence. Use explicit method calls instead." + ) class ScalarMatMulExpr(ScalarInfixExpr): @@ -471,7 +527,7 @@ def __init__(self, left, right, *, nrows, ncols): utils._output_types[MatrixMatMulExpr] = Matrix -def _dummy(obj, obj_type): +def _dummy(obj): with recorder.skip_record: return output_type(obj)(BOOL, *obj.shape, name="") @@ -484,6 +540,33 @@ def _mismatched(left, right, method, op): raise DimensionMismatch # pragma: no cover +def _check_literal(scalar, value, *, within, argname): + """Raise TypeError now if ``value`` cannot be a literal beside ``scalar``, as ``+`` does. + + How a literal is typed depends on the op (``literals._literal_operand``), which + comes later, so here it is typed as ``==`` types it, the op that takes the most + literals, and dropped. A Python number is always a literal, so it is not checked. + Nor is a literal beside an infix expression, which has no dtype until an op is + applied: ``(s & t) & object()`` raises only then. + """ + if type(value) in _LITERAL_KINDS or isinstance(scalar, ScalarInfixExpr): + return + value, dtype = _literal_operand(scalar.dtype, value, binary.eq) + try: + # A float too large for a float32 leaf warns once, when the op types the literal + with np.errstate(over="ignore"): + Scalar.from_value(value, dtype, is_cscalar=False, name="") + except TypeError: + _expect_type( + scalar, + value, + Scalar, + within=within, + argname=argname, + extra_message="Literal scalars also accepted.", + ) + + def _ewise_infix_expr(left, right, *, method, within): left_type = output_type(left) right_type = output_type(right) @@ -516,14 +599,15 @@ def _ewise_infix_expr(left, right, *, method, within): left._expect_type(right, tuple(types), within=within, argname="right") if right_type in types: right._expect_type(left, tuple(types), within=within, argname="left") - if left_type is Scalar: + if left_type is Scalar or right_type is Scalar: + # A literal operand stays a literal, on its own side, until an op types it + if left_type is not Scalar: + _check_literal(right, left, within=within, argname="left") + elif right_type is not Scalar: + _check_literal(left, right, within=within, argname="right") if method == "ewise_mult": return ScalarEwiseMultExpr(left, right) return ScalarEwiseAddExpr(left, right) - if right_type is Scalar: - if method == "ewise_mult": - return ScalarEwiseMultExpr(right, left) - return ScalarEwiseAddExpr(right, left) raise TypeError( # pragma: no cover (sanity) f"Bad types for ewise infix: {type(left).__name__}, {type(right).__name__}" ) diff --git a/graphblas/core/infixmethods.py b/graphblas/core/infixmethods.py index 71201ad9c..dfe539468 100644 --- a/graphblas/core/infixmethods.py +++ b/graphblas/core/infixmethods.py @@ -2,7 +2,8 @@ from ..dtypes import BOOL from .infix import MatrixInfixExpr, ScalarInfixExpr, VectorInfixExpr from .matrix import Matrix, MatrixExpression, MatrixIndexExpr, TransposedMatrix -from .scalar import Scalar, ScalarExpression, ScalarIndexExpr +from .scalar import Scalar, ScalarExpression, ScalarIndexExpr, _literal_scalar +from .udt import _ewise_add_casts_udt, _plus_zero from .utils import output_type from .vector import Vector, VectorExpression, VectorIndexExpr @@ -14,6 +15,16 @@ def call_op(self, other, op, *, outer=False, union=False): if type1 in types: if type2 in types: if outer: + if _ewise_add_casts_udt(self, other, op): + # ewise_add copies an entry present in only one input into + # the result type, which GraphBLAS cannot do for a UDT, so + # infix takes the union with each operand's zero for the + # missing value, as ``-`` takes 0. Its float fields are + # -0.0, which plus adds without changing any value; it is + # false for lor and lxor too. ``ewise_add`` itself raises. + return self.ewise_union( + other, op, _plus_zero(self.dtype), _plus_zero(other.dtype) + ) return self.ewise_add(other, op) if union: return self.ewise_union(other, op, False, False) @@ -21,18 +32,16 @@ def call_op(self, other, op, *, outer=False, union=False): return op(self, other) if type2 in types: return op(self, other) - # Scalars + # Scalars. A literal on the left stays there: `1 - s` is minus(1, s), typed as + # `binary.minus(1 & s)` is, since rminus may not exist for s's dtype (a UDT). + if type1 is not Scalar: + method = "ewise_add" if outer else "ewise_union" if union else "ewise_mult" + self = _literal_scalar(other, self, op, within=method, argname="left") if outer: - if type1 is Scalar: - return self.ewise_add(other, op) - return other.ewise_add(self, op.commutes_to) + return self.ewise_add(other, op) if union: - if type1 is Scalar: - return self.ewise_union(other, op, False, False) - return other.ewise_union(self, op.commutes_to, False, False) - if type1 is Scalar: - return self.ewise_mult(other, op) - return other.ewise_mult(self, op.commutes_to) + return self.ewise_union(other, op, False, False) + return self.ewise_mult(other, op) def __divmod__(self, other): @@ -65,7 +74,7 @@ def __xor__(self, other): if expr.dtype != BOOL: raise TypeError( f"The __xor__ infix operator, `x ^ y`, is not supported for {expr.dtype.name} dtype." - " It is only supported for BOOL dtype (and it uses ewise_add--the union)." + " It is only supported for BOOL dtype (and it uses ewise_add, the union)." ) return expr @@ -75,7 +84,7 @@ def __rxor__(self, other): if expr.dtype != BOOL: raise TypeError( f"The __xor__ infix operator, `x ^ y`, is not supported for {expr.dtype.name} dtype." - " It is only supported for BOOL dtype (and it uses ewise_add--the union)." + " It is only supported for BOOL dtype (and it uses ewise_add, the union)." ) return expr @@ -92,7 +101,7 @@ def __ixor__(self, other): elif other.dtype != BOOL: raise TypeError( f"The __ixor__ infix operator, `x ^= y`, is not supported for {other.dtype.name} " - "dtype. It is only supported for BOOL dtype (and it uses ewise_add--the union)." + "dtype. It is only supported for BOOL dtype (and it uses ewise_add, the union)." ) else: self(binary.lxor) << other @@ -111,13 +120,13 @@ def __ior__(self, other): if expr.dtype != BOOL: raise TypeError( f"The __ior__ infix operator, `x |= y`, is not supported for {expr.dtype.name} " - "dtype. It is only supported for BOOL dtype (and it uses ewise_add--the union)." + "dtype. It is only supported for BOOL dtype (and it uses ewise_add, the union)." ) self << expr elif other.dtype != BOOL: raise TypeError( f"The __ior__ infix operator, `x |= y`, is not supported for {other.dtype.name} " - "dtype. It is only supported for BOOL dtype (and it uses ewise_add--the union)." + "dtype. It is only supported for BOOL dtype (and it uses ewise_add, the union)." ) else: self(binary.lor) << other @@ -129,7 +138,7 @@ def __iand__(self, other): if expr.dtype != BOOL: raise TypeError( f"The __iand__ infix operator, `x &= y`, is not supported for {expr.dtype.name} dtype." - " It is only supported for BOOL dtype (and it uses ewise_mult--the intersection)." + " It is only supported for BOOL dtype (and it uses ewise_mult, the intersection)." ) self << expr return self diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py new file mode 100644 index 000000000..d43c97aa8 --- /dev/null +++ b/graphblas/core/literals.py @@ -0,0 +1,860 @@ +"""How python-graphblas types a literal operand, such as the ``1`` in ``v + 1``. + +A Python number is weak, as numpy 2 types it (NEP 50): it takes the dtype of +the operand beside it when that dtype's kind holds it, so ``int8_vec + 1`` is +INT8 and ``int8_vec * 0.5`` is FP64. A numpy value keeps its own dtype. Beside +a UDT, the same rules apply leaf by leaf. :func:`_literal_operand` decides for +each operand; a literal that must become an element of a UDT is checked by +:func:`_check_literal_fits`. + +The module also types a literal stored into an object +(:func:`_literal_store_dtype`) and a fill value +(:func:`_weak_builtin_literal_dtype` with no op). A literal that converts into +a UDT without fitting it is deprecated in :func:`_refuse_unless_exact` (MAINT). +""" + +import datetime +import math +import numbers +import struct +import warnings +from functools import lru_cache + +import numpy as np + +from .. import _STANDARD_OPERATOR_NAMES +from ..dtypes import FP64, lookup_dtype +from . import utils +from .dtypes import _nats_in_leaf_units, _view_if_same_layout +from .operator import IndexUnaryOp, ParameterizedUdf, SelectOp, TypedOpBase, binary_from_string +from .operator.binary import BinaryOp +from .operator.monoid import Monoid +from .udt import ( + BUILTIN_UDT_BINARY_OPS, + _casts_safely_into, + _dict_to_record, + _leaf_arrays, + _promoted_record, + _strip_leading_ones, +) + + +def _literal_operand(dtype, value, op, *, left=False): + """Return literal ``value`` and the dtype to make it, as ``op``'s operand beside ``dtype``. + + A dtype of ``None`` means the value's own dtype. This is where + python-graphblas decides how a literal operand is typed; beside a built-in + dtype, see :func:`_weak_builtin_literal_dtype`. + + A typed op (``op[T]``) takes each operand in its own input type, so the + literal is typed beside that instead of beside ``dtype``: the op's first + input type when the literal is its left operand (``left``), else its + second input or thunk type. GraphBLAS converts a value to a built-in input + type; a UDT input type it cannot, so the literal becomes an element of + that UDT. + + Beside a UDT, a literal becomes an element of that UDT, as it must for a + user's op, when its type allows it (:func:`_check_literal_fits`). With a + built-in op that lifts to UDTs, a literal is an operand in its own right, + typed as numpy 2 types it (NEP 50), by :func:`_literal_type`: + + - a numpy scalar or array keeps its dtype (strong); an array's type is the + array layout of its dtype (:func:`_array_literal_udt`); + - a Python ``bool``, ``int``, ``float`` or ``complex`` is weak: in each + element or leaf it takes that element's dtype when it fits the kind, so + ``int8_udt + 1`` stays ``int8_udt`` and ``int8_udt += 1`` can be stored, + while ``int8_udt * 2.5`` has float64 elements (see + :func:`_weak_literal_udt`), where converting it to the UDT would + multiply by 2; + - a tuple, list or dict of Python numbers, laid out like an element, is + weak in the same way, leaf by leaf: ``int8_udt + (1, 2, 3)`` stays + ``int8_udt`` and ``+ (0.5, 1.5, 2.5)`` has float64 elements. + + ``eq`` and ``ne`` type a literal the same way, as numpy compares, so + ``fp32_udt == 0.1`` compares in float32, as ``fp32_vec == 0.1`` does. + + A comparison keeps no result type, so a Python int beyond an integer + element's range is compared exactly, as numpy 2 compares it, as infinity + of its sign: every element compares with that infinity as with the int, + while a type that held both would round them (in FP64, the INT64 + ``2**63 - 1`` and the int ``2**63`` are equal). So ``int64_vec < 2**63`` is + all True and ``int8_vec == 300`` all False. Beside a UDT, where only ``eq`` + and ``ne`` compare, a literal with such an int in an integer leaf equals no + element: ``int8_udt == 300`` is False, not an error. + """ + if getattr(op, "is_positional", False): + return value, None + if isinstance(op, TypedOpBase): + dtype = op.type if left or op._type2 is None else op._type2 + if not dtype._is_udt: + literal_dtype = _weak_builtin_literal_dtype(dtype, value, op) + if ( + literal_dtype is not dtype + and literal_dtype is not None + and type(value) is int + and dtype.np_type.kind in "iu" + ): + # Only an int beyond dtype's range takes another type here, typed so + # by a comparison or by an op that gives another type (truediv). + if isinstance(op, str): + op = binary_from_string(op) + if op.name in _COMPARISONS and not op._anonymous: + return (math.inf if value > 0 else -math.inf), FP64 + return value, literal_dtype + if getattr(op, "_udt_types", True) is None: + # An op that does not take UDTs at all (``lt``) reports that itself. + return value, dtype + if isinstance(op, Monoid) and not op._anonymous: + # Element-wise, a Monoid is its BinaryOp (``get_typed_op`` types it so). + op = op.binaryop + if type(op) is BinaryOp and not op._anonymous: + is_comparison = op.name in ("eq", "ne") + if is_comparison or op.name in BUILTIN_UDT_BINARY_OPS: + # truediv gives float leaves whatever the int, so, as a comparison, it + # types an int out of an integer leaf's range exactly instead of raising. + exact = is_comparison or op.name == "truediv" + if is_comparison and not isinstance(value, (np.generic, np.ndarray)): + np_type = dtype.np_type + parts = value + if isinstance(value, dict) and np_type.names is not None: + parts = _dict_to_record(np_type, value) + # First: making an int beyond its leaf infinity reshapes the part. + _refuse_misfit_fields(dtype, np_type, parts) + if _int_beyond_leaf(np_type, parts) is not None: + # As infinity in its leaf, it still equals no element, and it + # stays out of an array numpy makes of numpy values beside it. + value = _beyond_as_infinity(np_type, parts) + literal_type = _literal_type(dtype, value, exact=exact) + if literal_type is not None: + return value, literal_type + _check_literal_fits(dtype, value) + return value, dtype + + +# Built-in ops that compare their operands, and so keep no result type. +_COMPARISONS = {"eq", "ne", "lt", "le", "gt", "ge"} +_COMPARISONS.update([f"value{name}" for name in list(_COMPARISONS)]) + + +# The types a Python float or complex can overflow. +_SINGLE_PRECISION = {np.dtype(np.float32), np.dtype(np.complex64)} + + +def _weak_builtin_literal_dtype(dtype, value, op=None): + """Return the dtype of Python number ``value`` beside built-in ``dtype``, or ``None``. + + A Python ``bool``, ``int``, ``float`` or ``complex`` is weak, as numpy 2 + types it (NEP 50), by the table UDT leaves use + (:func:`_weak_literal_leaf`): ``int8_vec + 1`` is INT8, + ``fp32_vec * 0.5`` is FP32, ``int8_vec * 0.5`` is FP64, and + ``int8_vec + 300`` raises ``OverflowError``. python-graphblas applies the + table itself, so numpy 1 gives the same answers. ``None`` keeps the value's + own dtype (strong): for anything but a Python number, for a typed op, which + fixes the operand types itself, and for the thunk of an index op that does + not compare it with the values, such as a row index. + + A comparison keeps no result type, so an int beyond ``dtype``'s range + takes a type that holds it rather than raise (:func:`_literal_operand` then + compares it as infinity). ``op`` of ``None``, for a fill value, does the + same, so the value is not lost. + """ + kind = _LITERAL_KINDS.get(type(value)) + if kind is None: + return None + if isinstance(op, str): + try: + op = binary_from_string(op) + except Exception: + return None + exact = op is None + if op is not None: + if isinstance(op, TypedOpBase): + return None + exact = getattr(op, "name", None) in _COMPARISONS and not op._anonymous + if isinstance(op, (IndexUnaryOp, SelectOp)) and not exact: + return None + try: + np_type = _weak_literal_leaf([value], kind, dtype.np_type, exact=exact) + except OverflowError: + if exact or not _gives_another_type(op, dtype): + raise + # An op whose result is not the operand's type has no narrow result to + # keep, so an int out of range is typed exactly, as for a comparison: + # int8_vec / 300 is FP64, as numpy divides integers in float64. + np_type = _weak_literal_leaf([value], kind, dtype.np_type, exact=True) + if np.result_type(dtype.np_type, np_type).kind not in "iu" and not _is_builtin_op(op): + # The two then unify to a float type, which rounds an int beyond 64 + # bits: in FP64, the INT64 2**63 - 1 and the int 2**63 are equal. A + # built-in op that gives another type computes in float anyway (a + # comparison is given infinity instead), but a user's op may compare. + raise + # Where the backend has no complex types (vanilla), lookup_dtype registers a + # complex literal's dtype as a UDT, so the literal keeps its own type. + rv = dtype if np_type == dtype.np_type else lookup_dtype(np_type) + if kind != "b" and rv.np_type in _SINGLE_PRECISION: + parts = (value.real, value.imag) if kind == "c" else (float(value),) + try: + for part in parts: + struct.pack(" 0 else -math.inf) if type(x) is int and not lo <= x <= hi else x + for x in values + ] + return np.array(values, dtype=object).reshape(shape).tolist() if shape else values[0] + + +def _array_literal_udt(dtype, value): + """Return the UDT a numpy array ``value`` is beside UDT ``dtype``, or ``None`` to convert it. + + An array has a dtype, so it is strong, as numpy 2 treats it: + ``int8_udt + np.array([0.5, 1, 2])`` has float64 elements and + ``+ np.array([1, 2, 3])`` int64 ones. Its type is the array layout of its + dtype: in the UDT's shape when it broadcasts to that (``np.array([1])`` + beside ``INT8[3]``), which a literal that must become an element needs, + else in its own shape when the two broadcast together, so that an op's + result takes the shape numpy's would (``INT8[3]`` and a 2-by-3 array give + ``[2, 3]``). A 0-d structured array or ``np.void`` beside a record is its + record dtype. Anything else, such as a numeric array beside a record, is + converted into the UDT. + """ + np_type = dtype.np_type + if value.dtype.names is not None: + return lookup_dtype(value.dtype) if value.ndim == 0 else None + # A bool array converts as it is (a BOOL array UDT cannot be registered), + # and bool casts safely into every numeric element. + if np_type.names is not None or value.dtype.kind not in "iufc": + return None + shape = np_type.subdtype[1] + if _strip_leading_ones(value.shape) == _strip_leading_ones(shape): + shape = value.shape + else: + try: + np.broadcast_to(value, shape) + except ValueError: + try: + np.broadcast_shapes(value.shape, shape) + except ValueError: + return None + shape = value.shape + return lookup_dtype(np.dtype((value.dtype, shape))) diff --git a/graphblas/core/mask.py b/graphblas/core/mask.py index 3bda2188a..f812e2119 100644 --- a/graphblas/core/mask.py +++ b/graphblas/core/mask.py @@ -18,6 +18,8 @@ def __init__(self, mask): def __eq__(self, other): raise TypeError(f"__eq__ not defined for objects of type {type(self)}.") + __hash__ = None + def __bool__(self): raise TypeError(f"__bool__ not defined for objects of type {type(self)}.") @@ -73,8 +75,6 @@ def new(self, dtype=None, *, complement=False, mask=None, name=None, **opts): val(self, **opts) << True return val - from .base import _check_mask - mask = _check_mask(mask) d = _COMPLEMENT_MASKS if complement else _COMBINE_MASKS func = d[type(self), type(mask)] @@ -93,8 +93,6 @@ def __and__(self, other, **opts): This uses faster recipes than the above for all combinations of input mask types, and aims to be memory efficient when operating on complemented masks. """ - from .base import _check_mask - other = _check_mask(other) complement = self.complement or other.complement d = _COMPLEMENT_MASKS if complement else _COMBINE_MASKS @@ -119,8 +117,6 @@ def __or__(self, other, **opts): This uses faster recipes than the above for all combinations of input mask types, and aims to be memory efficient when operating on complemented masks. """ - from .base import _check_mask - other = _check_mask(other) func = _MASK_OR[type(self), type(other)] return func(self, other, opts) @@ -200,6 +196,31 @@ def _name_html(self): return f"~{self.parent._name_html}.V" +def _check_mask(mask, output=None, strict_kind=False): + if not isinstance(mask, Mask): + # Convert bool objects to value masks + if utils.output_type(mask).__name__ in {"Vector", "Matrix"}: + if mask.dtype != BOOL: + raise TypeError( + f"Mask must be boolean objects (got {mask.dtype}) " + "or indicate values (M.V) or structure (M.S)" + ) + mask = mask.V # auto-compute (will raise if disabled) + else: + raise TypeError(f"Invalid mask: {type(mask)}") + if output is not None: + if output.ndim == 1 and mask.parent.ndim != 1: + raise TypeError(f"Mask object must be type Vector; got {type(mask.parent)}") + # A full-tensor op (ewise, mxm, apply, extract, ...) into a Matrix needs + # a Matrix mask. Assignment is exempt (`strict_kind` stays False for it): + # a Vector mask on a Matrix row/column assign is valid and is validated + # separately in Matrix.__setitem__. Without this, a Vector mask on a + # full-Matrix op leaked a raw cffi "struct GB_Matrix_opaque" error. + if strict_kind and output.ndim == 2 and mask.parent.ndim != 2: + raise TypeError(f"Mask object must be type Matrix; got {type(mask.parent)}") + return mask + + # Recipes to combine two masks. # Legend: # A: any diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index bf20cc953..299d98d8c 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -7,10 +7,12 @@ from ..dtypes import _INDEX, FP64, INT64, lookup_dtype, unify from ..exceptions import DimensionMismatch, InvalidValue, NoValue, check_status from . import _supports_udfs, automethods, ffi, lib, utils -from .base import BaseExpression, BaseType, _check_mask, call +from .base import BaseExpression, BaseType, call from .descriptor import lookup as descriptor_lookup +from .dtypes import _raise_dtype_or_arraylike from .expr import _ALL_INDICES, AmbiguousAssignOrExtract, IndexerResolver, InfixExprBase, Updater -from .mask import Mask, StructuralMask, ValueMask +from .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype +from .mask import Mask, StructuralMask, ValueMask, _check_mask from .operator import ( UNKNOWN_OPCLASS, _get_typed_op_from_exprs, @@ -19,6 +21,7 @@ get_typed_op, op_from_string, ) +from .operator.select import _value_select_in_own_type from .scalar import ( _COMPLETE, _MATERIALIZE, @@ -27,10 +30,20 @@ ScalarIndexExpr, _as_scalar, _scalar_index, + _union_default, +) +from .udt import ( + _cast_for_store, + _element_shapes_differ, + _ewise_add_cast_error, + _ewise_add_needs_cast, + _fused_store_expr, + _store_cast_op, ) from .utils import ( _CArray, _Pointer, + check_element_shape, class_property, get_order, ints_to_numpy_buffer, @@ -100,6 +113,15 @@ def _reposition(updater, indices, chunk): def _power(updater, A, n, op): opts = updater.opts + + def times(left, right): + # ``op`` is typed for A, and GraphBLAS casts a built-in operand of another + # type to it, but not a UDT: an int8 UDT squares to float64 under truediv. + # So a product of UDT operands of other types takes ``op`` typed for them. + if A.dtype._is_udt and (left.dtype != op.type or right.dtype != op.type2): + return get_typed_op(op.parent, left.dtype, right.dtype, kind="semiring")(left @ right) + return op(left @ right) + if n == 0: v = Vector.from_scalar(op.binaryop.monoid.identity, A._nrows, A.dtype, name="v_diag") updater << v.diag(name="M_diag") @@ -107,6 +129,14 @@ def _power(updater, A, n, op): if n == 1: updater << A return + if (out_dtype := updater.parent.dtype) is not A.dtype and ( + out_dtype._is_udt or A.dtype._is_udt + ): + # The last product is stored by mxm, which cannot cast a UDT as it writes; + # raise before computing the others. + expr = times(A, A) + if _store_cast_op(expr.dtype, out_dtype) is not None: + _fused_store_expr(expr, out_dtype) # Use repeated squaring: compute A^2, A^4, A^8, etc., and combine terms as needed. # See `numpy.linalg.matrix_power` for a simpler implementation to understand how this works. # We reuse `result` and `square` outputs, and use `square_expr` so masks can be applied. @@ -139,14 +169,14 @@ def _power(updater, A, n, op): result = square.dup(name="Power", **opts) elif n == 0: # All done! No more terms to compute - updater << op(result @ square) + updater << times(result, square) return elif result is A: # Now we need to create a new matrix for the intermediate result - result = op(result @ square).new(name="Power", **opts) + result = times(result, square).new(name="Power", **opts) else: # Main branch: multiply `square` into `result` - result(**opts) << op(result @ square) + result(**opts) << times(result, square) n, bit = divmod(n, 2) if square_expr is not None: # We need to perform another squaring, so evaluate current `square_expr` first @@ -158,10 +188,10 @@ def _power(updater, A, n, op): square << square_expr if square is None: # First iteration! Create expression for first square - square_expr = op(A @ A) + square_expr = times(A, A) else: # Expression for repeated squaring - square_expr = op(square @ square) + square_expr = times(square, square) class Matrix(BaseType): @@ -180,7 +210,7 @@ class Matrix(BaseType): """ - __slots__ = "_nrows", "_ncols", "_parent", "ss" + __slots__ = "_nrows", "_ncols", "_parent" ndim = 2 _is_transposed = False _name_counter = itertools.count() @@ -189,7 +219,10 @@ class Matrix(BaseType): def __new__(cls, dtype=FP64, nrows=0, ncols=0, *, name=None): self = object.__new__(cls) - self.dtype = lookup_dtype(dtype) + try: + self.dtype = lookup_dtype(dtype) + except (ValueError, TypeError) as exc: + _raise_dtype_or_arraylike("Matrix", dtype, exc) nrows = _as_scalar(nrows, _INDEX, is_cscalar=True) ncols = _as_scalar(ncols, _INDEX, is_cscalar=True) self.name = f"M_{next(Matrix._name_counter)}" if name is None else name @@ -198,8 +231,6 @@ def __new__(cls, dtype=FP64, nrows=0, ncols=0, *, name=None): self._nrows = nrows.value self._ncols = ncols.value self._parent = None - if backend == "suitesparse": - self.ss = ss(self) return self @classmethod @@ -211,8 +242,6 @@ def _from_obj(cls, gb_obj, dtype, nrows, ncols, *, parent=None, name=None): self._nrows = nrows self._ncols = ncols self._parent = parent - if backend == "suitesparse": - self.ss = ss(self) return self def __del__(self): @@ -228,13 +257,19 @@ def _as_vector(self, *, name=None): """Cast this Matrix with one column to a Vector. This is SuiteSparse-specific and may change in the future. - This does not copy the matrix. + This does not copy the matrix, but converts it to be held by column if it + is held by row. """ if self._ncols != 1: raise ValueError( f"Matrix must have a single column (not {self._ncols}) to be cast to a Vector" ) if backend == "suitesparse": + # SuiteSparse holds a GrB_Vector as a one-column matrix by column. Cast + # from a matrix held by row (a reshape of a by-row Matrix is one), the + # Vector has size 1 to SuiteSparse whatever its Python-side size. + if self.ss.config["format"] != "by_col": + self.ss.config["format"] = "by_col" return Vector._from_obj( ffi.cast("GrB_Vector*", self.gb_obj), self.dtype, @@ -368,7 +403,7 @@ def __sizeof__(self): size = ffi_new("size_t*") check_status(lib.GxB_Matrix_memoryUsage(size, self.gb_obj[0]), self) return size[0] + object.__sizeof__(self) - raise TypeError("Unable to get size of Matrix with backend: {backend}") + raise TypeError(f"Unable to get size of Matrix with backend: {backend}") def isequal(self, other, *, check_dtype=False, **opts): """Check for exact equality (same size, same structure). @@ -377,13 +412,18 @@ def isequal(self, other, *, check_dtype=False, **opts): ---------- other : Matrix The matrix to compare against - check_dtypes : bool + check_dtype : bool If True, also checks that dtypes match Returns ------- bool + Notes + ----- + Elements of array UDTs must have the same shape, apart from leading axes + of length 1, as in ``np.array_equal``; ``==`` broadcasts instead. + See Also -------- :meth:`isclose` : For equality check of floating point dtypes @@ -394,6 +434,8 @@ def isequal(self, other, *, check_dtype=False, **opts): ) if check_dtype and self.dtype != other.dtype: return False + if _element_shapes_differ(self.dtype, other.dtype): + return False if self._nrows != other._nrows: return False if self._ncols != other._ncols: @@ -417,7 +459,8 @@ def isequal(self, other, *, check_dtype=False, **opts): def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False, **opts): """Check for approximate equality (including same size and same structure). - Equivalent to: ``abs(a-b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)``. + Equivalent to: ``abs(a-b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)``, + except that, as in ``math.isclose``, an infinity is close only to itself. Parameters ---------- @@ -532,7 +575,7 @@ def to_coo(self, dtype=None, *, rows=True, columns=True, values=True, sort=True) Requested dtype for the output values array. rows : bool, default=True Whether to return rows; will return ``None`` for rows if ``False`` - columns :bool, default=True + columns : bool, default=True Whether to return columns; will return ``None`` for columns if ``False`` values : bool, default=True Whether to return values; will return ``None`` for values if ``False`` @@ -636,13 +679,14 @@ def build(self, rows, columns, values, *, dup_op=None, clear=False, nrows=None, # TODO: accept `dtype` keyword to match the dtype of `values`? rows = ints_to_numpy_buffer(rows, np.uint64, name="row indices") columns = ints_to_numpy_buffer(columns, np.uint64, name="column indices") - values, dtype = values_to_numpy_buffer(values, self.dtype) + values, _dtype = values_to_numpy_buffer(values, self.dtype) n = values.shape[0] if rows.size != n or columns.size != n: raise ValueError( "`rows` and `columns` and `values` lengths must match: " - f"{rows.size}, {columns.size}, {values.size}" + f"{rows.size}, {columns.size}, {n}" ) + check_element_shape(values, self.dtype, 1) if clear: self.clear() if nrows is not None or ncols is not None: @@ -707,7 +751,11 @@ def dup(self, dtype=None, *, clear=False, mask=None, name=None, **opts): dtype = self.dtype rv = Matrix(dtype, nrows=self._nrows, ncols=self._ncols, name=name) if not clear: - rv(mask=mask, **opts)[...] = self + if rv.dtype._is_udt or self.dtype._is_udt: + # ``<<``, not assign: a UDT of another type casts in one pass. + rv(mask=mask, **opts) << self + else: + rv(mask=mask, **opts)[...] = self else: if opts: # Ignore opts for now @@ -775,7 +823,11 @@ def wait(self, how="materialize"): the computations may be delayed and not yet safe to use by multiple threads. Use wait to force completion of the Matrix. - Has no effect in `blocking mode <../user_guide/init.html#graphblas-modes>`__. + In `blocking mode <../user_guide/init.html#graphblas-modes>`__, computations are + already complete, but with SuiteSparse:GraphBLAS, "materialize" may still build the + hyper_hash of a hypersparse Matrix that lacks one, such as one from ``from_coo`` or + an import (on recent versions, only with more non-empty rows or columns than the + ``GxB_HYPER_HASH`` setting, 1024 by default); see ``Matrix.ss.unpack_hyperhash``. """ how = how.lower() @@ -829,6 +881,12 @@ def from_coo( ): """Create a new Matrix from row and column indices and values. + .. warning:: + When ``nrows`` or ``ncols`` is omitted, the shape is inferred from + the largest row or column index, so trailing all-empty rows or + columns are dropped. Pass ``nrows`` and ``ncols`` explicitly to pin + the shape. + Parameters ---------- rows : list or np.ndarray @@ -886,7 +944,8 @@ def from_coo( if backend == "suitesparse": C.ss.build_scalar(rows, columns, values.tolist()) else: - C.build(rows, columns, np.broadcast_to(values, rows.size), dup_op=binary.any) + values = np.broadcast_to(values, (rows.size, *dtype.np_type.shape)) + C.build(rows, columns, values, dup_op=binary.any) else: # Add the data # This needs to be the original data to get proper error messages @@ -1019,6 +1078,7 @@ def _from_csx(cls, fmt, indptr, indices, values, dtype, num, check_num, name): f"expected {check_num}, got {ncols}" ) if values.ndim == 0: + values = np.broadcast_to(values, dtype.np_type.shape) # One whole element if backend == "suitesparse": # SuiteSparse GxB can handle iso-value if fmt is _CSR_FORMAT: @@ -1042,7 +1102,8 @@ def _from_csx(cls, fmt, indptr, indices, values, dtype, num, check_num, name): dtype=dtype, name=name, ) - values = np.broadcast_to(values, indices.size) + values = np.broadcast_to(values, (indices.size, *dtype.np_type.shape)) + check_element_shape(values, dtype, 1) new_mat = ffi_new("GrB_Matrix*") rv = cls._from_obj(new_mat, dtype, nrows, ncols, name=name) if dtype._is_udt: @@ -1131,15 +1192,15 @@ def from_csc( ---------- indptr : list or np.ndarray Pointers for each column into row_indices and values; ``indptr.size == ncols + 1``. - col_indices : list or np.ndarray - Column indices. + row_indices : list or np.ndarray + Row indices. values : list or np.ndarray or scalar, default 1.0 List of values. If a scalar is provided, all values will be set to this single value. dtype : Data type of the Matrix. If not provided, the values will be inspected to choose an appropriate dtype. nrows : int, optional - Number of rows in the Matrix. If not provided, ``ncols`` is computed + Number of rows in the Matrix. If not provided, ``nrows`` is computed from the maximum row index found in ``row_indices``. ncols : int, optional Number of columns in the Matrix. ncols is computed as ``len(indptr) - 1``. @@ -1223,11 +1284,16 @@ def from_dcsr( """ if backend == "suitesparse": + values, dtype = values_to_numpy_buffer(values, dtype, subarray_after=1) + is_iso = values.ndim == 0 + if is_iso: + values = np.broadcast_to(values, dtype.np_type.shape) # One whole element return cls.ss.import_hypercsr( rows=compressed_rows, indptr=indptr, col_indices=col_indices, values=values, + is_iso=is_iso, nrows=nrows, ncols=ncols, dtype=dtype, @@ -1308,11 +1374,16 @@ def from_dcsc( """ if backend == "suitesparse": + values, dtype = values_to_numpy_buffer(values, dtype, subarray_after=1) + is_iso = values.ndim == 0 + if is_iso: + values = np.broadcast_to(values, dtype.np_type.shape) # One whole element return cls.ss.import_hypercsc( cols=compressed_cols, indptr=indptr, row_indices=row_indices, values=values, + is_iso=is_iso, nrows=nrows, ncols=ncols, dtype=dtype, @@ -1435,6 +1506,7 @@ def from_dense(cls, values, missing_value=None, *, dtype=None, name=None, **opts raise ValueError("A >2d array is required to create a dense Matrix with subdtype") if values.ndim > 2 and dtype.np_type.subdtype is None: raise ValueError(f"values array must be 2d to create dense Matrix with dtype {dtype}") + check_element_shape(values, dtype, 2) if backend == "suitesparse": # Should we try to handle F-contiguous data w/o a copy? rv = cls.ss.import_fullr(values, dtype=dtype, name=name) @@ -1499,7 +1571,8 @@ def to_dense(self, fill_value=None, dtype=None, **opts): # dtype of fill_value can upcast the dtype if type(fill_value) is not Scalar: try: - fill_value = Scalar.from_value(fill_value, is_cscalar=None, name="") + dtype = _weak_builtin_literal_dtype(self.dtype, fill_value) + fill_value = Scalar.from_value(fill_value, dtype, is_cscalar=None, name="") except TypeError: fill_value = self._expect_type( fill_value, @@ -1529,7 +1602,7 @@ def from_dicts( Parameters ---------- - d : Mapping or Sequence + nested_dicts : Mapping or Sequence The dict-like object to convert. The keys will be cast to uint64 for the indices. dtype : Data type of the Matrix. If not provided, the values will be inspected @@ -1590,10 +1663,7 @@ def from_dicts( else: # If we know the dtype, then using `np.fromiter` is much faster dtype = lookup_dtype(dtype) - if dtype.np_type.subdtype is not None and np.__version__[:5] in {"1.21.", "1.22."}: - values, dtype = values_to_numpy_buffer(list(iter_values), dtype) # FLAKY COVERAGE - else: - values = np.fromiter(iter_values, dtype.np_type) + values = np.fromiter(iter_values, dtype.np_type) return getattr(cls, methodname)( *args, indptr, col_indices, values, dtype, nrows=nrows, ncols=ncols, name=name ) @@ -1826,18 +1896,17 @@ def to_dicts(self, order="rowwise"): # This is pretty fast, but can anybody make it faster? ;) cols = cols.tolist() values = values.tolist() + # Consecutive pairs of indptr, without sliding_window_view, which refuses + # the one-element indptr of a Matrix with no entries. + indptr = indptr.tolist() return { row: dict(zip(cols[start:stop], values[start:stop], strict=True)) - for row, (start, stop) in zip( - compressed_rows.tolist(), - np.lib.stride_tricks.sliding_window_view(indptr, 2).tolist(), - strict=True, + for row, start, stop in zip( + compressed_rows.tolist(), indptr[:-1], indptr[1:], strict=True ) } # Alternative - # slices = list( - # itertools.starmap(slice, np.lib.stride_tricks.sliding_window_view(indptr, 2).tolist()) - # ) + # slices = list(itertools.starmap(slice, zip(indptr[:-1], indptr[1:]))) # return { # r: dict(zip(c, v)) # for r, c, v in zip( @@ -1921,14 +1990,22 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): # Per the spec, op may be a semiring, but this is weird, so don't. self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") - if other.ndim == 1: + if other.ndim == 1 and self._ncols != other._size: # Broadcast rowwise from the right - if self._ncols != other._size: + raise DimensionMismatch( + "Dimensions not compatible for broadcasting Vector from the right " + f"to rows of Matrix in {method_name}. Matrix.ncols (={self._ncols}) " + f"must equal Vector.size (={other._size})." + ) + if _ewise_add_needs_cast(op, self.dtype, other.dtype): + # GraphBLAS would report the UDT cast before the shapes. + if other.ndim == 2 and self.shape != other.shape: raise DimensionMismatch( - "Dimensions not compatible for broadcasting Vector from the right " - f"to rows of Matrix in {method_name}. Matrix.ncols (={self._ncols}) " - f"must equal Vector.size (={other._size})." + f"Matrix shapes must match in {method_name}: {self.shape} and {other.shape}" ) + raise _ewise_add_cast_error(op, self.dtype, other.dtype) + + if other.ndim == 1: return MatrixExpression( method_name, None, @@ -2099,42 +2176,8 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): ) temp_op = get_typed_op(op, self.dtype, other.dtype, kind="binary") - left_dtype = temp_op.type - dtype = left_dtype if left_dtype._is_udt else None - if type(left_default) is not Scalar: - try: - left = Scalar.from_value( - left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - left = self._expect_type( - left_default, - Scalar, - within=method_name, - keyword_name="left_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar - right_dtype = temp_op.type2 - dtype = right_dtype if right_dtype._is_udt else None - if type(right_default) is not Scalar: - try: - right = Scalar.from_value( - right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - right = self._expect_type( - right_default, - Scalar, - within=method_name, - keyword_name="right_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar + left = _union_default(self, left_default, temp_op.type, op, "left_default") + right = _union_default(self, right_default, temp_op.type2, op, "right_default") if is_infix: op1 = _get_typed_op_from_exprs(op, self, right, kind="binary") @@ -2412,6 +2455,7 @@ def apply(self, op, right=None, *, left=None): """ method_name = "apply" + right_literal = None extra_message = ( "apply only accepts UnaryOp with no scalars or BinaryOp with `left` or `right` scalar " "or IndexUnaryOp with `right` thunk." @@ -2419,6 +2463,11 @@ def apply(self, op, right=None, *, left=None): if isinstance(op, str): op = op_from_string(op) op, opclass = find_opclass(op) + if opclass == "Monoid" and self.dtype._is_udt: + # Applying a Monoid applies its BinaryOp. Typing it as that op lets a + # literal promote as it would with the BinaryOp, where a Monoid takes + # one type and would need the literal converted into the UDT. + op, opclass = op.binaryop, "BinaryOp" if opclass in {"IndexUnaryOp", "SelectOp"}: # Provide default value for index unary if right is None: @@ -2440,7 +2489,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = None elif right is None: if type(left) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + left, dtype = _literal_operand(self.dtype, left, op, left=True) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: @@ -2476,7 +2525,8 @@ def apply(self, op, right=None, *, left=None): expr_repr = "{1.name}.apply({op}, left={0._expr_name})" elif left is None: if type(right) is not Scalar: - dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None + right_literal = right + right, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -2519,7 +2569,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = "{0.name}.apply({op}, right={1._expr_name})" else: raise TypeError("Cannot provide both `left` and `right` to apply") - return MatrixExpression( + expr = MatrixExpression( method_name, cfunc_name, args, @@ -2530,6 +2580,8 @@ def apply(self, op, right=None, *, left=None): at=self._is_transposed, bt=self._is_transposed, ) + expr._right_literal = right_literal + return expr def select(self, op, thunk=None): """Create a new Matrix by applying ``op`` to each element of the Matrix @@ -2596,8 +2648,10 @@ def select(self, op, thunk=None): if thunk is None: thunk = False # most basic form of 0 when unifying dtypes + given = thunk if type(thunk) is not Scalar: - dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None + # A thunk beside a UDT becomes an element of it, so it must fit, as in apply. + thunk, dtype = _literal_operand(self.dtype, thunk, op) try: thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: @@ -2610,6 +2664,7 @@ def select(self, op, thunk=None): op=op, ) op = get_typed_op(op, self.dtype, thunk.dtype, is_right_scalar=True, kind="select") + op, thunk = _value_select_in_own_type(op, self.dtype, thunk, given) self._expect_op(op, ("SelectOp", "IndexUnaryOp"), within=method_name, argname="op") if thunk._is_cscalar: if thunk.dtype._is_udt: @@ -2949,7 +3004,7 @@ def setdiag(self, values, k=0, *, mask=None, accum=None, **opts): elif type(values) is Scalar: is_scalar = True else: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, values) try: # Try to make it a Scalar values = Scalar.from_value(values, dtype, is_cscalar=None, name="") @@ -2974,7 +3029,7 @@ def setdiag(self, values, k=0, *, mask=None, accum=None, **opts): elif v._size != size: raise DimensionMismatch( f"Dimensions not compatible for assigning length {v._size} Vector " - f"to {k}'th diagonal of Matrix with shape {self._nrows}x{self._ncols}." + f"to {k}'th diagonal of Matrix with shape {self._nrows}x{self._ncols}. " f"The Vector should be size {size}." ) @@ -3026,6 +3081,13 @@ def _extract_element( dtype = self.dtype else: dtype = lookup_dtype(dtype) + if dtype is not self.dtype and (dtype._is_udt or self.dtype._is_udt): + # See Vector._extract_element: extract in this type, then cast. + value = self._extract_element(resolved_indexes, None, opts, is_cscalar=is_cscalar) + if result is None: + result = Scalar(dtype, is_cscalar=is_cscalar, name=name) + result << value + return result rowidx, colidx = resolved_indexes.indices if self._is_transposed: rowidx, colidx = colidx, rowidx @@ -3088,7 +3150,7 @@ def _prep_for_extract(self, resolved_indexes): def _assign_element(self, resolved_indexes, value): rowidx, colidx = resolved_indexes.indices if type(value) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, value) try: value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except TypeError: @@ -3099,6 +3161,10 @@ def _assign_element(self, resolved_indexes, value): argname="value", extra_message="Literal scalars also accepted.", ) + if value.dtype._is_udt or self.dtype._is_udt: + # After the conversion: an autocomputed element of another UDT type + # arrives here as a Scalar too. + value = _cast_for_store(value, self.dtype) if value._is_cscalar: if value._empty: call("GrB_Matrix_removeElement", [self, rowidx.index, colidx.index]) @@ -3115,6 +3181,8 @@ def _assign_element(self, resolved_indexes, value): def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, opts): method_name = "__setitem__" + if output_type(value) in {Scalar, Vector, Matrix}: + value = _cast_for_store(value, self.dtype) rowidx, colidx = resolved_indexes.indices rowsize = rowidx.size rows = rowidx.index @@ -3282,6 +3350,13 @@ def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, o within=method_name, extra_message=extra_message, ) + if mask is not None and mask.parent.ndim != 2: + # Matrix value, Vector mask, Matrix index + # C(m)[I, J] << A + # C[I, J](m) << A + # This also catches whole-object updates such as `C(m) << A` + # and `C.dup(mask=m)`, which assign with `C(m)[...] = A`. + raise TypeError("Unable to use Vector mask on Matrix assignment to a Matrix") if is_submask: # C[I, J](M) << A expr_repr = ( @@ -3311,7 +3386,7 @@ def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, o ) else: if type(value) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, value) try: value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except (TypeError, ValueError): @@ -3532,10 +3607,20 @@ def _delete_element(self, resolved_indexes): if backend == "suitesparse": - Matrix.ss = class_property(Matrix.ss, ss) + # `.ss` is built lazily per access rather than stored on the instance: a + # stored `ss(self)` holds `_parent` back to this Matrix (and so does its + # config), forming a reference cycle that keeps the object (and its C-side + # GrB buffer) alive until the cyclic gc runs instead of dying by refcount. + # See gh-559. Class-level `Matrix.ss` still resolves to the `ss` class so + # the `import_*` classmethods keep working. + def _ss(self): + return ss(self) + + _ss.__name__ = _ss.__qualname__ = "ss" + Matrix.ss = class_property(property(_ss), ss) else: Matrix.ss = class_property( - Matrix.ss, 'ss attribute is only available with "suitesparse" backend', exceptional=True + property(), 'ss attribute is only available with "suitesparse" backend', exceptional=True ) @@ -3827,6 +3912,7 @@ class TransposedMatrix: ndim = 2 _is_scalar = False _is_transposed = True + __array_priority__ = BaseType.__array_priority__ # numpy's operators defer to ours __networkx_backend__ = "graphblas" __networkx_plugin__ = "graphblas" @@ -3887,6 +3973,11 @@ def diag(self, k=0, dtype=None, *, name=None, **opts): def _carg(self): return self._matrix.gb_obj[0] + @property + def _exc_arg(self): + # check_status reports an error against the Matrix: GraphBLAS has no transposed type + return self._matrix + @property def name(self): return f"{self._matrix.name}.T" @@ -3979,6 +4070,7 @@ def to_dicts(self, order="rowwise"): _extract_element = Matrix._extract_element _prep_for_extract = Matrix._prep_for_extract __eq__ = Matrix.__eq__ + __hash__ = None __bool__ = Matrix.__bool__ __getitem__ = Matrix.__getitem__ __contains__ = Matrix.__contains__ diff --git a/graphblas/core/operator/__init__.py b/graphblas/core/operator/__init__.py index d59c835b3..f40de919b 100644 --- a/graphblas/core/operator/__init__.py +++ b/graphblas/core/operator/__init__.py @@ -1,5 +1,6 @@ from .base import UNKNOWN_OPCLASS, OpBase, OpPath, ParameterizedUdf, TypedOpBase, find_opclass from .binary import BinaryOp, ParameterizedBinaryOp +from .indexbinary import IndexBinaryOp, ParameterizedIndexBinaryOp from .indexunary import IndexUnaryOp, ParameterizedIndexUnaryOp from .monoid import Monoid, ParameterizedMonoid from .select import ParameterizedSelectOp, SelectOp diff --git a/graphblas/core/operator/agg.py b/graphblas/core/operator/agg.py index 6b463a8a6..574f4ec1b 100644 --- a/graphblas/core/operator/agg.py +++ b/graphblas/core/operator/agg.py @@ -5,6 +5,7 @@ from ... import agg, backend, binary, monoid, semiring, unary from ...dtypes import INT64, lookup_dtype +from ...exceptions import UdfParseError from .. import _supports_udfs from ..utils import output_type @@ -28,6 +29,27 @@ def _get_types(ops, initdtype): class Aggregator: + """A reduction operator that collapses the values of a Matrix or Vector. + + An Aggregator is used with the ``reduce`` family of methods: ``Vector.reduce``, + ``Matrix.reduce_rowwise``, ``Matrix.reduce_columnwise``, and ``Matrix.reduce_scalar``. + Some aggregators return a summary value (``sum``, ``mean``, ``max``); others return a + position (``ss.argmin``, ``ss.argmax``, ``ss.first_index``). + + Built-in aggregators live in the ``graphblas.agg`` namespace, such as ``agg.sum``, + ``agg.mean``, and ``agg.count``. The position aggregators are SuiteSparse-specific and + live under ``agg.ss``; the bare ``agg.argmin`` spellings are deprecated. Unlike the + other operators, an Aggregator is not a single GraphBLAS object; many are built from a + monoid or semiring plus an optional finalize step, composed per dtype on first use. + + Examples + -------- + >>> import graphblas as gb + >>> v = gb.Vector.from_coo([0, 1, 2], [1, 2, 3]) + >>> int(v.reduce(gb.agg.sum).new()) + 6 + """ + opclass = "Aggregator" def __init__( @@ -88,14 +110,48 @@ def types(self): def __getitem__(self, dtype): dtype = lookup_dtype(dtype) if not self._any_dtype and dtype not in self.types: - raise KeyError(f"{self.name} does not work with {dtype}") + # When this is a monoid-based aggregator and ``dtype`` is a UDT, + # try to compile the underlying monoid for it. Built-in monoids + # like ``plus``, ``min``, and ``max`` auto-generate field-by-field + # UDT versions, which lets ``agg.sum[udt]``, ``agg.min[udt]``, + # and friends work transparently. Composite or semiring-based + # aggregators (e.g., ``hypot``) aren't handled here; their + # pipelines don't lift to UDTs cleanly. + if ( + dtype._is_udt + and self._monoid is not None + and self._semiring is None + and self._composite is None + and self._custom is None + and self._finalize is None + and self._applybegin is None + ): + # ``_compile_udt`` can raise ``KeyError`` (no field-wise + # identity for this op/dtype) or ``UdfParseError`` (wrapped + # Numba compile failure). Treat both as "does not work + # with" rather than leaking the underlying error. + try: + typed_monoid = self._monoid[dtype] + except (KeyError, TypeError, UdfParseError): + pass + else: + self.types[dtype] = typed_monoid.return_type + if dtype not in self.types: + raise KeyError(f"{self.name} does not work with {dtype}") if dtype not in self._typed_ops: self._typed_ops[dtype] = TypedAggregator(self, dtype) return self._typed_ops[dtype] def __contains__(self, dtype): - dtype = lookup_dtype(dtype) - return self._any_dtype or dtype in self.types + # Routes through ``__getitem__`` (which can lift the underlying + # monoid for a UDT) so ``dtype in agg.sum`` reports True for any + # UDT the monoid can auto-compile. Catches the same exception set + # ``__getitem__`` raises. + try: + self[dtype] + except (KeyError, TypeError, UdfParseError): + return False + return True def __repr__(self): if self.name in agg._deprecated: @@ -103,6 +159,8 @@ def __repr__(self): return f"agg.{self.name}" def __reduce__(self): + # Aggregator has no register_anonymous/register_new, so every + # instance is a module-level singleton; pickle by name. if self.name in agg._deprecated: return f"agg.ss.{self.name}" return f"agg.{self.name}" @@ -153,6 +211,26 @@ def __init__(self, agg, dtype): def __repr__(self): return f"agg.{self.name}[{self.type}]" + @property + def jit_c_source(self): + """C source of the underlying monoid kernel for monoid-based aggs, else ``None``. + + Composite and custom aggregators (``agg.hypot``, ``agg.ss.first``, + etc.) don't have a single underlying kernel, so they return ``None``. + """ + mon = getattr(self.parent, "_monoid", None) + if mon is None or self.type not in mon: + return None + return mon[self.type].jit_c_source + + @property + def jit_c_name(self): + """C name of the underlying monoid kernel for monoid-based aggs, else ``None``.""" + mon = getattr(self.parent, "_monoid", None) + if mon is None or self.type not in mon: + return None + return mon[self.type].jit_c_name + def _new(self, updater, expr, *, in_composite=False): agg = self.parent opts = updater.opts diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 97b2c9fbd..b9cd34e52 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -3,13 +3,13 @@ from types import BuiltinFunctionType, ModuleType from ... import _STANDARD_OPERATOR_NAMES, backend, op -from ...dtypes import BOOL, INT8, UINT64, _supports_complex, lookup_dtype -from .. import _has_numba, _supports_udfs, lib +from ...dtypes import BOOL, UINT64, _supports_complex, lookup_dtype +from ...exceptions import UdfParseError +from .. import _has_numba, _supports_udfs, ffi, lib from ..expr import InfixExprBase from ..utils import output_type if _has_numba: - import numba from numba import NumbaError else: NumbaError = TypeError @@ -111,9 +111,7 @@ def _call_op(op, left, right=None, thunk=None, **kwargs): if right is None and thunk is None: if isinstance(left, InfixExprBase): # op(A & B), op(A | B), op(A @ B) - return getattr(left.left, f"_{left.method_name}")( - left.right, op, is_infix=True, **kwargs - ) + return left._call_left(f"_{left.method_name}", op, is_infix=True, **kwargs) if find_opclass(op)[1] == "Semiring": raise TypeError( f"Bad type when calling {op!r}. Got type: {type(left)}.\n" @@ -163,69 +161,22 @@ def _call_op(op, left, right=None, thunk=None, **kwargs): ) -if _has_numba: - - def _get_udt_wrapper(numba_func, return_type, dtype, dtype2=None, *, include_indexes=False): - ztype = INT8 if return_type == BOOL else return_type - xtype = INT8 if dtype == BOOL else dtype - nt = numba.types - wrapper_args = [nt.CPointer(ztype.numba_type), nt.CPointer(xtype.numba_type)] - if include_indexes: - wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) - if dtype2 is not None: - ytype = INT8 if dtype2 == BOOL else dtype2 - wrapper_args.append(nt.CPointer(ytype.numba_type)) - wrapper_sig = nt.void(*wrapper_args) - - zarray = xarray = yarray = BL = BR = yarg = yname = rcidx = "" - if return_type._is_udt: - if return_type.np_type.subdtype is None: - zarray = " z = numba.carray(z_ptr, 1)\n" - zname = "z[0]" - else: - zname = "z_ptr[0]" - BR = "[0]" - else: - zname = "z_ptr[0]" - if return_type == BOOL: - BL = "bool(" - BR = ")" - - if dtype._is_udt: - if dtype.np_type.subdtype is None: - xarray = " x = numba.carray(x_ptr, 1)\n" - xname = "x[0]" - else: - xname = "x_ptr" - elif dtype == BOOL: - xname = "bool(x_ptr[0])" - else: - xname = "x_ptr[0]" - - if dtype2 is not None: - yarg = ", y_ptr" - if dtype2._is_udt: - if dtype2.np_type.subdtype is None: - yarray = " y = numba.carray(y_ptr, 1)\n" - yname = ", y[0]" - else: - yname = ", y_ptr" - elif dtype2 == BOOL: - yname = ", bool(y_ptr[0])" - else: - yname = ", y_ptr[0]" - - if include_indexes: - rcidx = ", row, col" - - d = {"numba": numba, "numba_func": numba_func} - text = ( - f"def wrapper(z_ptr, x_ptr{rcidx}{yarg}):\n" - f"{zarray}{xarray}{yarray}" - f" {zname} = {BL}numba_func({xname}{rcidx}{yname}){BR}\n" - ) - exec(text, d) # pylint: disable=exec-used - return d["wrapper"], wrapper_sig +# Maps ``opclass`` to the SuiteSparse C type name SS allocates for it. +# ``IndexBinaryOp`` is in the GxB_ namespace (SS-specific, added in 9.4); +# the rest are GrB_. ``SelectOp`` is implemented on top of +# ``GrB_IndexUnaryOp`` (BOOL-returning), so its handle frees through the +# same ``GrB_IndexUnaryOp_free``. Used by ``TypedOpBase.__del__`` to +# synthesize the pointer cell for ``_free``, and by +# ``udf._finalize_udt_op`` to call ``_new``. +_GB_OBJ_C_TYPENAME = { + "UnaryOp": "GrB_UnaryOp", + "BinaryOp": "GrB_BinaryOp", + "IndexUnaryOp": "GrB_IndexUnaryOp", + "SelectOp": "GrB_IndexUnaryOp", + "IndexBinaryOp": "GxB_IndexBinaryOp", + "Monoid": "GrB_Monoid", + "Semiring": "GrB_Semiring", +} class TypedOpBase: @@ -237,8 +188,15 @@ class TypedOpBase: "gb_obj", "gb_name", "_type2", + "_jit_c_info", + "_gb_obj_owner", "__weakref__", ) + # Subclasses whose ``gb_obj`` was allocated via ``GrB__new`` / + # ``GxB__new`` (TypedUser*Op, _BoundIndexBinaryOp) override this so + # ``__del__`` frees the SuiteSparse handle. Built-in typed ops point at + # SuiteSparse's permanent built-in singletons and must never free. + _owns_gb_obj = False def __init__(self, parent, name, type_, return_type, gb_obj, gb_name, dtype2=None): self.parent = parent @@ -248,6 +206,26 @@ def __init__(self, parent, name, type_, return_type, gb_obj, gb_name, dtype2=Non self.gb_obj = gb_obj self.gb_name = gb_name self._type2 = dtype2 + # ``(c_name, c_definition)`` when SuiteSparse JIT-compiled a kernel + # for this typed op; ``None`` for built-in ops and for UDT ops with + # no JIT path. + self._jit_c_info = None + # Set when ``gb_obj`` is borrowed from another typed op rather than + # allocated here (``SelectOp._from_indexunary``). Storing the owner + # both suppresses our free and keeps the handle alive as long as we + # point at it; disclaiming ownership without naming an owner would + # leave this op holding a dangling handle. + self._gb_obj_owner = None + + @property + def jit_c_name(self): + """The C symbol name SuiteSparse uses for this op's JIT kernel, or ``None``.""" + return self._jit_c_info[0] if self._jit_c_info is not None else None + + @property + def jit_c_source(self): + """C source SuiteSparse JIT-compiles for this op, or ``None`` when no JIT kernel exists.""" + return self._jit_c_info[1] if self._jit_c_info is not None else None def __repr__(self): classname = self.opclass.lower() @@ -268,6 +246,55 @@ def __reduce__(self): return (getitem, (self.parent, self.type)) return (getitem, (self.parent, (self.type, self._type2))) + def __del__(self): + # Free the SuiteSparse handle we allocated. Built-in typed ops alias + # SuiteSparse's permanent built-in singletons and must never free. + # Mirrors the ``Matrix.__del__`` / ``Vector.__del__`` pattern. + if not type(self)._owns_gb_obj: + return + # A borrowed handle belongs to ``_gb_obj_owner``, which we keep alive. + # ``getattr`` guards the case where ``__init__`` raised before the slot + # was set. + if getattr(self, "_gb_obj_owner", None) is not None: + return + gb_obj = getattr(self, "gb_obj", None) + if gb_obj is None or lib is None or ffi is None: + # Interpreter shutdown can clear ``lib`` / ``ffi`` before + # finalizers run; SS will clean up the handles at process exit. + return + c_type_name = _GB_OBJ_C_TYPENAME.get(self.opclass) + if c_type_name is None: # pragma: no cover (defensive) + return + free_fn = getattr(lib, f"{c_type_name}_free", None) + if free_fn is None: + # ``GxB_IndexBinaryOp_free`` is absent on SS < 9.4; that build + # also can't allocate one in the first place, so this path is + # unreachable in practice but guarded for safety. + return + # ``GrB__free`` takes a pointer-to-pointer (sets ``*p = NULL`` + # after free). Synthesize a cell pointing at our handle and call it. + free_fn(ffi.new(f"{c_type_name}*", gb_obj)) + + +class _BinaryopJitDelegate: + """Mixin for ops that don't own a JIT kernel; defer introspection to ``binaryop``. + + Monoids and semirings reuse the binary op's JIT kernel. The setter + accepts ``None`` so ``TypedOpBase.__init__``'s slot init is a no-op. + """ + + __slots__ = () + + @property + def _jit_c_info(self): + return self.binaryop._jit_c_info + + @_jit_c_info.setter + def _jit_c_info(self, value): + # No-op so ``TypedOpBase.__init__``'s ``self._jit_c_info = None`` + # slot-init succeeds. The kernel lives on ``binaryop``. + pass + def _deserialize_parameterized(parameterized_op, args, kwargs): return parameterized_op(*args, **kwargs) @@ -277,6 +304,10 @@ class ParameterizedUdf: __slots__ = "name", "__call__", "_anonymous", "__weakref__" is_positional = False _custom_dtype = None + # Subclasses set this to the OpBase subclass they parameterize (e.g., + # ``ParameterizedUnaryOp._op_class = UnaryOp``). Assigned after the + # OpBase subclass is defined to avoid an import-order cycle. + _op_class = None def __init__(self, name, anonymous): self.name = name @@ -288,6 +319,30 @@ def __init__(self, name, anonymous): def _call(self, *args, **kwargs): raise NotImplementedError + def __reduce__(self): + # The namespace prefix (``unary``, ``binary``, ...) comes from the + # OpBase subclass each parameterized op wraps. Standard ops pickle by + # name; user-registered ones pickle the reduce tuple and re-register + # on load via ``_deserialize`` (which dispatches through ``_op_class``). + name = f"{self._op_class._modname}.{self.name}" + if not self._anonymous and name in _STANDARD_OPERATOR_NAMES: + return name + return (self._deserialize, (self.name, self.func, self._anonymous, self._is_udt)) + + @classmethod + def _deserialize(cls, name, func, anonymous, is_udt=False): + """Re-register a parameterized UDF on unpickle, or reuse if already present. + + Shared by the five ``Parameterized*Op`` subclasses; each sets + ``_op_class`` to the matching OpBase subclass for the dispatch below. + """ + op_cls = cls._op_class + if anonymous: + return op_cls.register_anonymous(func, name, parameterized=True, is_udt=is_udt) + if (rv := op_cls._find(name)) is not None: + return rv + return op_cls.register_new(name, func, parameterized=True, is_udt=is_udt) + _VARNAMES = tuple(x for x in dir(lib) if x[0] != "_") @@ -337,7 +392,7 @@ def __getitem__(self, type_): raise KeyError(f"{self.name} does not work with {type_}") else: return self._typed_ops[type_] - # This is a UDT or is able to operate on UDTs such as `first` any `any` + # This is a UDT or is able to operate on UDTs such as `first` and `any` dtype = lookup_dtype(type_) return self._compile_udt(dtype, dtype) @@ -362,7 +417,7 @@ def __delitem__(self, type_): def __contains__(self, type_): try: self[type_] - except (TypeError, KeyError, NumbaError): + except (TypeError, KeyError, NumbaError, UdfParseError): return False return True @@ -407,7 +462,7 @@ def _initialize(cls, include_in_ops=True): """Initialize operators for this operator type. include_in_ops determines whether the operators are included in the - ``gb.ops`` namespace in addition to the defined module. + ``gb.op`` namespace in addition to the defined module. """ if cls._initialized: # pragma: no cover (safety) return @@ -500,6 +555,39 @@ def _deserialize(cls, name, *args): return rv # Should we verify this is what the user expects? return cls.register_new(name, *args) + @classmethod + def _deserialize_udf(cls, name, orig_func, is_udt): + """Re-register a named UDF on unpickle, or reuse if already present. + + Shared by the five UDF-capable subclasses (UnaryOp, BinaryOp, + IndexUnaryOp, SelectOp, IndexBinaryOp), all of which use the + default ``__reduce__`` below. + """ + if (rv := cls._find(name)) is not None: + return rv + return cls.register_new(name, orig_func, is_udt=is_udt) + + @classmethod + def _deserialize_anon_udf(cls, func, name, is_udt): + """Re-register an anonymous UDF on unpickle.""" + return cls.register_anonymous(func, name, is_udt=is_udt) + + def __reduce__(self): + """Default ``__reduce__`` for UDF-capable subclasses. + + Assumes the instance has ``orig_func`` and ``_is_udt`` attributes (all + five UDF-capable subclasses do). ``Monoid``, ``Semiring``, and + ``Aggregator`` define their own ``__reduce__`` because their pickle + shape differs (they hold a binary op + identity, etc.). + """ + if self._anonymous: + if hasattr(self.orig_func, "_parameterized_info"): + return (_deserialize_parameterized, self.orig_func._parameterized_info) + return (type(self)._deserialize_anon_udf, (self.orig_func, self.name, self._is_udt)) + if (name := f"{self._modname}.{self.name}") in _STANDARD_OPERATOR_NAMES: + return name + return (type(self)._deserialize_udf, (self.name, self.orig_func, self._is_udt)) + @classmethod def _check_supports_udf(cls, method_name): if not _supports_udfs: @@ -513,7 +601,8 @@ def _check_supports_udf(cls, method_name): def find_opclass(gb_op): if isinstance(gb_op, OpBase): - opclass = type(gb_op).__name__ + # A subclass may name the class it is (``ibo(theta)`` is a BinaryOp). + opclass = getattr(gb_op, "opclass", type(gb_op).__name__) elif isinstance(gb_op, TypedOpBase): opclass = gb_op.opclass elif isinstance(gb_op, ParameterizedUdf): diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 3ee089fe4..6abccb39e 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -26,52 +26,161 @@ from .. import _has_numba, _supports_udfs, ffi, lib from ..dtypes import _sample_values from ..expr import InfixExprBase + +# Imported unconditionally (plain dict, no numba): ``_compile_udt`` consults it +# even when numba is absent. +from ..udt import BUILTIN_UDT_BINARY_OPS as _BUILTIN_UDT_BINARY_OPS from .base import ( _SS_OPERATORS, OpBase, ParameterizedUdf, TypedOpBase, _call_op, - _deserialize_parameterized, _hasop, ) if _has_numba: import numba - from .base import _get_udt_wrapper + from .udf import ( + _bool_to_int8, + _compile_udf_for_udt, + _finalize_udt_op, + _get_udt_wrapper, + _resolve_udt_return_type, + ) if _supports_complex: from ...dtypes import FC32, FC64 ffi_new = ffi.new if _has_numba: - _udt_mask_cache = {} - - def _udt_mask(dtype): - """Create mask to determine which bytes of UDTs to use for equality check.""" - if dtype in _udt_mask_cache: - return _udt_mask_cache[dtype] - if dtype.subdtype is not None: - mask = _udt_mask(dtype.subdtype[0]) - N = reduce(mul, dtype.subdtype[1]) - rv = np.concatenate([mask] * N) - elif dtype.names is not None: - prev_offset = mask = None - masks = [] - for name in dtype.names: - dtype2, offset = dtype.fields[name] - if mask is not None: - masks.append(np.pad(mask, (0, offset - prev_offset - mask.size))) - mask = _udt_mask(dtype2) - prev_offset = offset - masks.append(np.pad(mask, (0, dtype.itemsize - prev_offset - mask.size))) - rv = np.concatenate(masks) - else: - rv = np.ones(dtype.itemsize, dtype=bool) - # assert rv.size == dtype.itemsize - _udt_mask_cache[dtype] = rv - return rv + + def _make_udt_comparison(dtype, dtype2, *, is_eq): + """Build a cfunc-ready wrapper for UDT eq or ne comparison. + + Compares per-leaf with IEEE semantics (NaN compares unequal even + to itself); scalar broadcasts to every leaf when paired with a + UDT. Mismatched-shape UDT pairs raise ``KeyError``. A UDT that is + neither a record nor an array (bytes, say) compares only with its + own type, as one value. Returns ``(wrapper_func, wrapper_sig)``. + """ + from ..udt import ( + _check_leaves_compare, + _check_udt_pair, + _get_udt_info, + _iter_record_leaves, + ) + from .udt_codegen import _array_operand_source + + info_x = _get_udt_info(dtype) + info_y = _get_udt_info(dtype2) + op_name = "eq" if is_eq else "ne" + _check_udt_pair(op_name, dtype, dtype2, info_x, info_y) + _check_leaves_compare(op_name, dtype, dtype2) + + nt = numba.types + cmp_op = "==" if is_eq else "!=" + join = " and " if is_eq else " or " + + if info_x is None and info_y is None: + # Bytes, a datetime, or a complex number where the backend has no + # complex types: an element is one value, compared as numpy compares + # it, as a record field of that type is. ``_check_udt_pair`` has made + # sure both are the one UDT, text of one kind, or dates or durations. + wrapper = _compile_codegen( + f"def wrapper(z_ptr, x_ptr, y_ptr):\n z_ptr[0] = x_ptr[0] {cmp_op} y_ptr[0]\n", + func_name="wrapper", + source_label=f"", + ) + x_ptr_type = nt.CPointer(dtype.numba_type) + y_ptr_type = nt.CPointer(dtype2.numba_type) + return wrapper, nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type) + + x_is_scalar = info_x is None + y_is_scalar = info_y is None + udt_dtype = dtype2 if x_is_scalar else dtype + udt_info = info_y if x_is_scalar else info_x + + if udt_info[0] == "array": + # Array UDT: a scalar side is read through its pointer, with no + # ``numba.carray``. The wrapper sees an array side as a flat + # pointer-to-element, not pointer-to-Record; passing the Record + # numba_type here makes ``numba.carray`` try to read N + # record-sized chunks. Two array UDTs may differ in base dtype, + # so each side gets its own; typing both as the left one would + # reinterpret the bytes. Their shapes broadcast, as for the + # arithmetic ops, so eq is True when every pair of elements numpy + # would compare is equal. + shape_x = None if x_is_scalar else dtype.np_type.subdtype[1] + shape_y = None if y_is_scalar else dtype2.np_type.subdtype[1] + shape = np.broadcast_shapes(*(s for s in (shape_x, shape_y) if s is not None)) + N = reduce(mul, shape) + x_ptr_type = nt.CPointer(numba.from_dtype(dtype.np_type.base)) + y_ptr_type = nt.CPointer(numba.from_dtype(dtype2.np_type.base)) + wrapper_sig = nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type) + ns = {} + x_setup, x_ref = _array_operand_source("x", shape_x, shape, ns) + y_setup, y_ref = _array_operand_source("y", shape_y, shape, ns) + # A loop, not a chain of N terms, whose typing time grows faster + # than linearly with N (see ``_make_array_wrapper``). It reads every + # element rather than stopping at the first that decides, which + # lets the compiler vectorize it: faster on the equal values that + # need every element read anyway. + accumulate = "&=" if is_eq else "|=" + src = ( + f"def wrapper(z_ptr, x_ptr, y_ptr):\n{x_setup}{y_setup}" + f" z = {is_eq}\n" + f" for i in range({N}):\n" + f" z {accumulate} {x_ref} {cmp_op} {y_ref}\n" + f" z_ptr[0] = z\n" + ) + wrapper = _compile_codegen( + src, + func_name="wrapper", + source_label=f"", + extra_ns=ns, + ) + return wrapper, wrapper_sig + + # Record UDT (possibly nested): chain leaf comparisons. An + # array-valued sub-field compares as a whole, against the other + # operand's field or a scalar, and reduces to one bool. + np_type = udt_dtype.np_type + terms = [] + for py_path, _c, leaf_dtype in _iter_record_leaves(np_type): + x_expr = "x_ptr[0]" if x_is_scalar else f"x[0]{py_path}" + y_expr = "y_ptr[0]" if y_is_scalar else f"y[0]{py_path}" + if leaf_dtype.subdtype is not None: + reducer = ".all()" if is_eq else ".any()" + terms.append(f"({x_expr} {cmp_op} {y_expr}){reducer}") + else: + terms.append(f"({x_expr} {cmp_op} {y_expr})") + expr = join.join(terms) if terms else ("True" if is_eq else "False") + lines = ["def wrapper(z_ptr, x_ptr, y_ptr):"] + if not x_is_scalar: + lines.append(" x = numba.carray(x_ptr, 1)") + if not y_is_scalar: + lines.append(" y = numba.carray(y_ptr, 1)") + lines.append(f" z_ptr[0] = {expr}") + src = "\n".join(lines) + "\n" + wrapper = _compile_codegen( + src, + func_name="wrapper", + source_label=f"", + ) + x_ptr_type = ( + nt.CPointer(numba.from_dtype(dtype.np_type)) + if x_is_scalar + else nt.CPointer(dtype.numba_type) + ) + y_ptr_type = ( + nt.CPointer(numba.from_dtype(dtype2.np_type)) + if y_is_scalar + else nt.CPointer(dtype2.numba_type) + ) + wrapper_sig = nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type) + return wrapper, wrapper_sig class TypedBuiltinBinaryOp(TypedOpBase): @@ -94,9 +203,7 @@ def __call__(self, left, right=None, *, left_default=None, right_default=None): f">>> {self}(x | y, left_default=0, right_default=0)\n\nwhere x and y " "are Vectors or Matrices, and left_default and right_default are scalars." ) - return left.left._ewise_union( - left.right, self, left_default, right_default, is_infix=True - ) + return left._call_left("_ewise_union", self, left_default, right_default, is_infix=True) return _call_op(self, left, right) @property @@ -129,6 +236,7 @@ def type2(self): class TypedUserBinaryOp(TypedOpBase): __slots__ = "_monoid" opclass = "BinaryOp" + _owns_gb_obj = True def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None): super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2) @@ -208,20 +316,6 @@ def commutes_to(self): is_commutative = TypedBuiltinBinaryOp.is_commutative - def __reduce__(self): - name = f"binary.{self.name}" - if not self._anonymous and name in _STANDARD_OPERATOR_NAMES: - return name - return (self._deserialize, (self.name, self.func, self._anonymous)) - - @staticmethod - def _deserialize(name, func, anonymous): - if anonymous: - return BinaryOp.register_anonymous(func, name, parameterized=True) - if (rv := BinaryOp._find(name)) is not None: - return rv - return BinaryOp.register_new(name, func, parameterized=True) - def _floordiv(x, y): return x // y # pragma: no cover (numba) @@ -245,11 +339,51 @@ def _rpow(x, y): def _isclose(rel_tol=1e-7, abs_tol=0.0): def inner(x, y): # pragma: no cover (numba) - return x == y or abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol) + if x == y: + return True + # As in math.isclose, an infinity is close only to itself + if abs(x) == np.inf or abs(y) == np.inf: + return False + return _abs_diff(x, y) <= max(rel_tol * max(_magnitude(x), _magnitude(y)), abs_tol) return inner +# Integer arithmetic in Numba wraps: x - y in int64 or uint64 (Numba widens smaller +# unsigned types to uint64, so 5 - 6 is 2**64 - 1), and abs at a signed type's +# minimum. Python ints do not, so for them these are plain Python; the overloads +# below compile them for Numba's integers so that they do not wrap either. +def _abs_diff(x, y): + return abs(x - y) + + +def _magnitude(x): + return abs(x) + + +if _has_numba: + from numba.extending import overload + + @overload(_abs_diff) + def _abs_diff_overload(x, y): + if isinstance(x, numba.types.Integer) and x == y: + + def abs_diff(x, y): # pragma: no cover (numba) + # The difference fits uint64, where both convert modulo 2**64 + if x >= y: + return np.uint64(x) - np.uint64(y) + return np.uint64(y) - np.uint64(x) + + return abs_diff + return lambda x, y: abs(x - y) + + @overload(_magnitude) + def _magnitude_overload(x): + if isinstance(x, numba.types.Integer): + return lambda x: abs(np.float64(x)) + return lambda x: abs(x) # noqa: PLW0108 (an overload returns a Python function) + + _MAX_INT64 = np.iinfo(np.int64).max @@ -297,12 +431,8 @@ def _pair(x, y): return 1 # pragma: no cover (numba) -def _first_dtype(op, dtype, dtype2): - if dtype._is_udt or dtype2._is_udt: - return op._compile_udt(dtype, dtype2) - - -def _second_dtype(op, dtype, dtype2): +def _udt_dtype(op, dtype, dtype2): + """Custom dtype handler that compiles an operator for UDTs on first use.""" if dtype._is_udt or dtype2._is_udt: return op._compile_udt(dtype, dtype2) @@ -311,11 +441,22 @@ def _pair_dtype(op, dtype, dtype2): return op[INT64] +if _has_numba: + from .udt_codegen import ( + _compile_codegen, + _has_jit_set, + _maybe_warn_jit_skipped, + compile_udt_binary_wrapper, + set_jit_c_mixed_on_op, + set_jit_c_on_op, + ) + + class BinaryOp(OpBase): """Takes two inputs and returns one output, possibly of a different data type. Built-in and registered BinaryOps are located in the ``graphblas.binary`` namespace - as well as in the ``graphblas.ops`` combined namespace. + as well as in the ``graphblas.op`` combined namespace. """ __slots__ = ( @@ -432,13 +573,34 @@ class BinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") success = False - binary_udf = numba.njit(func) + # The error model has to be set here, not only on the cfunc wrapper + # below: ``.compile(sig)`` further down builds the specialization the + # wrapper then reuses, and a Dispatcher keeps one compilation per + # signature. Whichever compile happens first fixes the model, so + # setting it only on the wrapper leaves ``x // 0`` raising + # ZeroDivisionError inside a cfunc, where Numba prints the traceback + # and returns, handing GraphBLAS an element it never wrote. + # + # ``cache`` is only True for the module-level built-in UDFs. Numba can + # only key a cache entry on a stable on-disk source, which lambdas and + # interactively defined functions do not have, so ops registered by + # users stay uncached. + if cache: + try: + binary_udf = numba.njit(func, error_model="numpy", cache=True) + except RuntimeError: + # Numba has nowhere to write the cache ("no locator available"): a + # read-only install whose user has no writable cache directory + # either. Compile uncached, as for an op a user registers. + binary_udf = numba.njit(func, error_model="numpy") + else: + binary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) return_types = {} nt = numba.types @@ -463,11 +625,8 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): elif type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: ret_type = INT8 - # Numba is unable to handle BOOL correctly right now, but we have a workaround - # See: https://github.com/numba/numba/issues/5395 - # We're relying on coercion behaving correctly here - input_type = INT8 if type_ == BOOL else type_ - return_type = INT8 if ret_type == BOOL else ret_type + input_type = _bool_to_int8(type_) + return_type = _bool_to_int8(ret_type) # Build wrapper because GraphBLAS wants pointers and void return wrapper_sig = nt.void( @@ -497,7 +656,9 @@ def binary_wrapper(z, x, y): # pragma: no cover (numba) def binary_wrapper(z, x, y): # pragma: no cover (numba) z[0] = binary_udf(x[0], y[0]) - binary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(binary_wrapper) + binary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + binary_wrapper + ) new_binary = ffi_new("GrB_BinaryOp*") check_status_carg( lib.GrB_BinaryOp_new( @@ -524,110 +685,111 @@ def _compile_udt(self, dtype, dtype2): dtypes = (dtype, dtype2) if dtypes in self._udt_types: return self._udt_ops[dtypes] + from ..udt import _refuse_half - if self.name == "eq" and not self._anonymous and _has_numba: - nt = numba.types - # assert dtype.np_type == dtype2.np_type - itemsize = dtype.np_type.itemsize - mask = _udt_mask(dtype.np_type) - ret_type = BOOL - wrapper_sig = nt.void( - nt.CPointer(INT8.numba_type), - nt.CPointer(UINT8.numba_type), - nt.CPointer(UINT8.numba_type), - ) - # PERF: we can probably make this faster - if mask.all(): - - def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover (numba) - x = numba.carray(x_ptr, itemsize) - y = numba.carray(y_ptr, itemsize) - # for i in range(itemsize): - # if x[i] != y[i]: - # z_ptr[0] = False - # break - # else: - # z_ptr[0] = True - z_ptr[0] = (x == y).all() + _refuse_half(self, dtype, dtype2) - else: + # Built-in ops lift to UDTs by name. A UDF registered under one of those + # names (``BinaryOp.register_anonymous(func, "plus", is_udt=True)``) is + # anonymous, and is compiled from its own function below. + is_lifted = not self._anonymous and ( + self.name in _BUILTIN_UDT_BINARY_OPS or self.name in ("eq", "ne") + ) - def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover (numba) - x = numba.carray(x_ptr, itemsize) - y = numba.carray(y_ptr, itemsize) - # for i in range(itemsize): - # if mask[i] and x[i] != y[i]: - # z_ptr[0] = False - # break - # else: - # z_ptr[0] = True - z_ptr[0] = (x[mask] == y[mask]).all() - - elif self.name == "ne" and not self._anonymous and _has_numba: - nt = numba.types - # assert dtype.np_type == dtype2.np_type - itemsize = dtype.np_type.itemsize - mask = _udt_mask(dtype.np_type) - ret_type = BOOL - wrapper_sig = nt.void( - nt.CPointer(INT8.numba_type), - nt.CPointer(UINT8.numba_type), - nt.CPointer(UINT8.numba_type), + # Built-in arithmetic ops set ``_udt_types = {}`` to enable UDT + # dispatch, which also routes plain-scalar misses (e.g. ``plus[BOOL]``) + # here. Re-raise with the legacy single-dtype message so callers like + # ``mapnumpy`` see the same error as before UDT auto-lift. + if ( + is_lifted + and self.name in _BUILTIN_UDT_BINARY_OPS + and not dtype._is_udt + and not dtype2._is_udt + ): + raise KeyError(f"{self.name} does not work with {dtype}") + + if is_lifted and self.name in ("eq", "ne") and _has_numba: + binary_wrapper, wrapper_sig = _make_udt_comparison( + dtype, dtype2, is_eq=(self.name == "eq") ) - if mask.all(): - - def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover (numba) - x = numba.carray(x_ptr, itemsize) - y = numba.carray(y_ptr, itemsize) - # for i in range(itemsize): - # if x[i] != y[i]: - # z_ptr[0] = True - # break - # else: - # z_ptr[0] = False - z_ptr[0] = (x != y).any() - - else: - - def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover (numba) - x = numba.carray(x_ptr, itemsize) - y = numba.carray(y_ptr, itemsize) - # for i in range(itemsize): - # if mask[i] and x[i] != y[i]: - # z_ptr[0] = True - # break - # else: - # z_ptr[0] = False - z_ptr[0] = (x[mask] != y[mask]).any() + ret_type = BOOL + elif _has_numba and is_lifted: + # Auto-generate a per-leaf wrapper for built-in arithmetic. Most + # of these names (plus, minus, times, truediv, min, max) have no + # ``_numba_func`` and would fall through to the KeyError below + # without this branch. ``floordiv`` does have a ``_numba_func``, + # which Numba won't compile against records; the per-leaf wrapper + # is the only path that works for it on UDTs. + py_op = _BUILTIN_UDT_BINARY_OPS[self.name] + binary_wrapper, wrapper_sig, ret_type = compile_udt_binary_wrapper( + self.name, py_op, dtype, dtype2 + ) elif self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtypes} types") else: numba_func = self._numba_func sig = (dtype.numba_type, dtype2.numba_type) - numba_func.compile(sig) # Should we catch and give additional error message? - ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) - binary_wrapper, wrapper_sig = _get_udt_wrapper(numba_func, ret_type, dtype, dtype2) - - binary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(binary_wrapper) - new_binary = ffi_new("GrB_BinaryOp*") - check_status_carg( - lib.GrB_BinaryOp_new( - new_binary, binary_wrapper.cffi, ret_type._carg, dtype._carg, dtype2._carg - ), - "BinaryOp", - new_binary[0], - ) - op = TypedUserBinaryOp( - self, - self.name, - dtype, - ret_type, - new_binary[0], - dtype2=dtype2, + _compile_udf_for_udt( + numba_func, sig, op_kind="binary", op_name=self.name, dtypes=(dtype, dtype2) + ) + numba_ret_type = numba_func.overloads[sig].signature.return_type + ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + binary_wrapper, wrapper_sig = _get_udt_wrapper( + numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type + ) + + op = _finalize_udt_op( + self, dtype, dtype2, ret_type, binary_wrapper, wrapper_sig, TypedUserBinaryOp ) - self._udt_types[dtypes] = ret_type - self._udt_ops[dtypes] = op + # Set JIT C definition for auto-generated built-in UDT ops. A same-type + # pair whose result is that type too gets the one-type kernel below; a + # mixed pair (a scalar with a UDT, or two UDTs of one layout with other + # element types) gets a kernel with a C type for each of z, x and y, + # when it has a C form (``set_jit_c_mixed_on_op``); eq/ne on a mixed + # pair keep the cfunc. ``_has_jit_set`` gates access to + # ``lib.GrB_BinaryOp_set_String``, which is absent on SS < 9. + if is_lifted and self.name in _BUILTIN_UDT_BINARY_OPS: + is_jittable = ret_type == dtype + else: + is_jittable = is_lifted + if _has_numba and _has_jit_set and dtype == dtype2 and is_jittable: + if self.name in _BUILTIN_UDT_BINARY_OPS: + op._jit_c_info = set_jit_c_on_op( + op.gb_obj, + self.name, + _BUILTIN_UDT_BINARY_OPS[self.name], + dtype, + lib.GrB_BinaryOp_set_String, + arity=2, + ) + else: + # eq/ne return BOOL, not the UDT; the kernel is a chain of + # per-leaf scalar comparisons. + from .udt_codegen import set_jit_c_comparison_on_op + + op._jit_c_info = set_jit_c_comparison_on_op( + op.gb_obj, + self.name, + dtype, + lib.GrB_BinaryOp_set_String, + is_eq=(self.name == "eq"), + ) + _maybe_warn_jit_skipped(op._jit_c_info, self.name, dtype.name) + elif _has_numba and _has_jit_set and is_lifted and self.name in _BUILTIN_UDT_BINARY_OPS: + # (no warning when there is none: the cfunc is this op's documented path) + from ..udt import _udt_op_types + + _, input_types = _udt_op_types(self.name, dtype, dtype2) + op._jit_c_info = set_jit_c_mixed_on_op( + op.gb_obj, + self.name, + _BUILTIN_UDT_BINARY_OPS[self.name], + dtype, + dtype2, + ret_type, + input_types, + ) return op @classmethod @@ -671,7 +833,9 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals return cls._build(name, func, anonymous=True, is_udt=is_udt) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, _cache=False + ): """Register a new BinaryOp and save it to ``graphblas.binary`` namespace. Parameters @@ -719,11 +883,16 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal >>> dir(gb.binary) [..., 'max_zero', ...] - This is how ``gb.binary.isclose`` is defined: + ``gb.binary.isclose`` is defined like this, except that it computes ``x - y`` + and ``abs(x)`` for integers without wrapping: >>> def user_isclose(rel_tol=1e-7, abs_tol=0.0): >>> def inner(x, y): - >>> return x == y or abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol) + >>> if x == y: + >>> return True + >>> if abs(x) == np.inf or abs(y) == np.inf: + >>> return False + >>> return abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol) >>> return inner >>> gb.binary.register_new("user_isclose", user_isclose, parameterized=True) @@ -733,13 +902,19 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "_cache": _cache, + }, ) elif parameterized: binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) else: - binary_op = cls._build(name, func, is_udt=is_udt) + binary_op = cls._build(name, func, is_udt=is_udt, cache=_cache) setattr(module, funcname, binary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -792,13 +967,20 @@ def _initialize(cls): if _supports_udfs: # Add floordiv # cdiv truncates towards 0, while floordiv truncates towards -inf - BinaryOp.register_new("floordiv", _floordiv, lazy=True) # cast to integer - BinaryOp.register_new("rfloordiv", _rfloordiv, lazy=True) # cast to integer + # cache=True persists the numba compilation of these built-in UDFs to + # disk so it is paid once per machine, not once per process. It shaves + # ~20% off a single op's first-touch build and ~45% when several are + # used (the per-process object-link and LLVM-init cost is not cached). + # Only module-level built-ins get cache=True (see _build). + BinaryOp.register_new("floordiv", _floordiv, lazy=True, _cache=True) # cast to integer + BinaryOp.register_new( + "rfloordiv", _rfloordiv, lazy=True, _cache=True + ) # cast to integer # For aggregators - BinaryOp.register_new("absfirst", _absfirst, lazy=True) - BinaryOp.register_new("abssecond", _abssecond, lazy=True) - BinaryOp.register_new("rpow", _rpow, lazy=True) + BinaryOp.register_new("absfirst", _absfirst, lazy=True, _cache=True) + BinaryOp.register_new("abssecond", _abssecond, lazy=True, _cache=True) + BinaryOp.register_new("rpow", _rpow, lazy=True, _cache=True) # For algorithms binary._delayed["binom"] = (_register_binom, {}) # Lazy with custom creation @@ -915,29 +1097,37 @@ def _initialize(cls): right = getattr(binary, right_name) left._semiring_commutes_to = right right._semiring_commutes_to = left - # Allow some functions to work on UDTs + # ``any`` uses ``_second`` so ``monoid.any[udt].reduce(v)`` folds as + # ``acc = v_i`` and returns an actual element. With ``_first`` the + # accumulator would never advance, so the reduce would always return + # the identity. for binop, func in [ (binary.first, _first), (binary.second, _second), (binary.pair, _pair), - (binary.any, _first), + (binary.any, _second), ]: binop.orig_func = func - if _has_numba: - binop._numba_func = numba.njit(func) - else: - binop._numba_func = None + binop._numba_func = numba.njit(func, error_model="numpy") if _has_numba else None binop._udt_types = {} binop._udt_ops = {} - binary.any._numba_func = binary.first._numba_func + binary.any._numba_func = binary.second._numba_func + binary.first._custom_dtype = _udt_dtype + binary.second._custom_dtype = _udt_dtype + binary.pair._custom_dtype = _pair_dtype + # eq and ne walk the UDT leaf-by-leaf (special-cased in _compile_udt + # via _make_udt_comparison) so NaN propagation matches IEEE semantics + # rather than bit equality. binary.eq._udt_types = {} binary.eq._udt_ops = {} binary.ne._udt_types = {} binary.ne._udt_ops = {} - # Set custom dtype handling - binary.first._custom_dtype = _first_dtype - binary.second._custom_dtype = _second_dtype - binary.pair._custom_dtype = _pair_dtype + # Element-wise arithmetic ops on UDTs are auto-generated from + # per-field / per-element scalar ops; the attributes that enable this + # are seeded in ``__init__`` (keyed on ``_BUILTIN_UDT_BINARY_OPS``). + # Do not seed them here with ``getattr(binary, op_name)``: that would + # force lazily-registered UDF ops like ``floordiv`` to numba-compile + # for every dtype at import time (about half a second of startup here). cls._initialized = True def __init__( @@ -962,15 +1152,14 @@ def __init__( if is_udt: self._udt_types = {} # {(dtype, dtype): DataType} self._udt_ops = {} # {(dtype, dtype): TypedUserBinaryOp} - - def __reduce__(self): - if self._anonymous: - if hasattr(self.orig_func, "_parameterized_info"): - return (_deserialize_parameterized, self.orig_func._parameterized_info) - return (self.register_anonymous, (self.orig_func, self.name)) - if (name := f"binary.{self.name}") in _STANDARD_OPERATOR_NAMES: - return name - return (self._deserialize, (self.name, self.orig_func)) + elif _has_numba and not anonymous and name in _BUILTIN_UDT_BINARY_OPS: + # Built-in arithmetic ops auto-lift to UDTs field-by-field. + # Seeded here rather than in ``_initialize`` so that delayed ops + # (e.g. ``floordiv``) are not force-built at import time. An + # anonymous op that shares a built-in's name is a user's op. + self._udt_types = {} + self._udt_ops = {} + self._custom_dtype = _udt_dtype __call__ = TypedBuiltinBinaryOp.__call__ is_commutative = TypedBuiltinBinaryOp.is_commutative @@ -983,3 +1172,6 @@ def monoid(self): self._monoid = Monoid._find(self.name) return self._monoid + + +ParameterizedBinaryOp._op_class = BinaryOp diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py new file mode 100644 index 000000000..4539387a9 --- /dev/null +++ b/graphblas/core/operator/indexbinary.py @@ -0,0 +1,712 @@ +import inspect +from operator import getitem +from types import FunctionType + +import numpy as np + +from ... import _STANDARD_OPERATOR_NAMES, indexbinary +from ...dtypes import BOOL, INT8, UINT64, lookup_dtype +from ...exceptions import EmptyObject, UdfParseError, check_status_carg +from .. import _has_numba, ffi, lib +from ..dtypes import _sample_values, unify +from .base import OpBase, ParameterizedUdf, TypedOpBase +from .binary import BinaryOp + +_has_idxbinop = hasattr(lib, "GxB_IndexBinaryOp_new") + +if _has_numba: + import numba + + from .udf import ( + _bool_to_int8, + _compile_udf_for_udt, + _get_udt_wrapper_indexbinary, + _resolve_udt_return_type, + ) + +ffi_new = ffi.new + + +def _rebind_indexbinaryop(parent, type_, theta): + # ``type_`` is one type, or ``(x, y, theta)`` types when they differ. + # ``theta`` was stored as a raw Python/numpy value. For UDT thunks, + # ``TypedBuiltinIndexBinaryOp.__call__`` would route a raw value through + # ``Scalar.from_value(theta)`` with no dtype, which then fails inside the + # multi-dim ndarray ``Scalar.value`` setter for array UDT values. Wrap + # as a typed Scalar here so the bind uses an explicit dtype. + typed = parent[type_] + thunk_type = typed.thunk_type + if thunk_type._is_udt and not isinstance(theta, (int, float, bool, complex)): + from ..scalar import Scalar + + theta = Scalar.from_value(theta, dtype=thunk_type, is_cscalar=False, name="") + return typed(theta) + + +def _theta_to_scalar(theta, dtype): + """Wrap a raw (non-Scalar) theta as a GrB Scalar for binding. + + ``dtype`` is the thunk dtype when known (a UDT thunk or an explicit + ``dtype=``), else ``None`` to type theta by its own value: a Python + number by its default type, a numpy scalar or ``np.void`` by its dtype, + and a numpy array by its array layout (``INT32[3]`` for an int32 array + of shape ``(3,)``). A tuple, list or dict has no type of its own, so + raise a clear error that names the value and the working alternatives. + """ + from ..scalar import Scalar + + if dtype is None and isinstance(theta, np.ndarray) and theta.ndim > 0: + try: + dtype = lookup_dtype(np.dtype((theta.dtype, theta.shape))) + except (TypeError, ValueError): + pass # a bool array, for one, is no UDT; ``from_value`` says so + try: + return Scalar.from_value(theta, dtype, is_cscalar=False, name="") # pragma: is_grbscalar + except TypeError as exc: + if dtype is None: + raise TypeError( + f"Cannot infer a dtype for theta from {theta!r}; give theta as a Scalar of " + "its type (Scalar.from_value(theta, dtype)), or name the op's type " + "(op[dtype](theta), or op(theta, dtype=...))." + ) from exc + raise + + +def _refuse_empty_theta(theta): + # SuiteSparse binds theta without checking that it holds a value: binding a new + # empty Scalar crashes, and a copy of one binds whatever its value array holds. + # ``apply`` with an IndexUnaryOp raises EmptyObject for an empty thunk. + if theta._is_empty: + raise EmptyObject("theta is an empty Scalar; bind a Scalar that holds a value") + + +class _BoundIndexBinaryOp(TypedOpBase): + """A BinaryOp produced by binding a theta value to a typed IndexBinaryOp. + + ``ibo[dtype](theta)`` returns one, and so does ``ibo(theta)`` each time it + runs, for the type it runs in (:class:`_UntypedBoundIndexBinaryOp`). + Under the hood this is a single typed ``GrB_BinaryOp`` (constructed via + ``GxB_BinaryOp_new_IndexOp``), so SuiteSparse accepts it anywhere a + typed ``GrB_BinaryOp`` is expected: ``ewise_mult`` / ``ewise_add`` and, + via :func:`Semiring.register_anonymous`, as the multiplier of a + Semiring that then works in ``mxm`` / ``mxv`` / ``vxm``. + """ + + __slots__ = "_theta", "_theta_type" + opclass = "BinaryOp" + # Each ``iop(theta)`` call allocates a fresh ``GrB_BinaryOp`` via + # ``GxB_BinaryOp_new_IndexOp`` and bound IBOs are never cached, so this + # is the per-call allocation that ``TypedOpBase.__del__`` must release. + # Without this, every ``iop(theta)`` (including each pickle round-trip) + # leaks one ``GrB_BinaryOp`` for the life of the process. + _owns_gb_obj = True + # A bound IBO is monomorphic in its operand types (one ``(x, y, z, + # theta)`` combination). Expose the polymorphism / commutativity + # markers that ``Semiring.__init__``, ``Semiring.is_positional``, + # ``Semiring._custom_dtype``, and ``TypedUserSemiring`` probe as + # class-level defaults, so they get sane answers (and no + # ``AttributeError``) when a bound IBO is passed in place of a + # ``BinaryOp`` parent. + _is_udt = False + _udt_types = None + _udt_ops = None + _semiring_commutes_to = None + commutes_to = None + is_commutative = False + # No Monoid is made of it, so reduce refuses it as it refuses a BinaryOp + # that is not part of one. + monoid = None + # ``is_positional`` and ``_custom_dtype`` are BinaryOp-specific dispatch + # hints: when a polymorphic BinaryOp is queried for a type it doesn't + # support, the resolver in ``utils.get_typed_op`` falls back through + # them. A bound IBO is already type-monomorphic (one ``GrB_BinaryOp`` + # for one (x, y, z, theta) combination), so neither fallback applies. + is_positional = False + _custom_dtype = None + + def __reduce__(self): + return ( + _rebind_indexbinaryop, + (self.parent, _ibo_types_key(self), self._theta), + ) + + +class _UntypedBoundIndexBinaryOp(BinaryOp): + """A BinaryOp produced by binding a theta value to an IndexBinaryOp, typed when it runs. + + ``ibo(theta)`` returns one. As with any other BinaryOp, its inputs choose + its type each time it runs. Theta is typed when it is given, as an + IndexUnaryOp's thunk is typed by its value: a Python int is INT64, a float + FP64 and a bool BOOL; a numpy scalar, a numpy array (as an array UDT) or a + Scalar keeps its own type. The op runs at the type that holds both inputs + and theta, and the default theta, False, changes no type. An IndexBinaryOp + registered with ``is_udt=True`` instead runs as ``ibo[X, Y, T]`` for inputs + of types X and Y and theta of type T, compiled for those three types, so + theta is never converted. + + It runs the typed op ``ibo[...](theta)``, built once per pair of input types. + """ + + __slots__ = "_ibo", "_theta" + # For ``find_opclass``, where the class name would not say BinaryOp. + opclass = "BinaryOp" + + def __init__(self, ibo, theta): + # ``is_udt=True`` sends every lookup (``get_typed_op``, a Semiring's) to + # ``_compile_udt`` below, UDT or not. + super().__init__(ibo.name, anonymous=True, is_udt=True) + self._ibo = ibo + self._theta = theta + + def _compile_udt(self, dtype, dtype2): + if dtype2 is None: # pragma: no cover (safety) + dtype2 = dtype + dtypes = (dtype, dtype2) + if dtypes in self._udt_ops: + return self._udt_ops[dtypes] + op = self._ibo._get_typed(dtype, dtype2, self._theta.dtype)(self._theta) + self._udt_types[dtypes] = op.return_type + self._udt_ops[dtypes] = op + return op + + def __getitem__(self, type_): + # A type named here is the op's type, as in ``ibo[type_](theta)``. + return self._ibo[type_](self._theta) + + def __reduce__(self): + return (self._ibo, (self._theta,)) + + def __repr__(self): + return f"{self._ibo!r}({self._theta.value!r})" + + +def _ibo_types_key(op): + """The type, or ``(x, y, theta)`` types when they differ, that index ``op``'s parent.""" + dtype2 = op.type if op._type2 is None else op._type2 + if op.type == dtype2 == op._theta_type: + return op.type + return (op.type, dtype2, op._theta_type) + + +class _TypedIndexBinaryOpMixin: + """The theta type, repr and pickle of a typed IndexBinaryOp, whose theta has its own type.""" + + __slots__ = () + + @property + def thunk_type(self): + return self._theta_type + + def __repr__(self): + dtype2 = self.type if self._type2 is None else self._type2 + if self._theta_type == dtype2: + return super().__repr__() + return f"indexbinary.{self.name}[{self.type.name}, {dtype2.name}, {self._theta_type.name}]" + + def __reduce__(self): + return (getitem, (self.parent, _ibo_types_key(self))) + + +class TypedBuiltinIndexBinaryOp(TypedOpBase): + __slots__ = () + opclass = "IndexBinaryOp" + + def __call__(self, theta=None): + """Bind a theta value, returning a BinaryOp. + + The result is usable directly in ``ewise_mult`` and ``ewise_add``, + or as the multiplier of a Semiring (see + :meth:`Semiring.register_anonymous`) for ``mxm`` / ``mxv`` / ``vxm``. + + Parameters + ---------- + theta : scalar, optional + The theta parameter to bind. Defaults to 0 (False). Its value is bound + here, so changing a Scalar theta afterwards does not change the BinaryOp. + + Returns + ------- + TypedOpBase + A BinaryOp built from this IndexBinaryOp with the given theta. + + """ + from ..scalar import Scalar + + if theta is None: + theta = False + theta_value = theta.value if isinstance(theta, Scalar) else theta + if not isinstance(theta, Scalar): + # A raw value for a UDT thunk carries the thunk dtype; otherwise the + # value is inferred (and a raw UDT value raises a clear error). + tt = self.thunk_type + theta = _theta_to_scalar(theta, tt if tt._is_udt else None) + elif theta._is_cscalar: + # fmt: off + # Pass dtype explicitly; an array-UDT value would otherwise infer wrong. + val, dt = theta.value, theta.dtype + theta = Scalar.from_value(val, dt, is_cscalar=False, name="") # pragma: is_grbscalar + # fmt: on + _refuse_empty_theta(theta) + new_binop = ffi_new("GrB_BinaryOp*") + check_status_carg( + lib.GxB_BinaryOp_new_IndexOp(new_binop, self.gb_obj, theta._carg), + "BinaryOp", + new_binop[0], + ) + rv = _BoundIndexBinaryOp( + self.parent, + self.name, + self.type, + self.return_type, + new_binop[0], + f"{self.name}_bound", + dtype2=self._type2, + ) + rv._theta = theta_value + rv._theta_type = self.thunk_type + return rv + + +class TypedUserIndexBinaryOp(_TypedIndexBinaryOpMixin, TypedOpBase): + __slots__ = ("_theta_type",) + opclass = "IndexBinaryOp" + _owns_gb_obj = True + + def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None, theta_type=None): + super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2) + if theta_type is None: + theta_type = type_ if dtype2 is None else dtype2 + self._theta_type = theta_type + + @property + def orig_func(self): + return self.parent.orig_func + + @property + def _numba_func(self): + return self.parent._numba_func + + def __call__(self, theta=None): + return TypedBuiltinIndexBinaryOp.__call__(self, theta) + + __call__.__doc__ = TypedBuiltinIndexBinaryOp.__call__.__doc__ + + +class ParameterizedIndexBinaryOp(ParameterizedUdf): + __slots__ = "func", "__signature__", "_is_udt" + + def __init__(self, name, func, *, anonymous=False, is_udt=False): + self.func = func + self.__signature__ = inspect.signature(func) + self._is_udt = is_udt + if name is None: + name = getattr(func, "__name__", name) + super().__init__(name, anonymous) + + def _call(self, *args, **kwargs): + idxbinop = self.func(*args, **kwargs) + idxbinop._parameterized_info = (self, args, kwargs) + return IndexBinaryOp.register_anonymous(idxbinop, self.name, is_udt=self._is_udt) + + +class IndexBinaryOp(OpBase): + """Takes two inputs with their indices plus a thunk, and returns one output. + + The function has the signature ``f(x, ix, jx, y, iy, jy, theta) -> z``, + where ``ix, jx`` are the row and column indices of ``x``, and ``iy, jy`` + are the row and column indices of ``y``. + + Binding a theta value (``ibo(theta)``) produces a BinaryOp usable + directly in ``ewise_mult`` and ``ewise_add``. Its type is chosen when it + runs, from both inputs and theta; ``ibo[dtype](theta)`` binds to one type. + With ``is_udt=True``, x, y and theta each keep their own type, and + ``ibo[X, Y, T](theta)`` names all three. + To use it as a multiplier in ``mxm`` / ``mxv`` / ``vxm``, wrap it in a + Semiring via :meth:`Semiring.register_anonymous`; per SuiteSparse, the + additive monoid itself cannot be IndexBinaryOp-based. + + IndexBinaryOps live in the ``graphblas.indexbinary`` namespace. There are + no built-ins; all IndexBinaryOps are user-defined. + """ + + __slots__ = "orig_func", "_is_udt", "_numba_func" + _module = indexbinary + _modname = "indexbinary" + _custom_dtype = None + # Read by ``OpBase.__getitem__`` for a type the op does not have; no + # IndexBinaryOp is positional. + is_positional = False + # Read by the error for a bound op where a Monoid is expected (reduce). + monoid = None + _typed_class = TypedBuiltinIndexBinaryOp + _typed_user_class = TypedUserIndexBinaryOp + # No built-in IndexBinaryOps; no parse config needed + _parse_config = { + "trim_from_front": 4, + "num_underscores": 1, + "re_exprs": [], + } + + @classmethod + def _build(cls, name, func, *, is_udt=False, anonymous=False): + if not _has_idxbinop: + raise RuntimeError( + "IndexBinaryOp requires SuiteSparse:GraphBLAS 9.4+ " + "(python-suitesparse-graphblas 9.3.1+)" + ) + if not isinstance(func, FunctionType): + raise TypeError(f"UDF argument must be a function, not {type(func)}") + if name is None: + name = getattr(func, "__name__", "") + success = False + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + indexbinary_udf = numba.njit(func, error_model="numpy") + new_type_obj = cls( + name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexbinary_udf + ) + return_types = {} + nt = numba.types + if not is_udt: + for type_ in _sample_values: + sig = ( + type_.numba_type, + UINT64.numba_type, + UINT64.numba_type, + type_.numba_type, + UINT64.numba_type, + UINT64.numba_type, + type_.numba_type, + ) + try: + indexbinary_udf.compile(sig) + except numba.TypingError: + continue + ret_type = lookup_dtype(indexbinary_udf.overloads[sig].signature.return_type) + if ret_type != type_ and ( + ("INT" in ret_type.name and "INT" in type_.name) + or ("FP" in ret_type.name and "FP" in type_.name) + or ("FC" in ret_type.name and "FC" in type_.name) + or ( + type_ == UINT64 + and ret_type.name == "FP64" + and return_types.get(lookup_dtype("INT64")) == lookup_dtype("INT64") + ) + ): + ret_type = type_ + elif type_ == BOOL and ret_type.name == "INT64" and return_types.get(INT8) == INT8: + ret_type = INT8 + + input_type = _bool_to_int8(type_) + return_type = _bool_to_int8(ret_type) + + # Build a wrapper that calls z = f(x, ix, jx, y, iy, jy, theta). + # C signature: void(z*, x*, ix, jx, y*, iy, jy, theta*). + wrapper_sig = nt.void( + nt.CPointer(return_type.numba_type), + nt.CPointer(input_type.numba_type), + UINT64.numba_type, + UINT64.numba_type, + nt.CPointer(input_type.numba_type), + UINT64.numba_type, + UINT64.numba_type, + nt.CPointer(input_type.numba_type), + ) + + if type_ == BOOL: + if ret_type == BOOL: + + def indexbinary_wrapper( + z, x, ix, jx, y, iy, jy, theta + ): # pragma: no cover (numba) + z[0] = bool( + indexbinary_udf( + bool(x[0]), ix, jx, bool(y[0]), iy, jy, bool(theta[0]) + ) + ) + + else: + + def indexbinary_wrapper( + z, x, ix, jx, y, iy, jy, theta + ): # pragma: no cover (numba) + z[0] = indexbinary_udf( + bool(x[0]), ix, jx, bool(y[0]), iy, jy, bool(theta[0]) + ) + + elif ret_type == BOOL: + + def indexbinary_wrapper( + z, x, ix, jx, y, iy, jy, theta + ): # pragma: no cover (numba) + z[0] = bool(indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0])) + + else: + + def indexbinary_wrapper( + z, x, ix, jx, y, iy, jy, theta + ): # pragma: no cover (numba) + z[0] = indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0]) + + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexbinary_wrapper + ) + new_idxbinop = ffi_new("GxB_IndexBinaryOp*") + check_status_carg( + lib.GxB_IndexBinaryOp_new( + new_idxbinop, + indexbinary_wrapper.cffi, + ret_type.gb_obj, + type_.gb_obj, + type_.gb_obj, + type_.gb_obj, + ffi_new("char[]", name.encode()), + ffi.NULL, + ), + "IndexBinaryOp", + new_idxbinop[0], + ) + op = cls._typed_user_class(new_type_obj, name, type_, ret_type, new_idxbinop[0]) + new_type_obj._add(op) + success = True + return_types[type_] = ret_type + if success or is_udt: + return new_type_obj + raise UdfParseError("Unable to parse function using Numba") + + def _compile_udt(self, dtype, dtype2, theta_type=None): + if not _has_idxbinop: + # KeyError (not RuntimeError) so ``udt in op`` and the resolver + # chain in ``__contains__`` / ``[dtype]`` lookups return cleanly. + raise KeyError( + "IndexBinaryOp requires SuiteSparse:GraphBLAS 9.4+ " + "(python-suitesparse-graphblas 9.3.1+)" + ) + if dtype2 is None: + dtype2 = dtype + if theta_type is None: + # ``ibo[X, Y]`` and ``ibo[T]``: theta takes y's type, and casts to it + theta_type = dtype2 + dtypes = (dtype, dtype2, theta_type) + if dtypes in self._udt_types: + return self._udt_ops[dtypes] + from ..udt import _refuse_half + + _refuse_half(self, *dtypes) + if self._numba_func is None: + raise KeyError(f"{self.name} does not work with {dtypes} types") + + numba_func = self._numba_func + sig = ( + dtype.numba_type, + UINT64.numba_type, + UINT64.numba_type, + dtype2.numba_type, + UINT64.numba_type, + UINT64.numba_type, + theta_type.numba_type, + ) + _compile_udf_for_udt( + numba_func, sig, op_kind="indexbinary", op_name=self.name, dtypes=dtypes + ) + numba_ret_type = numba_func.overloads[sig].signature.return_type + ret_type = _resolve_udt_return_type(numba_ret_type, *dtypes) + indexbinary_wrapper, wrapper_sig = _get_udt_wrapper_indexbinary( + numba_func, ret_type, dtype, dtype2, theta_type, numba_ret_type=numba_ret_type + ) + + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexbinary_wrapper + ) + new_idxbinop = ffi_new("GxB_IndexBinaryOp*") + check_status_carg( + lib.GxB_IndexBinaryOp_new( + new_idxbinop, + indexbinary_wrapper.cffi, + ret_type._carg, + dtype._carg, + dtype2._carg, + theta_type._carg, + ffi_new("char[]", self.name.encode()), + ffi.NULL, + ), + "IndexBinaryOp", + new_idxbinop[0], + ) + op = TypedUserIndexBinaryOp( + self, + self.name, + dtype, + ret_type, + new_idxbinop[0], + dtype2=dtype2, + theta_type=theta_type, + ) + self._udt_types[dtypes] = ret_type + self._udt_ops[dtypes] = op + return op + + def _get_typed(self, dtype, dtype2, theta_type): + """The typed op for x, y and theta of these types, as ``ibo(theta)`` runs it. + + With ``is_udt``, it is compiled for the three types as they are. Without + it, x, y and theta share one type, the one that holds all three. + """ + if self._is_udt: + return self._compile_udt(dtype, dtype2, theta_type) + if dtype._is_udt or dtype2._is_udt or theta_type._is_udt: + raise KeyError(f"{self.name} does not work with ({dtype}, {dtype2}, {theta_type})") + return self[unify(unify(dtype, dtype2), theta_type)] + + def __getitem__(self, type_): + if type(type_) is tuple and len(type_) == 3: + return self._get_typed(*(lookup_dtype(dtype) for dtype in type_)) + return super().__getitem__(type_) + + def _add(self, op, *, is_jit=False): + if not is_jit: + return super()._add(op) + # A JIT op is registered for its x, y and theta types + dtypes = (op.type, op._type2, op.thunk_type) + self.types[dtypes] = op.return_type + self._udt_types[dtypes] = op.return_type + self._udt_ops[dtypes] = op + + @classmethod + def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + """Register an IndexBinaryOp without adding it to the ``indexbinary`` namespace. + + Because it is not registered in the namespace, the name is optional. + + Parameters + ---------- + func : FunctionType + The function to compile. For all current backends, this must be + compilable with ``numba.njit``. The function takes seven inputs + ``(x, ix, jx, y, iy, jy, theta)``: ``x`` and ``y`` are element + values, ``ix, jx`` and ``iy, jy`` are their row and column indices + (int64), and ``theta`` is a scalar parameter. + name : str, optional + The name of the operator. This *does not* show up as ``gb.indexbinary.{name}``. + parameterized : bool, default False + When True, create a parameterized user-defined operator, so that + additional parameters can be "baked into" the operator when used. + is_udt : bool, default False + Whether the operator is intended to operate on user-defined types. + + Returns + ------- + IndexBinaryOp or ParameterizedIndexBinaryOp + + """ + cls._check_supports_udf("register_anonymous") + if parameterized: + return ParameterizedIndexBinaryOp(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt) + + @classmethod + def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + """Register a new IndexBinaryOp under the ``graphblas.indexbinary`` namespace. + + Parameters + ---------- + name : str + The name of the operator. Available afterwards as ``gb.indexbinary.{name}``. + func : FunctionType + The function to compile. For all current backends, this must be + compilable with ``numba.njit``. The function takes seven inputs + ``(x, ix, jx, y, iy, jy, theta)``: ``x`` and ``y`` are element + values, ``ix, jx`` and ``iy, jy`` are their row and column indices + (int64), and ``theta`` is a scalar parameter. + parameterized : bool, default False + When True, create a parameterized user-defined operator. + is_udt : bool, default False + Whether the operator is intended to operate on user-defined types. + lazy : bool, default False + When True, defer compilation until the operator is first used. + + Examples + -------- + >>> gb.indexbinary.register_new("index_dist", lambda x, ix, jx, y, iy, jy, t: abs(ix - iy)) + + """ + cls._check_supports_udf("register_new") + module, funcname = cls._remove_nesting(name) + if lazy: + module._delayed[funcname] = ( + cls.register_new, + {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + ) + elif parameterized: + idxbinop = ParameterizedIndexBinaryOp(name, func, is_udt=is_udt) + setattr(module, funcname, idxbinop) + else: + idxbinop = cls._build(name, func, is_udt=is_udt) + setattr(module, funcname, idxbinop) + + if not cls._initialized: # pragma: no cover (safety) + _STANDARD_OPERATOR_NAMES.add(f"{cls._modname}.{name}") + if not lazy: + return idxbinop + + @classmethod + def _initialize(cls): + if cls._initialized: + return + super()._initialize(include_in_ops=False) + # No built-in IndexBinaryOps to register. + cls._initialized = True + + def __init__(self, name, func=None, *, anonymous=False, is_udt=False, numba_func=None): + super().__init__(name, anonymous=anonymous) + self.orig_func = func + self._numba_func = numba_func + self._is_udt = is_udt + if is_udt: + self._udt_types = {} # {(dtype, dtype2, theta_type): DataType} + self._udt_ops = {} # {(dtype, dtype2, theta_type): TypedUserIndexBinaryOp} + + def __call__(self, theta=None, dtype=None): + """Bind a theta value to produce a BinaryOp. + + Parameters + ---------- + theta : scalar, optional + The theta parameter to bind. Defaults to 0 (False). Its value is bound + here, so changing a Scalar theta afterwards does not change the BinaryOp. + dtype : dtype, optional + The type of the BinaryOp, as in ``op[dtype](theta)``. Without it, the + BinaryOp's type is chosen each time it runs, from the types of both + inputs and theta, as other BinaryOps choose theirs. Theta has the type + of its value: a Python int is INT64 and a float FP64, and a numpy + scalar, a numpy array or a Scalar keeps its own type. + + Returns + ------- + BinaryOp + A BinaryOp built from this IndexBinaryOp with the given theta. + + """ + from ...dtypes import lookup_dtype as _lookup_dtype + from ..scalar import Scalar + + if theta is None: + theta = False + if dtype is None: + if isinstance(theta, Scalar): + _refuse_empty_theta(theta) + # Bound when it runs, so keep the value theta has now. + theta = theta.dup(name="") + else: + theta = _theta_to_scalar(theta, None) + if not self._is_udt and theta.dtype._is_udt: + raise KeyError(f"{self.name} does not work with {theta.dtype}") + return _UntypedBoundIndexBinaryOp(self, theta) + dtype = _lookup_dtype(dtype) + if not isinstance(theta, Scalar): + # The named dtype types the value, a UDT's included. + theta = _theta_to_scalar(theta, dtype) + return self[dtype](theta) + + +ParameterizedIndexBinaryOp._op_class = IndexBinaryOp diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 6fdacbcc1..a8eca2460 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -7,12 +7,18 @@ from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, ffi, lib from ..dtypes import _sample_values -from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op, _deserialize_parameterized +from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op if _has_numba: import numba - from .base import _get_udt_wrapper + from .udf import ( + _bool_to_int8, + _compile_udf_for_udt, + _finalize_udt_op, + _get_udt_wrapper, + _resolve_udt_return_type, + ) ffi_new = ffi.new @@ -33,6 +39,7 @@ def thunk_type(self): class TypedUserIndexUnaryOp(TypedOpBase): __slots__ = () opclass = "IndexUnaryOp" + _owns_gb_obj = True def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None): super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2) @@ -65,22 +72,6 @@ def _call(self, *args, **kwargs): indexunary._parameterized_info = (self, args, kwargs) return IndexUnaryOp.register_anonymous(indexunary, self.name, is_udt=self._is_udt) - def __reduce__(self): - # NOT COVERED - name = f"indexunary.{self.name}" - if not self._anonymous and name in _STANDARD_OPERATOR_NAMES: - return name - return (self._deserialize, (self.name, self.func, self._anonymous)) - - @staticmethod - def _deserialize(name, func, anonymous): - # NOT COVERED - if anonymous: - return IndexUnaryOp.register_anonymous(func, name, parameterized=True) - if (rv := IndexUnaryOp._find(name)) is not None: - return rv - return IndexUnaryOp.register_new(name, func, parameterized=True) - class IndexUnaryOp(OpBase): """Takes one input and a thunk and returns one output, possibly of a different data type. @@ -125,7 +116,9 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if name is None: name = getattr(func, "__name__", "") success = False - indexunary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + indexunary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexunary_udf ) @@ -152,11 +145,8 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): elif type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: ret_type = INT8 - # Numba is unable to handle BOOL correctly right now, but we have a workaround - # See: https://github.com/numba/numba/issues/5395 - # We're relying on coercion behaving correctly here - input_type = INT8 if type_ == BOOL else type_ - return_type = INT8 if ret_type == BOOL else ret_type + input_type = _bool_to_int8(type_) + return_type = _bool_to_int8(ret_type) # Build wrapper because GraphBLAS wants pointers and void return wrapper_sig = nt.void( @@ -188,7 +178,9 @@ def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) z[0] = indexunary_udf(x[0], row, col, y[0]) - indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexunary_wrapper) + indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexunary_wrapper + ) new_indexunary = ffi_new("GrB_IndexUnaryOp*") check_status_carg( lib.GrB_IndexUnaryOp_new( @@ -215,37 +207,25 @@ def _compile_udt(self, dtype, dtype2): dtypes = (dtype, dtype2) if dtypes in self._udt_types: return self._udt_ops[dtypes] + from ..udt import _refuse_half + + _refuse_half(self, dtype, dtype2) if self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtypes} types") numba_func = self._numba_func sig = (dtype.numba_type, UINT64.numba_type, UINT64.numba_type, dtype2.numba_type) - numba_func.compile(sig) # Should we catch and give additional error message? - ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) - indexunary_wrapper, wrapper_sig = _get_udt_wrapper( - numba_func, ret_type, dtype, dtype2, include_indexes=True + _compile_udf_for_udt( + numba_func, sig, op_kind="indexunary", op_name=self.name, dtypes=(dtype, dtype2) ) - - indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexunary_wrapper) - new_indexunary = ffi_new("GrB_IndexUnaryOp*") - check_status_carg( - lib.GrB_IndexUnaryOp_new( - new_indexunary, indexunary_wrapper.cffi, ret_type._carg, dtype._carg, dtype2._carg - ), - "IndexUnaryOp", - new_indexunary[0], + numba_ret_type = numba_func.overloads[sig].signature.return_type + ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + indexunary_wrapper, wrapper_sig = _get_udt_wrapper( + numba_func, ret_type, dtype, dtype2, include_indexes=True, numba_ret_type=numba_ret_type ) - op = TypedUserIndexUnaryOp( - self, - self.name, - dtype, - ret_type, - new_indexunary[0], - dtype2=dtype2, + return _finalize_udt_op( + self, dtype, dtype2, ret_type, indexunary_wrapper, wrapper_sig, TypedUserIndexUnaryOp ) - self._udt_types[dtypes] = ret_type - self._udt_ops[dtypes] = op - return op @classmethod def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): @@ -258,8 +238,8 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals func : FunctionType The function to compile. For all current backends, this must be able to be compiled with ``numba.njit``. - ``func`` takes four input parameters--any dtype, int64, int64, - any dtype and returns any dtype. The first argument (any dtype) is + ``func`` takes four input parameters (any dtype, int64, int64, + any dtype) and returns any dtype. The first argument (any dtype) is the value of the input Matrix or Vector, the second argument (int64) is the row index of the Matrix or the index of the Vector, the third argument (int64) is the column index of the Matrix or 0 for a Vector, @@ -308,8 +288,8 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal func : FunctionType The function to compile. For all current backends, this must be able to be compiled with ``numba.njit``. - ``func`` takes four input parameters--any dtype, int64, int64, - any dtype and returns any dtype. The first argument (any dtype) is + ``func`` takes four input parameters (any dtype, int64, int64, + any dtype) and returns any dtype. The first argument (any dtype) is the value of the input Matrix or Vector, the second argument (int64) is the row index of the Matrix or the index of the Vector, the third argument (int64) is the column index of the Matrix or 0 for a Vector, @@ -348,7 +328,7 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized}, + {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, ) elif parameterized: indexunary_op = ParameterizedIndexUnaryOp(name, func, is_udt=is_udt) @@ -425,18 +405,10 @@ def __init__( self.is_positional = is_positional self._is_udt = is_udt if is_udt: - self._udt_types = {} # {dtype: DataType} - self._udt_ops = {} # {dtype: TypedUserIndexUnaryOp} - - def __reduce__(self): - if self._anonymous: - if hasattr(self.orig_func, "_parameterized_info"): - # NOT COVERED - return (_deserialize_parameterized, self.orig_func._parameterized_info) - return (self.register_anonymous, (self.orig_func, self.name)) - if (name := f"indexunary.{self.name}") in _STANDARD_OPERATOR_NAMES: - return name - # NOT COVERED - return (self._deserialize, (self.name, self.orig_func)) + self._udt_types = {} # {(dtype, thunk_type): DataType} + self._udt_ops = {} # {(dtype, thunk_type): TypedUserIndexUnaryOp} __call__ = TypedBuiltinIndexUnaryOp.__call__ + + +ParameterizedIndexUnaryOp._op_class = IndexUnaryOp diff --git a/graphblas/core/operator/monoid.py b/graphblas/core/operator/monoid.py index e3f218a90..e6d2ce07e 100644 --- a/graphblas/core/operator/monoid.py +++ b/graphblas/core/operator/monoid.py @@ -2,6 +2,8 @@ import re from collections.abc import Mapping +import numpy as np + from ... import _STANDARD_OPERATOR_NAMES, binary, monoid, op from ...dtypes import ( BOOL, @@ -18,11 +20,106 @@ lookup_dtype, ) from ...exceptions import check_status_carg -from .. import ffi, lib +from .. import _has_numba, ffi, lib from ..utils import libget -from .base import OpBase, ParameterizedUdf, TypedOpBase, _hasop +from .base import OpBase, ParameterizedUdf, TypedOpBase, _BinaryopJitDelegate, _hasop from .binary import BinaryOp, ParameterizedBinaryOp, TypedBuiltinBinaryOp +# Monoid names for which we can auto-generate UDT identities. +_BUILTIN_UDT_MONOIDS = {"plus", "times", "min", "max"} + + +def _is_builtin_binaryop(binaryop): + """True if ``binaryop`` wraps built-in GraphBLAS operators rather than a UDF. + + Most built-ins are constructed with neither a Python func nor a numba func. + That covers ``binary.numpy.float_power`` too, which is assembled out of + ``binary.pow``'s typed ops. A few built-ins (``any``, ``first``, ``pair``, + ``second``) carry a Python ``orig_func`` as their UDT fallback, so the func + test alone misses them; their typed ops are still ``TypedBuiltinBinaryOp``, + which is the authoritative signal. A UDF's typed ops are never + ``TypedBuiltinBinaryOp``. Under ``mapnumpy=True`` the mapped numpy binary + ops hold the built-in's typed ops and no funcs, so both tests agree they + are built-in; the numpy monoids themselves are mapped without passing + through ``Monoid._build``. + """ + return ( + binaryop.orig_func is None + and binaryop._numba_func is None + or any(type(typed) is TypedBuiltinBinaryOp for typed in binaryop._typed_ops.values()) + ) + + +def _scalar_identity(monoid_name, scalar_dtype): + """Return the identity for a single numeric (or bool) numpy dtype, or ``None``. + + ``None`` signals that we don't have an identity for this combination, so + the caller raises a clear error. Bool fields are treated as 0/1: ``min`` + is logical AND (identity ``True``), ``max`` is logical OR (identity + ``False``). + """ + if monoid_name == "plus": + return scalar_dtype.type(0) + if monoid_name == "times": + return scalar_dtype.type(1) + if monoid_name not in {"min", "max"}: + return None + # Not ``np.issubdtype(scalar_dtype, np.integer)``: numpy files timedelta64 + # under ``np.integer``, and ``np.iinfo`` rejects it. + if scalar_dtype.kind in "iu": + info = np.iinfo(scalar_dtype) + return info.max if monoid_name == "min" else info.min + if np.issubdtype(scalar_dtype, np.floating): + return scalar_dtype.type(np.inf if monoid_name == "min" else -np.inf) + if scalar_dtype == np.dtype(np.bool_): + return monoid_name == "min" + return None + + +def _udt_identity(monoid_name, dtype): + """Generate an identity value for a built-in monoid on a UDT. + + Returns a value suitable for assignment to ``Scalar.value``: a numpy + array for array UDTs, or a tuple of per-field identities for record + UDTs. Nested records recurse, producing a nested tuple that matches the + dtype's structure (which is what ``Scalar.value =`` expects). + """ + np_type = dtype.np_type if hasattr(dtype, "np_type") else dtype + return _udt_identity_np(monoid_name, np_type, dtype) + + +def _udt_identity_np(monoid_name, np_type, top_dtype): + if np_type.subdtype is not None: + base_dtype, shape = np_type.subdtype + val = _scalar_identity(monoid_name, base_dtype) + # The binary op compiled first, and min and max compile for arrays of real numbers only + if val is None: # pragma: no cover (safety) + raise KeyError( + f"monoid.{monoid_name} does not work with {top_dtype}: " + f"base dtype {base_dtype} is not numeric" + ) + return np.full(shape, val, dtype=base_dtype) + if np_type.names is not None: + vals = [] + for field_name in np_type.names: + field_dtype = np_type.fields[field_name][0] + if field_dtype.names is not None or field_dtype.subdtype is not None: + val = _udt_identity_np(monoid_name, field_dtype, top_dtype) + else: + val = _scalar_identity(monoid_name, field_dtype) + if val is None: + raise KeyError( + f"monoid.{monoid_name} does not work with {top_dtype}: " + f"field {field_name!r} has unsupported dtype {field_dtype}" + ) + vals.append(val) + return tuple(vals) + # The binary op compiled first, and lifted ops refuse a UDT neither a record nor an array + raise KeyError( # pragma: no cover (safety) + f"monoid.{monoid_name} does not work with {top_dtype}" + ) + + ffi_new = ffi.new @@ -67,10 +164,15 @@ def is_idempotent(self): __call__ = TypedBuiltinBinaryOp.__call__ -class TypedUserMonoid(TypedOpBase): +class TypedUserMonoid(_BinaryopJitDelegate, TypedOpBase): __slots__ = "binaryop", "identity" opclass = "Monoid" is_commutative = True + # Deliberately not ``_owns_gb_obj = True``: a ``GrB_Monoid`` holds a + # pointer into its underlying ``GrB_BinaryOp``, and Python's cyclic GC + # makes no guarantee about which side is finalized first. Freeing the + # monoid after its binary op is gone aborts SuiteSparse. The leak is + # bounded by the cache in ``Monoid.{_typed_ops,_udt_ops}``. def __init__(self, parent, name, type_, return_type, gb_obj, binaryop, identity): super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}") @@ -133,22 +235,25 @@ def __reduce__(self): name = f"monoid.{self.name}" if not self._anonymous and name in _STANDARD_OPERATOR_NAMES: # pragma: no cover return name - return (self._deserialize, (self.name, self.binaryop, self.identity, self._anonymous)) + return ( + self._deserialize, + (self.name, self.binaryop, self.identity, self._anonymous, self._is_idempotent), + ) @staticmethod - def _deserialize(name, binaryop, identity, anonymous): + def _deserialize(name, binaryop, identity, anonymous, is_idempotent=False): if anonymous: - return Monoid.register_anonymous(binaryop, identity, name) + return Monoid.register_anonymous(binaryop, identity, name, is_idempotent=is_idempotent) if (rv := Monoid._find(name)) is not None: return rv - return Monoid.register_new(name, binaryop, identity) + return Monoid.register_new(name, binaryop, identity, is_idempotent=is_idempotent) class Monoid(OpBase): """Takes two inputs and returns one output, all of the same data type. Built-in and registered Monoids are located in the ``graphblas.monoid`` namespace - as well as in the ``graphblas.ops`` combined namespace. + as well as in the ``graphblas.op`` combined namespace. """ __slots__ = "_binaryop", "_identity", "_is_idempotent" @@ -180,6 +285,24 @@ class Monoid(OpBase): def _build(cls, name, binaryop, identity, *, is_idempotent=False, anonymous=False): if type(binaryop) is not BinaryOp: raise TypeError(f"binaryop must be a BinaryOp, not {type(binaryop)}") + if _is_builtin_binaryop(binaryop): + # Rejected up front, before any object is created: a built-in binaryop + # cannot back a user monoid. Its typed ops are ``TypedBuiltinBinaryOp`` + # with no ``_monoid`` slot, so ``TypedUserMonoid`` would raise an opaque + # AttributeError partway through, having already repointed + # ``binary.._monoid`` at the half-built monoid. Even if that were + # patched over, SuiteSparse ignores the identity handed to + # ``GrB_Monoid_new`` whenever the built-in op already has a built-in + # monoid (max, min, plus, times, any), so the result would silently + # disagree with ``monoid[dtype].identity``. + raise TypeError( + "binaryop must be a user-defined BinaryOp, not the built-in " + f"binary.{binaryop.name}. Built-in binary operators already own their " + "monoid association (binary.plus resolves to monoid.plus, and so on). " + "To build a monoid with a different identity, register the function as " + "a new BinaryOp (BinaryOp.register_new or BinaryOp.register_anonymous) " + "and pass that." + ) if name is None: name = binaryop.name new_type_obj = cls( @@ -221,7 +344,7 @@ def _compile_udt(self, dtype, dtype2): dtype2 = dtype elif dtype != dtype2: raise TypeError( - "Monoid inputs must be the same dtype (got {dtype} and {dtype2}); " + f"Monoid inputs must be the same dtype (got {dtype} and {dtype2}); " "unable to coerce when using UDTs." ) if dtype in self._udt_types: @@ -230,12 +353,34 @@ def _compile_udt(self, dtype, dtype2): from ..scalar import Scalar ret_type = binaryop.return_type - identity = Scalar.from_value(self._identity, dtype=ret_type, is_cscalar=True) + if ret_type != dtype: + # Checked here because GrB_Monoid_new_UDT reports it as a DomainMismatch + # with no message. + raise TypeError( + f"monoid.{self.name} does not work with {dtype}: a monoid's binary operator " + f"must return its input type, {dtype}, not {ret_type}" + ) + identity_val = self._identity + if identity_val is None and self._binaryop is None and self.name in _BUILTIN_UDT_MONOIDS: + # Auto-generate the identity for built-in monoids on UDTs. Only built-in + # monoids have no ``_binaryop``: a user monoid named like one (after its + # op's function, or by ``name=``) must not get that identity. + identity_val = _udt_identity(self.name, ret_type) + if identity_val is None: + raise KeyError( + f"monoid.{self.name} does not work with {dtype}: " + "no identity value (provide one via Monoid.register_anonymous)" + ) + if ret_type._is_udt: + identity = Scalar(ret_type, is_cscalar=True) + identity.value = identity_val + else: + identity = Scalar.from_value(identity_val, dtype=ret_type, is_cscalar=True) new_monoid = ffi_new("GrB_Monoid*") status = lib.GrB_Monoid_new_UDT(new_monoid, binaryop.gb_obj, identity.gb_obj) check_status_carg(status, "Monoid", new_monoid[0]) op = TypedUserMonoid( - new_monoid, + self, self.name, dtype, ret_type, @@ -243,7 +388,7 @@ def _compile_udt(self, dtype, dtype2): binaryop, identity, ) - self._udt_types[dtype] = dtype + self._udt_types[dtype] = ret_type self._udt_ops[dtype] = op return op @@ -256,10 +401,10 @@ def register_anonymous(cls, binaryop, identity, name=None, *, is_idempotent=Fals Parameters ---------- - binaryop: BinaryOp or ParameterizedBinaryOp + binaryop : BinaryOp or ParameterizedBinaryOp The binary operator of the monoid, which should be able to use the same dtype for both inputs and the output. - identity: scalar or Mapping + identity : scalar or Mapping The identity of the monoid such that ``op(x, identity) == x`` for any x. ``identity`` may also be a mapping from dtype to scalar. name : str, optional @@ -290,10 +435,10 @@ def register_new(cls, name, binaryop, identity, *, is_idempotent=False, lazy=Fal The name of the operator. This will show up as ``gb.monoid.{name}``. The name may contain periods, ".", which will result in nested objects such as ``gb.monoid.x.y.z`` for name ``"x.y.z"``. - binaryop: BinaryOp or ParameterizedBinaryOp + binaryop : BinaryOp or ParameterizedBinaryOp The binary operator of the monoid, which should be able to use the same dtype for both inputs and the output. - identity: scalar or Mapping + identity : scalar or Mapping The identity of the monoid such that ``op(x, identity) == x`` for any x. ``identity`` may also be a mapping from dtype to scalar. is_idempotent : bool, default False @@ -316,7 +461,12 @@ def register_new(cls, name, binaryop, identity, *, is_idempotent=False, lazy=Fal if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "binaryop": binaryop, "identity": identity}, + { + "name": name, + "binaryop": binaryop, + "identity": identity, + "is_idempotent": is_idempotent, + }, ) elif type(binaryop) is ParameterizedBinaryOp: monoid = ParameterizedMonoid(name, binaryop, identity, is_idempotent=is_idempotent) @@ -348,11 +498,30 @@ def __init__(self, name, binaryop=None, identity=None, *, is_idempotent=False, a self._udt_ops = {} # {dtype: TypedUserMonoid} def __reduce__(self): - if self._anonymous: - return (self.register_anonymous, (self._binaryop, self._identity, self.name)) - if (name := f"monoid.{self.name}") in _STANDARD_OPERATOR_NAMES: + if not self._anonymous and (name := f"monoid.{self.name}") in _STANDARD_OPERATOR_NAMES: return name - return (self._deserialize, (self.name, self._binaryop, self._identity)) + # Carry ``is_idempotent`` through pickle: ``register_anonymous`` / + # ``register_new`` take it as keyword-only, and the inherited + # ``OpBase._deserialize`` can't pass keyword args, so route through a + # dedicated deserializer here. + return ( + Monoid._deserialize_named_monoid, + ( + self.name, + self._binaryop, + self._identity, + self._anonymous, + self._is_idempotent, + ), + ) + + @staticmethod + def _deserialize_named_monoid(name, binaryop, identity, anonymous, is_idempotent): + if anonymous: + return Monoid.register_anonymous(binaryop, identity, name, is_idempotent=is_idempotent) + if (rv := Monoid._find(name)) is not None: + return rv + return Monoid.register_new(name, binaryop, identity, is_idempotent=is_idempotent) @property def binaryop(self): @@ -422,6 +591,13 @@ def _initialize(cls): any_._identity = 0 any_._udt_types = {} any_._udt_ops = {} + # Enable element-wise monoids on UDTs (``plus``, ``times``, ``min``, ``max``). + if _has_numba: + for name in _BUILTIN_UDT_MONOIDS: + mon = getattr(monoid, name, None) + if mon is not None: + mon._udt_types = {} + mon._udt_ops = {} cls._initialized = True commutes_to = TypedBuiltinMonoid.commutes_to diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 6de4fa89a..1c1138ae5 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -1,17 +1,144 @@ import inspect -from ... import _STANDARD_OPERATOR_NAMES, select +from ... import select from ...dtypes import BOOL, UINT64 -from ...exceptions import check_status_carg -from .. import _has_numba, ffi, lib -from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op, _deserialize_parameterized +from .. import _has_numba +from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op from .indexunary import IndexUnaryOp, TypedBuiltinIndexUnaryOp if _has_numba: - import numba + from .udf import _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper - from .base import _get_udt_wrapper -ffi_new = ffi.new + +_VALUE_OPS = {"valueeq", "valuene", "valuelt", "valuele", "valuegt", "valuege"} + + +def _value_select_in_own_type(op, dtype, thunk, given): + """Return ``op`` and ``thunk`` for ``select`` on built-in ``dtype``, as a value op of ``dtype``. + + SuiteSparse:GraphBLAS 8.0 to 10.5.1 runs ``select`` with a value + op typed for another type than the input's by reading the thunk's bytes as + the input's type, where the spec casts both to the op's type: on an INT8 + vector, ``valueeq`` with ``1.0`` compared with 0, and ``valuelt`` with + ``300`` compared with 44. ``apply`` with the same op is right. So the + comparison is made in the input's own type, with a thunk that selects the + same elements: for integers, ``x < 2.5`` is ``x < 3``, and a thunk no + element can equal, or that every element passes, selects nothing or + everything; for floats, the thunk is rounded toward the side that keeps + the comparison exact; for complex numbers, which have no order, a thunk + that does not convert exactly equals no element. ``given`` is the thunk as + the caller gave it, whose value is exact where ``thunk`` may have rounded + it (an int beyond 64 bits is FP64). Other ops and types are returned + unchanged. + """ + # MAINT 2026-10-28: works around the SuiteSparse:GraphBLAS bug above (seen in + # 8.0.2 through 10.5.1; fixed in 10.5.2 by DrTimothyAldenDavis/GraphBLAS#463). + # The fix casts both to the op's type, which can round the thunk: beside FP32 + # values, an INT64 thunk of 2**53 + 1 is FP64 2**53 and equals 2**53. Once + # every supported version has the fix, keep only what makes the comparison exact. + if op.type is dtype or dtype._is_udt or op.parent.name not in _VALUE_OPS: + return op, thunk + import math + + import numpy as np + + from ... import indexunary + from ..scalar import Scalar + + parent = op.parent + name = parent.name + if ( + parent is not getattr(select, name) + and parent is not getattr(indexunary, name) + or dtype.np_type.kind not in "biufc" + or thunk._is_empty + ): + return op, thunk + value = given.value if type(given) is Scalar else given + if isinstance(value, (np.generic, np.ndarray)): + value = value.item() + np_type = dtype.np_type + if isinstance(value, complex) and np_type.kind != "c": + if value.imag != 0: + if name not in {"valueeq", "valuene"}: # pragma: no cover (no complex order) + return op, thunk + value = math.nan # equal to nothing, as a complex number is to a real one + else: + value = value.real + if np_type.kind == "c": + # Only valueeq and valuene take complex values. The thunk is compared + # exactly, as for floats: one that does not convert exactly (2**53 + 1 + # beside complex64, or a NaN part) equals no element, as NaN does. + with np.errstate(over="ignore", invalid="ignore"): + near = np_type.type(value) + real, imag = (value.real, value.imag) if isinstance(value, complex) else (value, 0) + exact = float(near.real) == real and float(near.imag) == imag + new_name, new_value = name, near if exact else np_type.type(complex(math.nan, math.nan)) + elif np_type.kind == "f": + nan, ftype = np_type.type(np.nan), np_type.type + if value != value: # NaN: only ``!=`` holds, for every element + new_name, new_value = name, nan + else: + # A finite value too large for ftype is infinity there, which still + # compares right with every element, so numpy need not warn. + with np.errstate(over="ignore"): + near = ftype(value) + above = near if float(near) >= value else np.nextafter(near, ftype(math.inf)) + below = near if float(near) <= value else np.nextafter(near, ftype(-math.inf)) + new_name = name + if name in {"valueeq", "valuene"}: + # No element equals NaN: == selects nothing and != everything. + new_value = near if float(near) == value else nan + else: + new_value = above if name in {"valuelt", "valuege"} else below + else: + if np_type.kind == "b": + lo, hi = 0, 1 + else: + info = np.iinfo(np_type) + lo, hi = int(info.min), int(info.max) + keep_all, keep_none = ("valuege", lo), ("valuelt", lo) + if value != value: + new_name, new_value = keep_all if name == "valuene" else keep_none + elif name in {"valueeq", "valuene"}: + # An int is whole, and math.isfinite cannot take one too large for a float. + exact = type(value) is int or (math.isfinite(value) and value == math.floor(value)) + if exact and lo <= value <= hi: + new_name, new_value = name, int(value) + else: + new_name, new_value = keep_all if name == "valuene" else keep_none + else: + # Integers: x < t is x < ceil(t), x <= t is x <= floor(t), and so on. + if type(value) is float and math.isinf(value): + bound = value + elif name in {"valuelt", "valuege"}: + bound = math.ceil(value) + else: + bound = math.floor(value) + passes_all = { + "valuelt": bound > hi, + "valuele": bound >= hi, + "valuegt": bound < lo, + "valuege": bound <= lo, + }[name] + passes_none = { + "valuelt": bound <= lo, + "valuele": bound < lo, + "valuegt": bound >= hi, + "valuege": bound > hi, + }[name] + if passes_all: + new_name, new_value = keep_all + elif passes_none: + new_name, new_value = keep_none + else: + new_name, new_value = name, int(bound) + if np_type.kind == "b": + new_value = bool(new_value) + module = select if isinstance(parent, SelectOp) else indexunary + new_op = getattr(module, new_name)[dtype] + new_thunk = Scalar.from_value(new_value, dtype, is_cscalar=thunk._is_cscalar, name="") + return new_op, new_thunk class TypedBuiltinSelectOp(TypedOpBase): @@ -29,6 +156,11 @@ def __call__(self, val, thunk=None): class TypedUserSelectOp(TypedOpBase): __slots__ = () opclass = "SelectOp" + # The underlying object is a ``GrB_IndexUnaryOp``. Instances built by + # ``_from_indexunary`` borrow that handle and set ``_gb_obj_owner``, which + # suppresses the free here. Freeing from both sides is a use-after-free that + # no test catches: the freed handle usually still reads as valid. + _owns_gb_obj = True def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None): super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2) @@ -61,22 +193,6 @@ def _call(self, *args, **kwargs): sel._parameterized_info = (self, args, kwargs) return SelectOp.register_anonymous(sel, self.name, is_udt=self._is_udt) - def __reduce__(self): - # NOT COVERED - name = f"select.{self.name}" - if not self._anonymous and name in _STANDARD_OPERATOR_NAMES: - return name - return (self._deserialize, (self.name, self.func, self._anonymous)) - - @staticmethod - def _deserialize(name, func, anonymous): - # NOT COVERED - if anonymous: - return SelectOp.register_anonymous(func, name, parameterized=True) - if (rv := SelectOp._find(name)) is not None: - return rv - return SelectOp.register_new(name, func, parameterized=True) - class SelectOp(OpBase): """Identical to an :class:`IndexUnaryOp `, @@ -116,6 +232,13 @@ def _from_indexunary(cls, iop): t.return_type, t.gb_obj, ) + # Borrow the IndexUnaryOp's allocation instead of making a + # second one. Holding ``t`` keeps that handle alive for as long + # as this SelectOp can use it: ``iop`` is a temporary in + # ``register_anonymous``, so without this the handle is freed + # the moment it is collected and every call raises + # UninitializedObject. + op._gb_obj_owner = t else: op = cls._typed_class( obj, @@ -137,37 +260,24 @@ def _compile_udt(self, dtype, dtype2): dtypes = (dtype, dtype2) if dtypes in self._udt_types: return self._udt_ops[dtypes] + from ..udt import _refuse_half + + _refuse_half(self, dtype, dtype2) if self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtypes} types") # It would be nice if we could reuse compiling done for IndexUnaryOp numba_func = self._numba_func sig = (dtype.numba_type, UINT64.numba_type, UINT64.numba_type, dtype2.numba_type) - numba_func.compile(sig) # Should we catch and give additional error message? + _compile_udf_for_udt( + numba_func, sig, op_kind="select", op_name=self.name, dtypes=(dtype, dtype2) + ) select_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, BOOL, dtype, dtype2, include_indexes=True ) - - select_wrapper = numba.cfunc(wrapper_sig, nopython=True)(select_wrapper) - new_select = ffi_new("GrB_IndexUnaryOp*") - check_status_carg( - lib.GrB_IndexUnaryOp_new( - new_select, select_wrapper.cffi, BOOL._carg, dtype._carg, dtype2._carg - ), - "IndexUnaryOp", - new_select[0], - ) - op = TypedUserSelectOp( - self, - self.name, - dtype, - BOOL, - new_select[0], - dtype2=dtype2, + return _finalize_udt_op( + self, dtype, dtype2, BOOL, select_wrapper, wrapper_sig, TypedUserSelectOp ) - self._udt_types[dtypes] = BOOL - self._udt_ops[dtypes] = op - return op @classmethod def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): @@ -181,8 +291,8 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals func : FunctionType The function to compile. For all current backends, this must be able to be compiled with ``numba.njit``. - ``func`` takes four input parameters--any dtype, int64, int64, - any dtype and returns boolean. The first argument (any dtype) is + ``func`` takes four input parameters (any dtype, int64, int64, + any dtype) and returns boolean. The first argument (any dtype) is the value of the input Matrix or Vector, the second argument (int64) is the row index of the Matrix or the index of the Vector, the third argument (int64) is the column index of the Matrix or 0 for a Vector, @@ -232,8 +342,8 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal func : FunctionType The function to compile. For all current backends, this must be able to be compiled with ``numba.njit``. - ``func`` takes four input parameters--any dtype, int64, int64, - any dtype and returns boolean. The first argument (any dtype) is + ``func`` takes four input parameters (any dtype, int64, int64, + any dtype) and returns boolean. The first argument (any dtype) is the value of the input Matrix or Vector, the second argument (int64) is the row index of the Matrix or the index of the Vector, the third argument (int64) is the column index of the Matrix or 0 for a Vector, @@ -322,19 +432,10 @@ def __init__( self.is_positional = is_positional self._is_udt = is_udt if is_udt: - # NOT COVERED - self._udt_types = {} # {dtype: DataType} - self._udt_ops = {} # {dtype: TypedUserIndexUnaryOp} - - def __reduce__(self): - if self._anonymous: - if hasattr(self.orig_func, "_parameterized_info"): - # NOT COVERED - return (_deserialize_parameterized, self.orig_func._parameterized_info) - return (self.register_anonymous, (self.orig_func, self.name)) - if (name := f"select.{self.name}") in _STANDARD_OPERATOR_NAMES: - return name - # NOT COVERED - return (self._deserialize, (self.name, self.orig_func)) + self._udt_types = {} # {(dtype, thunk_type): DataType} + self._udt_ops = {} # {(dtype, thunk_type): TypedUserSelectOp} __call__ = TypedBuiltinSelectOp.__call__ + + +ParameterizedSelectOp._op_class = SelectOp diff --git a/graphblas/core/operator/semiring.py b/graphblas/core/operator/semiring.py index a8d18f1bf..f9d8603db 100644 --- a/graphblas/core/operator/semiring.py +++ b/graphblas/core/operator/semiring.py @@ -18,7 +18,15 @@ ) from ...exceptions import check_status_carg from .. import _supports_udfs, ffi, lib -from .base import _SS_OPERATORS, OpBase, ParameterizedUdf, TypedOpBase, _call_op, _hasop +from .base import ( + _SS_OPERATORS, + OpBase, + ParameterizedUdf, + TypedOpBase, + _BinaryopJitDelegate, + _call_op, + _hasop, +) from .binary import BinaryOp, ParameterizedBinaryOp from .monoid import Monoid, ParameterizedMonoid @@ -81,9 +89,16 @@ def type2(self): return self.type if self._type2 is None else self._type2 -class TypedUserSemiring(TypedOpBase): +class TypedUserSemiring(_BinaryopJitDelegate, TypedOpBase): + # ``_jit_c_info`` delegates to ``self.binaryop`` (the multiplier inside + # mxm). To inspect the additive monoid's kernel directly, use + # ``semi.monoid.jit_c_source``. __slots__ = "monoid", "binaryop" opclass = "Semiring" + # Deliberately not ``_owns_gb_obj = True``: see ``TypedUserMonoid``. + # Same ordering hazard: a ``GrB_Semiring`` holds pointers into its + # ``GrB_Monoid`` and ``GrB_BinaryOp``, and the cache in + # ``Semiring.{_typed_ops,_udt_ops}`` bounds the leak. def __init__(self, parent, name, type_, return_type, gb_obj, monoid, binaryop, dtype2=None): super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2) @@ -160,7 +175,7 @@ class Semiring(OpBase): and the Monoid taking the place of the standard addition operator. Built-in and registered Semirings are located in the ``graphblas.semiring`` namespace - as well as in the ``graphblas.ops`` combined namespace. + as well as in the ``graphblas.op`` combined namespace. """ __slots__ = "_monoid", "_binaryop" @@ -207,8 +222,24 @@ class Semiring(OpBase): def _build(cls, name, monoid, binaryop, *, anonymous=False): if type(monoid) is not Monoid: raise TypeError(f"monoid must be a Monoid, not {type(monoid)}") - if type(binaryop) is not BinaryOp: - raise TypeError(f"binaryop must be a BinaryOp, not {type(binaryop)}") + # SuiteSparse:GraphBLAS allows the multiplier of a Semiring to be + # any GrB_BinaryOp, "including those based on a GxB_IndexBinaryOp" + # (GraphBLAS.h, GrB_Semiring section). A ``_BoundIndexBinaryOp`` + # (``ibo[dtype](theta)``) is already exactly such a typed + # GrB_BinaryOp (built via ``GxB_BinaryOp_new_IndexOp``), so accept it + # directly and build one Semiring for its specific input/output type + # pair, rather than the regular per-typed-op iteration that + # polymorphic ``BinaryOp`` parents go through. ``ibo(theta)`` is such + # a parent, a BinaryOp subclass that binds theta per type as it runs. + from .indexbinary import _BoundIndexBinaryOp + + if isinstance(binaryop, _BoundIndexBinaryOp): + return cls._build_from_bound_indexbinary(name, monoid, binaryop, anonymous=anonymous) + if not isinstance(binaryop, BinaryOp): + raise TypeError( + f"binaryop must be a BinaryOp or a bound IndexBinaryOp " + f"(``ibo[dtype](theta)``), not {type(binaryop)}" + ) if name is None: name = f"{monoid.name}_{binaryop.name}".replace(".", "_") new_type_obj = cls(name, monoid, binaryop, anonymous=anonymous) @@ -243,6 +274,46 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): new_type_obj._add(op) return new_type_obj + @classmethod + def _build_from_bound_indexbinary(cls, name, monoid, bound_binop, *, anonymous=False): + """Build a Semiring whose multiplier is a bound IndexBinaryOp. + + ``bound_binop`` is already a single typed ``GrB_BinaryOp``, so this + creates exactly one ``GrB_Semiring`` (for the bound op's input + type), not one per polymorphic dispatch entry. The additive monoid + must accept the bound op's return type (auto-lifts for UDT return + types via ``monoid[udt]``). + """ + if name is None: + name = f"{monoid.name}_{bound_binop.parent.name}".replace(".", "_") + return_type = bound_binop.return_type + if return_type not in monoid: + raise TypeError( + f"Monoid {monoid.name!r} doesn't accept the bound IndexBinaryOp's " + f"return type {return_type!r}. The additive monoid and the " + f"multiplier must agree on the type the semiring reduces over." + ) + typed_monoid = monoid[return_type] + new_semiring = ffi_new("GrB_Semiring*") + check_status_carg( + lib.GrB_Semiring_new(new_semiring, typed_monoid.gb_obj, bound_binop.gb_obj), + "Semiring", + new_semiring[0], + ) + new_type_obj = cls(name, monoid, bound_binop, anonymous=anonymous) + typed_semiring = TypedUserSemiring( + new_type_obj, + name, + bound_binop.type, + typed_monoid.return_type, + new_semiring[0], + typed_monoid, + bound_binop, + dtype2=bound_binop._type2, + ) + new_type_obj._add(typed_semiring) + return new_type_obj + def _compile_udt(self, dtype, dtype2): if dtype2 is None: dtype2 = dtype @@ -256,7 +327,7 @@ def _compile_udt(self, dtype, dtype2): status = lib.GrB_Semiring_new(new_semiring, monoid.gb_obj, binaryop.gb_obj) check_status_carg(status, "Semiring", new_semiring[0]) op = TypedUserSemiring( - new_semiring, + self, self.name, dtype, ret_type, @@ -265,7 +336,7 @@ def _compile_udt(self, dtype, dtype2): binaryop, dtype2=dtype2, ) - self._udt_types[dtypes] = dtype + self._udt_types[dtypes] = ret_type self._udt_ops[dtypes] = op return op @@ -279,8 +350,12 @@ def register_anonymous(cls, monoid, binaryop, name=None): ---------- monoid : Monoid or ParameterizedMonoid The monoid of the semiring (like "plus" in the default "plus_times" semiring). - binaryop : BinaryOp or ParameterizedBinaryOp - The binaryop of the semiring (like "times" in the default "plus_times" semiring). + binaryop : BinaryOp, ParameterizedBinaryOp, or bound IndexBinaryOp + The multiplier of the semiring (like "times" in the default + "plus_times" semiring). An IndexBinaryOp bound to theta, + ``ibo(theta)``, is a BinaryOp. One bound at a type, + ``ibo[dtype](theta)``, is accepted as well; the resulting semiring + is monomorphic in the bound op's input/output types. name : str, optional The name of the operator. This *does not* show up as ``gb.semiring.{name}``. @@ -305,8 +380,11 @@ def register_new(cls, name, monoid, binaryop, *, lazy=False): such as ``gb.semiring.x.y.z`` for name ``"x.y.z"``. monoid : Monoid or ParameterizedMonoid The monoid of the semiring (like "plus" in the default "plus_times" semiring). - binaryop : BinaryOp or ParameterizedBinaryOp - The binaryop of the semiring (like "times" in the default "plus_times" semiring). + binaryop : BinaryOp, ParameterizedBinaryOp, or bound IndexBinaryOp + The multiplier of the semiring. ``ibo(theta)`` is a BinaryOp, and + an IndexBinaryOp bound at a type, ``ibo[dtype](theta)``, is also + accepted; the resulting semiring is monomorphic in the bound op's + input/output types. lazy : bool, default False If False (the default), then the function will be automatically compiled for builtin types (unless ``is_udt`` is True). diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py new file mode 100644 index 000000000..30b3e7ff4 --- /dev/null +++ b/graphblas/core/operator/udf.py @@ -0,0 +1,682 @@ +"""Compiling a user's function (UDF) on UDTs into a GraphBLAS operator. + +GraphBLAS calls an operator as a C function on pointers to raw elements. A UDF +written for a record or array UDT takes Numba records and arrays instead, and +may return a tuple, an array or a number to fill the result's element. The +functions here generate the Numba cfunc between the two: it reads each operand +through its pointer, calls the UDF, and writes what it returns into the output +element. A return that cannot fill that element is an error when the op is +typed, not inside the cfunc. ``_finalize_udt_op`` then registers the cfunc as +a GraphBLAS operator. + +Built-in operators that lift to UDTs are generated by ``udt_codegen``, which uses +these too. ``_bool_to_int8``, the numba#5395 workaround, is here as well: every +Numba UDF on BOOL uses it, not only those on UDTs. + +Importing this module needs Numba. +""" + +import re + +import numba +import numpy as np +from numba import NumbaError + +from ...dtypes import BOOL, INT8, UINT64, lookup_dtype +from ...exceptions import UdfParseError, check_status_carg +from .. import ffi, lib +from ..dtypes import _sample_values +from ..udt import _iter_record_leaves +from .base import _GB_OBJ_C_TYPENAME +from .udt_codegen import _compile_codegen + + +def _bool_to_int8(dtype): + """Return INT8 if ``dtype`` is BOOL, else the dtype unchanged. + + Numba can't compile cfuncs that read or write ``CPointer(boolean)`` + (errors like ``cannot store i1 to i8*`` / ``cond is not i1: i8``); see + numba/numba#5395. Routing BOOL through INT8 sidesteps that; GraphBLAS + coerces back at the cfunc boundary. + + MAINT 2026-10-18: still hits on Numba 0.68. Re-test periodically and + drop the INT8 routing when upstream is fixed. + """ + return INT8 if dtype == BOOL else dtype + + +def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, typed_user_cls): + """Compile the cfunc, allocate the ``GrB`` op, wrap it, and cache it. + + Shared tail for ``_compile_udt`` in UnaryOp / BinaryOp / IndexUnaryOp / + SelectOp. Looks up the SuiteSparse handle type and ``_new`` symbol + from ``typed_user_cls.opclass``. ``dtype2`` is ``None`` for unary + ops; the rest pass both. Returns the cached ``TypedUser*Op``. + """ + try: + wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) + except NumbaError as exc: + dtypes_str = str(dtype) if dtype2 is None else f"{dtype}, {dtype2}" + raise UdfParseError( + f"{parent_op._modname}.{parent_op.name} does not work with ({dtypes_str}): " + f"{_summarize_numba_typing_error(exc)}" + ) from exc + c_typename = _GB_OBJ_C_TYPENAME[typed_user_cls.opclass] + error_label = c_typename.removeprefix("GrB_").removeprefix("GxB_") + gb_obj = ffi.new(f"{c_typename}*") + new_func = getattr(lib, f"{c_typename}_new") + if dtype2 is None: + check_status_carg( + new_func(gb_obj, wrapper.cffi, ret_type._carg, dtype._carg), + error_label, + gb_obj[0], + ) + op = typed_user_cls(parent_op, parent_op.name, dtype, ret_type, gb_obj[0]) + key = dtype + else: + check_status_carg( + new_func(gb_obj, wrapper.cffi, ret_type._carg, dtype._carg, dtype2._carg), + error_label, + gb_obj[0], + ) + op = typed_user_cls(parent_op, parent_op.name, dtype, ret_type, gb_obj[0], dtype2=dtype2) + key = (dtype, dtype2) + parent_op._udt_types[key] = ret_type + parent_op._udt_ops[key] = op + return op + + +def _compile_udf_for_udt(numba_func, sig, *, op_kind, op_name, dtypes): + """Compile ``sig`` and re-raise Numba compilation errors as ``UdfParseError``. + + Catches the full ``NumbaError`` hierarchy (TypingError, LoweringError, + UnsupportedError, ...) so any compilation failure produces a + single-line UDT diagnostic. + """ + try: + numba_func.compile(sig) + except NumbaError as exc: + dtypes_str = ", ".join(str(d) for d in dtypes) + snippet = _summarize_numba_typing_error(exc) + raise UdfParseError( + f"{op_kind}.{op_name} does not work with ({dtypes_str}): {snippet}" + ) from exc + + +# Numba prefixes its actionable diagnostic line with one of these. +_NUMBA_DIAG_PREFIXES = ( + "No implementation of function", + "No conversion from", + "Field ", + "Cannot infer", + "Operator Overload", + "Invalid use of", + "use of undeclared", + "Untyped global", +) + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") + + +def _summarize_numba_typing_error(exc): + """Pull the most actionable lines out of a Numba TypingError. + + Numba's "No implementation of function ... found for signature:" + diagnostic puts the signature on the line *below* the prefix; the + signature is the actionable part. When the matched line ends with a + ``:``, append the next non-empty line so the user sees both. + """ + lines = [_ANSI_RE.sub("", line).strip() for line in str(exc).splitlines()] + for i, line in enumerate(lines): + if line.startswith(_NUMBA_DIAG_PREFIXES): + if line.endswith(":"): + for follow in lines[i + 1 :]: + if follow: + return f"{line} {follow}" + return line + for line in lines: + if line: + return line + return "Numba could not compile the function for these input types" + + +def _resolve_udt_return_type(numba_ret_type, *dtypes): + """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. + + When a UDF returns a tuple, Numba infers ``Tuple(...)`` rather than a + Record type. Match by field count, preferring a candidate whose field + types align with the Tuple's element types. An array the UDF builds + (``x + y``) is typed without its extents, so it takes the type of the + array-UDT operand with the same element type and rank. + """ + try: + return lookup_dtype(numba_ret_type) + except (ValueError, TypeError): + pass + for d in dtypes: + # Bytes and the like have a Numba type no DataType is registered for; + # an operand of that type, returned as it is, keeps its UDT. + if ( + d._is_udt + and d.np_type.names is None + and d.np_type.subdtype is None + and d.numba_type == numba_ret_type + ): + return d + if isinstance(numba_ret_type, numba.core.types.BaseTuple): + n = len(numba_ret_type.types) + + def _leaves(d): + """Return ``(python_path, c_path, leaf_dtype)`` for each leaf of a record UDT.""" + return list(_iter_record_leaves(d.np_type)) + + # Match by total leaf count (flat for shallow records, total + # leaves across nesting for nested records). + same_arity = [ + d for d in dtypes if d._is_udt and d.np_type.names is not None and len(_leaves(d)) == n + ] + # Prefer a UDT whose leaf types match the tuple elements element-wise. + for d in same_arity: + leaf_dtypes = [leaf for _py, _c, leaf in _leaves(d)] + try: + expected = [lookup_dtype(t).numba_type for t in leaf_dtypes] + except (ValueError, TypeError): # pragma: no cover (safety) + continue # no record UDT has such a leaf: a void field fails to register + if all(et == tt for et, tt in zip(expected, numba_ret_type.types, strict=True)): + return d + # No perfect match; fall back to the first arity-compatible UDT. + if same_arity: + return same_arity[0] + # Tuple return whose arity matches no input UDT: most likely the + # user is returning the wrong number of fields. List the candidate + # arities so the fix is obvious. + record_arities = sorted( + {len(_leaves(d)) for d in dtypes if d._is_udt and d.np_type.names is not None} + ) + if record_arities: + expected = " or ".join(str(a) for a in record_arities) + raise UdfParseError( + f"UDT UDF returned a tuple of length {n}; expected {expected} " + f"to match one of the input record UDTs." + ) + # All UDT inputs are array UDTs. Tuples don't map to those: the + # function should return a numpy array of the right shape (or a + # scalar) instead. + array_inputs = [d for d in dtypes if d._is_udt and d.np_type.subdtype is not None] + if array_inputs: + shape = array_inputs[0].np_type.subdtype[1] + raise UdfParseError( + f"UDT UDF returned a tuple of length {n}, but inputs are array UDTs of " + f"shape {shape}. Return a numpy array (e.g., ``np.array(...)``) or a " + f"scalar; tuple returns are only matched to record UDTs." + ) + elif isinstance(numba_ret_type, numba.core.types.Array): + # Numba types an array the UDF builds as a plain ``Array``: element + # type and rank, but no extents, which only an operand can supply. + # (An operand returned as-is keeps its ``NestedArray`` type, extents + # and all, and ``lookup_dtype`` resolved it above.) + array_inputs = list( + dict.fromkeys( # dedup: ``x`` and ``y`` are often the same UDT + d for d in dtypes if d._is_udt and d.np_type.subdtype is not None + ) + ) + if isinstance(numba_ret_type, numba.core.types.NestedArray): + # Extents survived, so ``lookup_dtype`` above named a UDT for + # them unless SuiteSparse refused to create it: builds without + # variable-length arrays cap a UDT at 128 bytes before SS 9.0 + # and 1024 from then on. Matching those extents to a same-rank + # operand would hand SuiteSparse an element of the wrong size. + shaped_like = [] + else: + shaped_like = [ + d + for d in array_inputs + if d.numba_type.dtype == numba_ret_type.dtype + and d.numba_type.ndim == numba_ret_type.ndim + ] + if len(shaped_like) == 1: + return shaped_like[0] + if shaped_like: + # Operands that differ only in length are indistinguishable + # here. Guessing would hand SuiteSparse an element of the + # wrong size, so say so. + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches more than one " + f"input array UDT ({', '.join(str(d) for d in shaped_like)}). " + f"Return one of the operands, or make the operands the same type." + ) + # An array UDT went in and an array came out, but not one that + # fits: name the mismatch rather than fall through to the generic + # "unsupported type", whose advice the user already followed. + if array_inputs: + described = "; ".join( + f"{d} elements are {d.numba_type.dtype} with shape {d.numba_type.shape}" + for d in array_inputs + ) + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches no input array " + f"UDT: {described}. Return an array matching one of those, or one of " + f"the operands." + ) + raise UdfParseError( + f"UDT UDF returned an unsupported type {numba_ret_type!r}. " + f"Return a scalar, a tuple matching a record UDT's fields, or a numpy array " + f"matching an array UDT's shape." + ) + + +def _array_udt_view(dtype): + """Return ``(base_element_numba_type, shape)`` for an array UDT. + + Array UDTs are addressed as a ``carray`` over their base elements, in + the UDT's declared shape, rather than as Numba's ``NestedArray``. Numba + models a ``NestedArray`` *value* as an array descriptor (data pointer, + shape, strides, ...), so loading or storing one through a ``CPointer`` + moves the descriptor rather than the element payload, corrupting + whatever follows it (and overrunning the element outright once the + descriptor is the wider of the two). + """ + nested = dtype.numba_type + return nested.dtype, nested.shape + + +def _input_operand(dtype, var): + """Return ``(setup_line, deref_expr, ptr_arg_type)`` for one input operand. + + - ``setup_line`` is the optional ``var = numba.carray(var_ptr, n)`` line + to add to the wrapper body, where ``n`` is 1 for a record UDT and the + element's shape for an array UDT (empty for non-UDT cases). + - ``deref_expr`` is the value to pass to ``numba_func`` for this operand. + - ``ptr_arg_type`` is the Numba ``CPointer(...)`` type for the wrapper signature. + """ + nt = numba.types + if dtype._is_udt: + if dtype.np_type.subdtype is None: + return ( + f" {var} = numba.carray({var}_ptr, 1)\n", + f"{var}[0]", + nt.CPointer(dtype.numba_type), + ) + base_numba, shape = _array_udt_view(dtype) + return ( + f" {var} = numba.carray({var}_ptr, {shape})\n", + var, + nt.CPointer(base_numba), + ) + if dtype == BOOL: + # Numba can't compile bool ptrs (numba/numba#5395); expose them + # as int8 and cast on deref. + # MAINT 2026-10-18: still hits on Numba 0.68; re-test periodically. + return "", f"bool({var}_ptr[0])", nt.CPointer(INT8.numba_type) + return "", f"{var}_ptr[0]", nt.CPointer(dtype.numba_type) + + +def _output_handler(return_type, numba_ret_type): + """Return ``(setup_line, ret_ptr_type, write_kind, write_info)``. + + ``write_kind`` is ``"record_fields"`` when the UDF returns a Tuple to be + unpacked into a record output, ``"array_elements"`` for an array UDT + output, otherwise ``"direct"``. ``write_info`` is one + ``(access_path, is_array)`` pair per leaf for ``"record_fields"`` and a + ``(BL, BR, zname)`` 3-tuple for ``"direct"``; ``"array_elements"`` needs + none, since ``setup_line`` already binds ``z`` to the right shape. + """ + nt = numba.types + ztype = INT8 if return_type == BOOL else return_type + ret_ptr_type = nt.CPointer(ztype.numba_type) + + if ( + numba_ret_type is not None + and isinstance(numba_ret_type, numba.core.types.BaseTuple) + and return_type._is_udt + and return_type.np_type.names is not None + ): + # ``write_info`` pairs each leaf field's Python access path + # (``"['a']"`` for flat, ``"['outer']['inner_a']"`` for nested) + # with whether that leaf is array-typed. The wrapper iterates + # these to write the flat tuple ``_result`` back leaf-by-leaf. + leaves = tuple( + (py, d.subdtype is not None) for py, _c, d in _iter_record_leaves(return_type.np_type) + ) + return ( + " z = numba.carray(z_ptr, 1)\n", + ret_ptr_type, + "record_fields", + leaves, + ) + if return_type._is_udt: + if return_type.np_type.subdtype is None: + return ( + " z = numba.carray(z_ptr, 1)\n", + ret_ptr_type, + "direct", + ("", "", "z[0]"), + ) + base_numba, shape = _array_udt_view(return_type) + return ( + f" z = numba.carray(z_ptr, {shape})\n", + nt.CPointer(base_numba), + "array_elements", + None, + ) + if return_type == BOOL: + return "", ret_ptr_type, "direct", ("bool(", ")", "z_ptr[0]") + return "", ret_ptr_type, "direct", ("", "", "z_ptr[0]") + + +def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): + """Assemble the Python source for a UDT cfunc wrapper. + + For record returns, the wrapper writes leaf-by-leaf; Numba does not + compile a nested-tuple assignment to an outer record field. Array + returns and array-typed record leaves slice-assign into the caller's + buffer, so a wrong-shape return cannot overrun it. The error does not reach + the caller either: the resulting ``ValueError`` is raised inside a + cfunc, which Numba prints and swallows, so the write stops there and + every leaf from that point on keeps whatever SuiteSparse had in the + buffer. :func:`_check_udf_fills_output` rejects such a return when the + op is typed, but it cannot always learn the shape, so this stays the + backstop. + """ + if zkind == "array_elements": + return ( + f"{signature_line}\n" + f"{body_setup}" + f" _result = {call_expr}\n" + f" z[:] = _result\n" + ) + if zkind == "record_fields": + # Array-typed leaves slice-assign. Numba's record-field setitem + # copies the destination's extent regardless of the source's, so a + # short source is read past its end and a long one is silently + # truncated; ``[:]`` makes both a shape error instead. + field_assigns = "".join( + f" z[0]{path}{'[:]' if is_array else ''} = _result[{i}]\n" + for i, (path, is_array) in enumerate(zinfo) + ) + return f"{signature_line}\n{body_setup} _result = {call_expr}\n{field_assigns}" + BL, BR, zname = zinfo + return f"{signature_line}\n{body_setup} {zname} = {BL}{call_expr}{BR}\n" + + +def _fits_by_broadcast(shape, expected): + """Whether a return of ``shape`` fills an ``expected``-shaped destination. + + The wrapper slice-assigns (``z[:] = ...``), and numpy broadcasts on + assignment, so an exact match is not the requirement: a ``(1,)`` return + legitimately fills a ``(6,)`` element, and a ``(1, 3)`` return fills + every row of a ``(2, 3)`` one. Rejecting those would refuse code that + works today. + + numpy's assignment rule is broadcasting plus a leading-``1`` strip when + the source has the higher rank, so ``(1, 6)`` fits ``(6,)`` while + ``(6, 1)`` does not. ``test_udt_broadcast_matches_numba_slice_assign`` + pins this against Numba's own slice-assign for both. + """ + shape = tuple(shape) + expected = tuple(expected) + while len(shape) > len(expected) and shape[0] == 1: + shape = shape[1:] + try: + return np.broadcast_shapes(shape, expected) == expected + except ValueError: + return False + + +def _udf_probe_value(dtype): + """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`. + + Only the shape of the UDF's result is read, so any values it accepts + will do. + """ + if dtype._is_udt: + np_type = dtype.np_type + if np_type.subdtype is None: + return np.ones(1, dtype=np_type)[0] + base_np_type, shape = np_type.subdtype + return np.ones(shape, dtype=base_np_type) + return _sample_values[dtype] + + +def _run_udf_probe(numba_func, arg_dtypes): + """Call the UDF on stand-in operands. Returns ``(result,)``, or ``None`` if it raised. + + A UDF that raises on the stand-in values is not checked, so one that + also returns a shape that does not fit is left to the wrapper, which + fails without telling the caller (see :func:`_compose_wrapper_body`). + Only the call is guarded, so a mistake in building the operands still + surfaces. + + This is where typing an op for a UDT runs the user's function rather + than only compiling it, and ``OpBase.__contains__`` is a typed lookup, + so ``udt in some_op`` can run it too. It costs no extra compile: Numba + types the stand-ins as it types the wrapper's operands (a C-contiguous + array, a record, a scalar), so the call compiles the specialization the + wrapper needs anyway. + """ + args = [_udf_probe_value(d) for d in arg_dtypes] + try: + return (numba_func(*args),) + except NumbaError as exc: # pragma: no cover (safety) + # The wrapper's operands have these types, so it would fail to compile too. + raise UdfParseError(_summarize_numba_typing_error(exc)) from exc + except Exception: + return None + + +def _raise_unless_fits(shape, expected, return_type, path, *, probed): + """Raise ``UdfParseError`` unless a return of ``shape`` fills an ``expected`` slot. + + ``path`` is the record field, or ``None`` for an array UDT's element. + """ + if _fits_by_broadcast(shape, expected): + return + field = "" if path is None else f" for field {path} of {return_type}" + # A probed shape holds for the values the probe used. Saying so helps + # the author of a UDF whose output shape depends on its inputs' values. + how = " when run on sample values" if probed else "" + slot = f"{return_type} elements are" if path is None else "that field holds" + raise UdfParseError( + f"UDT UDF returned an array of shape {tuple(shape)}{field}{how}, but {slot} " + f"{tuple(expected)}. Return an array of that shape, or one that broadcasts to it." + ) + + +def _raise_unless_sequence_fits(kind, length, expected, return_type, path, *, probed): + """Raise ``UdfParseError`` unless a ``kind`` of ``length`` fills record field ``path``. + + Numba copies a tuple or list into an array field element by element, + which works only for a one-dimensional field and only at its exact + length. Unlike an array, a sequence does not broadcast, so a 1-tuple + does not fill a 3-element field. A wrong length raises inside the cfunc + just as a wrong-shape array does; a field of higher rank fails the + wrapper compile. + """ + if expected == (length,): + return + how = " when run on sample values" if probed else "" + if len(expected) == 1: + fix = f"Return {expected[0]} values, or an array that broadcasts to {expected}." + else: + fix = f"A {kind} only fills a one-dimensional field; return an array." + raise UdfParseError( + f"UDT UDF returned a {kind} of length {length} for field {path} of {return_type}" + f"{how}, but that field holds {expected}. {fix}" + ) + + +# Numba types a record UDF might return for a scalar field, none of which the +# wrapper can write there, each with the kind the error names. ``Array`` +# covers ``NestedArray``. +_NON_SCALAR_KINDS = ( + ("an array", numba.core.types.Array), + ("a tuple", numba.core.types.BaseTuple), + ("a list", numba.core.types.List), +) + + +def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes): + """Reject a UDF whose return cannot fill the output element. + + The wrapper slice-assigns an array return, and each array-typed record + leaf, into the element, so a shape that does not fit raises inside the + cfunc, where the caller never sees it (see + :func:`_compose_wrapper_body`). So does a record leaf the UDF can + return as ``None`` (``v if cond else None``, typed ``Optional``), array + or scalar, on the path that returns it; a leaf that is always ``None`` + does not compile at all, and Numba's traceback says why at length. Nor + does an array, tuple or list returned for a scalar leaf: Numba has no + record setitem for them. Checking here makes each an error when the op + is typed. + + An ``Optional`` leaf is rejected from its type, whichever branch the + sample values would take. Numba also keeps the extents of an operand, + or operand field, returned as-is (a ``NestedArray``), so those are + checked from the type alone. An array the UDF builds (``x + y``, + ``x[:2]``) is a plain ``Array``, which has a rank but no extents, so + only running the UDF, on stand-ins built from ``arg_dtypes``, can say + what shape it has. A tuple or list for an array leaf must match its + length exactly (:func:`_raise_unless_sequence_fits`); a tuple's length + is in its type, and a list's is learned the same way as a built + array's. A UDF that returns no built array or list is never run. + """ + if zkind == "array_elements": + # (record field path, Numba type, slot shape, index into the tuple) + slots = [(None, numba_ret_type, return_type.np_type.shape, None)] + elif zkind == "record_fields": + # The tuple is flat over the leaves, one element per leaf. + leaves = list(_iter_record_leaves(return_type.np_type)) + if len(numba_ret_type.types) != len(leaves): + # _resolve_udt_return_type matched the record to the tuple's length, + # so this would be a wrong field count, not a wrong shape. + return # pragma: no cover (safety) + slots = [ + (path, numba_ret_type.types[i], leaf.shape, i) + for i, (path, _c, leaf) in enumerate(leaves) + ] + else: + return + unknown = [] + for path, nb_type, expected, i in slots: + if isinstance(nb_type, (numba.core.types.Optional, numba.core.types.NoneType)): + field = "" if path is None else f" for field {path}" + raise UdfParseError( + f"UDT UDF can return None{field}, but {return_type} cannot hold None. " + f"Return a value on every path." + ) + if not expected: + for kind, cls in _NON_SCALAR_KINDS: + if isinstance(nb_type, cls): + raise UdfParseError( + f"UDT UDF returned {kind} for field {path} of {return_type}, but " + f"that field holds a scalar. Return a scalar for it." + ) + continue # a scalar leaf: nothing to fit + if isinstance(nb_type, numba.core.types.NestedArray): + _raise_unless_fits(nb_type.shape, expected, return_type, path, probed=False) + elif isinstance(nb_type, numba.core.types.Array) or nb_type is None: + # No extents to read (or, with no return type given, no type + # at all), so learn the shape by running the UDF. + unknown.append((path, expected, i)) + elif isinstance(nb_type, numba.core.types.BaseTuple): + _raise_unless_sequence_fits( + "tuple", len(nb_type.types), expected, return_type, path, probed=False + ) + elif isinstance(nb_type, numba.core.types.List): + unknown.append((path, expected, i)) # a list's length is not in its type + if unknown and (ran := _run_udf_probe(numba_func, arg_dtypes)) is not None: + for path, expected, i in unknown: + value = ran[0] if i is None else ran[0][i] + if isinstance(value, list): + _raise_unless_sequence_fits( + "list", len(value), expected, return_type, path, probed=True + ) + else: + _raise_unless_fits(np.shape(value), expected, return_type, path, probed=True) + + +def _get_udt_wrapper( + numba_func, return_type, dtype, dtype2=None, *, include_indexes=False, numba_ret_type=None +): + """Build a Numba cfunc wrapper for unary, binary, indexunary, or select UDFs on UDTs. + + ``include_indexes=True`` inserts ``(row, col)`` between ``x`` and + ``y`` in both the wrapper signature and the call to ``numba_func``, + matching the IndexUnaryOp and SelectOp shape. The IndexBinaryOp path + (four indices plus a theta operand) uses + :func:`_get_udt_wrapper_indexbinary`. + """ + nt = numba.types + zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) + xsetup, xderef, xptr_type = _input_operand(dtype, "x") + wrapper_args = [zptr_type, xptr_type] + arg_dtypes = [dtype] + if include_indexes: + wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) + arg_dtypes.extend([UINT64, UINT64]) + ysetup, yderef_expr, yarg = "", "", "" + if dtype2 is not None: + ysetup, yderef, yptr_type = _input_operand(dtype2, "y") + wrapper_args.append(yptr_type) + arg_dtypes.append(dtype2) + yarg = ", y_ptr" + yderef_expr = f", {yderef}" + wrapper_sig = nt.void(*wrapper_args) + _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) + + rcidx = ", row, col" if include_indexes else "" + signature_line = f"def wrapper(z_ptr, x_ptr{rcidx}{yarg}):" + body_setup = f"{zsetup}{xsetup}{ysetup}" + call_expr = f"numba_func({xderef}{rcidx}{yderef_expr})" + + text = _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr) + kind = "indexunary" if include_indexes else ("binary" if dtype2 is not None else "unary") + wrapper = _compile_codegen( + text, + func_name="wrapper", + source_label=f"", + extra_ns={"numba_func": numba_func}, + ) + return wrapper, wrapper_sig + + +def _get_udt_wrapper_indexbinary( + numba_func, return_type, dtype, dtype2, theta_type, *, numba_ret_type=None +): + """Build a Numba cfunc wrapper for IndexBinaryOp UDFs on UDTs. + + Signature: ``f(x, ix, jx, y, iy, jy, theta) -> z``, where ``x``, ``y`` + and ``theta`` have types ``dtype``, ``dtype2`` and ``theta_type``. + """ + nt = numba.types + zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) + xsetup, xderef, xptr_type = _input_operand(dtype, "x") + ysetup, yderef, yptr_type = _input_operand(dtype2, "y") + tsetup, tderef, tptr_type = _input_operand(theta_type, "t") + wrapper_sig = nt.void( + zptr_type, + xptr_type, + UINT64.numba_type, + UINT64.numba_type, + yptr_type, + UINT64.numba_type, + UINT64.numba_type, + tptr_type, + ) + arg_dtypes = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, theta_type] + _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) + + signature_line = "def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):" + body_setup = f"{zsetup}{xsetup}{ysetup}{tsetup}" + call_expr = f"numba_func({xderef}, ix, jx, {yderef}, iy, jy, {tderef})" + + text = _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr) + wrapper = _compile_codegen( + text, + func_name="wrapper", + source_label=( + f"" + ), + extra_ns={"numba_func": numba_func}, + ) + return wrapper, wrapper_sig diff --git a/graphblas/core/operator/udt_codegen.py b/graphblas/core/operator/udt_codegen.py new file mode 100644 index 000000000..0ac54404a --- /dev/null +++ b/graphblas/core/operator/udt_codegen.py @@ -0,0 +1,1757 @@ +"""Code generation for the operators that compute on UDTs. + +``core/udt.py`` decides which operand pairs an op takes and what type its +result has; this module writes the code that computes it, for record UDTs +(C structs) and array UDTs (such as ``FP64[3]``): + +- ``_udt_c_typedef``, a UDT's C struct, set when the UDT is registered; +- under ``if _has_numba:``, the Numba cfuncs: ``compile_udt_binary_wrapper``, + ``compile_udt_unary_wrapper``, ``compile_udt_cast_wrapper`` and + ``fused_store_op``; +- the JIT C definitions SuiteSparse compiles in place of a cfunc + (``set_jit_c_on_op``, ``set_jit_c_comparison_on_op``). + +``_compile_codegen`` compiles the generated Python, for this module, +``operator.udf`` and the comparisons in ``operator.binary``. An op is compiled +when it is first typed for a UDT. +""" + +import ast +import hashlib +import itertools +import linecache +from functools import lru_cache, reduce +from operator import mul + +import numpy as np + +from ... import backend +from .. import _has_numba, ffi, lib +from ..udt import ( + BUILTIN_UDT_BINARY_OPS, + BUILTIN_UDT_UNARY_OPS, + _check_udt_pair, + _get_udt_info, + _iter_record_leaves, + _udt_op_types, +) + +if _has_numba: + import numba + + +_codegen_counter = itertools.count(1) + + +def _compile_codegen(src, *, func_name, source_label, extra_ns=None): + """Compile a generated Python source string and return the named function. + + Centralizes the small amount of ``exec``-based code generation that the + UDT operator wrappers need (Numba has to see a literal function body, + not a closure, to type-check shape-specialized arithmetic). Three things + this helper does that a bare ``exec`` doesn't: + + 1. ``ast.parse`` runs first so a codegen typo raises a clear + ``RuntimeError`` with the offending source attached, at the call + site, rather than as a cryptic ``SyntaxError`` from ``exec`` or a + confusing ``TypingError`` from Numba's first compile. + 2. ``compile(..., filename, "exec")`` uses a human-readable synthetic + filename like ``" #7"``, and the + generated source is registered with ``linecache`` so any later + traceback shows real lines instead of ``:??``. + 3. The execution namespace is constructed here, so the names visible + to the generated code are auditable in one place. Extra entries + (e.g., the user's compiled ``numba_func`` for wrapper bodies) come + in via ``extra_ns``. + + Returns ``namespace[func_name]``. + """ + try: + ast.parse(src) + except SyntaxError as exc: + # Codegen bug, not user input; surface the source so a future + # maintainer can see exactly what was generated. + raise RuntimeError( + f"Generated code for {source_label!r} is not valid Python " + f"(parse error: {exc}). Source:\n{src}" + ) from exc + # Counter suffix so two codegens with the same label still get distinct + # cache keys (e.g., the same op compiled for two different UDTs that + # happen to print the same shape summary). + filename = f"{source_label} #{next(_codegen_counter)}" + linecache.cache[filename] = ( + len(src), + None, + src.splitlines(keepends=True), + filename, + ) + code = compile(src, filename, "exec") + # Generated code never calls Python's ``min`` / ``max`` (still reachable as + # builtins): they order a NaN by argument position, where SuiteSparse's + # ``fmin`` ignores it, so :func:`_minmax_expr` spells the comparison out, + # with ``signbit`` for the signed-zero tie. Complex ``*`` and ``/`` call + # :func:`_complex_times` and :func:`_complex_truediv` (see _expr_binary). + namespace = {"abs": abs, "np": np, "signbit": np.signbit} + if _has_numba: + namespace["numba"] = numba + namespace["_complex_times"] = _complex_times + namespace["_complex_truediv"] = _complex_truediv + if extra_ns: + namespace.update(extra_ns) + exec(code, namespace) + return namespace[func_name] + + +# Ops that use function-call syntax rather than infix (e.g., min(a, b), not "a min b") +_FUNC_BINARY_OPS = {"min", "max"} +_FUNC_UNARY_OPS = {"abs"} + +NP_TO_C_TYPES = { + np.dtype(np.bool_): "_Bool", + np.dtype(np.int8): "int8_t", + np.dtype(np.int16): "int16_t", + np.dtype(np.int32): "int32_t", + np.dtype(np.int64): "int64_t", + np.dtype(np.uint8): "uint8_t", + np.dtype(np.uint16): "uint16_t", + np.dtype(np.uint32): "uint32_t", + np.dtype(np.uint64): "uint64_t", + np.dtype(np.float32): "float", + np.dtype(np.float64): "double", + # Complex types are layout-compatible with numpy's; SS's GraphBLAS.h + # typedefs ``GxB_FC32_t`` / ``GxB_FC64_t`` to C99 ``float _Complex`` / + # ``double _Complex`` (or MSVC's ``_Fcomplex`` / ``_Dcomplex``). The + # JIT include chain pulls in GraphBLAS.h, so these names are in scope. + # Plus, minus, times, truediv, ainv compile natively on _Complex; abs + # uses cabs / cabsf (see _c_expr_unary). min, max, floordiv don't make + # sense on complex; ``udt._element_types`` rejects them before any codegen. + np.dtype(np.complex64): "GxB_FC32_t", + np.dtype(np.complex128): "GxB_FC64_t", +} + + +# ``_MIN`` spellings from , keyed by itemsize. Signed integer +# division traps on ``MIN / -1``; the generated C tests for it explicitly. +_C_INT_MIN = {1: "INT8_MIN", 2: "INT16_MIN", 4: "INT32_MIN", 8: "INT64_MIN"} + + +def _is_float(np_dtype): + return np_dtype.kind == "f" + + +def _is_signed_int(np_dtype): + return np_dtype.kind == "i" + + +# Vanilla strips GxB callables but keeps GxB constants, so the bare +# ``hasattr`` would lie; gate on the backend too. +_has_jit_set = backend == "suitesparse" and hasattr(lib, "GxB_JIT_C_NAME") + +# Names that can't be used as record field names in the JIT C source: C +# keywords, and the macros of the JIT include chain (``GraphBLAS.h`` +# transitively pulls in ````, ````, ````, +# ````, ````, ```` and more). The preprocessor +# expands a macro inside any declarator or expression it appears in, so a field +# declared ``double M_PI ;`` becomes ``double 3.14159... ;`` and won't compile; +# the separate namespace for struct members doesn't help. (A UDT's own name +# never collides: its C name has a prefix, ``_C_TYPE_NAME_PREFIX``. The typedef +# names listed here mattered when the type's C name was its Python name.) +# +# Either way the JIT compile fails. SuiteSparse 9.4 and later then raise +# JitError for that call and set ``jit_c_control`` to 'load', so the process +# compiles no more kernels. Block these eagerly: the op runs through the +# slower Numba cfunc path, and the user sees the "without JIT" warning. +_C_RESERVED = frozenset( + { + # C keywords + "auto", + "break", + "case", + "char", + "const", + "continue", + "default", + "do", + "double", + "else", + "enum", + "extern", + "float", + "for", + "goto", + "if", + "inline", + "int", + "long", + "register", + "restrict", + "return", + "short", + "signed", + "sizeof", + "static", + "struct", + "switch", + "typedef", + "union", + "unsigned", + "void", + "volatile", + "while", + "_Alignas", + "_Alignof", + "_Atomic", + "_Bool", + "_Complex", + "_Generic", + "_Imaginary", + "_Noreturn", + "_Static_assert", + "_Thread_local", + # C++ keywords that some compilers also reserve; cheap insurance + "bool", + "class", + "new", + "delete", + "template", + "namespace", + "this", + "true", + "false", + "nullptr", + # / : macros + the FILE typedef + "NULL", + "EOF", + "offsetof", + "stdin", + "stdout", + "stderr", + "FILE", + # / : typedefs at risk in the outer position + "size_t", + "ptrdiff_t", + "wchar_t", + "intptr_t", + "uintptr_t", + "intmax_t", + "uintmax_t", + "int8_t", + "int16_t", + "int32_t", + "int64_t", + "uint8_t", + "uint16_t", + "uint32_t", + "uint64_t", + # numeric macros + "INFINITY", + "NAN", + "HUGE_VAL", + "HUGE_VALF", + "HUGE_VALL", + "M_E", + "M_LOG2E", + "M_LOG10E", + "M_LN2", + "M_LN10", + "M_PI", + "M_PI_2", + "M_PI_4", + "M_1_PI", + "M_2_PI", + "M_2_SQRTPI", + "M_SQRT2", + "M_SQRT1_2", + # : ``complex`` and ``I`` in particular expand inside + # any field declarator. + "complex", + "imaginary", + "I", + "_Complex_I", + "_Imaginary_I", + "CMPLX", + "CMPLXF", + "CMPLXL", + # + "errno", + } +) | frozenset( + # Object-like macros the C standard defines in the other headers of the JIT + # include chain. A platform adds its own (POSIX errno codes, ...), which + # still fail to compile. + [ + # + *("CHAR_BIT", "SCHAR_MIN", "SCHAR_MAX", "UCHAR_MAX", "CHAR_MIN", "CHAR_MAX"), + *("MB_LEN_MAX", "SHRT_MIN", "SHRT_MAX", "USHRT_MAX", "INT_MIN", "INT_MAX"), + *("UINT_MAX", "LONG_MIN", "LONG_MAX", "ULONG_MAX", "LLONG_MIN", "LLONG_MAX"), + "ULLONG_MAX", + # + *( + f"{u}INT{kind}{bits}_{end}" + for kind in ("", "_LEAST", "_FAST") + for bits in (8, 16, 32, 64) + for u, end in (("", "MIN"), ("", "MAX"), ("U", "MAX")) + ), + *("INTPTR_MIN", "INTPTR_MAX", "UINTPTR_MAX", "INTMAX_MIN", "INTMAX_MAX"), + *("UINTMAX_MAX", "PTRDIFF_MIN", "PTRDIFF_MAX", "SIG_ATOMIC_MIN", "SIG_ATOMIC_MAX"), + *("SIZE_MAX", "WCHAR_MIN", "WCHAR_MAX", "WINT_MIN", "WINT_MAX"), + # + *( + f"{prefix}{conv}{size}" + for prefix, convs in (("PRI", "diouxX"), ("SCN", "dioux")) + for conv in convs + for size in ( + *(f"{kind}{bits}" for kind in ("", "LEAST", "FAST") for bits in (8, 16, 32, 64)), + "MAX", + "PTR", + ) + ), + # + *("BUFSIZ", "FILENAME_MAX", "FOPEN_MAX", "L_tmpnam", "SEEK_CUR", "SEEK_END"), + *("SEEK_SET", "TMP_MAX", "_IOFBF", "_IOLBF", "_IONBF"), + # + *("EXIT_FAILURE", "EXIT_SUCCESS", "RAND_MAX", "MB_CUR_MAX"), + # + *("EDOM", "ERANGE", "EILSEQ"), + # + *("FP_INFINITE", "FP_NAN", "FP_NORMAL", "FP_SUBNORMAL", "FP_ZERO", "FP_FAST_FMA"), + *("FP_FAST_FMAF", "FP_FAST_FMAL", "FP_ILOGB0", "FP_ILOGBNAN", "MATH_ERRNO"), + *("MATH_ERREXCEPT", "math_errhandling"), + # + *( + f"ATOMIC_{kind}_LOCK_FREE" + for kind in ("BOOL", "CHAR", "CHAR16_T", "CHAR32_T", "WCHAR_T", "SHORT", "INT") + ), + *("ATOMIC_LONG_LOCK_FREE", "ATOMIC_LLONG_LOCK_FREE", "ATOMIC_POINTER_LOCK_FREE"), + "ATOMIC_FLAG_INIT", + # , which the C library headers bring in on some platforms + *("SIGABRT", "SIGFPE", "SIGILL", "SIGINT", "SIGSEGV", "SIGTERM"), + *("SIG_DFL", "SIG_ERR", "SIG_IGN"), + ] +) + +# GraphBLAS's own names, many of them macros of ``GraphBLAS.h`` and its JIT headers +# (``GrB_INDEX_MAX``, ``GxB_IMPLEMENTATION_MAJOR``), and more with each version, so a +# field may not start with any of these. +_C_RESERVED_PREFIXES = ("GB_", "GrB_", "GxB_") + + +def _is_valid_c_identifier(name): + """True if ``name`` is a valid C identifier and not a reserved word or GraphBLAS's name.""" + return ( + isinstance(name, str) + and name.isidentifier() + and name not in _C_RESERVED + and not name.startswith(_C_RESERVED_PREFIXES) + ) + + +# Process-global counter for synthesizing C names when the user-supplied +# Python name doesn't make a valid C identifier of 127 characters or fewer (or +# the UDT was registered without ``name=`` at all). The Python-side ``DataType.name`` +# is left alone; this only affects ``GxB_JIT_C_NAME``. ``itertools.count`` is +# atomic enough for our purposes; collisions across processes are harmless +# because the SS JIT cache files are keyed by ``(c_name, content_hash)`` per +# process state. +_synthetic_udt_counter = itertools.count(1) + +# The C typedef given to SuiteSparse under each ``GxB_JIT_C_NAME`` in this +# process. SuiteSparse's JIT knows a type by that name alone (its hash is the +# name's), so two layouts under one name would share compiled kernels, and a +# kernel compiled for one would read and write the other's element size. +_jit_c_type_defs = {} + + +# The C name of a UDT is its name after this prefix. SuiteSparse's JIT puts the type's name in +# a kernel beside its own identifiers, so a UDT named like one of them did not compile: a +# local of its templates (``A``, ``C``, ``p``, ``i``, ``result``), a parameter of the op +# definitions made here (``z``, ``x``), or anything the C headers declare (``div``, +# ``index``, ``GB_void``). No identifier there starts with the prefix. +_C_TYPE_NAME_PREFIX = "gbudt_" + + +# The ops whose C name is ``{op}_{type}`` (``set_jit_c_on_op``, +# ``set_jit_c_comparison_on_op``): the lifted arithmetic, eq and ne. A mixed op +# shortens its own name (``_make_jit_c_mixed_definition``). +_ONE_TYPE_OP_NAMES = (*BUILTIN_UDT_BINARY_OPS, *BUILTIN_UDT_UNARY_OPS, "eq", "ne") + + +def _max_type_c_name_len(): + """The longest C name a UDT may have, so every ``{op}_{type}`` built from it fits. + + SuiteSparse refuses a name of ``GxB_MAX_NAME_LEN`` (128) characters or more; the + longest op prefix is ``floordiv_``, so this is 118. + """ + return lib.GxB_MAX_NAME_LEN - 1 - max(len(f"{name}_") for name in _ONE_TYPE_OP_NAMES) + + +def _pick_c_type_name(python_name): + """Return a valid C identifier for use as ``GxB_JIT_C_NAME`` on a UDT named ``python_name``. + + ``gbudt_`` and the name (see ``_C_TYPE_NAME_PREFIX``), which stays readable in + introspection and JIT cache file names. A name starting with ``_gb``, which + python-graphblas keeps for the types it makes itself (a synthetic ``_gbudt_NNN`` or + ``_gbnest_NNN``, the tensor engine's ``_gbtensor_...``), and which no kernel identifier + starts with, is kept as it is. When the result is not a C identifier (a name with dots, + or none given), or is too long for each op built from it to fit SuiteSparse's + ``GxB_MAX_NAME_LEN`` (``_max_type_c_name_len``), mint a fresh ``_gbudt_NNN`` name + instead, so the UDT and its ops still take the JIT path. A name SuiteSparse already has + for a type goes to ``_udt_c_typedef`` with ``is_c_name`` instead. + """ + if not isinstance(python_name, str): + c_name = "" + elif python_name.startswith("_gb"): + c_name = python_name + else: + c_name = _C_TYPE_NAME_PREFIX + python_name + if c_name.isidentifier() and len(c_name) <= _max_type_c_name_len(): + return c_name + return f"_gbudt_{next(_synthetic_udt_counter)}" + + +def _udt_c_typedef(python_name, np_type, *, is_c_name=False): + """Build the JIT C typedef for a record or array UDT, or return ``None``. + + Returns ``(type_name, typedef)`` on success, where ``type_name`` is the + C identifier we register with SuiteSparse (``gbudt_`` and ``python_name``, or + a synthetic name; see ``_pick_c_type_name``). With ``is_c_name``, + ``python_name`` is already that C name: one SuiteSparse has for the type, or a + synthetic one made here. Returns ``None`` when the UDT still can't be + expressed in C after synthesizing the top-level name (a field name + collides with a C reserved word, or a leaf field type isn't in + ``NP_TO_C_TYPES``, e.g. object or datetime). + + Nested record UDTs (structured field whose dtype is itself a record) + are supported: each inner struct is emitted with a synthesized + ``_gbnest_NNN`` C name before the outer typedef in the same + ``GxB_JIT_C_DEFINITION`` string, so SuiteSparse's JIT C file ends up + with the full chain of typedefs in scope. Array UDTs are flattened to + a single-dimension C array so that the generated per-element operator + code can use flat indices like ``x->v[i]``, matching the + (row-major contiguous) numpy memory layout. + """ + # If the numpy layout disagrees with what a C compiler would produce, + # the JIT kernel would read fields at the wrong offsets. Skip rather + # than register a broken typedef; the cfunc path is still correct. + if np_type.names is not None and not _is_c_compatible_layout(np_type): + return None + type_name = python_name if is_c_name else _pick_c_type_name(python_name) + if np_type.names is not None: + # Inner typedefs (for nested structured fields) are prepended to the + # final definition string so the outer struct's field types are in + # scope when SuiteSparse compiles the kernel. + inner_typedefs = [] + fields = [] + for field_name in np_type.names: + # Field names still have to be valid C; synthesizing per-field + # names would mean tracking a name map for the codegen and would + # break the human-readable JIT source (``z->_f0 = ...``). Not + # worth it for the corner case where a user named a field + # ``"class"``; fall back to cfunc there. + if not _is_valid_c_identifier(field_name): + return None + field_dtype = np_type.fields[field_name][0] + if field_dtype.names is not None: + # Nested struct field. Recurse to build the inner typedef. + # The inner C name is synthesized to avoid collisions with + # user-supplied type names elsewhere in the process. + inner_name = f"_gbnest_{next(_synthetic_udt_counter)}" + inner_info = _udt_c_typedef(inner_name, field_dtype, is_c_name=True) + if inner_info is None: + return None + inner_resolved_name, inner_typedef = inner_info + inner_typedefs.append(inner_typedef) + fields.append(f"{inner_resolved_name} {field_name}") + continue + c_type = NP_TO_C_TYPES.get(field_dtype) + if c_type is None: + return None + fields.append(f"{c_type} {field_name}") + outer_typedef = f"typedef struct {{ {' ; '.join(fields)} ; }} {type_name} ;" + if inner_typedefs: + typedef = " ".join([*inner_typedefs, outer_typedef]) + else: + typedef = outer_typedef + return type_name, typedef + if np_type.subdtype is not None: + base_dtype, shape = np_type.subdtype + base_c_type = NP_TO_C_TYPES.get(base_dtype) + if base_c_type is None: + return None + size = reduce(mul, shape) + typedef = f"typedef struct {{ {base_c_type} v [{size}] ; }} {type_name} ;" + return type_name, typedef + return None + + +def _c_aligned_version(np_type): + """Return an ``align=True`` version of ``np_type`` with every nested + record also aligned recursively. + + ``np.dtype([..., (name, inner_struct)], align=True)`` respects the + *inner_struct*'s own declared alignment, so a packed inner struct stays + packed at its natural-1-byte alignment and the outer's ``coord`` field + lands at the outer's next-after-flag offset (4 for ``int32``), not at + the 8-byte boundary a C compiler would pick. To compare against what C + would emit, we need to rebuild the inner as ``align=True`` first, then + use that rebuilt type as the outer's field type. + """ + fields = [] + for name in np_type.names: + f = np_type.fields[name][0] + if f.names is not None: + f = _c_aligned_version(f) + fields.append((name, f)) + return np.dtype(fields, align=True) + + +def _is_c_compatible_layout(np_type): + """Return True iff ``np_type``'s layout matches what a C compiler will + pick for the corresponding ``typedef struct``. + + A user-supplied ``np.dtype([(name, dtype), ...])`` is *packed* by default + (no padding between fields). A C struct *aligns* fields to their natural + boundary. For ``[(int32, float64)]`` numpy uses offsets ``0, 4`` and + itemsize 12; C uses offsets ``0, 8`` and itemsize 16. The JIT-compiled + kernel reads fields at C offsets but the numpy buffer holds them at the + packed offsets, so the JIT would read garbage. Detect and refuse JIT in + that case; the user can either re-register with ``align=True``, use the + dict / dataclass form (which already does), or accept the cfunc path. + + Array UDTs are always C-compatible (single contiguous run of one type). + """ + if np_type.names is None: + return True + aligned = _c_aligned_version(np_type) + if aligned.itemsize != np_type.itemsize: + return False + for name in np_type.names: + if np_type.fields[name][1] != aligned.fields[name][1]: + return False + if not _is_c_compatible_layout(np_type.fields[name][0]): + return False + return True + + +# Numba function generators, called lazily from each op's ``_compile_udt``. +if _has_numba: + + def _expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype): + """Python-source builder; sibling of :func:`_c_expr_binary` for the C JIT. + + ``x_expr`` and ``y_expr`` are scalars, never arrays. ``x_dtype`` and + ``y_dtype`` are the numpy dtypes the op computes in for each operand: + the built-in op's input types, which the operands are converted to + first. These ops consult them: + + - ``min`` / ``max`` follow SuiteSparse's own ``GrB_MIN_FP64``, which + is C99 ``fmin`` and ignores a NaN operand rather than ordering or + propagating it; see :func:`_minmax_expr`. Only floating-point + operands carry the NaN and signed-zero rules, so the dtype decides + which of the two forms is emitted. + - ``//`` on two signed integers. Numba deliberately returns 0 for + ``INT64_MIN // -1`` to dodge the SIGFPE that x86 raises on the + unrepresentable quotient, while numpy wraps to ``INT64_MIN``. Since + ``a // -1`` is exactly ``-a``, routing that divisor through + negation (which wraps) reaches numpy's answer without the trap. + - ``*`` and ``/`` where either side is complex go through + :func:`_complex_times` and :func:`_complex_truediv`, the formulas + the C kernel spells out. Numba's own operators are CPython's, so + its ``/`` raises ``ZeroDivisionError`` on a zero divisor, and from + Python 3.14 (Numba 0.68) both recover infinities that numpy and the + C kernel leave as ``nan``. + + ``/`` on two integers needs nothing: the built-in typing makes its + result float64, as ``truediv`` on two INT64 vectors is FP64, so + ``x / 0`` is an infinity, as there. + """ + if py_op in _FUNC_BINARY_OPS: + return _minmax_expr(py_op, x_expr, y_expr, x_dtype, y_dtype) + if py_op == "//" and _is_signed_int(x_dtype) and _is_signed_int(y_dtype): + return f"(-{x_expr} if {y_expr} == -1 else {x_expr} // {y_expr})" + if py_op == "*" and "c" in (x_dtype.kind, y_dtype.kind): + return f"_complex_times({x_expr}, {y_expr})" + if py_op == "/" and "c" in (x_dtype.kind, y_dtype.kind): + return f"_complex_truediv({x_expr}, {y_expr})" + return f"{x_expr} {py_op} {y_expr}" + + @numba.njit(error_model="numpy") + def _complex_times(x, y): # pragma: no cover (numba) + """Multiply two complex numbers by the textbook formula, as numpy does. + + Numba's own ``*`` is CPython's, which from Python 3.14 (Numba 0.68) + adds C99 Annex G's recovery of infinities: ``1j * (inf+infj)`` is + ``-inf+infj`` there and ``nan+nanj`` in numpy. This is + :func:`_c_complex_times_stmt`, product for product, so the cfunc and + the C kernel agree whatever the Python and Numba versions. The parts + keep the operands' precision; the complex128 result holds them + exactly, for the caller to store in the field's type. + """ + ar = x.real + ai = x.imag + br = y.real + bi = y.imag + return complex(ar * br - ai * bi, ar * bi + ai * br) + + @numba.njit(error_model="numpy") + def _complex_truediv(x, y): # pragma: no cover (numba) + """Divide two complex numbers by Smith's method, as the C kernel does. + + This is CPython's ``_Py_c_quot`` before Python 3.14 (scale by the + larger part of the divisor), statement for statement as + :func:`_c_complex_truediv_stmt` spells it. Numba's own ``/`` raises + ``ZeroDivisionError`` on a zero divisor, outside the error model's + control, and from Python 3.14 (Numba 0.68) recovers infinities: + ``(inf+nanj) / (1+1j)`` is ``inf-infj`` there and ``nan+nanj`` in + numpy. A zero divisor divides the parts by ``0.0`` as reals, for + numpy's infinities. The result is complex128, as in + :func:`_complex_times`. + """ + ar = x.real + ai = x.imag + br = y.real + bi = y.imag + if br == 0 and bi == 0: + return complex(ar / 0.0, ai / 0.0) + if abs(br) >= abs(bi): + r = bi / br + d = br + bi * r + return complex((ar + ai * r) / d, (ai - ar * r) / d) + if bi == 0: # the real part is NaN + return complex(np.nan, np.nan) + r = br / bi + d = br * r + bi + return complex((ar * r + ai) / d, (ai * r - ar) / d) + + def _minmax_expr(py_op, x_expr, y_expr, x_dtype, y_dtype): + """Return a Python expression for ``min`` / ``max`` matching C ``fmin``. + + ``GrB_MIN_FP64`` is C99 ``fmin``, so this builder and + :func:`_c_minmax_expr` both reproduce it or ``binary.min`` means one + thing on FP64 and another on a UDT: a NaN operand is ignored unless + both are NaN. C99 leaves the tie between ``-0.0`` and ``0.0`` to the + C library, which differs by platform, so both builders fix it: it + resolves to ``-0.0`` for ``min``, ``0.0`` for ``max``, whichever side + each sits on (Apple's libm answers so; glibc's x86-64 ``fmin`` answers + by position). + + Neither ``min`` nor ``np.fmin`` gets both rules right under Numba. + Python's builtin returns its first argument whenever the comparison + is false, so it keeps a NaN on the left and drops one on the right. + ``np.fmin`` fixes the NaN rule but Numba lowers it without the + signed-zero tie-break. Spelling the comparison out, here and in the C + kernel alike, is what keeps the two execution paths equal on every + platform. + + Integers have neither NaN nor a signed zero, so they take the same + plain comparison the C side emits. + """ + cmp_op = "<" if py_op == "min" else ">" + if not _is_float(x_dtype) and not _is_float(y_dtype): + return f"({x_expr} if {x_expr} {cmp_op} {y_expr} else {y_expr})" + # ``y != y`` is the NaN test on the right operand: when it holds, the + # left one wins whatever it is, which is fmin's "ignore the NaN" rule + # and also (correctly) returns NaN when both are NaN. + tie = f"signbit({x_expr})" if py_op == "min" else f"not signbit({x_expr})" + return ( + f"({x_expr} if ({x_expr} {cmp_op} {y_expr} or {y_expr} != {y_expr}" + f" or ({x_expr} == {y_expr} and {tie})) else {y_expr})" + ) + + def _expr_unary(py_op, operand): + """Python-source builder; sibling of :func:`_c_expr_unary` for JIT C.""" + if py_op in _FUNC_UNARY_OPS or py_op == "np.abs": + return f"{py_op}({operand})" + return f"{py_op}{operand}" + + def _cast_expr(expr, to_dtype, ns): + """Return Numba source converting scalar ``expr`` to numpy dtype ``to_dtype``. + + ``to_dtype`` of ``None`` leaves ``expr`` alone. The conversion function + is added to ``ns``, the generated code's namespace. + """ + if to_dtype is None: + return expr + name = f"_to_{to_dtype.name}" + ns[name] = numba.from_dtype(to_dtype) + return f"{name}({expr})" + + def _make_record_func( + leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scalar=False, leaf_types=None + ): + """Build a Numba njit function for a record UDT. + + ``leaf_paths`` is a sequence of Python access strings ``"['a']"`` or, + for nested records, ``"['outer']['inner_a']"``. The generated function + always returns a *flat* tuple of leaf values regardless of nesting + depth; the wrapper (``udf._compose_wrapper_body``) walks the same leaf paths when + writing the result back, so nested-record outputs land at the + correct depth without nested tuple construction (which Numba can't + ``setitem``-assign to a record field). + + When ``x_is_scalar`` or ``y_is_scalar`` is True, that argument is a + plain scalar (not a record), so it is used directly for all leaves. + A binary op needs ``leaf_types``, one ``(x_cast, y_cast, x_dtype, + y_dtype, result)`` tuple per leaf: the numpy dtype to convert each + operand to before the op (``None`` for none), the dtypes the op then + computes in, which :func:`_expr_binary` keys its semantics on, and the + result leaf's dtype. + + An array-valued leaf is computed element by element, with the same + scalar expression as a scalar leaf or an array UDT element, into an + array of the result leaf's dtype. As a whole-array expression it ran + through Numba's ufuncs, which pick their own loop (``int8 / int8`` in + float32), raise on a complex zero divisor, and take none of the scalar + forms of ``min``, ``max`` and ``//``. + """ + ns = {} + lines = [] + if arity == 2: + parts = [] + for i, path in enumerate(leaf_paths): + x_cast, y_cast, x_dtype, y_dtype, result = leaf_types[i] + if not result.shape: + x_expr = _cast_expr("x" if x_is_scalar else f"x{path}", x_cast, ns) + y_expr = _cast_expr("y" if y_is_scalar else f"y{path}", y_cast, ns) + parts.append(_expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype)) + continue + # Two array leaves have the same shape (``_check_udt_pair``). + ns[f"_r{i}"] = result.base + x_expr = _cast_expr("x" if x_is_scalar else f"x{path}[k]", x_cast, ns) + y_expr = _cast_expr("y" if y_is_scalar else f"y{path}[k]", y_cast, ns) + lines += [ + f" r{i} = np.empty({result.shape}, _r{i})", + f" for k in np.ndindex({result.shape}):", + f" xv = {x_expr}", + f" yv = {y_expr}", + f" r{i}[k] = {_expr_binary(py_op, 'xv', 'yv', x_dtype, y_dtype)}", + ] + parts.append(f"r{i}") + sig = "x, y" + else: + # Python's ``abs`` does not take an array field in Numba; ``np.abs`` + # takes both. + unary_op = "np.abs" if py_op == "abs" else py_op + parts = [_expr_unary(unary_op, f"x{path}") for path in leaf_paths] + sig = "x" + body = ", ".join(parts) + # Single-leaf tuple needs the trailing comma to remain a tuple. + ret = f"({body},)" if len(leaf_paths) == 1 else f"({body})" + lines.append(f" return {ret}") + src = f"def _op({sig}):\n" + "\n".join(lines) + "\n" + op_func = _compile_codegen( + src, + func_name="_op", + source_label=f"", + extra_ns=ns, + ) + return numba.njit(op_func, error_model="numpy") + + def _array_operand_source(side, operand_shape, shape, ns): + """Return ``(setup, ref)``, Numba source that reads operand ``side`` at element ``i``. + + ``i`` is a flat position in a result of ``shape``. ``operand_shape`` is + the operand's shape, or ``None`` for a plain scalar, read for every + element. An operand that broadcasts to ``shape`` reads through a + constant flat index map, added to ``ns``. + """ + if operand_shape is None: + return "", f"{side}_ptr[0]" + size = reduce(mul, operand_shape) + setup = f" {side} = numba.carray({side}_ptr, {size})\n" + # Broadcasting repeats no element when the sizes agree, so the shapes + # differ at most by leading axes of length 1 and the flat order is the same. + if size == reduce(mul, shape): + return setup, f"{side}[i]" + ns[f"{side}_index"] = np.broadcast_to(np.arange(size).reshape(operand_shape), shape).ravel() + return setup, f"{side}[{side}_index[i]]" + + def _make_array_wrapper( + shape, + base_numba_type, + arity, + py_op, + *, + x_type, + y_type=None, + x_shape=None, + y_shape=None, + x_cast=None, + y_cast=None, + store_cast=None, + ): + """Build a cfunc-ready wrapper for an array UDT (element-by-element). + + ``shape`` and ``base_numba_type`` are the shape and element type of the + result. ``x_type`` and ``y_type`` are the operands' element types. For + a binary op, ``x_shape`` and ``y_shape`` are the operands' shapes, which + broadcast to the result's; ``None`` makes that side a plain scalar of + its type, broadcast to all elements. + ``x_cast`` and ``y_cast`` are numpy dtypes to convert each operand to + before the op, if any. ``store_cast``, a ``(computed, stored)`` pair of + numpy dtypes, makes each element computed as ``computed`` and stored as + ``stored``, cast as a store casts (:func:`_cast_leaf_source`); then + ``base_numba_type`` is the stored element type. + + Returns (wrapper_func, wrapper_sig). + """ + nt = numba.types + ns = {} + + def store(expr): + if store_cast is None: + return expr + computed, stored = store_cast + return _cast_leaf_source(computed, stored, _cast_expr(expr, computed, ns), ns) + + size = reduce(mul, shape) + # A loop, not one statement per element: Numba's typing time grows + # faster than linearly with the number of statements, so unrolling took + # seconds for a 32 by 32 array and minutes for a 64 by 64 one. + arrays = f" z = numba.carray(z_ptr, {size})\n" + if arity == 2: + # Each operand is read once per element, into ``xv`` and ``yv``, + # so an op that names an operand more than once reads it once. + body = "" + computed = [] + for side, operand_shape, operand_type, cast in [ + ("x", x_shape, x_type, x_cast), + ("y", y_shape, y_type, y_cast), + ]: + setup, ref = _array_operand_source(side, operand_shape, shape, ns) + arrays += setup + body += f" {side}v = {_cast_expr(ref, cast, ns)}\n" + computed.append( + numba.np.numpy_support.as_dtype(operand_type) if cast is None else cast + ) + body += f" z[i] = {store(_expr_binary(py_op, 'xv', 'yv', *computed))}\n" + params = "z_ptr, x_ptr, y_ptr" + sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(x_type), nt.CPointer(y_type)) + else: + body = f" z[i] = {store(_expr_unary(py_op, 'x[i]'))}\n" + params = "z_ptr, x_ptr" + arrays += f" x = numba.carray(x_ptr, {size})\n" + sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(x_type)) + src = f"def _op({params}):\n{arrays} for i in range({size}):\n{body}" + op_func = _compile_codegen( + src, + func_name="_op", + source_label=f"", + extra_ns=ns, + ) + return op_func, sig + + # MAINT 2026-05-21: this function, ``compile_udt_unary_wrapper`` below, + # and ``_make_jit_c_definition`` all branch on ``kind == "record"`` vs the + # array path with near-identical scaffolding. When a third op family + # (ternary, indexbinary, ...) needs the same treatment, fold the three + # into one shape-parametrized helper instead of pasting a third copy. + def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2, out=None): + """Compile a built-in element-wise binary op for a UDT. + + Handles two cases: + + - Both sides are UDTs that pair (``udt._check_udt_pair``): a + field-by-field (record) or element-by-element (array) op. + - One side is a UDT and the other is a scalar type: the scalar is + broadcast to all fields or elements (e.g., ``Point + int`` adds + the int to every field). + + Each element is computed as the op computes it on the operands' + element dtypes (see ``udt._udt_op_types``). + + ``out``, a UDT of the result's layout, makes the op store each element + into that type instead, cast as a store casts (:func:`fused_store_op`). + + Returns ``(wrapper_func, wrapper_sig, ret_type)``. Raises ``KeyError`` + when the dtype combination is not supported, with a clear message for + common mistakes like passing two record UDTs with different field + names. + """ + from .udf import _get_udt_wrapper + + info_x = _get_udt_info(dtype) + info_y = _get_udt_info(dtype2) + if info_x is None and info_y is None: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}). " + f"Element-wise UDT ops require a record dtype (named fields) " + f"or an array dtype (e.g., FP64[3])." + ) + _check_udt_pair(op_name, dtype, dtype2, info_x, info_y) + ret_type, input_types = _udt_op_types(op_name, dtype, dtype2) + + x_is_scalar = info_x is None # left side is a plain scalar type + y_is_scalar = info_y is None # right side is a plain scalar type + udt_dtype = dtype2 if x_is_scalar else dtype + + # Numba promotes mixed operands its own way (uint64 with int64 is + # int64, where GraphBLAS uses float64) and only then stores the value + # into the result's type, so each operand is converted first to the + # type the built-in op computes in. Same-type operands are left alone. + def cast(in_dtype, elem): + return None if in_dtype == elem.base else in_dtype + + if udt_dtype.np_type.names is not None: + from .udf import _compile_udf_for_udt + + # Use leaf paths so the same codegen handles nested-record UDTs + # uniformly. A non-nested record's leaves are its top-level + # fields, with paths like ``"['a']"``. + leaves = list(_iter_record_leaves(udt_dtype.np_type)) + leaves_x = None if x_is_scalar else list(_iter_record_leaves(dtype.np_type)) + leaves_y = None if y_is_scalar else list(_iter_record_leaves(dtype2.np_type)) + ret_leaves = list(_iter_record_leaves(ret_type.np_type)) + leaf_types = [ + ( + cast(in_x, dtype.np_type if x_is_scalar else leaves_x[i][2]), + cast(in_y, dtype2.np_type if y_is_scalar else leaves_y[i][2]), + in_x, + in_y, + ret_leaves[i][2], + ) + for i, (in_x, in_y) in enumerate(input_types) + ] + func = _make_record_func( + [py for py, _c, _d in leaves], + 2, + py_op, + x_is_scalar=x_is_scalar, + y_is_scalar=y_is_scalar, + leaf_types=leaf_types, + ) + sig = (dtype.numba_type, dtype2.numba_type) + _compile_udf_for_udt( + func, sig, op_kind="binary", op_name=op_name, dtypes=(dtype, dtype2) + ) + if out is not None: + wrapper, wrapper_sig = _make_record_store_wrapper( + func, ret_type, out, dtype, dtype2 + ) + return wrapper, wrapper_sig, out + numba_ret_type = func.overloads[sig].signature.return_type + # Numba can still return a leaf wider than ``ret_type``'s (int8 + # plus int8 is int64 in Numba, INT8 here); the wrapper casts each + # leaf as it stores it, which for those ops gives the same value. + wrapper, wrapper_sig = _get_udt_wrapper( + func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type + ) + else: + # Each side has its own element type, since the bases may differ. + x_type = dtype.np_type.base + y_type = dtype2.np_type.base + ((in_x, in_y),) = input_types + base, shape = ret_type.np_type.subdtype + stored = base if out is None else out.np_type.subdtype[0] + wrapper, wrapper_sig = _make_array_wrapper( + shape, + numba.from_dtype(stored), + 2, + py_op, + x_type=numba.from_dtype(x_type), + y_type=numba.from_dtype(y_type), + x_shape=None if x_is_scalar else dtype.np_type.shape, + y_shape=None if y_is_scalar else dtype2.np_type.shape, + x_cast=cast(in_x, x_type), + y_cast=cast(in_y, y_type), + store_cast=None if out is None else (base, stored), + ) + return wrapper, wrapper_sig, ret_type if out is None else out + + def _make_record_store_wrapper(func, result_type, out_type, dtype, dtype2=None): + """Build a cfunc-ready wrapper storing ``func``'s result as UDT ``out_type``. + + ``func`` returns one value per leaf of ``result_type``, the type the op + computes in. Each value becomes that leaf's dtype and is then cast into + the matching leaf of ``out_type`` as a store casts it + (:func:`_cast_leaf_source`), so storing into an object of another type + needs no temporary of the result's own type. + + Returns (wrapper_func, wrapper_sig). + """ + from .udf import _input_operand + + nt = numba.types + ns = {"numba_func": func} + lines = [" z = numba.carray(z_ptr, 1)"] + params, args, sig_args = ["z_ptr"], [], [nt.CPointer(out_type.numba_type)] + for var, operand in [("x", dtype), ("y", dtype2)]: + if operand is None: + continue + setup, deref, ptr_type = _input_operand(operand, var) + if setup: + lines.append(setup.rstrip("\n")) + params.append(f"{var}_ptr") + args.append(deref) + sig_args.append(ptr_type) + lines.append(f" r = numba_func({', '.join(args)})") + leaves = zip( + _iter_record_leaves(result_type.np_type), + _iter_record_leaves(out_type.np_type), + strict=True, + ) + for k, ((path, _c, computed), (_p, _cz, stored)) in enumerate(leaves): + if computed.shape: + value = _cast_expr(f"r[{k}][i]", computed.base, ns) + value = _cast_leaf_source(computed.base, stored.base, value, ns) + lines.append(f" for i in np.ndindex({computed.shape}):") + lines.append(f" z[0]{path}[i] = {value}") + else: + value = _cast_leaf_source(computed, stored, _cast_expr(f"r[{k}]", computed, ns), ns) + lines.append(f" z[0]{path} = {value}") + src = f"def _store({', '.join(params)}):\n" + "\n".join(lines) + "\n" + wrapper = _compile_codegen( + src, + func_name="_store", + source_label=f"", + extra_ns=ns, + ) + return wrapper, nt.void(*sig_args) + + @lru_cache + def _saturating_cast(dst): + """Return a Numba function converting a float to integer dtype ``dst`` as GraphBLAS does. + + NaN is 0, a value beyond the type's range is its nearest bound, and + anything else truncates toward zero. Numba's own conversion is LLVM's, + which leaves the first two undefined, and numpy's differs by platform. + """ + info = np.iinfo(dst) + lo = "0.0" if dst.kind == "u" else repr(float(info.min)) + src = ( + "def _sat(x):\n" + " x = _f64(x)\n" + f" if x != x or x <= {lo}:\n" + f" return _to(0 if x != x else {info.min})\n" + f" if x >= {float(info.max)!r}:\n" + f" return _to({info.max})\n" + " return _to(x)\n" + ) + func = _compile_codegen( + src, + func_name="_sat", + source_label=f"", + extra_ns={"_f64": numba.float64, "_to": numba.from_dtype(dst)}, + ) + return numba.njit(func) + + def _cast_leaf_source(src, dst, expr, ns): + """Return Numba source converting ``expr`` from numpy dtype ``src`` to ``dst``. + + It converts as GraphBLAS casts built-in types on store: anything + nonzero is True, a complex number loses its imaginary part, a float + stored as an integer saturates (:func:`_saturating_cast`), and + integers wrap. Helpers are added to ``ns``, the generated code's + namespace. + """ + if src == dst: + return expr + if dst.kind == "b": + return f"({expr} != 0)" + if src.kind == "c" and dst.kind != "c": + expr = f"({expr}).real" + src = np.dtype(f"f{src.itemsize // 2}") + if src.kind == "f" and dst.kind in "iu": + ns[f"_sat_{dst.name}"] = _saturating_cast(dst) + return f"_sat_{dst.name}({expr})" + ns[f"_to_{dst.name}"] = numba.from_dtype(dst) + return f"_to_{dst.name}({expr})" + + def compile_udt_cast_wrapper(dtype, to_dtype): + """Compile a cfunc wrapper that stores an element of UDT ``dtype`` as UDT ``to_dtype``. + + ``udt._udt_cast_error`` has accepted the pair. Each element of an array + UDT, and each record leaf, converts by :func:`_cast_leaf_source`. + + Returns (wrapper_func, wrapper_sig). + """ + nt = numba.types + ns = {} + np_x, np_z = dtype.np_type, to_dtype.np_type + if np_x.names is not None: + body = [" x = numba.carray(x_ptr, 1)", " z = numba.carray(z_ptr, 1)"] + for (path, _c, leaf_x), (_path, _cz, leaf_z) in zip( + _iter_record_leaves(np_x), _iter_record_leaves(np_z), strict=True + ): + if leaf_x.shape: + value = _cast_leaf_source(leaf_x.base, leaf_z.base, f"x[0]{path}[i]", ns) + body.append(f" for i in np.ndindex({leaf_x.shape}):") + body.append(f" z[0]{path}[i] = {value}") + else: + value = _cast_leaf_source(leaf_x, leaf_z, f"x[0]{path}", ns) + body.append(f" z[0]{path} = {value}") + sig = nt.void(nt.CPointer(to_dtype.numba_type), nt.CPointer(dtype.numba_type)) + else: + (base_x, shape), base_z = np_x.subdtype, np_z.subdtype[0] + size = reduce(mul, shape) + body = [ + f" x = numba.carray(x_ptr, {size})", + f" z = numba.carray(z_ptr, {size})", + f" for i in range({size}):", + f" z[i] = {_cast_leaf_source(base_x, base_z, 'x[i]', ns)}", + ] + sig = nt.void( + nt.CPointer(numba.from_dtype(base_z)), nt.CPointer(numba.from_dtype(base_x)) + ) + src = "def _cast(z_ptr, x_ptr):\n" + "\n".join(body) + "\n" + wrapper = _compile_codegen( + src, + func_name="_cast", + source_label=f"", + extra_ns=ns, + ) + return wrapper, sig + + def compile_udt_unary_wrapper(op_name, py_op, dtype, out=None): + """Compile a built-in element-wise unary op for a UDT. + + ``out`` stores the result as another UDT, as in + :func:`compile_udt_binary_wrapper`. + + Returns (wrapper_func, wrapper_sig, ret_type). + Raises KeyError if the dtype is not supported. + """ + from .udf import _get_udt_wrapper, _resolve_udt_return_type + + info = _get_udt_info(dtype) + if info is None: + raise KeyError( + f"unary.{op_name} does not work with {dtype}. " + f"Element-wise UDT ops require a record dtype (named fields) " + f"or an array dtype (e.g., FP64[3])." + ) + + kind, _detail = info + if kind == "record": + from .udf import _compile_udf_for_udt + + leaf_paths = [py for py, _c, _d in _iter_record_leaves(dtype.np_type)] + func = _make_record_func(leaf_paths, 1, py_op) + sig = (dtype.numba_type,) + _compile_udf_for_udt(func, sig, op_kind="unary", op_name=op_name, dtypes=(dtype,)) + numba_ret_type = func.overloads[sig].signature.return_type + ret_type = _resolve_udt_return_type(numba_ret_type, dtype) + if out is not None: + wrapper, wrapper_sig = _make_record_store_wrapper(func, ret_type, out, dtype) + return wrapper, wrapper_sig, out + wrapper, wrapper_sig = _get_udt_wrapper( + func, ret_type, dtype, numba_ret_type=numba_ret_type + ) + else: + base_dtype, shape = dtype.np_type.subdtype + ret_type = dtype + stored = base_dtype if out is None else out.np_type.subdtype[0] + wrapper, wrapper_sig = _make_array_wrapper( + shape, + numba.from_dtype(stored), + 1, + py_op, + x_type=numba.from_dtype(base_dtype), + store_cast=None if out is None else (base_dtype, stored), + ) + return wrapper, wrapper_sig, ret_type if out is None else out + + _fused_store_parents = {} + + def fused_store_op(typed_op, out_type): + """Return ``typed_op`` storing its result as ``out_type``, or ``None`` if it cannot. + + The returned op computes what ``typed_op`` computes, in the same type, + and casts each element or leaf into ``out_type`` as it writes it, as a + store casts (``udt._store_cast_op``). So a result goes into an object + of another UDT type in the one GraphBLAS call that computes it, with no + temporary. Only the built-in ops that lift to UDTs have such a form; + ``udt._store_cast_op`` has already accepted the pair of types. + """ + from .udf import _finalize_udt_op + + parent = getattr(typed_op, "parent", None) + if parent is None or parent._anonymous: + return None + if typed_op.opclass == "BinaryOp" and parent.name in BUILTIN_UDT_BINARY_OPS: + from .binary import BinaryOp, TypedUserBinaryOp + + cls, typed_cls = BinaryOp, TypedUserBinaryOp + dtype2 = typed_op.type2 + py_op = BUILTIN_UDT_BINARY_OPS[parent.name] + elif typed_op.opclass == "UnaryOp" and parent.name in BUILTIN_UDT_UNARY_OPS: + from .unary import TypedUserUnaryOp, UnaryOp + + cls, typed_cls = UnaryOp, TypedUserUnaryOp + dtype2 = None + py_op = BUILTIN_UDT_UNARY_OPS[parent.name] + else: + return None + key = (cls, parent.name, out_type) + fused = _fused_store_parents.get(key) + if fused is None: + fused = _fused_store_parents[key] = cls( + f"{parent.name}_as_{out_type.name}", anonymous=True, is_udt=True + ) + dtype = typed_op.type + if (types := dtype if dtype2 is None else (dtype, dtype2)) in fused._udt_ops: + return fused._udt_ops[types] + if dtype2 is None: + wrapper, sig, _ = compile_udt_unary_wrapper(parent.name, py_op, dtype, out=out_type) + else: + wrapper, sig, _ = compile_udt_binary_wrapper( + parent.name, py_op, dtype, dtype2, out=out_type + ) + return _finalize_udt_op(fused, dtype, dtype2, out_type, wrapper, sig, typed_cls) + + +# JIT C code generators below. + + +def _c_assign_binary(py_op, target, lhs, rhs, field_dtype): + """Return a C *statement* assigning ``lhs py_op rhs`` to ``target``. + + Almost every op is a single expression, but floating-point ``//`` and + complex ``/`` and ``*`` need temporaries (see :func:`_c_float_floordiv_stmt`, + :func:`_c_complex_truediv_stmt` and :func:`_c_complex_times_stmt`), so this + is the entry point the kernel builder uses rather than :func:`_c_expr_binary`. + """ + if py_op == "//" and _is_float(field_dtype): + return _c_float_floordiv_stmt(target, lhs, rhs, field_dtype) + if py_op == "/" and field_dtype.kind == "c": + return _c_complex_truediv_stmt(target, lhs, rhs, field_dtype) + if py_op == "*" and field_dtype.kind == "c": + return _c_complex_times_stmt(target, lhs, rhs, field_dtype) + return f"{target} = {_c_expr_binary(py_op, lhs, rhs, field_dtype)} ;" + + +def _c_expr_binary(py_op, lhs, rhs, field_dtype): + """Return a C expression for a binary op: e.g., ``(x->a) + (y->a)``. + + ``field_dtype`` is the numpy dtype of the operands and of the element + being written: a C kernel exists only for an op whose result has the + operands' type. Two of the ops need it because the C spelling that looks + obvious disagrees with what the Numba cfunc computes: ``//`` is not C + ``/``, and ``min`` / ``max`` have to ignore NaN and break a signed-zero tie + as :func:`_minmax_expr` does. It is required rather than optional so a new + caller can't silently get the wrong kernel. :func:`_expr_binary` is the + sibling for the Numba cfunc; SuiteSparse picks between the two on its + own, so they must agree. + + Floating-point ``//`` and complex ``/`` and ``*`` have no single-expression + form; go through :func:`_c_assign_binary`. + """ + if py_op in _FUNC_BINARY_OPS: + return _c_minmax_expr(py_op, lhs, rhs, field_dtype) + if py_op == "/": + return _c_truediv_expr(lhs, rhs, field_dtype) + if py_op == "//": + return _c_int_floordiv_expr(lhs, rhs, field_dtype) + # Everything left (``+``, ``-``, and ``*`` on real fields) spells the same + # in C as in Python and needs no dtype-specific handling. + return f"({lhs}) {py_op} ({rhs})" + + +def _c_minmax_expr(py_op, lhs, rhs, field_dtype): + """Return a C expression for ``min`` / ``max`` on a field, as the Numba cfunc computes it. + + ``GrB_MIN_FP64`` is C99 ``fmin``: it ignores a NaN operand rather than + ordering it, and a UDT field does the same. C99 leaves the tie between + ``-0.0`` and ``0.0`` to the implementation, though: Apple's libm returns + ``-0.0`` for ``fmin``, while glibc's x86-64 ``fmin`` (one ``minsd``) + returns an operand by position. So the comparison is spelled out, the same + as :func:`_minmax_expr` spells it for the cfunc: a NaN on the right loses, + and a tie resolves to ``-0.0`` for ``min`` and ``0.0`` for ``max``, + whichever side each sits on. The C kernel and the cfunc then agree on + every platform; only SuiteSparse's built-in FP64 ``min`` may break that + tie the other way, where the C library does. ``signbit`` comes from + ````, which reaches the C JIT kernel through ``GraphBLAS.h``, the + same route :func:`_c_float_floordiv_stmt` relies on for ``floor`` and + ``fmod``. + + Integers have neither NaN nor a signed zero, so a plain comparison is + exact for them. + """ + cmp_op = "<" if py_op == "min" else ">" + if not _is_float(field_dtype): + return f"(({lhs}) {cmp_op} ({rhs}) ? ({lhs}) : ({rhs}))" + tie = f"signbit({lhs})" if py_op == "min" else f"!signbit({lhs})" + return ( + f"((({lhs}) {cmp_op} ({rhs}) || ({rhs}) != ({rhs}) " + f"|| (({lhs}) == ({rhs}) && {tie})) ? ({lhs}) : ({rhs}))" + ) + + +def _c_truediv_expr(lhs, rhs, field_dtype): + """Return a C expression for true division on a float field. + + Only float and complex fields reach a C kernel. True division of integers + or bools gives float64, as on built-in vectors, so its result never has + the operands' type, which a C kernel needs (``BinaryOp._compile_udt``). + Complex division is a statement (:func:`_c_complex_truediv_stmt`). + """ + if not _is_float(field_dtype): + raise ValueError(f"no C expression divides {field_dtype}") # pragma: no cover (safety) + return f"({lhs}) / ({rhs})" + + +def _c_complex_truediv_stmt(target, lhs, rhs, field_dtype): + """Return a C block dividing two complex fields into ``target``, as the Numba cfunc does. + + C's own ``/`` on ``_Complex`` cannot be used. SuiteSparse's C JIT compiles + with ``-fcx-limited-range`` under GCC, which replaces C99 Annex G division + with the naive formula: it squares the divisor's parts, so ``1e20 / 1e20`` + in complex64 gave ``nan`` and ``(1+1j) / (1e160+1e160j)`` gave ``0`` on + Linux. This block is the cfunc's :func:`_complex_truediv`, statement for + statement: CPython's ``_Py_c_quot`` before Python 3.14 (Smith's method: + scale by the larger part of the divisor), computed in the field's own + precision with real arithmetic only, so no flag changes it. Each product + is its own statement: clang contracts a multiply and an add within one + expression into a fused multiply-add, which rounds once where Numba + rounds twice, and GCC under ``-std=c11`` does not contract at all. A zero + divisor divides the parts by ``0.0`` as reals, for numpy's infinities. + + The constructor is SuiteSparse's ``GB_CMPLX32``/``GB_CMPLX64`` rather + than C11's ``CMPLXF``/``CMPLX``, which macOS does not define. + ``GB_complex.h`` is on the C JIT's include path, names that exact gap, and + its fallback builds the value through a two-element array, so it keeps + infinities and NaNs as given. + """ + ctype, cmplx, creal, cimag, fabs, zero = _c_complex_names(field_dtype) + return ( + f"{{ {ctype} gb_ar = {creal} ({lhs}) ; {ctype} gb_ai = {cimag} ({lhs}) ; " + f"{ctype} gb_br = {creal} ({rhs}) ; {ctype} gb_bi = {cimag} ({rhs}) ; " + f"if (gb_br == {zero} && gb_bi == {zero}) " + f"{{ {target} = {cmplx} (gb_ar / {zero}, gb_ai / {zero}) ; }} " + f"else if ({fabs} (gb_br) >= {fabs} (gb_bi)) {{ " + f"{ctype} gb_r = gb_bi / gb_br ; {ctype} gb_t = gb_bi * gb_r ; " + f"{ctype} gb_d = gb_br + gb_t ; {ctype} gb_u = gb_ai * gb_r ; " + f"{ctype} gb_v = gb_ar * gb_r ; " + f"{target} = {cmplx} ((gb_ar + gb_u) / gb_d, (gb_ai - gb_v) / gb_d) ; }} " + f"else if (gb_bi == {zero}) {{ {target} = {cmplx} (NAN, NAN) ; }} " + f"else {{ {ctype} gb_r = gb_br / gb_bi ; {ctype} gb_t = gb_br * gb_r ; " + f"{ctype} gb_d = gb_t + gb_bi ; {ctype} gb_u = gb_ar * gb_r ; " + f"{ctype} gb_v = gb_ai * gb_r ; " + f"{target} = {cmplx} ((gb_u + gb_ai) / gb_d, (gb_v - gb_ar) / gb_d) ; }} }}" + ) + + +def _c_complex_times_stmt(target, lhs, rhs, field_dtype): + """Return a C block multiplying two complex fields into ``target``, as the Numba cfunc does. + + numpy multiplies complex numbers by the textbook formula. C99 Annex G + adds a recovery step that turns some of its ``nan`` results into + infinities, and whether C's ``*`` on ``_Complex`` takes it depends on the + compiler: GCC under SuiteSparse's ``-fcx-limited-range`` does not, clang + does. So ``1j * (inf+infj)`` was ``nan+nanj`` on Linux and the cfunc but + ``-inf+infj`` from a C kernel on macOS. CPython takes it too from 3.14, + and Numba follows CPython, so the cfunc does not use Numba's ``*`` + either: this block is :func:`_complex_times`, the textbook formula in + real arithmetic, one product per statement so clang cannot contract it + into a fused multiply-add (see :func:`_c_complex_truediv_stmt`). + """ + ctype, cmplx, creal, cimag, _, _ = _c_complex_names(field_dtype) + return ( + f"{{ {ctype} gb_ar = {creal} ({lhs}) ; {ctype} gb_ai = {cimag} ({lhs}) ; " + f"{ctype} gb_br = {creal} ({rhs}) ; {ctype} gb_bi = {cimag} ({rhs}) ; " + f"{ctype} gb_rr = gb_ar * gb_br ; {ctype} gb_ii = gb_ai * gb_bi ; " + f"{ctype} gb_ri = gb_ar * gb_bi ; {ctype} gb_ir = gb_ai * gb_br ; " + f"{target} = {cmplx} (gb_rr - gb_ii, gb_ri + gb_ir) ; }}" + ) + + +def _c_complex_names(field_dtype): + """Return the C type, constructor and part functions for a complex field.""" + if field_dtype.itemsize == 8: + return "float", "GB_CMPLX32", "crealf", "cimagf", "fabsf", "0.0f" + return "double", "GB_CMPLX64", "creal", "cimag", "fabs", "0.0" + + +def _c_int_floordiv_expr(lhs, rhs, field_dtype): + """Return a C expression for integer floor division. + + Three ways C ``/`` differs from what the cfunc computes: + + - It truncates toward zero; Python ``//`` floors. Subtract one when the + operands have different signs and the remainder is non-zero. + - It raises SIGFPE (process death, not an exception) when the divisor is + zero. Numba's numpy error model returns 0, as does ``np.floor_divide``. + - It raises SIGFPE on ``INT_MIN / -1``, whose true quotient is not + representable. numpy wraps to ``INT_MIN``, which is what the cfunc + reaches through negation (see ``_expr_binary``). + """ + if field_dtype.kind in ("u", "b"): + # Both operands are non-negative, so truncation already floors. + return f"(({rhs}) == 0 ? 0 : (({lhs}) / ({rhs})))" + type_min = _C_INT_MIN[field_dtype.itemsize] + floored = f"({lhs}) / ({rhs}) - ((({lhs}) % ({rhs}) != 0) && ((({lhs}) < 0) != (({rhs}) < 0)))" + return ( + f"(({rhs}) == 0 ? 0 : " + f"((({lhs}) == {type_min} && ({rhs}) == -1) ? {type_min} : ({floored})))" + ) + + +def _c_float_floordiv_stmt(target, lhs, rhs, field_dtype): + """Return a C block computing floating-point ``//`` into ``target``. + + ``floor(a / b)`` is not floor division. numpy and CPython both compute + the quotient from the remainder (``(a - fmod(a, b)) / b``) and snap it, + which is more accurate and handles infinities differently: ``1.0 // 0.1`` + is 9.0 but ``floor(1.0 / 0.1)`` is 10.0, and ``inf // 2.0`` is NaN but + ``floor(inf / 2.0)`` is ``inf``. This mirrors numpy's ``npy_divmod`` + (which CPython's ``float_divmod`` matches), so the C JIT kernel, the + Numba cfunc, and ``np.floor_divide`` all agree. + + The sign fix-up needs the remainder twice and the quotient twice, so + this emits a statement with temporaries rather than one expression. + """ + is_f32 = field_dtype.itemsize == 4 + ctype = "float" if is_f32 else "double" + sfx = "f" if is_f32 else "" + half = "0.5f" if is_f32 else "0.5" + one = "1.0f" if is_f32 else "1.0" + zero = "0.0f" if is_f32 else "0.0" + return ( + f"{{ {ctype} gb_a = ({lhs}) ; {ctype} gb_b = ({rhs}) ; {ctype} gb_q ; " + # A zero divisor is the one case numpy answers straight from the + # division: +-inf, or NaN for 0/0. + f"if (gb_b == {zero}) {{ gb_q = gb_a / gb_b ; }} " + f"else {{ {ctype} gb_m = fmod{sfx} (gb_a, gb_b) ; " + f"{ctype} gb_d = (gb_a - gb_m) / gb_b ; " + # NaN compares false both ways, so a NaN remainder skips the + # adjustment and carries through to the quotient. + f"if (gb_m != {zero} && ((gb_b < {zero}) != (gb_m < {zero}))) {{ gb_d -= {one} ; }} " + f"if (gb_d != {zero}) {{ gb_q = floor{sfx} (gb_d) ; " + f"if (gb_d - gb_q > {half}) {{ gb_q += {one} ; }} }} " + f"else {{ gb_q = copysign{sfx} ({zero}, gb_a / gb_b) ; }} }} " + f"{target} = gb_q ; }}" + ) + + +def _c_expr_unary(py_op, operand, field_dtype): + """Return a C expression for a unary op: e.g., ``-(x->a)``. + + ``field_dtype`` is only consulted for ``abs``: + + - Float fields use ``fabs`` / ``fabsf`` so ``abs(-0.0)`` returns + ``+0.0`` (matching Python). The naive ternary preserves the sign bit. + - Complex fields use ``cabs`` / ``cabsf``: the magnitude (a real). It's + assigned to a complex field via implicit ``double -> _Complex`` + conversion (``imag = 0``), matching Numba's behavior when Python + ``abs`` on a complex value is written back to a complex record field. + - Integer fields keep the ternary; ``abs`` of unsigned is a no-op and + ``abs`` of signed int wraps on INT_MIN, both matching the cfunc. + """ + if py_op == "abs": + if field_dtype.kind == "f": + fn = "fabsf" if field_dtype.itemsize == 4 else "fabs" + return f"{fn}({operand})" + if field_dtype.kind == "c": + fn = "cabsf" if field_dtype.itemsize == 8 else "cabs" + return f"{fn}({operand})" + return f"(({operand}) < 0 ? -({operand}) : ({operand}))" + if py_op == "-": + return f"-({operand})" + raise ValueError(f"Unknown unary C op: {py_op}") # pragma: no cover (safety) + + +def _make_jit_c_definition(op_name, py_op, dtype, arity): + """Generate a JIT C function definition for a UDT operator. + + Returns ``(c_name, c_defn)``, or ``None`` if the dtype can't be expressed + in C (unsupported field types, or names that collide with C reserved words). + An op not defined on the field types (``min`` on a complex field) never + gets here: ``udt._element_types`` rejects it when the op is typed. + """ + np_type = dtype.np_type + # Prefer the C name SuiteSparse already has for this type. ``GxB_JIT_C_NAME`` + # on a ``GrB_Type`` is one-shot, so a later ``dtype.name`` rename does not + # propagate to SS. If the op's signature referenced ``dtype.name`` it would + # use an undefined struct name, and SS would silently fall back to the + # Numba cfunc instead of JIT-compiling. ``jit_c_name`` is also where the + # ``gbudt_`` or synthetic ``_gbudt_NNN`` name (see ``_pick_c_type_name``) is + # recorded; using it keeps op codegen consistent. + typedef_info = _dtype_c_typedef(dtype) + if typedef_info is None: + return None + type_name, _typedef = typedef_info + c_name = f"{op_name}_{type_name}" + + if arity == 2: + params = f"{type_name} *z, const {type_name} *x, const {type_name} *y" + else: + params = f"{type_name} *z, const {type_name} *x" + + if np_type.names is not None: + # Use leaf C paths so nested records emit ``z->outer.inner_a = ...`` + # alongside the inner+outer typedefs. Non-nested records degenerate + # to plain ``z->name``. + leaves = list(_iter_record_leaves(np_type)) + if arity == 2: + # Pass the leaf dtype to the binary statement builder so + # type-sensitive ops (min, max, and the divisions) emit correct C. + assigns = " ".join( + _c_assign_binary(py_op, f"z->{c}", f"x->{c}", f"y->{c}", leaf_dtype) + for _py, c, leaf_dtype in leaves + ) + else: + assigns = " ".join( + f"z->{c} = {_c_expr_unary(py_op, f'x->{c}', leaf_dtype)} ;" + for _py, c, leaf_dtype in leaves + ) + else: # array UDT, flattened to v[size] + base_dtype, shape = np_type.subdtype + size = reduce(mul, shape) + # A loop, as in the Numba wrapper: one statement per element made the + # C compiler take seconds on a 64 by 64 array, once per JIT cache. + if arity == 2: + # A statement, not an expression: floating-point ``//`` needs a + # block with temporaries (see ``_c_assign_binary``). + assign = _c_assign_binary(py_op, "z->v[i]", "x->v[i]", "y->v[i]", base_dtype) + else: + assign = f"z->v[i] = {_c_expr_unary(py_op, 'x->v[i]', base_dtype)} ;" + assigns = f"for (int64_t i = 0 ; i < {size} ; i++) {{ {assign} }}" + return c_name, f"void {c_name} ({params}) {{ {assigns} }}" + + +def _c_type_name(dtype): + """The C name SuiteSparse knows ``dtype`` by: a UDT's pinned JIT name, or the C type of + a built-in scalar dtype. ``None`` when it has none. + """ + if dtype.np_type.names is None and dtype.np_type.subdtype is None: + return NP_TO_C_TYPES.get(dtype.np_type) + typedef_info = _dtype_c_typedef(dtype) + return None if typedef_info is None else typedef_info[0] + + +def _dtype_c_typedef(dtype): + """``_udt_c_typedef`` for UDT ``dtype`` under the C name SuiteSparse has for it. + + That is its ``jit_c_name``: pinned when a UDT from numpy registers, or the name of a + type registered from C source. A UDT with none (not expressible in C) gets the name + ``_pick_c_type_name`` would give it. + """ + if (c_name := dtype.jit_c_name) is not None: + return _udt_c_typedef(c_name, dtype.np_type, is_c_name=True) + return _udt_c_typedef(dtype.name, dtype.np_type) + + +def _c_operand(side, path, stored, computed): + """C source reading operand ``side`` (``*x``, or the element at ``path``) in the type the + op computes in. A conversion in C (integer to floating, a real to complex, a narrower + type to a wider one) gives what numpy's ``astype`` gives, which the cfunc uses. + """ + expr = f"(*{side})" if path is None else f"{side}->{path}" + if stored == computed: + return expr + return f"(({NP_TO_C_TYPES[computed]}) ({expr}))" + + +def _make_jit_c_mixed_definition(op_name, py_op, dtype, dtype2, ret_type, input_types): + """Generate a JIT C function definition for a lifted op whose operand types differ, or + whose result type differs from them: a scalar with an array or record UDT (either + order), two array UDTs of one shape with different element types, two records with the + same fields and different leaf types. + + The signature has a C type for each of z, x and y, so the C name carries all three. A + scalar operand is read as ``(*x)`` for every element or leaf. Returns ``(c_name, + c_defn)``, or ``None`` when one of the types has no C name, when an array operand + broadcasts (another shape than the result's), when a record has an array leaf, or when + an element is not computed in the result's element type (integer ``truediv`` computes + in its operands' type and returns float); those keep the Numba cfunc. + + ``input_types`` is what :func:`udt._udt_op_types` returns with ``ret_type``: each + element is computed as the built-in op on the operands' element types computes it, + after converting both operands to that type, as the cfunc does. + """ + names = [_c_type_name(t) for t in (ret_type, dtype, dtype2)] + if None in names: + return None + z_name, x_name, y_name = names + np_x, np_y, np_z = dtype.np_type, dtype2.np_type, ret_type.np_type + + def assign(target, x_path, y_path, x_stored, y_stored, types, out): + in_x, in_y = types + if not in_x == in_y == out: + return None + lhs = _c_operand("x", x_path, x_stored, in_x) + rhs = _c_operand("y", y_path, y_stored, in_y) + return _c_assign_binary(py_op, target, lhs, rhs, out) + + if np_z.names is not None: + leaves_z = list(_iter_record_leaves(np_z)) + leaves_x = list(_iter_record_leaves(np_x)) if np_x.names is not None else None + leaves_y = list(_iter_record_leaves(np_y)) if np_y.names is not None else None + stmts = [] + for i, (_py, c, leaf) in enumerate(leaves_z): + # A record with an array leaf has no C type name, so it returned above + if leaf.shape: # pragma: no cover (safety) + return None + x_path, x_stored = (None, np_x) if leaves_x is None else (c, leaves_x[i][2]) + y_path, y_stored = (None, np_y) if leaves_y is None else (c, leaves_y[i][2]) + stmt = assign(f"z->{c}", x_path, y_path, x_stored, y_stored, input_types[i], leaf) + if stmt is None: + return None + stmts.append(stmt) + body = " ".join(stmts) + else: + base_z, shape_z = np_z.subdtype + operands = [] + for np_t in (np_x, np_y): + if np_t.subdtype is None: + operands.append((None, np_t)) + elif np_t.subdtype[1] != shape_z: + return None # (broadcasting reads through an index map) + else: + operands.append(("v[i]", np_t.subdtype[0])) + (x_path, x_stored), (y_path, y_stored) = operands + stmt = assign("z->v[i]", x_path, y_path, x_stored, y_stored, input_types[0], base_z) + if stmt is None: + return None + body = f"for (int64_t i = 0 ; i < {reduce(mul, shape_z)} ; i++) {{ {stmt} }}" + # The length of x's name keeps two pairs from sharing a C name ("A" and "B_C", "A_B" and + # "C"), and keeps it from any one-type name (``{op_name}_{type}``, whose type name cannot + # start with a digit). SuiteSparse finds a compiled kernel by its op's name, so two ops + # with one name would run each other's kernel. + c_name = f"{op_name}_{len(x_name)}_{x_name}_{y_name}" + if len(c_name) >= lib.GxB_MAX_NAME_LEN: + # SuiteSparse keeps at most 127 characters (see ``_set_jit_c_strings``). A length of 0 + # cannot come from the readable form, so this name is still the pair's own. + c_name = f"{op_name}_0_{hashlib.sha256(c_name.encode()).hexdigest()[:32]}" + params = f"{z_name} *z, const {x_name} *x, const {y_name} *y" + return c_name, f"void {c_name} ({params}) {{ {body} }}" + + +def _make_jit_c_comparison_definition(op_name, dtype, *, is_eq): + """Generate a JIT C function definition for ``binary.eq[udt]`` / + ``binary.ne[udt]``. + + Returns ``(c_name, c_defn)``, or ``None`` when the dtype can't be + expressed in C. The kernel signature is + ``void op(_Bool *z, const Udt *x, const Udt *y)``: each leaf field + contributes a scalar ``==`` (or ``!=``) comparison; record-UDT leaves + are chained with ``&&`` (eq) or ``||`` (ne). Top-level array UDTs loop + over their elements. + + Array-valued sub-fields inside a record (e.g. + ``[("weights", (float64, 3))]``) aren't supported: ``_udt_c_typedef`` + rejects records with array sub-fields (``NP_TO_C_TYPES`` has no entry + for sub-array dtypes), so this function returns ``None`` before + iterating leaves in that case. Adding support would mean extending + ``_udt_c_typedef`` to emit ``double name[N]`` and ``_make_jit_c_definition`` + to unroll array sub-fields in arithmetic codegen too; both legs need + to land together or the eq/ne path is alone in supporting it. + + IEEE NaN propagation comes for free: C ``a == b`` is false when either + side is NaN, so two records both carrying NaN compare unequal under + ``eq`` (and equal under ``ne``), matching the cfunc path's leaf-wise + semantic comparison. + """ + np_type = dtype.np_type + # ``_udt_c_typedef`` returns None when the numpy layout doesn't match a + # C compiler's layout for the same struct. + typedef_info = _dtype_c_typedef(dtype) + if typedef_info is None: + return None + type_name, _typedef = typedef_info + c_name = f"{op_name}_{type_name}" + op = "==" if is_eq else "!=" + join = " && " if is_eq else " || " + if np_type.subdtype is not None: + # Array UDT: a loop over every element (see ``_make_jit_c_definition`` + # and the Numba wrapper in ``binary._make_udt_comparison``). + _base, shape = np_type.subdtype + size = reduce(mul, shape) + accumulate = "&=" if is_eq else "|=" + body = ( + f"_Bool r = {int(is_eq)} ; for (int64_t i = 0 ; i < {size} ; i++) " + f"{{ r {accumulate} ((x->v[i]) {op} (y->v[i])) ; }} *z = r ;" + ) + else: + terms = [f"((x->{c}) {op} (y->{c}))" for _py, c, _d in _iter_record_leaves(np_type)] + body = f"*z = {join.join(terms) if terms else int(is_eq)} ;" + params = f"_Bool *z, const {type_name} *x, const {type_name} *y" + return c_name, f"void {c_name} ({params}) {{ {body} }}" + + +def _maybe_warn_jit_skipped(jit_info, op_name, dtype_name): + """Emit a ``NoJITWarning`` when JIT setup was skipped despite SS supporting JIT. + + Called by op-class ``_compile_udt`` paths after ``set_jit_c_*_on_op`` + returns ``None``. The ``_has_jit_set`` gate keeps the warning silent on + SS < 8 where JIT is fundamentally absent. + """ + if jit_info is None and _has_jit_set: + from ..ss.jit_config import _maybe_warn_no_jit + + _maybe_warn_no_jit(op_name=op_name, dtype_name=dtype_name) + + +def _set_jit_c_strings(gb_obj, c_name, c_defn, set_string_func): + """Pin ``GxB_JIT_C_NAME`` + ``GxB_JIT_C_DEFINITION`` on a fresh op handle. + + Both strings are one-shot on SuiteSparse: a second set returns + ``GrB_ALREADY_SET`` silently. The return is not checked here because + callers always pass a freshly-allocated handle. + + Arming an op with C source is the point at which this process has + committed to wanting a JIT compiler, so it is also where SuiteSparse's + non-compiling default gets raised. + """ + from ..ss.jit_config import _enable_jit_for_udt + + _enable_jit_for_udt() + # SuiteSparse refuses a name of GxB_MAX_NAME_LEN (128) characters or more, then takes the + # definition under the empty name: its kernels call no function and leave their outputs + # uninitialized. So the definition goes only with the name, and without it the op keeps + # its cfunc. + info = set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) + if info != lib.GrB_SUCCESS: + return False + set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) + return True + + +def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq): + """Generate and set JIT C name+definition for ``eq``/``ne`` on a UDT. + + Sibling of :func:`set_jit_c_on_op` for the structurally different + comparison kernel (BOOL output, chained leaf comparisons). Returns + ``(c_name, c_defn)`` on success, ``None`` when the dtype can't be + expressed in C; callers cache the returned strings on the op for + introspection (see ``TypedUserBinaryOp.jit_c_source``). Callers check + ``_has_jit_set`` first, as for :func:`set_jit_c_on_op`. + """ + result = _make_jit_c_comparison_definition(op_name, dtype, is_eq=is_eq) + if result is None: + return None + if not _set_jit_c_strings(gb_obj, *result, set_string_func): + return None + return result + + +def set_jit_c_mixed_on_op(gb_obj, op_name, py_op, dtype, dtype2, ret_type, input_types): + """Generate and set the JIT C name and definition of a lifted binary op on two types + (:func:`_make_jit_c_mixed_definition`). Returns ``(c_name, c_defn)``, or ``None`` when + the op has no C form (it keeps the Numba cfunc). Callers check ``_has_jit_set`` first, + as for :func:`set_jit_c_on_op`: without it, SuiteSparse has no + ``GrB_BinaryOp_set_String``. + """ + result = _make_jit_c_mixed_definition(op_name, py_op, dtype, dtype2, ret_type, input_types) + if result is None: + return None + if not _set_jit_c_strings(gb_obj, *result, lib.GrB_BinaryOp_set_String): + return None + return result + + +def set_jit_c_on_op(gb_obj, op_name, py_op, dtype, set_string_func, arity=2): + """Generate and set JIT C name+definition on a GrB operator, if possible. + + Returns ``(c_name, c_defn)`` when the dtype is expressible in C, ``None`` + otherwise. Callers may cache the returned strings for introspection (see + ``TypedUserUnaryOp.jit_c_source`` and ``TypedUserBinaryOp.jit_c_source``). + Callers check ``_has_jit_set`` first: without it, SuiteSparse has no + ``set_string_func`` to pass. + """ + result = _make_jit_c_definition(op_name, py_op, dtype, arity) + if result is None: + return None + if not _set_jit_c_strings(gb_obj, *result, set_string_func): + return None + return result diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index 26e0ca61c..e8f0870e0 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -27,7 +27,6 @@ OpBase, ParameterizedUdf, TypedOpBase, - _deserialize_parameterized, _hasop, ) @@ -36,7 +35,13 @@ if _has_numba: import numba - from .base import _get_udt_wrapper + from .udf import ( + _bool_to_int8, + _compile_udf_for_udt, + _finalize_udt_op, + _get_udt_wrapper, + _resolve_udt_return_type, + ) ffi_new = ffi.new @@ -73,6 +78,7 @@ def __call__(self, val): class TypedUserUnaryOp(TypedOpBase): __slots__ = () opclass = "UnaryOp" + _owns_gb_obj = True def __init__(self, parent, name, type_, return_type, gb_obj): super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}") @@ -104,20 +110,6 @@ def _call(self, *args, **kwargs): unary._parameterized_info = (self, args, kwargs) return UnaryOp.register_anonymous(unary, self.name, is_udt=self._is_udt) - def __reduce__(self): - name = f"unary.{self.name}" - if not self._anonymous and name in _STANDARD_OPERATOR_NAMES: # pragma: no cover - return name - return (self._deserialize, (self.name, self.func, self._anonymous)) - - @staticmethod - def _deserialize(name, func, anonymous): - if anonymous: - return UnaryOp.register_anonymous(func, name, parameterized=True) - if (rv := UnaryOp._find(name)) is not None: - return rv - return UnaryOp.register_new(name, func, parameterized=True) - def _identity(x): return x # pragma: no cover (numba) @@ -127,11 +119,22 @@ def _one(x): return 1 # pragma: no cover (numba) +if _has_numba: + from ..udt import BUILTIN_UDT_UNARY_OPS as _BUILTIN_UDT_UNARY_OPS + from .udt_codegen import ( + _has_jit_set, + _maybe_warn_jit_skipped, + compile_udt_cast_wrapper, + compile_udt_unary_wrapper, + set_jit_c_on_op, + ) + + class UnaryOp(OpBase): """Takes one input and returns one output, possibly of a different data type. Built-in and registered UnaryOps are located in the ``graphblas.unary`` namespace - as well as in the ``graphblas.ops`` combined namespace. + as well as in the ``graphblas.op`` combined namespace. """ __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" @@ -174,7 +177,9 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): if name is None: name = getattr(func, "__name__", "") success = False - unary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + unary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=unary_udf) return_types = {} nt = numba.types @@ -199,11 +204,8 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): elif type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: ret_type = INT8 - # Numba is unable to handle BOOL correctly right now, but we have a workaround - # See: https://github.com/numba/numba/issues/5395 - # We're relying on coercion behaving correctly here - input_type = INT8 if type_ == BOOL else type_ - return_type = INT8 if ret_type == BOOL else ret_type + input_type = _bool_to_int8(type_) + return_type = _bool_to_int8(ret_type) # Build wrapper because GraphBLAS wants pointers and void return wrapper_sig = nt.void( @@ -232,7 +234,9 @@ def unary_wrapper(z, x): def unary_wrapper(z, x): z[0] = unary_udf(x[0]) # pragma: no cover (numba) - unary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(unary_wrapper) + unary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + unary_wrapper + ) new_unary = ffi_new("GrB_UnaryOp*") check_status_carg( lib.GrB_UnaryOp_new( @@ -252,26 +256,46 @@ def unary_wrapper(z, x): def _compile_udt(self, dtype, dtype2): if dtype in self._udt_types: return self._udt_ops[dtype] + from ..udt import _refuse_half + + _refuse_half(self, dtype) + # See the matching note in ``BinaryOp._compile_udt``: plain-scalar + # misses on UDT-dispatch-enabled ops get the legacy message. + if not dtype._is_udt and self._numba_func is None: + raise KeyError(f"{self.name} does not work with {dtype}") + if self._numba_func is None and _has_numba and self.name in _BUILTIN_UDT_UNARY_OPS: + # Auto-generate an element-wise UDT function for the built-in + # unary ops (``ainv``, ``abs``). + py_op = _BUILTIN_UDT_UNARY_OPS[self.name] + unary_wrapper, wrapper_sig, ret_type = compile_udt_unary_wrapper( + self.name, py_op, dtype + ) + op = _finalize_udt_op( + self, dtype, None, ret_type, unary_wrapper, wrapper_sig, TypedUserUnaryOp + ) + # ``_has_jit_set`` gates access to ``lib.GrB_UnaryOp_set_String``, + # which is absent on SS < 9; the attribute lookup itself raises + # there, so it must not be evaluated as an argument first. + if _has_jit_set: + op._jit_c_info = set_jit_c_on_op( + op.gb_obj, self.name, py_op, dtype, lib.GrB_UnaryOp_set_String, arity=1 + ) + _maybe_warn_jit_skipped(op._jit_c_info, self.name, dtype.name) + return op if self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtype}") numba_func = self._numba_func sig = (dtype.numba_type,) - numba_func.compile(sig) # Should we catch and give additional error message? - ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) - - unary_wrapper, wrapper_sig = _get_udt_wrapper(numba_func, ret_type, dtype) - unary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(unary_wrapper) - new_unary = ffi_new("GrB_UnaryOp*") - check_status_carg( - lib.GrB_UnaryOp_new(new_unary, unary_wrapper.cffi, ret_type._carg, dtype._carg), - "UnaryOp", - new_unary[0], + _compile_udf_for_udt(numba_func, sig, op_kind="unary", op_name=self.name, dtypes=(dtype,)) + numba_ret_type = numba_func.overloads[sig].signature.return_type + ret_type = _resolve_udt_return_type(numba_ret_type, dtype) + unary_wrapper, wrapper_sig = _get_udt_wrapper( + numba_func, ret_type, dtype, numba_ret_type=numba_ret_type + ) + return _finalize_udt_op( + self, dtype, None, ret_type, unary_wrapper, wrapper_sig, TypedUserUnaryOp ) - op = TypedUserUnaryOp(self, self.name, dtype, ret_type, new_unary[0]) - self._udt_types[dtype] = ret_type - self._udt_ops[dtype] = op - return op @classmethod def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): @@ -295,7 +319,8 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals can be used as: ``A.ewise_mult(B, gb.binary.isclose(rel_tol=1e-5))``. When creating a parameterized user-defined operator, the ``func`` parameter must be a callable that *returns* a function that will - then get compiled. See the ``user_isclose`` example below. + then get compiled. See the ``user_isclose`` example in + ``BinaryOp.register_new``. is_udt : bool, default False Whether the operator is intended to operate on user-defined types. If True, then the function will not be automatically compiled for @@ -333,7 +358,8 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal can be used as: ``A.ewise_mult(B, gb.binary.isclose(rel_tol=1e-5))``. When creating a parameterized user-defined operator, the ``func`` parameter must be a callable that *returns* a function that will - then get compiled. See the ``user_isclose`` example below. + then get compiled. See the ``user_isclose`` example in + ``BinaryOp.register_new``. is_udt : bool, default False Whether the operator is intended to operate on user-defined types. If True, then the function will not be automatically compiled for @@ -357,7 +383,7 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized}, + {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, ) elif parameterized: unary_op = ParameterizedUnaryOp(name, func, is_udt=is_udt) @@ -437,11 +463,18 @@ def _initialize(cls): ]: unop.orig_func = func if _has_numba: - unop._numba_func = numba.njit(func) + unop._numba_func = numba.njit(func, error_model="numpy") else: unop._numba_func = None unop._udt_types = {} unop._udt_ops = {} + # Enable element-wise unary ops on UDTs (``ainv``, ``abs``). + if _has_numba: + for op_name in _BUILTIN_UDT_UNARY_OPS: + unop = getattr(unary, op_name, None) + if unop is not None: + unop._udt_types = {} + unop._udt_ops = {} cls._initialized = True def __init__( @@ -463,13 +496,27 @@ def __init__( self._udt_types = {} # {dtype: DataType} self._udt_ops = {} # {dtype: TypedUserUnaryOp} - def __reduce__(self): - if self._anonymous: - if hasattr(self.orig_func, "_parameterized_info"): - return (_deserialize_parameterized, self.orig_func._parameterized_info) - return (self.register_anonymous, (self.orig_func, self.name)) - if (name := f"unary.{self.name}") in _STANDARD_OPERATOR_NAMES: - return name - return (self._deserialize, (self.name, self.orig_func)) - __call__ = TypedBuiltinUnaryOp.__call__ + + +ParameterizedUnaryOp._op_class = UnaryOp + + +# {to_dtype: UnaryOp}; each one's typed ops are keyed by the type they cast from. +_udt_casts = {} + + +def _udt_cast_op(dtype, to_dtype): + """Return the UnaryOp storing an element of UDT ``dtype`` as UDT ``to_dtype``. + + ``udt._store_cast_op`` decides when to use it, after + ``udt._udt_cast_error`` has accepted the pair. It is compiled once + per pair and has no JIT kernel. + """ + parent = _udt_casts.get(to_dtype) + if parent is None: + parent = _udt_casts[to_dtype] = UnaryOp(f"cast_to_{to_dtype}", anonymous=True, is_udt=True) + if dtype in parent._udt_ops: + return parent._udt_ops[dtype] + wrapper, wrapper_sig = compile_udt_cast_wrapper(dtype, to_dtype) + return _finalize_udt_op(parent, dtype, None, to_dtype, wrapper, wrapper_sig, TypedUserUnaryOp) diff --git a/graphblas/core/operator/utils.py b/graphblas/core/operator/utils.py index 1442a9b5e..c3a8f5b8f 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -1,8 +1,21 @@ from types import BuiltinFunctionType, FunctionType, ModuleType -from ... import backend, binary, config, indexunary, monoid, op, select, semiring, unary +from ... import ( + _STANDARD_OPERATOR_NAMES, + backend, + binary, + config, + indexbinary, + indexunary, + monoid, + op, + select, + semiring, + unary, +) from ...dtypes import UINT64, lookup_dtype, unify from ..expr import InfixExprBase +from ..udt import _monoid_pair_type from .base import ( _SS_OPERATORS, OpBase, @@ -14,6 +27,7 @@ find_opclass, ) from .binary import BinaryOp +from .indexbinary import IndexBinaryOp from .indexunary import IndexUnaryOp from .monoid import Monoid from .select import SelectOp @@ -28,6 +42,27 @@ BinaryOp._initialize() Monoid._initialize() Semiring._initialize() + # IndexBinaryOp has no built-ins, but ``_initialize`` still needs to + # set ``cls._initialized = True``. Without that, user ``register_new`` + # calls hit the "not initialized" branch and mistakenly add user op + # names to ``_STANDARD_OPERATOR_NAMES``. Pickle then emits a string + # reference that won't resolve in a fresh process. Cross-process pickle + # for user-registered IBOs (and bound IBOs derived from them) requires + # the tuple-form ``__reduce__`` that re-registers the function in the + # child. + IndexBinaryOp._initialize() + # Every op the namespaces hold now is python-graphblas's own, also ones made + # outside ``register_new`` (cdiv, truediv, rtruediv, the semirings on cdiv) and + # ones made on first use (binom). Record them all as standard: they pickle by + # name, and ``literals._is_builtin_op`` counts them built-in, with no list to + # keep in step. + for _module in [binary, unary, monoid, semiring, indexunary, select, indexbinary]: + _modname = _module.__name__.rsplit(".", 1)[-1] + _STANDARD_OPERATOR_NAMES.update(f"{_modname}.{key}" for key in _module._delayed) + for _val in list(vars(_module).values()): + if isinstance(_val, (OpBase, ParameterizedUdf)) and not _val._anonymous: + _STANDARD_OPERATOR_NAMES.add(f"{_modname}.{_val.name}") + del _module, _modname, _val except Exception: # pragma: no cover (debug) # Exceptions here can often get ignored by Python import traceback @@ -47,6 +82,25 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala # Handle special cases such as first and second (may have UDTs) if op._custom_dtype is not None and (rv := op._custom_dtype(op, dtype, dtype2)) is not None: return rv + # UDTs do not cast, so never unify a UDT with another type: numpy would + # promote two records to one of them (or a third), and GraphBLAS would + # reject the operand of the other type with GrB_DOMAIN_MISMATCH. An op + # that does not lift to UDTs says so here with a KeyError, where a UDT + # with a built-in type mostly failed to unify with numpy's TypeError. + if dtype is not dtype2 and (dtype._is_udt or dtype2._is_udt): + if op.is_positional: + return op[UINT64] + if op._udt_types is None: + raise KeyError(f"{op.name} does not work with ({dtype}, {dtype2})") + if isinstance(op, Monoid): + # A Monoid takes one type, but used element-wise it is its + # BinaryOp, which takes the pair: ``v.ewise_mult(w, monoid.plus)`` + # works where ``v * w`` does. + if kind == "binary": + return get_typed_op(op.binaryop, dtype, dtype2) + # Indexed with the pair, it takes the type the pair promotes to. + return op[_monoid_pair_type(op.binaryop.name, dtype, dtype2)] + return op._compile_udt(dtype, dtype2) # Generic case: try to unify the two dtypes try: return op[ @@ -198,9 +252,10 @@ def get_semiring(monoid, binaryop, name=None): monoid_type = monoid.type monoid = monoid.parent if isinstance(binaryop, BinaryOp): - binary_type = None + binary_type = binary_type2 = None else: binary_type = binaryop.type + binary_type2 = getattr(binaryop, "type2", binary_type) binaryop = binaryop.parent if monoid._anonymous or binaryop._anonymous: rv = Semiring.register_anonymous(monoid, binaryop, name=name) @@ -253,6 +308,9 @@ def get_semiring(monoid, binaryop, name=None): setattr(module, funcname, rv) if binary_type is not None: + if binary_type2 is not binary_type and (binary_type._is_udt or binary_type2._is_udt): + # A BinaryOp typed on two different UDTs gives a semiring typed on both. + return get_typed_op(rv, binary_type, binary_type2) return rv[binary_type] if monoid_type is not None: return rv[monoid_type] @@ -261,6 +319,8 @@ def get_semiring(monoid, binaryop, name=None): unary.register_new = UnaryOp.register_new unary.register_anonymous = UnaryOp.register_anonymous +indexbinary.register_new = IndexBinaryOp.register_new +indexbinary.register_anonymous = IndexBinaryOp.register_anonymous indexunary.register_new = IndexUnaryOp.register_new indexunary.register_anonymous = IndexUnaryOp.register_anonymous select.register_new = SelectOp.register_new diff --git a/graphblas/core/recorder.py b/graphblas/core/recorder.py index ca776f697..90f35bb82 100644 --- a/graphblas/core/recorder.py +++ b/graphblas/core/recorder.py @@ -179,4 +179,4 @@ def __repr__(self): skip_record = Recorder(start=False) -skip_record.data = collections.deque([], 0) +skip_record.data = collections.deque(maxlen=0) diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 25aef5743..0b711c258 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -8,13 +8,34 @@ from . import _has_numba, _supports_udfs, automethods, ffi, lib, utils from .base import BaseExpression, BaseType, call from .expr import AmbiguousAssignOrExtract +from .literals import ( + _SINGLE_PRECISION, + _check_literal_fits, + _check_scalar_fits, + _literal_operand, + _literal_store_dtype, + _literal_type, +) from .operator import get_typed_op +from .operator.binary import _isclose +from .udt import ( + _element_bytes, + _element_shapes_differ, + _element_value, + _ewise_add_needs_cast, + _get_udt_info, + _has_half_leaf, + _is_lifted_op, + _literal_shape_differs, + _scalar_ewise_add_converted, + _udt_eq_without_numba, +) from .utils import _Pointer, output_type, wrapdoc if _supports_udfs: from ..binary import isclose else: - from .operator.binary import _isclose as isclose + isclose = _isclose ffi_new = ffi.new @@ -131,13 +152,19 @@ def _repr_html_(self, collapse=False, expr=None): def __eq__(self, other): """Check equality. - Equality comparison uses :meth:`isequal`. Use that directly for finer control of - what is considered equal. + Compares with ``binary.eq``, as ``==`` on a Vector or Matrix does, so the + elements of array UDTs broadcast (an ``INT8[3]`` Scalar of ``[1, 1, 1]`` + equals ``1``), but returns a bool. Two empty Scalars are equal, and + ``s == None`` checks whether ``s`` is empty. :meth:`isequal` is stricter. + Without Numba, which ``binary.eq`` needs for UDTs, a UDT's values are + compared in numpy, which broadcasts them the same way. """ - return self.isequal(other) + return _scalar_eq(self, other, "__eq__") + + __hash__ = None def __ne__(self, other): - return not self.isequal(other) + return not _scalar_eq(self, other, "__ne__") def __bool__(self): """Truthiness check. @@ -178,7 +205,7 @@ def __sizeof__(self): size = ffi_new("size_t*") check_status(lib.GxB_Scalar_memoryUsage(size, self.gb_obj[0]), self) return base + size[0] - raise TypeError("Unable to get size of GrB_Scalar with backend: {backend}") + raise TypeError(f"Unable to get size of GrB_Scalar with backend: {backend}") def isequal(self, other, *, check_dtype=False): """Check for exact equality (including whether the value is missing). @@ -187,13 +214,18 @@ def isequal(self, other, *, check_dtype=False): ---------- other : Scalar Scalar to compare against - check_dtypes : bool, default=False + check_dtype : bool, default=False If True, also checks that dtypes match Returns ------- bool + Notes + ----- + Elements of array UDTs must have the same shape, apart from leading axes + of length 1, as in ``np.array_equal``; ``==`` broadcasts instead. + See Also -------- :meth:`isclose` : For equality check of floating point dtypes @@ -202,8 +234,15 @@ def isequal(self, other, *, check_dtype=False): if type(other) is not Scalar: if other is None: return self._is_empty + dtype = None + if self.dtype._is_udt: + # A literal is compared as given, as ``==`` types it, not converted + # into the UDT (0.5 would become 0 in an int field), and an array + # literal must have the element's shape, as in np.array_equal. + if _literal_shape_differs(self.dtype, other): + return False + dtype = _literal_type(self.dtype, other, exact=True) or self.dtype try: - dtype = self.dtype if self.dtype._is_udt else None other = Scalar.from_value(other, dtype, is_cscalar=None, name="s_isequal") except TypeError: other = self._expect_type( @@ -217,6 +256,8 @@ def isequal(self, other, *, check_dtype=False): check_dtype = False if check_dtype and self.dtype != other.dtype: return False + if _element_shapes_differ(self.dtype, other.dtype): + return False if self._is_empty: return other._is_empty if other._is_empty: @@ -231,7 +272,8 @@ def isequal(self, other, *, check_dtype=False): def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False): """Check for approximate equality (including whether the value is missing). - Equivalent to: ``abs(a-b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)``. + Equivalent to: ``abs(a-b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)``, + except that, as in ``math.isclose``, an infinity is close only to itself. Parameters ---------- @@ -249,6 +291,11 @@ def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False): bool """ + if self.dtype._is_udt: + raise TypeError( + f"Scalar.isclose is not defined for user-defined types (got {self.dtype.name!r}); " + "use isequal for exact comparison." + ) if type(other) is not Scalar: if other is None: return self._is_empty @@ -273,8 +320,9 @@ def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False): # We can't yet call a UDF on a scalar as part of the spec, so let's do it ourselves isclose_func = isclose(rel_tol, abs_tol) if not _has_numba: - # Check if types are compatible - get_typed_op( + # Check if types are compatible, and compare in the type they unify to, as + # with Numba (INT64 with UINT64 is in FP64). Python ints do not wrap in x - y. + eq = get_typed_op( binary.eq, self.dtype, other.dtype, @@ -282,7 +330,9 @@ def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False): is_right_scalar=True, kind="binary", ) - return isclose_func(self.value, other.value) + x = eq.type.np_type.type(self.value).item() + y = eq.type2.np_type.type(other.value).item() + return isclose_func(x, y) isclose_func = get_typed_op( isclose_func, self.dtype, @@ -291,7 +341,11 @@ def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False): is_right_scalar=True, kind="binary", ) - return isclose_func._numba_func(self.value, other.value) + # Pass the values in the op's input types, as ewise_mult does; Numba types a + # Python int as int64, which a UINT64 value need not fit + x = isclose_func.type.np_type.type(self.value) + y = isclose_func.type2.np_type.type(other.value) + return isclose_func._numba_func(x, y) def clear(self): """In-place operation which clears the value in the Scalar. @@ -344,11 +398,7 @@ def value(self): dtype_name = "UDT" if is_udt else self.dtype.name call(f"GrB_Scalar_extractElement_{dtype_name}", [_Pointer(scalar), self]) if is_udt: - np_type = self.dtype.np_type - rv = np.array(ffi.buffer(scalar.gb_obj[0 : np_type.itemsize])) - if np_type.subdtype is None: - return rv.view(np_type)[0] - return rv.view(np_type.subdtype[0]).reshape(np_type.subdtype[1]) + return _element_value(self.dtype, scalar.gb_obj) return scalar.gb_obj[0] @value.setter @@ -359,16 +409,7 @@ def value(self, val): if output_type(val) is Scalar: val = val.value # raise below if wrong type (as determined by cffi) if self.dtype._is_udt: - np_type = self.dtype.np_type - if np_type.subdtype is None: - arr = np.empty(1, dtype=np_type) - if isinstance(val, dict) and np_type.names is not None: - val = _dict_to_record(np_type, val) - else: - arr = np.empty(np_type.subdtype[1], dtype=np_type.subdtype[0]) - arr[:] = val - self.gb_obj[0 : self.dtype.np_type.itemsize] = arr.view(np.uint8) - # self.gb_obj[0:self.dtype.np_type.itemsize] = bytes(val) + self.gb_obj[0 : self.dtype.np_type.itemsize] = _element_bytes(self.dtype, val) else: self.gb_obj[0] = val self._empty = False @@ -377,7 +418,7 @@ def value(self, val): val = _Pointer(_as_scalar(val, self.dtype, is_cscalar=True)) dtype_name = "UDT" else: - val = _as_scalar(val, is_cscalar=True) + val = _as_scalar(val, _literal_store_dtype(self.dtype, val), is_cscalar=True) dtype_name = val.dtype.name call(f"GrB_Scalar_setElement_{dtype_name}", [self, val]) @@ -455,7 +496,11 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=None, name=None): else: new_scalar = Scalar(dtype, is_cscalar=is_cscalar, name=name) if not clear and not self._is_empty: - if new_scalar.is_cscalar and not new_scalar.dtype._is_udt: + if self.dtype._is_udt: + # Cast as storing it would (``udt._store_cast_op``), not as numpy + # does; a UDT does not cast to a built-in dtype at all. + new_scalar << self + elif new_scalar.is_cscalar and not new_scalar.dtype._is_udt: # Cast value so we don't raise given explicit dup with dtype new_scalar.value = new_scalar.dtype.np_type.type(self.value) else: @@ -653,22 +698,15 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): self = op(self).new() if isinstance(other, ScalarEwiseAddExpr): other = op(other).new() + if output_type(self) is not Scalar: # A literal on the left, as in `op(1 | s)` + self = _literal_scalar(other, self, op, within=method_name, argname="left") if type(other) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None - try: - other = Scalar.from_value(other, dtype, is_cscalar=False, name="") - except TypeError: - other = self._expect_type( - other, - Scalar, - within=method_name, - keyword_name="other", - extra_message="Literal scalars also accepted.", - op=op, - ) + other = _literal_scalar(self, other, op, within=method_name, keyword_name="other") op = get_typed_op(op, self.dtype, other.dtype, kind="binary") self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") + if _ewise_add_needs_cast(op, self.dtype, other.dtype): + self, other, op = _scalar_ewise_add_converted(self, other, op) return ScalarExpression( method_name, f"GrB_Vector_eWiseAdd_{op.opclass}", @@ -712,6 +750,7 @@ def ewise_mult(self, other, op=binary.times): def _ewise_mult(self, other, op=binary.times, is_infix=False): method_name = "ewise_mult" + given = other if is_infix: from .infix import ScalarEwiseMultExpr @@ -724,23 +763,14 @@ def _ewise_mult(self, other, op=binary.times, is_infix=False): self = op(self).new() if isinstance(other, ScalarEwiseMultExpr): other = op(other).new() + if output_type(self) is not Scalar: # A literal on the left, as in `op(1 & s)` + self = _literal_scalar(other, self, op, within=method_name, argname="left") if type(other) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None - try: - other = Scalar.from_value(other, dtype, is_cscalar=False, name="") - except TypeError: - other = self._expect_type( - other, - Scalar, - within=method_name, - keyword_name="other", - extra_message="Literal scalars also accepted.", - op=op, - ) + other = _literal_scalar(self, other, op, within=method_name, keyword_name="other") op = get_typed_op(op, self.dtype, other.dtype, kind="binary") self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") - return ScalarExpression( + expr = ScalarExpression( method_name, f"GrB_Vector_eWiseMult_{op.opclass}", [self._as_vector(), other._as_vector()], @@ -748,6 +778,9 @@ def _ewise_mult(self, other, op=binary.times, is_infix=False): is_cscalar=False, scalar_as_vector=True, ) + if output_type(given) is not Scalar: + expr._right_literal = given # a literal, as Scalar comparisons give it + return expr def ewise_union(self, other, op, left_default, right_default): """Perform element-wise computation on the union of sparse values, @@ -799,62 +832,25 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): self = op(self, left_default=left_default, right_default=right_default).new() if isinstance(other, ScalarEwiseAddExpr): other = op(other, left_default=left_default, right_default=right_default).new() + if output_type(self) is not Scalar: # A literal on the left, as in `op(1 | s)` + self = _literal_scalar(other, self, op, within=method_name, argname="left") right_dtype = self.dtype dtype = right_dtype if right_dtype._is_udt else None if type(other) is not Scalar: - try: - other = Scalar.from_value(other, dtype, is_cscalar=False, name="") - except TypeError: - other = self._expect_type( - other, - Scalar, - within=method_name, - keyword_name="other", - extra_message="Literal scalars also accepted.", - op=op, - ) + other = _literal_scalar(self, other, op, within=method_name, keyword_name="other") else: + # A typed Scalar keeps its dtype beside a lifted op, as it does for + # Vector.ewise_union: ``s8 - sf`` is float, not ``sf`` cut to int8. + if dtype is not None and other.dtype != dtype and _is_lifted_op(op): + dtype = None + _check_scalar_fits(dtype, other) other = _as_scalar(other, dtype, is_cscalar=False) # pragma: is_grbscalar temp_op = get_typed_op(op, self.dtype, other.dtype, kind="binary") - left_dtype = temp_op.type - dtype = left_dtype if left_dtype._is_udt else None - if type(left_default) is not Scalar: - try: - left = Scalar.from_value( - left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - left = self._expect_type( - left_default, - Scalar, - within=method_name, - keyword_name="left_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar - right_dtype = temp_op.type2 - dtype = right_dtype if right_dtype._is_udt else None - if type(right_default) is not Scalar: - try: - right = Scalar.from_value( - right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - right = self._expect_type( - right_default, - Scalar, - within=method_name, - keyword_name="right_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar + left = _union_default(self, left_default, temp_op.type, op, "left_default") + right = _union_default(self, right_default, temp_op.type2, op, "right_default") op1 = get_typed_op(op, self.dtype, right.dtype, kind="binary") op2 = get_typed_op(op, left.dtype, other.dtype, kind="binary") @@ -930,7 +926,7 @@ def apply(self, op, right=None, *, left=None): """ expr = self._as_vector().apply(op, right, left=left) - return ScalarExpression( + rv = ScalarExpression( expr.method_name, expr.cfunc_name, expr.args, @@ -940,6 +936,8 @@ def apply(self, op, right=None, *, left=None): is_cscalar=False, scalar_as_vector=True, ) + rv._right_literal = expr._right_literal + return rv def select(self, op, thunk=None): expr = self._as_vector().select(op, thunk) @@ -1001,6 +999,7 @@ def _repr_html_(self): is_cscalar = Scalar.is_cscalar is_grbscalar = Scalar.is_grbscalar + __hash__ = None # Begin auto-generated code: Scalar _get_value = automethods._get_value @@ -1075,6 +1074,7 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): is_cscalar = Scalar.is_cscalar is_grbscalar = Scalar.is_grbscalar + __hash__ = None # Begin auto-generated code: Scalar _get_value = automethods._get_value @@ -1125,6 +1125,63 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): # End auto-generated code: Scalar +def _scalar_eq(scalar, other, within): + """Return ``scalar == other`` as a bool, by ``binary.eq`` (see ``Scalar.__eq__``). + + ``within`` is the method an error names, ``__eq__`` or ``__ne__``. + """ + if other is None: + return scalar._is_empty + if type(other) is not Scalar and output_type(other) is Scalar: + other = other.new(name="s_eq_other") # an expression, such as ``-s`` or ``v[0]`` + if type(other) is Scalar and (scalar._is_empty or other._is_empty): + return scalar._is_empty and other._is_empty + if scalar._is_empty: + return False + if not scalar.dtype._is_udt: + # For built-in dtypes, compare the two values in the type ``eq`` is typed + # for, which is what it computes, without a GraphBLAS call: ``s == 0`` + # takes about 2 us instead of 11. + if type(other) is Scalar: + other_dtype, value = other.dtype, other.value + else: + value, other_dtype = _literal_operand(scalar.dtype, other, binary.eq) + if other_dtype is not None and not other_dtype._is_udt: + if other_dtype is not scalar.dtype: + other_dtype = get_typed_op(binary.eq, scalar.dtype, other_dtype).type + convert = other_dtype.np_type.type + if other_dtype.np_type in _SINGLE_PRECISION: + # 1e300 beside FP32 is infinity, and _literal_operand has + # warned at the caller's line, so numpy need not warn again. + with np.errstate(over="ignore"): + return (convert(scalar.value) == convert(value)).item() + return (convert(scalar.value) == convert(value)).item() + udts = [scalar.dtype, other.dtype] if type(other) is Scalar else [scalar.dtype] + udts = [dtype for dtype in udts if dtype._is_udt] + if udts and ( + not _has_numba + or any(_get_udt_info(dtype) is None or _has_half_leaf(dtype.np_type) for dtype in udts) + ): + # binary.eq needs Numba. A UDT that is neither a record nor an array (bytes, + # float16) holds one value, which numpy compares as binary.eq does, with no + # GraphBLAS call; and numpy compares float16, which Numba cannot. + return _udt_eq_without_numba(scalar, other) + if type(other) is not Scalar: + # As ewise_mult converts it, but an error names the method the user called. + value, dtype = _literal_operand(scalar.dtype, other, binary.eq) + try: + other = Scalar.from_value(value, dtype, is_cscalar=False, name="s_eq_other") + except TypeError: + other = scalar._expect_type( + other, + Scalar, + within=within, + argname="other", + extra_message="Literal scalars also accepted.", + ) + return bool(scalar.ewise_mult(other, binary.eq).new(name="s_eq").value) + + def _as_scalar(scalar, dtype=None, *, is_cscalar): if type(scalar) is not Scalar: return Scalar.from_value(scalar, dtype, is_cscalar=is_cscalar, name="") @@ -1133,16 +1190,63 @@ def _as_scalar(scalar, dtype=None, *, is_cscalar): return scalar -def _dict_to_record(np_type, d): - """Converts e.g. ``{"x": 1, "y": 2.3}`` to ``(1, 2.3)``.""" - rv = [] - for name, (dtype, _) in np_type.fields.items(): - val = d[name] - if dtype.names is not None and isinstance(val, dict): - rv.append(_dict_to_record(dtype, val)) - else: - rv.append(val) - return tuple(rv) +def _literal_scalar(scalar, value, op, *, within, argname=None, keyword_name=None): + """Return literal ``value``, ``op``'s operand beside ``scalar`` on either side, as a Scalar. + + It is typed as any literal operand is (``literals._literal_operand``), as the + left operand when ``argname`` is ``"left"``. A value that cannot be a Scalar + raises TypeError, naming ``argname`` or ``keyword_name``. + """ + value, dtype = _literal_operand(scalar.dtype, value, op, left=argname == "left") + try: + return Scalar.from_value(value, dtype, is_cscalar=False, name="") + except TypeError: + return scalar._expect_type( + value, + Scalar, + within=within, + argname=argname, + keyword_name=keyword_name, + extra_message="Literal scalars also accepted.", + op=op, + ) + + +def _union_default(obj, default, input_dtype, op, keyword_name): + """Return ``default`` for ``obj.ewise_union`` as a Scalar, where the op takes ``input_dtype``. + + GraphBLAS casts a default to the op's input type, and a UDT casts only to + itself, so where that is a UDT the default becomes an element of it and must + fit (``literals._check_literal_fits``). Otherwise a literal default is typed + as any literal operand is (``literals._literal_operand``), and a Scalar keeps + its type. ``keyword_name`` names the argument when ``default`` is not a scalar. + """ + dtype = input_dtype if input_dtype._is_udt else None + if type(default) is Scalar: + _check_scalar_fits(dtype, default) + return _as_scalar(default, dtype, is_cscalar=False) # pragma: is_grbscalar + if dtype is not None: + _check_literal_fits(dtype, default) + else: + default, dtype = _literal_operand( + input_dtype, default, op, left=keyword_name == "left_default" + ) + try: + return Scalar.from_value( + default, + dtype, + is_cscalar=False, # pragma: is_grbscalar + name="", + ) + except TypeError: + return obj._expect_type( + default, + Scalar, + within="ewise_union", + keyword_name=keyword_name, + extra_message="Literal scalars also accepted.", + op=op, + ) _MATERIALIZE = Scalar.from_value(lib.GrB_MATERIALIZE, is_cscalar=True, name="GrB_MATERIALIZE") diff --git a/graphblas/core/ss/__init__.py b/graphblas/core/ss/__init__.py index 10a6fed94..7ee1da469 100644 --- a/graphblas/core/ss/__init__.py +++ b/graphblas/core/ss/__init__.py @@ -1,5 +1,17 @@ import suitesparse_graphblas as _ssgb -(version_major, version_minor, version_bug) = map(int, _ssgb.__version__.split(".")[:3]) +# Version of the SuiteSparse:GraphBLAS C library, which is what every version +# gate in this package is really asking about. Read from the library's own +# constants rather than parsed out of ``_ssgb.__version__``: the wrapper's +# version normally tracks the library, but a development build between +# releases reports the previous one (10.4.1.0 against a 10.5.0 library), and a +# gate keyed on that takes the wrong branch without saying so. These are cffi +# ``#define`` constants, so they read at import time, before ``GrB_init``. +try: + version_major = _ssgb.lib.GxB_IMPLEMENTATION_MAJOR + version_minor = _ssgb.lib.GxB_IMPLEMENTATION_MINOR + version_bug = _ssgb.lib.GxB_IMPLEMENTATION_SUB +except AttributeError: # pragma: no cover (build without the GxB defines) + version_major, version_minor, version_bug = map(int, _ssgb.__version__.split(".")[:3]) _IS_SSGB7 = version_major == 7 diff --git a/graphblas/core/ss/_utils.py b/graphblas/core/ss/_utils.py new file mode 100644 index 000000000..ca4fc87d3 --- /dev/null +++ b/graphblas/core/ss/_utils.py @@ -0,0 +1,116 @@ +import hashlib +import itertools +import re + +from suitesparse_graphblas.utils import unclaim_buffer + +from ...exceptions import _error_code_lookup +from .. import ffi, lib +from ..dtypes import _string_to_dtype + + +def _unclaim_taken(*arrays_and_pointers): + """Unclaim each array whose pointer GraphBLAS cleared, as it does when it takes the data. + + Call it right after an import or pack, before checking the status. GraphBLAS + moves the arrays into the object first and only then runs a secure import's + check; when that fails, it frees the object and the arrays with it, so NumPy + must stop owning them although the call failed. A call that fails before the + move leaves every pointer as it was, and the arrays stay the caller's. + """ + for array, ptr in arrays_and_pointers: + if ptr[0] == ffi.NULL: + unclaim_buffer(array) + + +def _pack_securely(import_, obj, format, opts, **kwargs): + """Pack ``kwargs`` (data in ``format``, for ``import_``) into ``obj`` with a secure check. + + A secure pack whose check fails frees ``obj``'s GraphBLAS object itself + (SuiteSparse:GraphBLAS 10.5.2 and earlier), and the Python object keeps the + freed handle. A secure import that fails frees only the new object it made. + So the data is imported securely into a new object, and its content moves + into ``obj`` with an unpack and a fast pack, which take its arrays as they + are. A failure leaves ``obj`` unchanged. + """ + if obj.ndim == 1: + shape = {"size": obj._size} + else: + shape = {"nrows": obj._nrows, "ncols": obj._ncols} + new = import_( + **kwargs, **shape, dtype=obj.dtype, secure_import=True, method="import", opts=opts + ) + obj.ss.pack_any(**new.ss.unpack(format, **opts), take_ownership=True, **opts) + return obj + + +def _resolve_serialized_dtype(data_obj, data_nbytes, obj_kind): + """Resolve the dtype of a serialized Matrix/Vector buffer. + + ``GxB_deserialize_type_name`` returns the type's short name (a built-in + like ``'INT64'`` or, for a UDT, its ``GxB_JIT_C_NAME``). For UDTs the + short name may not match any Python-side registration (e.g., when an + anonymous UDT was registered by-shape on the producer side). Fall back + to ``GrB_NAME`` (the numpy repr that ``serialize`` writes into the + object) when the short name fails to resolve. + + ``obj_kind`` is ``"Matrix"`` or ``"Vector"`` and only shows up in error + messages. + """ + cname = ffi.new(f"char[{lib.GxB_MAX_NAME_LEN}]") + info = lib.GxB_deserialize_type_name(cname, data_obj, data_nbytes) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info](f"{obj_kind} deserialize failed to get the dtype name") + dtype_name = b"".join(itertools.takewhile(b"\x00".__ne__, cname)).decode() + orig_error = None + if dtype_name: + try: + return _string_to_dtype(dtype_name) + except (ValueError, TypeError, SyntaxError) as exc: + orig_error = exc + if hasattr(lib, "GxB_Serialized_get_String"): + dtype_size = ffi.new("size_t*") + info = lib.GxB_Serialized_get_SIZE(data_obj, dtype_size, lib.GrB_NAME, data_nbytes) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info](f"{obj_kind} deserialize failed to get the size of name") + dtype_char = ffi.new(f"char[{dtype_size[0]}]") + info = lib.GxB_Serialized_get_String(data_obj, dtype_char, lib.GrB_NAME, data_nbytes) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info](f"{obj_kind} deserialize failed to get the name") + return _string_to_dtype(ffi.string(dtype_char).decode()) + # No GrB_NAME fallback available (older SS); surface the original + # resolution failure with the actual dtype name attached. + raise ValueError( + f"{obj_kind} deserialize cannot resolve dtype {dtype_name!r} and " + f"GxB_Serialized_get_String is unavailable in this SuiteSparse build" + ) from orig_error + + +def _jit_c_source(name, funcname, jit_c_definition, types=()): + """Return the C name and definition to give SuiteSparse for an op registered from C. + + SuiteSparse finds a compiled kernel in memory by its op's name (and the types in the + kernel's encoding), and checks an op's definition only when it loads a kernel from disk. + So two ops of one C name with different definitions would run each other's kernel: an op + defined again under its name, or two of one last name in two namespaces (``ss.f`` and + ``ss.nested.f``). The C name is ``funcname``, then ``_{len(t)}_{t}`` for each of + ``types`` (built-in types by name, UDTs by their C name), then a digest of these, the + full dotted ``name`` and the definition, as ``f_3a1c...``. SuiteSparse keeps a name of at + most 127 characters, so a longer one is ``funcname``'s start and a longer digest. One + name and definition give one C name on purpose: registered again, they reuse the kernel + compiled for them. The function is renamed in the definition wherever it is named, so a + recursive call follows, but not a struct member of the same name (after ``.`` or ``->``). + """ + parts = [re.sub(r"\W", "_", t.jit_c_name or t.name) for t in types] + readable = funcname + "".join(f"_{len(part)}_{part}" for part in parts) + digest = hashlib.sha256(f"{readable}\0{name}\0{jit_c_definition}".encode()).hexdigest() + c_name = f"{readable}_{digest[:12]}" + if len(c_name) >= lib.GxB_MAX_NAME_LEN: + c_name = f"{funcname[:100]}_{digest[:16]}" + + def rename(match): + if jit_c_definition[: match.start()].rstrip().endswith((".", "->")): + return match.group() + return c_name + + return c_name, re.sub(rf"\b{re.escape(funcname)}\b", rename, jit_c_definition) diff --git a/graphblas/core/ss/binary.py b/graphblas/core/ss/binary.py index d53608818..d940ac342 100644 --- a/graphblas/core/ss/binary.py +++ b/graphblas/core/ss/binary.py @@ -5,6 +5,7 @@ from ..operator.base import TypedOpBase from ..operator.binary import BinaryOp, TypedUserBinaryOp from . import _IS_SSGB7 +from ._utils import _jit_c_source ffi_new = ffi.new @@ -44,7 +45,12 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. Parameters ---------- @@ -55,6 +61,8 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): jit_c_definition : str The C definition as a string of the user-defined function. For example: ``"void absdiff (double *z, double *x, double *y) { (*z) = fabs ((*x) - (*y)) ; }"``. + A UDT made from a numpy dtype is named in C by its ``jit_c_name`` (usually ``gbudt_`` + and its name), not its Python name. left_type : dtype The dtype of the left operand of the binary operator. right_type : dtype @@ -106,23 +114,28 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): else: # We use `is_udt=True` to make dtype handling flexible and explicit. rv = BinaryOp(name, is_udt=True) + from .jit_config import _compiling_c_source + + c_name, c_definition = _jit_c_source(name, funcname, jit_c_definition) gb_obj = ffi_new("GrB_BinaryOp*") - check_status_carg( - lib.GxB_BinaryOp_new( - gb_obj, - NULL, - ret_type._carg, - left_type._carg, - right_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), - ), - "BinaryOp", - gb_obj[0], - ) + with _compiling_c_source(f"gb.binary.{name}"): + check_status_carg( + lib.GxB_BinaryOp_new( + gb_obj, + NULL, + ret_type._carg, + left_type._carg, + right_type._carg, + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), + ), + "BinaryOp", + gb_obj[0], + ) op = TypedJitBinaryOp( rv, funcname, left_type, ret_type, gb_obj[0], jit_c_definition, dtype2=right_type ) + op._jit_c_info = (c_name, c_definition) rv._add(op, is_jit=True) setattr(module, funcname, rv) return rv diff --git a/graphblas/core/ss/config.py b/graphblas/core/ss/config.py index 70a7dd196..4d2a20573 100644 --- a/graphblas/core/ss/config.py +++ b/graphblas/core/ss/config.py @@ -52,8 +52,38 @@ def __init__(self, parent=None, context=None): if k not in rd: # pragma: no branch (safety) rd[k] = k cls._initialized = True + # Writes are unrestricted until __init__ finishes, and the instance + # attributes set by then become the allowlist for __setattr__. + # Subclasses that add instance attributes (such as Context) must set + # them before calling super().__init__(). + object.__setattr__(self, "_initializing", True) self._parent = parent self._context = context + object.__setattr__(self, "_internal_attrs", frozenset(self.__dict__) - {"_initializing"}) + object.__setattr__(self, "_initializing", False) + + def __setattr__(self, key, value): + # An instance built without running __init__ has no allowlist yet, so stay + # permissive rather than raising a confusing AttributeError from the guard. + # Properties (such as Context._context) go through so their setters run. + if ( + self.__dict__.get("_initializing", True) + or key in self._internal_attrs + or isinstance(getattr(type(self), key, None), property) + ): + object.__setattr__(self, key, value) + return + if (option := key.lower()) in self._options: + if option in self._read_only: + raise AttributeError(f"Config option {option!r} is read-only") + raise AttributeError( + f"Cannot set config option {option!r} by attribute assignment; " + f"this does not change the config. " + f"Use item assignment instead: config[{option!r}] = value" + ) + raise AttributeError( + f"Unknown config option {key!r}; known options are {sorted(self._options)}." + ) def __delitem__(self, key): raise TypeError("Configuration options can't be deleted.") @@ -210,5 +240,5 @@ def __repr__(self): + "})" ) - def _ipython_key_completions_(self): # pragma: no cover (ipython) + def _ipython_key_completions_(self): return list(self) diff --git a/graphblas/core/ss/context.py b/graphblas/core/ss/context.py index f93d1ec1c..99026e16f 100644 --- a/graphblas/core/ss/context.py +++ b/graphblas/core/ss/context.py @@ -20,9 +20,12 @@ class Context(BaseConfig): _context_keys = {"chunk", "gpu_id", "nthreads"} _options = { "chunk": (lib.GxB_CONTEXT_CHUNK, "double"), - "gpu_id": (lib.GxB_CONTEXT_GPU_ID, "int"), "nthreads": (lib.GxB_CONTEXT_NTHREADS, "int"), } + # GxB_CONTEXT_GPU_ID existed in SS:GraphBLAS <=10.1 as a scalar int. + # It was renamed to GxB_CONTEXT_GPU_IDS in 10.2 and changed to a list type. + if hasattr(lib, "GxB_CONTEXT_GPU_ID"): + _options["gpu_id"] = (lib.GxB_CONTEXT_GPU_ID, "int") _defaults = { "nthreads": 0, "chunk": 0, @@ -30,22 +33,25 @@ class Context(BaseConfig): } def __init__(self, engage=True, *, stack=True, nthreads=None, chunk=None, gpu_id=None): - super().__init__() + # Instance attributes must exist before super().__init__() finishes, + # because BaseConfig.__setattr__ allows only the attributes set by then. self.gb_obj = ffi_new("GxB_Context*") check_status_carg(lib.GxB_Context_new(self.gb_obj), "Context", self.gb_obj[0]) + self._prev_context = None + super().__init__() if stack: context = threadlocal.context self["nthreads"] = context["nthreads"] if nthreads is None else nthreads self["chunk"] = context["chunk"] if chunk is None else chunk - self["gpu_id"] = context["gpu_id"] if gpu_id is None else gpu_id + if "gpu_id" in self._options: + self["gpu_id"] = context["gpu_id"] if gpu_id is None else gpu_id else: if nthreads is not None: self["nthreads"] = nthreads if chunk is not None: self["chunk"] = chunk - if gpu_id is not None: + if gpu_id is not None and "gpu_id" in self._options: self["gpu_id"] = gpu_id - self._prev_context = None if engage: self.engage() @@ -66,7 +72,7 @@ def dup(self, engage=True, *, nthreads=None, chunk=None, gpu_id=None): nthreads = self["nthreads"] if chunk is None: chunk = self["chunk"] - if gpu_id is None: + if gpu_id is None and "gpu_id" in self._options: gpu_id = self["gpu_id"] return type(self)(engage, stack=False, nthreads=nthreads, chunk=chunk, gpu_id=gpu_id) diff --git a/graphblas/core/ss/descriptor.py b/graphblas/core/ss/descriptor.py index 781661b7b..c8a2fcade 100644 --- a/graphblas/core/ss/descriptor.py +++ b/graphblas/core/ss/descriptor.py @@ -44,10 +44,13 @@ class _DescriptorConfig(BaseConfig): _options.update( { "chunk": (lib.GxB_CONTEXT_CHUNK, "double"), - "gpu_id": (lib.GxB_CONTEXT_GPU_ID, "int"), "nthreads": (lib.GxB_CONTEXT_NTHREADS, "int"), } ) + # GxB_CONTEXT_GPU_ID existed in SS:GraphBLAS <=10.1 as a scalar int. + # It was renamed to GxB_CONTEXT_GPU_IDS in 10.2 and changed to a list type. + if hasattr(lib, "GxB_CONTEXT_GPU_ID"): + _options["gpu_id"] = (lib.GxB_CONTEXT_GPU_ID, "int") _enumerations = { # GrB "output_replace": { @@ -150,7 +153,7 @@ def get_descriptor(**opts): Whether and how to compress the data for serialization. The default is "zstd" with ``compression_level=1`` compression_level : int - 1-9 for "lz4hc" compression: (1) is the fastest, (9) is the most compact (default 1). + 1-9 for "lz4hc" compression: (1) is the fastest, (9) is the most compact (default 9). 1-19 for "zstd" compression: (1) is the fastest, (19) is the most compact (default 1). Ignored for other compression methods. @@ -209,7 +212,7 @@ def _set_compression(desc, compression, level): comp = _str_to_compression[compression] if level is not None: if compression not in {"lz4hc", "zstd"}: - raise TypeError('level argument is only valid when using "lz4hc" compression') + raise TypeError('level argument is only valid when using "lz4hc" or "zstd" compression') level = int(level) upper = 9 if compression == "lz4hc" else 19 default = 9 if compression == "lz4hc" else 1 diff --git a/graphblas/core/ss/dtypes.py b/graphblas/core/ss/dtypes.py index d2eb5b416..54d3fb812 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -29,22 +29,38 @@ def register_new(name, jit_c_definition, *, np_type=None): ) if not name.isidentifier(): raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}") - if name in core.dtypes._registry or hasattr(dtypes.ss, name): + from ..operator.udt_codegen import _jit_c_type_defs + + # SuiteSparse's JIT knows a type by its C name alone, so it must not be + # another type's C name either. + if name in core.dtypes._registry or hasattr(dtypes.ss, name) or name in _jit_c_type_defs: raise ValueError(f"{name!r} name for dtype is unavailable") - if len(name) > lib.GxB_MAX_NAME_LEN: + # GxB_MAX_NAME_LEN counts the terminating NUL, so SuiteSparse keeps one fewer + # character; a name that long would no longer match its typedef and would + # fail to compile. + if len(name) >= lib.GxB_MAX_NAME_LEN: raise ValueError( - f"`name` argument is too large. Max size is {lib.GxB_MAX_NAME_LEN}; got {len(name)}" + "`name` argument is too large. " + f"Max size is {lib.GxB_MAX_NAME_LEN - 1}; got {len(name)}" ) if name not in jit_c_definition: raise ValueError("`name` argument must be same name as the typedef in `jit_c_definition`") if "struct" not in jit_c_definition: raise ValueError("Only struct typedefs are currently allowed for JIT dtypes") + # Registering a type by its C typedef is a request to compile C, so it is + # a fair place to raise SuiteSparse's non-compiling default. + from .jit_config import _compiling_c_source + gb_obj = ffi.new("GrB_Type*") - status = lib.GxB_Type_new( - gb_obj, 0, ffi_new("char[]", name.encode()), ffi_new("char[]", jit_c_definition.encode()) - ) - check_status_carg(status, "Type", gb_obj[0]) + with _compiling_c_source(f"gb.dtypes.ss.{name}"): + status = lib.GxB_Type_new( + gb_obj, + 0, + ffi_new("char[]", name.encode()), + ffi_new("char[]", jit_c_definition.encode()), + ) + check_status_carg(status, "Type", gb_obj[0]) # Let SuiteSparse:GraphBLAS determine the size (we gave 0 as size above) size_ptr = ffi_new("size_t*") @@ -72,17 +88,22 @@ def register_new(name, jit_c_definition, *, np_type=None): np_type = np.dtype((np.uint8, size)) save_np_type = False if _has_numba: - numba_type = numba.typeof(np_type).dtype + numba_type = core.dtypes._numba_type(np_type) else: numba_type = None + np_type = core.dtypes._flatten_subarrays(np_type) + # For now, let's use "opaque" unsigned bytes for the c type. rv = core.dtypes.DataType(name, gb_obj, None, f"uint8_t[{size}]", numba_type, np_type) core.dtypes._registry[gb_obj] = rv if save_np_type or np_type not in core.dtypes._registry: core.dtypes._registry[np_type] = rv - if numba_type is not None and (save_np_type or numba_type not in core.dtypes._registry): - core.dtypes._registry[numba_type] = rv - core.dtypes._registry[numba_type.name] = rv setattr(dtypes.ss, name, rv) + # Only now, so that a registration that fails above (Numba cannot type an + # array of bool, say) leaves the name free to register again. The type it + # made is never used, so no JIT kernel knows its name. + _jit_c_type_defs[name] = jit_c_definition + # Its C name and definition are the user's, which ops lifted to it name it by. + rv._jit_c_info = (name, jit_c_definition) return rv diff --git a/graphblas/core/ss/indexbinary.py b/graphblas/core/ss/indexbinary.py new file mode 100644 index 000000000..8fd53a27e --- /dev/null +++ b/graphblas/core/ss/indexbinary.py @@ -0,0 +1,151 @@ +from ... import backend +from ...dtypes import lookup_dtype +from ...exceptions import check_status_carg +from .. import NULL, ffi, lib +from ..operator.base import TypedOpBase +from ..operator.indexbinary import ( + IndexBinaryOp, + TypedBuiltinIndexBinaryOp, + _has_idxbinop, + _TypedIndexBinaryOpMixin, +) +from ._utils import _jit_c_source + +ffi_new = ffi.new + + +class TypedJitIndexBinaryOp(_TypedIndexBinaryOpMixin, TypedOpBase): + __slots__ = "_jit_c_definition", "_theta_type" + opclass = "IndexBinaryOp" + + def __init__( + self, parent, name, type_, return_type, gb_obj, jit_c_definition, dtype2, theta_type + ): + super().__init__(parent, name, type_, return_type, gb_obj, name, dtype2=dtype2) + self._jit_c_definition = jit_c_definition + self._theta_type = theta_type + + @property + def jit_c_definition(self): + return self._jit_c_definition + + def __call__(self, theta=None): + return TypedBuiltinIndexBinaryOp.__call__(self, theta) + + __call__.__doc__ = TypedBuiltinIndexBinaryOp.__call__.__doc__ + + +def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): + """Register a new IndexBinaryOp using the SuiteSparse:GraphBLAS JIT compiler. + + This creates an IndexBinaryOp by compiling the C string definition of the function. + It requires a shell call to a C compiler. The resulting operator will be as + fast as if it were built-in to SuiteSparse:GraphBLAS and does not have the + overhead of additional function calls as when using ``gb.indexbinary.register_new``. + + This is an advanced feature that requires a C compiler and proper configuration. + Configuration is handled by ``gb.ss.config``; see its docstring for details. + By default, the JIT caches results in ``~/.SuiteSparse/``. For more information, + see the SuiteSparse:GraphBLAS user guide. + + Only one type signature may be registered at a time, but repeated calls using + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. + + Parameters + ---------- + name : str + The name of the operator. This will show up as ``gb.indexbinary.ss.{name}``. + The name may contain periods, ".", which will result in nested objects + such as ``gb.indexbinary.ss.x.y.z`` for name ``"x.y.z"``. + jit_c_definition : str + The C definition as a string of the user-defined function. For example: + ``"void my_idxbin (double *z, double *x, GrB_Index ix, GrB_Index jx, "`` + ``"double *y, GrB_Index iy, GrB_Index jy, double *theta) "`` + ``"{ (*z) = (*x) + (*y) + (*theta) ; }"`` + A UDT made from a numpy dtype is named in C by its ``jit_c_name`` (usually ``gbudt_`` + and its name), not its Python name. + x_type : dtype + The dtype of the first operand (x). + y_type : dtype + The dtype of the second operand (y). + theta_type : dtype + The dtype of the theta parameter. + ret_type : dtype + The dtype of the result. + + Returns + ------- + IndexBinaryOp + + See Also + -------- + gb.indexbinary.register_new + gb.indexbinary.register_anonymous + + """ + if backend != "suitesparse": # pragma: no cover (safety) + raise RuntimeError( + "`gb.indexbinary.ss.register_new` invalid when not using 'suitesparse' backend" + ) + if not _has_idxbinop: + import suitesparse_graphblas as ssgb + + raise RuntimeError( + "IndexBinaryOp requires SuiteSparse:GraphBLAS 9.4+; " + f"current version is {ssgb.__version__}" + ) + x_type = lookup_dtype(x_type) + y_type = lookup_dtype(y_type) + theta_type = lookup_dtype(theta_type) + ret_type = lookup_dtype(ret_type) + name = name if name.startswith("ss.") else f"ss.{name}" + module, funcname = IndexBinaryOp._remove_nesting(name, strict=False) + if hasattr(module, funcname): + rv = getattr(module, funcname) + if not isinstance(rv, IndexBinaryOp): + IndexBinaryOp._remove_nesting(name) + if ( + (x_type, y_type, theta_type) in rv.types + or rv._udt_types is not None + and (x_type, y_type, theta_type) in rv._udt_types + ): + raise TypeError( + f"IndexBinaryOp gb.indexbinary.{name} already defined for " + f"({x_type}, {y_type}, {theta_type}) input types" + ) + else: + rv = IndexBinaryOp(name, is_udt=True) + from .jit_config import _compiling_c_source + + # Theta's type is not in the kernel's encoding, so the C name carries the types + c_name, c_definition = _jit_c_source( + name, funcname, jit_c_definition, (x_type, y_type, theta_type) + ) + gb_obj = ffi_new("GxB_IndexBinaryOp*") + with _compiling_c_source(f"gb.indexbinary.{name}"): + check_status_carg( + lib.GxB_IndexBinaryOp_new( + gb_obj, + NULL, + ret_type._carg, + x_type._carg, + y_type._carg, + theta_type._carg, + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), + ), + "IndexBinaryOp", + gb_obj[0], + ) + op = TypedJitIndexBinaryOp( + rv, funcname, x_type, ret_type, gb_obj[0], jit_c_definition, y_type, theta_type + ) + op._jit_c_info = (c_name, c_definition) + rv._add(op, is_jit=True) + setattr(module, funcname, rv) + return rv diff --git a/graphblas/core/ss/indexunary.py b/graphblas/core/ss/indexunary.py index b60837acf..77eafe067 100644 --- a/graphblas/core/ss/indexunary.py +++ b/graphblas/core/ss/indexunary.py @@ -5,6 +5,7 @@ from ..operator.base import TypedOpBase from ..operator.indexunary import IndexUnaryOp, TypedUserIndexUnaryOp from . import _IS_SSGB7 +from ._utils import _jit_c_source ffi_new = ffi.new @@ -39,7 +40,12 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. This will also create a SelectOp operator under ``gb.select.ss`` if the return type is boolean. @@ -54,6 +60,8 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): The C definition as a string of the user-defined function. For example: ``"void diffy (double *z, double *x, GrB_Index i, GrB_Index j, double *y) "`` ``"{ (*z) = (i + j) * fabs ((*x) - (*y)) ; }"`` + A UDT made from a numpy dtype is named in C by its ``jit_c_name`` (usually ``gbudt_`` + and its name), not its Python name. input_type : dtype The dtype of the operand of the indexunary operator. thunk_type : dtype @@ -105,23 +113,28 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): else: # We use `is_udt=True` to make dtype handling flexible and explicit. rv = IndexUnaryOp(name, is_udt=True) + from .jit_config import _compiling_c_source + + c_name, c_definition = _jit_c_source(name, funcname, jit_c_definition) gb_obj = ffi_new("GrB_IndexUnaryOp*") - check_status_carg( - lib.GxB_IndexUnaryOp_new( - gb_obj, - NULL, - ret_type._carg, - input_type._carg, - thunk_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), - ), - "IndexUnaryOp", - gb_obj[0], - ) + with _compiling_c_source(f"gb.indexunary.{name}"): + check_status_carg( + lib.GxB_IndexUnaryOp_new( + gb_obj, + NULL, + ret_type._carg, + input_type._carg, + thunk_type._carg, + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), + ), + "IndexUnaryOp", + gb_obj[0], + ) op = TypedJitIndexUnaryOp( rv, funcname, input_type, ret_type, gb_obj[0], jit_c_definition, dtype2=thunk_type ) + op._jit_c_info = (c_name, c_definition) rv._add(op, is_jit=True) if ret_type == BOOL: from ..operator.select import SelectOp @@ -147,6 +160,7 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): op2 = TypedJitSelectOp( selectop, funcname, input_type, ret_type, gb_obj[0], jit_c_definition, dtype2=thunk_type ) + op2._jit_c_info = (c_name, c_definition) selectop._add(op2, is_jit=True) setattr(select_module, funcname, selectop) setattr(module, funcname, rv) diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py new file mode 100644 index 000000000..62d78a44a --- /dev/null +++ b/graphblas/core/ss/jit_config.py @@ -0,0 +1,575 @@ +"""Helpers for repairing the SuiteSparse:GraphBLAS JIT compiler configuration. + +The C library is built by conda-build with the build host's compiler paths +baked into ``GxB_JIT_C_COMPILER_NAME`` and ``GxB_JIT_C_COMPILER_FLAGS``. +Those paths (e.g., ``/Users/runner/...``) don't exist in the user's +environment. When SuiteSparse JIT tries to compile a UDT kernel the compile +step fails silently: the ``.c`` source is written but no ``.dylib`` or +``.so`` is produced, and SS falls back to the Numba function-pointer path. +The caller sees a working op, but the JIT speedup is lost, and there is +no error or warning to make the regression visible. + +``fix_jit_config()`` replaces the baked-in compiler with one from +``$CONDA_PREFIX/bin/`` and strips build-time-only flags (``-isysroot``, +``-fdebug-prefix-map``). For non-conda installs (pure pip), ``sysconfig`` +may expose a usable compiler. ``fix_jit_config(use_sysconfig=True)`` will +try that. + +``jit_compiler_is_usable()`` is a non-invasive check that returns True iff +the configured compiler exists on disk. Useful for emitting a one-time +warning at import or at first JIT use. + +``jit_c_control`` itself is set when the ``suitesparse`` backend initializes, +from ``graphblas.config["jit_c_control"]`` (``'run'`` unless configured +otherwise; see :func:`_set_jit_c_control_at_init`). UDT work raises that +default ``'run'`` to ``'on'`` when it needs C compiled, but keeps a ``'run'`` +the user set (:func:`_enable_jit_for_udt`). +The trade-offs of the JIT, and what to do when it misbehaves, are in the user +guide's initialization page (docs/user_guide/init.rst). +""" + +import contextlib +import hashlib +import os +import pathlib +import platform +import re +import subprocess + +# ``ffi`` and ``lib`` are imported in the functions that use them: importing them +# here initializes the backend when this module is imported first, and init +# imports this module for ``_set_jit_c_control_at_init``. + +# Mapping from the values that show up in baked-in ``-arch`` flags to what +# ``platform.machine()`` returns on the running host. +_ARCH_ALIASES = { + "x86_64": ("x86_64", "amd64"), + "arm64": ("arm64", "aarch64"), + "aarch64": ("arm64", "aarch64"), + "i386": ("i386", "x86"), + "ppc": ("ppc", "powerpc"), + "ppc64": ("ppc64", "powerpc64"), +} + + +def _ss_config(): + # Imported lazily to avoid pulling all of ``gb.ss`` into module-load time. + from ... import ss + + return ss.config + + +def jit_compiler_is_usable(): + """True iff the configured JIT compiler path exists on disk. + + Cheap (no compile attempt); use before suggesting ``fix_jit_config()``. + """ + return _compiler_exists(_ss_config()) + + +def _compiler_exists(cfg): + cc = cfg.get("jit_c_compiler_name", "") + return bool(cc) and pathlib.Path(cc).exists() + + +def fix_jit_config(*, use_sysconfig=True, probe=True): + """Repair the SuiteSparse:GraphBLAS JIT compiler configuration. + + Replaces the baked-in compiler path (which often points at a conda-build + host that doesn't exist in user environments) with one from + ``$CONDA_PREFIX/bin/``, strips build-time-only flags (``-isysroot``, + ``-fdebug-prefix-map``), and sets ``jit_c_control`` to ``'on'``. + + Parameters + ---------- + use_sysconfig : bool, default True + When ``$CONDA_PREFIX`` isn't set (pure pip install), try the + compiler from ``sysconfig.get_config_var("CC")``. Set to ``False`` + to restrict the repair to a conda environment only. + probe : bool, default True + After fixing the config, JIT-compile a trivial type to verify the + compiler actually works. A failed compile makes SuiteSparse lower + ``jit_c_control`` to ``'load'``, and from SuiteSparse 9.4 on it also + raises ``JitError`` that one time; the probe takes that failure, so + later operations use the generic kernels without raising. + + Returns + ------- + True + Fix applied and (if ``probe``) verified working. + False + Fix attempted but the probe failed. ``jit_c_control`` is now + whatever SuiteSparse left it at. + None + No environment available to fix from. There's no ``$CONDA_PREFIX``, + and either ``use_sysconfig=False`` or no sysconfig compiler is set. + """ + cfg = _ss_config() + rv = _repair_jit_compiler(cfg, use_sysconfig=use_sysconfig) + if rv is None: + return None + # Enabling compilation is what the caller asked for, whatever + # ``graphblas.config["jit_c_control"]`` set at init. + cfg["jit_c_control"] = "on" + # An explicit user-driven fix is a clean opportunity to re-arm + # ``NoJITWarning``: if the repair worked, the next UDT auto-lift that + # *still* falls back to cfunc (different cause: UDT layout, etc.) + # deserves a fresh notification rather than silent suppression. + _warned_no_jit_for.clear() + # Re-arm the cached probe answer for the same reason: a pre-repair False + # is stale now, so let the next UDT op re-derive it against the fixed + # toolchain instead of trusting the old verdict. + global _jit_enabled_for_udt + _jit_enabled_for_udt = None + if not probe: + return True + return _probe_jit(cfg) + + +def _repair_jit_compiler(cfg, *, use_sysconfig=True): + """Point the JIT compiler settings at something that exists on this host. + + Only rewrites the compiler name and flags; ``jit_c_control`` is left + alone, so on its own this changes nothing about what SuiteSparse will + do. Callers that want compilation enabled set the control themselves. + """ + if conda_prefix := os.environ.get("CONDA_PREFIX", ""): + return _fix_from_conda(cfg, conda_prefix) + if use_sysconfig: + return _fix_from_sysconfig(cfg) + return None + + +def _fix_from_conda(cfg, conda_prefix): + """Conda-aware fix; swap the compiler path with one from the environment's binaries: + ``$CONDA_PREFIX/bin/``, or ``$CONDA_PREFIX/Library/bin/`` with ``.exe`` names (Windows). + """ + jit_cc = cfg["jit_c_compiler_name"] + if not pathlib.Path(jit_cc).exists(): + # (a Windows path's basename, also when read on another system) + cc_basename = re.split(r"[\\/]", jit_cc)[-1] + prefix = pathlib.Path(conda_prefix) + names = [cc_basename, "cc", "clang", "gcc"] + names += [n + ".exe" for n in names if not n.lower().endswith(".exe")] + for bin_dir in [prefix / "bin", prefix / "Library" / "bin"]: + local_cc = next((bin_dir / n for n in names if (bin_dir / n).exists()), None) + if local_cc is not None: + cfg["jit_c_compiler_name"] = str(local_cc) + break + else: + return None # nothing usable + _fix_compiler_flags(cfg) + return True + + +def _fix_from_sysconfig(cfg): + """Pure-pip fallback; pull compiler info from ``sysconfig``.""" + import sysconfig + + cc = sysconfig.get_config_var("CC") + cflags = sysconfig.get_config_var("CFLAGS") + include = sysconfig.get_path("include") + if not (cc and cflags and include): + return None + cfg["jit_c_compiler_name"] = cc + cfg["jit_c_compiler_flags"] = f"{cflags} -I{include}" + if libs := sysconfig.get_config_var("LIBS"): + cfg["jit_c_libraries"] = libs + return True + + +def _fix_compiler_flags(cfg): + """Replace build-time-only paths in ``jit_c_compiler_flags``.""" + flags = cfg["jit_c_compiler_flags"] + isysroot_match = re.search(r"-isysroot\s+(\S+)", flags) + if isysroot_match and not pathlib.Path(isysroot_match.group(1)).exists(): + try: + sdk_path = subprocess.check_output( + ["xcrun", "--show-sdk-path"], text=True, stderr=subprocess.DEVNULL + ).strip() + flags = re.sub(r"-isysroot\s+\S+", f"-isysroot {sdk_path}", flags) + except (subprocess.CalledProcessError, FileNotFoundError): + # No Xcode SDK (Linux, or macOS without Xcode CLT). + flags = re.sub(r"-isysroot\s+\S+", "", flags) + flags = _strip_mismatched_arch(flags) + # Build-time debug path remapping is irrelevant in a user environment. + flags = re.sub(r"-fdebug-prefix-map=\S+", "", flags) + cfg["jit_c_compiler_flags"] = flags + + +def _strip_mismatched_arch(flags): + """Strip any ``-arch FOO`` that doesn't match the host architecture. + + conda-forge's ``python-suitesparse-graphblas`` bakes the build host's + ``-arch`` into ``jit_c_compiler_flags``. On a different host (e.g., an + arm64 Mac running an x86_64-built package) the JIT compile produces + objects for the wrong arch and the link fails. Leaving the flag out + lets the compiler default to the host arch. + """ + host = platform.machine().lower() + return re.sub( + r"-arch\s+(\S+)", + lambda m: "" if m.group(1).lower() not in _ARCH_ALIASES.get(host, (host,)) else m.group(0), + flags, + ) + + +# The settings a compile depends on. The probe's type is named after them; see +# ``_probe_jit``. +_TOOLCHAIN_KEYS = ( + "jit_c_compiler_name", + "jit_c_compiler_flags", + "jit_c_linker_flags", + "jit_c_libraries", +) + + +def _probe_jit(cfg): + """Compile a trivial type with the JIT, or load it from the cache, to check the config works. + + ``jit_c_control`` must be ``'on'``, and is left wherever SuiteSparse puts + it, which is not predictable from here. A library built without the JIT + clamps every write down to ``'run'`` (``GB_jitifyer_set_control``), so + ``'on'`` never takes on such a build no matter how good the compiler is; + separately, a compile failure drops it to ``'load'``, a load failure to + ``'run'``, and a failed hash insert to ``'pause'``. The contract here is + only "did this work?", so the caller gets a bool. + + It goes through the C API, so nothing is registered on the Python side, + and it can run while the backend initializes. Once a kernel is in the + cache, a later probe loads it (about 2 ms) instead of compiling (about + 0.3 s). The type is named after the toolchain settings, so a compiler or + flags this cache has not seen compile a probe of their own, rather than + load one that another toolchain left in a shared cache. + """ + from .. import ffi, lib + + toolchain = "\n".join(cfg.get(key, "") for key in _TOOLCHAIN_KEYS) + name = "_jit_probe_" + hashlib.sha256(toolchain.encode()).hexdigest()[:16] + defn = f"typedef struct {{ int _probe ; }} {name} ;" + type_ptr = ffi.new("GrB_Type*") + # Size 0 asks SuiteSparse to get the size from a compiled kernel. + if lib.GxB_Type_new(type_ptr, 0, name.encode(), defn.encode()) != lib.GrB_SUCCESS: + return False + lib.GrB_Type_free(type_ptr) + return cfg["jit_c_control"] == "on" + + +class _CConfig: + """The ``gb.ss.config`` keys :func:`_set_jit_c_control_at_init` needs, through the C API. + + That runs while the backend initializes, possibly in the middle of + importing a core module, and importing ``graphblas.ss`` (for its + ``config``) would import the core modules. + """ + + _strings = { + "jit_c_compiler_name": "GxB_JIT_C_COMPILER_NAME", + "jit_c_compiler_flags": "GxB_JIT_C_COMPILER_FLAGS", + "jit_c_linker_flags": "GxB_JIT_C_LINKER_FLAGS", + "jit_c_libraries": "GxB_JIT_C_LIBRARIES", + "jit_cache_path": "GxB_JIT_CACHE_PATH", + } + + def __init__(self): + from .. import lib + + self._controls = { + "off": lib.GxB_JIT_OFF, + "pause": lib.GxB_JIT_PAUSE, + "run": lib.GxB_JIT_RUN, + "load": lib.GxB_JIT_LOAD, + "on": lib.GxB_JIT_ON, + } + + def __getitem__(self, key): + from .. import ffi, lib + + if key == "jit_c_control": + val = ffi.new("int32_t*") + info = lib.GxB_Global_Option_get_INT32(lib.GxB_JIT_C_CONTROL, val) + else: + val = ffi.new("char**") + info = lib.GxB_Global_Option_get_CHAR(getattr(lib, self._strings[key]), val) + if info != lib.GrB_SUCCESS: # pragma: no cover (safety) + raise RuntimeError(f"Could not get SuiteSparse:GraphBLAS option {key!r}") + if key == "jit_c_control": + return next(k for k, v in self._controls.items() if v == val[0]) + return ffi.string(val[0]).decode() + + def __setitem__(self, key, value): + from .. import lib + + if key == "jit_c_control": + info = lib.GxB_Global_Option_set_INT32(lib.GxB_JIT_C_CONTROL, self._controls[value]) + else: + option = getattr(lib, self._strings[key]) + info = lib.GxB_Global_Option_set_CHAR(option, value.encode()) + if info != lib.GrB_SUCCESS: # pragma: no cover (safety) + raise RuntimeError(f"Could not set SuiteSparse:GraphBLAS option {key!r}") + + def get(self, key, default=None): + return self[key] if key == "jit_c_control" or key in self._strings else default + + +def _cache_is_writable(cfg): + """Whether SuiteSparse can write new kernels to its cache. + + A probe kernel already in the cache loads from a read-only one, and the + first compile after it would fail (with ``JitError`` from SuiteSparse 9.4 + on), so the probe alone does not answer this. + """ + path = cfg.get("jit_cache_path") + return bool(path) and all( + os.access(pathlib.Path(path, sub), os.W_OK | os.X_OK) for sub in ("c", "lib", "tmp") + ) + + +def _set_jit_c_control_at_init(control): + """Set ``jit_c_control`` as the backend initializes, from ``graphblas.config``. + + ``control`` is ``graphblas.config["jit_c_control"]``, already checked, or + None when it is not set: then the control is the default ``'run'``, which + UDT work raises (:func:`_enable_jit_for_udt`). A ``'run'`` that is set is + the user's choice, and UDT work keeps it (``_user_set_run``). + Every value but ``'on'`` is set as given. ``'on'`` is set only where it + works, which :func:`_probe_jit` checks by compiling a trivial type (or + loading it from the cache, after the first time). Where it does not, the + control is set to ``'run'`` (where the conda-forge and PyPI builds of + SuiteSparse start it), with no error or warning: a library built without + the JIT, which clamps every write to ``'run'``; no compiler on disk, even + after :func:`_repair_jit_compiler`; a cache directory it cannot write to; + or a compiler that fails the probe. + The probe is load-bearing: from SuiteSparse 9.4 on a failed compile raises + ``JitError`` (once; SuiteSparse then drops the control to ``'load'``), and + without the probe that would be the first user operation that needs a new + kernel, such as a slice, which SuiteSparse precompiles no kernel for. + """ + from . import _IS_SSGB7 + + if _IS_SSGB7: + return # no JIT before SuiteSparse:GraphBLAS 8 + global _jit_enabled_for_udt, _configured_jit_c_control, _user_set_run + cfg = _CConfig() + _configured_jit_c_control = control + _user_set_run = control == "run" + if control != "on": + cfg["jit_c_control"] = "run" if control is None else control + return + if not _compiler_exists(cfg): + _repair_jit_compiler(cfg) + _jit_enabled_for_udt = False + if _compiler_exists(cfg) and _cache_is_writable(cfg): + cfg["jit_c_control"] = "on" + _jit_enabled_for_udt = cfg["jit_c_control"] == "on" and _probe_jit(cfg) + if not _jit_enabled_for_udt: + # Not the 'load' a failed probe leaves: 'load' still loads kernels from + # the on-disk cache, which other builds of this version may have written. + cfg["jit_c_control"] = "run" + + +def _auto_fix_jit_at_import(): + """Repair the JIT compiler path at ``gb.ss`` import; designed not to raise. + + Called unguarded from ``graphblas/ss/__init__.py``, so any exception here + breaks ``import graphblas.ss``. The body sticks to dict ops and the + string rewriting in :func:`_repair_jit_compiler`. + + Deliberately leaves ``jit_c_control`` alone. An attribute access such as + ``gb.ss.about["library_version"]`` imports this submodule, and that must + not change what any later operation computes or which kernels SuiteSparse + is willing to load from its on-disk cache. The control is set when the + backend initializes (:func:`_set_jit_c_control_at_init`), which also + repairs the compiler when it sets ``'on'``. Otherwise compilation is + enabled later, by :func:`_enable_jit_for_udt`, when a UDT actually needs a + kernel built, or a type or op is registered from C source. + """ + cfg = _ss_config() + if "jit_c_control" not in cfg: + return + if not jit_compiler_is_usable(): + _repair_jit_compiler(cfg) + + +# Tri-state: ``None`` until the compiler is probed (at init when that sets +# ``'on'``, or at the first request to compile C source while the control +# allows compiling), then the answer to "can this process JIT-compile?" until +# ``fix_jit_config`` re-arms it. +_jit_enabled_for_udt = None + +# ``graphblas.config["jit_c_control"]`` as init found it, None when it was not set. +_configured_jit_c_control = None +# Whether a ``'run'`` in effect is one the user set, which ``_enable_jit_for_udt`` +# keeps: a ``'run'`` configured before init, or one written to ``gb.ss.config`` +# (whose setter keeps this current) while something else is configured. With +# nothing configured, a ``'run'`` written there is the default, as before it. +_user_set_run = False + + +def _enable_jit_for_udt(): + """Enable JIT compilation the first time C source needs compiling. + + ``jit_c_control`` defaults to ``'run'`` (``graphblas.config``, set at + init), which runs kernels already loaded but neither compiles nor loads + any. UDT auto-lift wants ``'on'``; without it every UDT op falls back to + the Numba function-pointer path (typically 2-3x slower for elementwise + ops). An op or type registered from C source has no fallback: under + ``'run'`` its creation fails outright. + + This is where that bump belongs, rather than at import: registering a + type or an op from C source, or arming an op with C source, is an act + that plainly involves compiling C, so enabling the compiler is not a + surprise. Reading ``gb.ss.about`` is not, so it leaves the setting alone. + + Only the default ``'run'`` is raised. A ``'run'`` the user set (through + ``graphblas.config["jit_c_control"]``, or ``gb.ss.config``; see + ``_user_set_run``) is honored, often to keep processes that run at the + same time from compiling into one cache (GraphBLAS #449). So are + ``'off'``, ``'pause'``, and ``'load'``, and ``'load'`` is also where + SuiteSparse leaves the control after a compile fails, so raising it would + retry that compile at every later UDT op (raising ``JitError`` each time on + SuiteSparse 9.4 and later). Returns True iff a kernel armed now will be + compiled. Whether this process can compile at all is probed once and + cached. + """ + global _jit_enabled_for_udt + cfg = _ss_config() + control = cfg.get("jit_c_control") + if control in (None, "off", "pause", "load") or (control == "run" and _user_set_run): + # ``None`` is a library with no JIT at all. The others are honored and + # not cached, so they decide nothing once the control is back at the + # default. + return False + if _jit_enabled_for_udt is None: + if not jit_compiler_is_usable(): + _repair_jit_compiler(cfg) + # A probe kernel already in a read-only cache loads, and the first new + # kernel would then fail to compile; see ``_cache_is_writable``. + if _jit_enabled_for_udt := jit_compiler_is_usable() and _cache_is_writable(cfg): + cfg["jit_c_control"] = "on" + # A library built without the JIT clamps that write back down to + # ``'run'``. Otherwise the probe is load-bearing: a failed compile + # drops the control to ``'load'`` and, from SuiteSparse 9.4 on, + # also raises ``JitError``, so without the probe that one error + # would reach the first user operation that needs a kernel. + _jit_enabled_for_udt = cfg["jit_c_control"] == "on" and _probe_jit(cfg) + elif _jit_enabled_for_udt and control == "run": + # The cache answers "can this process compile?", which is settled once. + # The control is separate state that anything may have lowered since, + # and a kernel armed now still needs it raised. + cfg["jit_c_control"] = "on" + return _jit_enabled_for_udt + + +@contextlib.contextmanager +def _compiling_c_source(name): + """Enable the JIT to create ``name`` from C source, and explain a failure. + + SuiteSparse compiles the C source, or loads the kernel from its cache, as + part of creating the type or op. When it cannot, it reports a bare + ``NullPointer`` (op) or ``InvalidValue`` (type), or from 9.4 a + ``JitError`` for a failed compile, none with a message. ``jit_c_control`` + before and after says which case it was, so the same exception type is + raised again with the reason, chained to the original. + """ + from ...exceptions import InvalidValue, JitError, NullPointer + + cfg = _ss_config() + initial = cfg.get("jit_c_control") + _enable_jit_for_udt() + before = cfg.get("jit_c_control") + try: + yield + except (NullPointer, InvalidValue, JitError) as exc: + after = cfg.get("jit_c_control") + if after in {"off", "pause"} and after == before: + reason = ( + f"jit_c_control is {after!r}, so SuiteSparse will not compile it. " + "Set gb.ss.config['jit_c_control'] = 'on' to allow that." + ) + elif after == "load" and before == "on": + reason = ( + "SuiteSparse failed to compile it and set jit_c_control to 'load', so " + "nothing else compiles in this process until jit_c_control is set back " + "to 'on'. SuiteSparse prints the compiler's output when " + "gb.ss.config['burble'] is True." + ) + elif after == before == "load" and initial != "load": + # ``_enable_jit_for_udt`` only ever writes 'on', so 'load' here means + # the probe it just ran failed to compile. + reason = ( + "the C compiler does not work here: SuiteSparse could not compile a " + "small test type, and set jit_c_control to 'load'. To see the " + "compiler's output, set gb.ss.config['burble'] = True and " + "jit_c_control back to 'on', then try again." + ) + elif after == before == "load": + reason = ( + "jit_c_control is 'load', so SuiteSparse only loads kernels already in " + "its cache, and this one is not there. SuiteSparse also sets 'load' " + "itself after a compile fails. Set gb.ss.config['jit_c_control'] = " + "'on' to compile it." + ) + elif after == before == "run" and _user_set_run: + reason = ( + "jit_c_control was set to 'run' (by graphblas.config['jit_c_control'], " + "the GRAPHBLAS_JIT_C_CONTROL environment variable, or gb.ss.config), " + "which neither compiles nor loads, and python-graphblas keeps a 'run' " + "that was set. To compile it, set gb.ss.config['jit_c_control'] = 'on', " + "or None for the default 'run', which python-graphblas raises to 'on' " + "for C source." + ) + elif after == before == "run": + reason = ( + "jit_c_control is still 'run', which neither compiles nor loads: " + "python-graphblas did not find the configured C compiler, or " + "SuiteSparse was built without the JIT. gb.ss.fix_jit_config() looks " + "for a compiler and turns the JIT on." + ) + else: + raise + raise type(exc)(f"Cannot create {name} from C source: {reason}") from exc + + +# Keyed by ``(op_name, dtype_name)`` so each distinct pair warns once. +# A user who registers several UDTs gets one warning per (op, dtype) pair +# rather than a single global swallow. +_warned_no_jit_for = set() + + +def _maybe_warn_no_jit(*, op_name="", dtype_name=""): + """Emit a ``NoJITWarning`` (once per ``(op_name, dtype_name)``) when UDT auto-lift falls back. + + The only caller is ``udt_codegen._maybe_warn_jit_skipped``, which fires + after C codegen returned nothing. Codegen reads the dtype and neither the + compiler path nor ``jit_c_control``, so the dtype is the cause here even + on a host that could not have compiled a kernel anyway. Do not re-derive + the cause from the live config; that is not the state codegen read. + + The introspection properties (``DataType.jit_c_definition`` and + ``TypedUserBinaryOp.jit_c_source``) show what was generated, or ``None`` + when codegen was skipped. + """ + key = (op_name, dtype_name) + if key in _warned_no_jit_for: + return + _warned_no_jit_for.add(key) + import warnings as _warnings + + from ...exceptions import NoJITWarning + + loc = f" (op={op_name!r}, dtype={dtype_name!r})" if op_name else "" + _warnings.warn( + f"UDT operator running without JIT compilation: this UDT is not " + f"expressible as a C struct{loc} (a field name is a C reserved word or " + f"stdlib macro, a field type isn't in the numpy-to-C map, a field is " + f"array-typed, or the record has a packed layout). The op still works " + f"through the Numba function-pointer fallback; only the JIT speedup is " + f"lost (typically 2-3x slower for elementwise ops, since SuiteSparse " + f"can't inline the kernel into its eWise and reduce templates). This " + f"warning fires once per (op, dtype) per process; silence with " + f"``warnings.filterwarnings('ignore', " + f"category=gb.exceptions.NoJITWarning)`` or by message match.", + NoJITWarning, + stacklevel=3, + ) diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index 509c56113..d39eb9145 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -1,22 +1,29 @@ import itertools import numpy as np -from suitesparse_graphblas.utils import claim_buffer, claim_buffer_2d, unclaim_buffer +from suitesparse_graphblas.utils import claim_buffer, claim_buffer_2d import graphblas as gb from ... import binary, monoid from ...dtypes import _INDEX, BOOL, INT64, UINT64, lookup_dtype -from ...exceptions import _error_code_lookup, check_status, check_status_carg +from ...exceptions import ( + IndexOutOfBound, + InvalidValue, + _error_code_lookup, + check_status, + check_status_carg, +) from .. import NULL, _has_numba, ffi, lib from ..base import call -from ..dtypes import _string_to_dtype from ..operator import get_typed_op from ..scalar import Scalar, _as_scalar, _scalar_index from ..utils import ( _CArray, + _distinct_buffers, _MatrixArray, _Pointer, + check_values_size, get_order, get_shape, ints_to_numpy_buffer, @@ -25,6 +32,7 @@ values_to_numpy_buffer, wrapdoc, ) +from ._utils import _pack_securely, _resolve_serialized_dtype, _unclaim_taken from .config import BaseConfig from .descriptor import get_descriptor @@ -201,9 +209,18 @@ def is_iso(self): @property def iso_value(self): + """The value every entry of an iso-valued Matrix holds, as a Scalar. + + SuiteSparse can mark a Matrix with no entries iso. It then holds no value, + and the Scalar is empty, as for an empty reduction. Raises ValueError if + the Matrix is not iso-valued. + """ if self.is_iso: + parent = self._parent + if parent.nvals == 0: + return Scalar(parent.dtype, name="") # This may not be thread-safe if the parent is being modified in another thread - return Scalar.from_value(next(self.itervalues()), dtype=self._parent.dtype, name="") + return Scalar.from_value(next(self.itervalues()), dtype=parent.dtype, name="") raise ValueError("Matrix is not iso-valued") @property @@ -376,7 +393,7 @@ def concat(self, tiles, **opts): graphblas.ss.concat """ - tiles, m, n, is_matrix = _concat_mn(tiles, is_matrix=True) + tiles, m, n, _is_matrix = _concat_mn(tiles, is_matrix=True) self._concat(tiles, m, n, opts) def build_scalar(self, rows, columns, value): @@ -852,7 +869,9 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m return rv if method == "export": - mhandle = ffi_new("GrB_Matrix*", parent._carg) + # GraphBLAS frees the Matrix as it hands the arrays over, and sets this + # handle to NULL, so an error after that cannot leave it dangling + mhandle = parent.gb_obj type_ = ffi_new("GrB_Type*") nrows_ = ffi_new("GrB_Index*") ncols_ = ffi_new("GrB_Index*") @@ -1117,6 +1136,12 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m rv["nrows"] = nrows rv["ncols"] = ncols elif format in {"fullr", "fullc"}: + if parent._nvals != nrows * ncols: + # SuiteSparse refuses with GrB_INVALID_VALUE and no message. + raise InvalidValue( + f"Matrix with {parent._nvals} of {nrows * ncols} entries is not full, " + f"so it cannot {method} to format {format!r}" + ) if format == "fullr": cfunc = getattr(lib, f"GxB_Matrix_{method}_FullR") else: @@ -1153,8 +1178,6 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m rv["is_iso"] = is_iso rv["format"] = format rv["values"] = values - if method == "export": - parent.gb_obj[0] = NULL if parent.dtype._is_udt: rv["dtype"] = parent.dtype return rv @@ -1199,12 +1222,17 @@ def import_csr( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (uint64 for indptr and col_indices) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Matrix. If not specified, this will be inferred from ``values``. @@ -1297,7 +1325,20 @@ def _import_csr( ): if format is not None and format.lower() != "csr": raise ValueError(f"Invalid format: {format!r}. Must be None or 'csr'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_csr, + matrix, + "csr", + opts, + indptr=indptr, + values=values, + col_indices=col_indices, + is_iso=is_iso, + sorted_cols=sorted_cols, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import indptr = ints_to_numpy_buffer( indptr, np.uint64, copy=copy, ownable=True, name="index pointers" ) @@ -1309,8 +1350,7 @@ def _import_csr( values, dtype = values_to_numpy_buffer( values, dtype, copy=copy, ownable=True, subarray_after=1 ) - if col_indices is values: - values = np.copy(values) + indptr, col_indices, values = _distinct_buffers(indptr, col_indices, values) Ap = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indptr)) Aj = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", col_indices)) Ax = ffi_new("void**", ffi.from_buffer("void*", values)) @@ -1320,6 +1360,8 @@ def _import_csr( else: mhandle = matrix._carg args = () + n = nrows if method == "import" else matrix._nrows + check_values_size(values, dtype, n, is_iso, indptr=indptr, indices=col_indices) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Matrix_{method}_CSR")( mhandle, @@ -1334,6 +1376,7 @@ def _import_csr( not sorted_cols, NULL if desc is None else desc._carg, ) + _unclaim_taken((indptr, Ap), (col_indices, Aj), (values, Ax)) if method == "import": check_status_carg( status, @@ -1343,9 +1386,6 @@ def _import_csr( matrix = gb.Matrix._from_obj(mhandle, dtype, nrows, ncols, name=name) else: check_status(status, matrix) - unclaim_buffer(indptr) - unclaim_buffer(col_indices) - unclaim_buffer(values) return matrix @classmethod @@ -1388,12 +1428,17 @@ def import_csc( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (uint64 for indptr and row_indices) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Matrix. If not specified, this will be inferred from ``values``. @@ -1486,7 +1531,20 @@ def _import_csc( ): if format is not None and format.lower() != "csc": raise ValueError(f"Invalid format: {format!r} Must be None or 'csc'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_csc, + matrix, + "csc", + opts, + indptr=indptr, + values=values, + row_indices=row_indices, + is_iso=is_iso, + sorted_rows=sorted_rows, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import indptr = ints_to_numpy_buffer( indptr, np.uint64, copy=copy, ownable=True, name="index pointers" ) @@ -1498,8 +1556,7 @@ def _import_csc( values, dtype = values_to_numpy_buffer( values, dtype, copy=copy, ownable=True, subarray_after=1 ) - if row_indices is values: - values = np.copy(values) + indptr, row_indices, values = _distinct_buffers(indptr, row_indices, values) Ap = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indptr)) Ai = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", row_indices)) Ax = ffi_new("void**", ffi.from_buffer("void*", values)) @@ -1509,6 +1566,8 @@ def _import_csc( else: mhandle = matrix._carg args = () + n = ncols if method == "import" else matrix._ncols + check_values_size(values, dtype, n, is_iso, indptr=indptr, indices=row_indices) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Matrix_{method}_CSC")( mhandle, @@ -1523,6 +1582,7 @@ def _import_csc( not sorted_rows, NULL if desc is None else desc._carg, ) + _unclaim_taken((indptr, Ap), (row_indices, Ai), (values, Ax)) if method == "import": check_status_carg( status, @@ -1532,9 +1592,6 @@ def _import_csc( matrix = gb.Matrix._from_obj(mhandle, dtype, nrows, ncols, name=name) else: check_status(status, matrix) - unclaim_buffer(indptr) - unclaim_buffer(row_indices) - unclaim_buffer(values) return matrix @classmethod @@ -1583,12 +1640,17 @@ def import_hypercsr( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (uint64 for rows, indptr, col_indices) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Matrix. If not specified, this will be inferred from ``values``. @@ -1689,7 +1751,22 @@ def _import_hypercsr( ): if format is not None and format.lower() != "hypercsr": raise ValueError(f"Invalid format: {format!r} Must be None or 'hypercsr'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_hypercsr, + matrix, + "hypercsr", + opts, + rows=rows, + indptr=indptr, + values=values, + col_indices=col_indices, + nvec=nvec, + is_iso=is_iso, + sorted_cols=sorted_cols, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import indptr = ints_to_numpy_buffer( indptr, np.uint64, copy=copy, ownable=True, name="index pointers" ) @@ -1704,8 +1781,7 @@ def _import_hypercsr( ) if not is_iso and values.ndim == 0: is_iso = True - if col_indices is values: - values = np.copy(values) + indptr, rows, col_indices, values = _distinct_buffers(indptr, rows, col_indices, values) Ap = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indptr)) Ah = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", rows)) Aj = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", col_indices)) @@ -1728,6 +1804,7 @@ def _import_hypercsr( else: mhandle = matrix._carg args = () + check_values_size(values, dtype, nvec, is_iso, indptr=indptr, indices=col_indices) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Matrix_{method}_HyperCSR")( mhandle, @@ -1745,6 +1822,7 @@ def _import_hypercsr( not sorted_cols, NULL if desc is None else desc._carg, ) + _unclaim_taken((indptr, Ap), (rows, Ah), (col_indices, Aj), (values, Ax)) if method == "import": check_status_carg( status, @@ -1754,10 +1832,6 @@ def _import_hypercsr( matrix = gb.Matrix._from_obj(mhandle, dtype, nrows, ncols, name=name) else: check_status(status, matrix) - unclaim_buffer(indptr) - unclaim_buffer(rows) - unclaim_buffer(col_indices) - unclaim_buffer(values) return matrix @classmethod @@ -1806,12 +1880,17 @@ def import_hypercsc( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (uint64 for indptr and row_indices) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Matrix. If not specified, this will be inferred from ``values``. @@ -1912,7 +1991,22 @@ def _import_hypercsc( ): if format is not None and format.lower() != "hypercsc": raise ValueError(f"Invalid format: {format!r} Must be None or 'hypercsc'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_hypercsc, + matrix, + "hypercsc", + opts, + cols=cols, + indptr=indptr, + values=values, + row_indices=row_indices, + nvec=nvec, + is_iso=is_iso, + sorted_rows=sorted_rows, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import indptr = ints_to_numpy_buffer( indptr, np.uint64, copy=copy, ownable=True, name="index pointers" ) @@ -1925,8 +2019,7 @@ def _import_hypercsc( values, dtype = values_to_numpy_buffer( values, dtype, copy=copy, ownable=True, subarray_after=1 ) - if row_indices is values: - values = np.copy(values) + indptr, cols, row_indices, values = _distinct_buffers(indptr, cols, row_indices, values) if not is_iso and values.ndim == 0: is_iso = True Ap = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indptr)) @@ -1951,6 +2044,7 @@ def _import_hypercsc( else: mhandle = matrix._carg args = () + check_values_size(values, dtype, nvec, is_iso, indptr=indptr, indices=row_indices) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Matrix_{method}_HyperCSC")( mhandle, @@ -1968,6 +2062,7 @@ def _import_hypercsc( not sorted_rows, NULL if desc is None else desc._carg, ) + _unclaim_taken((indptr, Ap), (cols, Ah), (row_indices, Ai), (values, Ax)) if method == "import": check_status_carg( status, @@ -1977,10 +2072,6 @@ def _import_hypercsc( matrix = gb.Matrix._from_obj(mhandle, dtype, nrows, ncols, name=name) else: check_status(status, matrix) - unclaim_buffer(indptr) - unclaim_buffer(cols) - unclaim_buffer(row_indices) - unclaim_buffer(values) return matrix @classmethod @@ -2028,12 +2119,17 @@ def import_bitmapr( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (bool for bitmap) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Matrix. If not specified, this will be inferred from ``values``. @@ -2122,7 +2218,19 @@ def _import_bitmapr( ): if format is not None and format.lower() != "bitmapr": raise ValueError(f"Invalid format: {format!r} Must be None or 'bitmapr'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_bitmapr, + matrix, + "bitmapr", + opts, + bitmap=bitmap, + values=values, + nvals=nvals, + is_iso=is_iso, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import bitmap = ints_to_numpy_buffer( bitmap, np.bool_, copy=copy, ownable=True, order="C", name="bitmap" ) @@ -2131,8 +2239,7 @@ def _import_bitmapr( values, dtype = values_to_numpy_buffer( values, dtype, copy=copy, ownable=True, order="C", subarray_after=2 ) - if bitmap is values: - values = np.copy(values) + bitmap, values = _distinct_buffers(bitmap, values) if method == "import": nrows, ncols = get_shape(nrows, ncols, dtype, bitmap=bitmap, values=values) else: @@ -2150,6 +2257,7 @@ def _import_bitmapr( else: mhandle = matrix._carg args = () + check_values_size(values, dtype, nrows * ncols, is_iso) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Matrix_{method}_BitmapR")( mhandle, @@ -2162,6 +2270,7 @@ def _import_bitmapr( nvals, NULL if desc is None else desc._carg, ) + _unclaim_taken((bitmap, Ab), (values, Ax)) if method == "import": check_status_carg( status, @@ -2171,8 +2280,6 @@ def _import_bitmapr( matrix = gb.Matrix._from_obj(mhandle, dtype, nrows, ncols, name=name) else: check_status(status, matrix) - unclaim_buffer(bitmap) - unclaim_buffer(values) return matrix @classmethod @@ -2220,12 +2327,17 @@ def import_bitmapc( memory buffers to GraphBLAS, the arrays must: - be FORTRAN contiguous - have the correct dtype (bool for bitmap) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Matrix. If not specified, this will be inferred from ``values``. @@ -2314,7 +2426,19 @@ def _import_bitmapc( ): if format is not None and format.lower() != "bitmapc": raise ValueError(f"Invalid format: {format!r} Must be None or 'bitmapc'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_bitmapc, + matrix, + "bitmapc", + opts, + bitmap=bitmap, + values=values, + nvals=nvals, + is_iso=is_iso, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import bitmap = ints_to_numpy_buffer( bitmap, np.bool_, copy=copy, ownable=True, order="F", name="bitmap" ) @@ -2323,8 +2447,7 @@ def _import_bitmapc( values, dtype = values_to_numpy_buffer( values, dtype, copy=copy, ownable=True, order="F", subarray_after=2 ) - if bitmap is values: - values = np.copy(values) + bitmap, values = _distinct_buffers(bitmap, values) if method == "import": nrows, ncols = get_shape(nrows, ncols, dtype, bitmap=bitmap, values=values) else: @@ -2342,6 +2465,7 @@ def _import_bitmapc( else: mhandle = matrix._carg args = () + check_values_size(values, dtype, nrows * ncols, is_iso) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Matrix_{method}_BitmapC")( mhandle, @@ -2354,6 +2478,7 @@ def _import_bitmapc( nvals, NULL if desc is None else desc._carg, ) + _unclaim_taken((bitmap, Ab), (values, Ax)) if method == "import": check_status_carg( status, @@ -2363,8 +2488,6 @@ def _import_bitmapc( matrix = gb.Matrix._from_obj(mhandle, dtype, nrows, ncols, name=name) else: check_status(status, matrix) - unclaim_buffer(bitmap) - unclaim_buffer(values) return matrix @classmethod @@ -2405,12 +2528,17 @@ def import_fullr( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Matrix. If not specified, this will be inferred from ``values``. @@ -2491,7 +2619,17 @@ def _import_fullr( ): if format is not None and format.lower() != "fullr": raise ValueError(f"Invalid format: {format!r} Must be None or 'fullr'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_fullr, + matrix, + "fullr", + opts, + values=values, + is_iso=is_iso, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import if method == "pack": dtype = matrix.dtype values, dtype = values_to_numpy_buffer( @@ -2509,10 +2647,12 @@ def _import_fullr( else: mhandle = matrix._carg args = () + check_values_size(values, dtype, nrows * ncols, is_iso) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Matrix_{method}_FullR")( mhandle, *args, Ax, values.nbytes, is_iso, NULL if desc is None else desc._carg ) + _unclaim_taken((values, Ax)) if method == "import": check_status_carg( status, @@ -2522,7 +2662,6 @@ def _import_fullr( matrix = gb.Matrix._from_obj(mhandle, dtype, nrows, ncols, name=name) else: check_status(status, matrix) - unclaim_buffer(values) return matrix @classmethod @@ -2563,12 +2702,17 @@ def import_fullc( memory buffers to GraphBLAS, the arrays must: - be FORTRAN contiguous - have the correct dtype - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Matrix. If not specified, this will be inferred from ``values``. @@ -2649,7 +2793,17 @@ def _import_fullc( ): if format is not None and format.lower() != "fullc": raise ValueError(f"Invalid format: {format!r}. Must be None or 'fullc'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_fullc, + matrix, + "fullc", + opts, + values=values, + is_iso=is_iso, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import if method == "pack": dtype = matrix.dtype values, dtype = values_to_numpy_buffer( @@ -2666,10 +2820,12 @@ def _import_fullc( else: mhandle = matrix._carg args = () + check_values_size(values, dtype, nrows * ncols, is_iso) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Matrix_{method}_FullC")( mhandle, *args, Ax, values.nbytes, is_iso, NULL if desc is None else desc._carg ) + _unclaim_taken((values, Ax)) if method == "import": check_status_carg( status, @@ -2679,7 +2835,6 @@ def _import_fullc( matrix = gb.Matrix._from_obj(mhandle, dtype, nrows, ncols, name=name) else: check_status(status, matrix) - unclaim_buffer(values) return matrix @classmethod @@ -2911,12 +3066,17 @@ def import_coor( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (uint64 for indptr and row_indices) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. For "coor", ownership of "rows" will never change. dtype : dtype, optional dtype of the new Matrix. @@ -3017,7 +3177,7 @@ def _import_coor( raise ValueError(f"Invalid format: {format!r}. Must be None or 'coor'.") if not sorted_rows: raise ValueError("sorted_rows must be True when importing 'coor' format") - indptr = indices_to_indptr(rows, nrows + 1) + indptr = _coo_indptr(rows, nrows, "rows", "nrows") return cls._import_csr( nrows=nrows, ncols=ncols, @@ -3079,12 +3239,17 @@ def import_cooc( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (uint64 for indptr and row_indices) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. For "cooc", ownership of "cols" will never change. dtype : dtype, optional dtype of the new Matrix. @@ -3185,7 +3350,7 @@ def _import_cooc( raise ValueError(f"Invalid format: {format!r}. Must be None or 'cooc'.") if not sorted_cols: raise ValueError("sorted_cols must be True when importing 'cooc' format") - indptr = indices_to_indptr(cols, ncols + 1) + indptr = _coo_indptr(cols, ncols, "cols", "ncols") return cls._import_csc( nrows=nrows, ncols=ncols, @@ -4059,10 +4224,10 @@ def serialize(self, compression="default", level=None, **opts): - "lz4hc": LZ4 compression that allows the compression level (1-9) to be set. Low compression level (1) is faster, high (9) is more compact. Default is 9. - "zstd": ZSTD compression, which allows compression level (1-19) to be set. - Low compression level (1) is faster, high (19) is more compact. Default is 19. + Low compression level (1) is faster, high (19) is more compact. Default is 1. - "none" or None: no compression - level : int [1-9], optional - The compression level, between 1 to 9, to use with "lz4hc" and "zstd" compression. + level : int, optional + The compression level to use with "lz4hc" (1-9) and "zstd" (1-19) compression. Level 1 is the fastest and largest, and is the default for "zstd" compression. Level 9 is the default when using "lz4hc" compression. nthreads : int, optional @@ -4088,13 +4253,22 @@ def serialize(self, compression="default", level=None, **opts): dtype_size = ffi_new("size_t*") status = lib.GrB_Type_get_SIZE(parent.dtype.gb_obj[0], dtype_size, lib.GrB_NAME) check_status_carg(status, "Type", parent.dtype.gb_obj[0]) - # Then get the name - dtype_char = ffi_new(f"char[{dtype_size[0]}]") - status = lib.GrB_Type_get_String(parent.dtype.gb_obj[0], dtype_char, lib.GrB_NAME) - check_status_carg(status, "Type", parent.dtype.gb_obj[0]) - # Then set the name - status = lib.GrB_Matrix_set_String(parent._carg, dtype_char, lib.GrB_NAME) - check_status_carg(status, "Matrix", parent._carg) + if dtype_size[0] >= lib.GxB_MAX_NAME_LEN: + # The dtype name is too long to safely store in the blob (GxB_Serialized_get_SIZE + # segfaults on names >= GxB_MAX_NAME_LEN). For named UDTs, use the short + # registered name instead; anonymous UDTs cannot round-trip without dtype=. + if not parent.dtype._is_anonymous: + val_obj = ffi.new("char[]", parent.dtype.name.encode()) + status = lib.GrB_Matrix_set_String(parent._carg, val_obj, lib.GrB_NAME) + check_status_carg(status, "Matrix", parent._carg) + else: + # Then get the name + dtype_char = ffi_new(f"char[{dtype_size[0]}]") + status = lib.GrB_Type_get_String(parent.dtype.gb_obj[0], dtype_char, lib.GrB_NAME) + check_status_carg(status, "Type", parent.dtype.gb_obj[0]) + # Then set the name + status = lib.GrB_Matrix_set_String(parent._carg, dtype_char, lib.GrB_NAME) + check_status_carg(status, "Matrix", parent._carg) check_status( lib.GxB_Matrix_serialize( @@ -4138,33 +4312,7 @@ def deserialize(cls, data, dtype=None, *, name=None, **opts): data = np.frombuffer(data, np.uint8) data_obj = ffi.from_buffer("void*", data) if dtype is None: - # Get the dtype name first (for non-UDTs) - cname = ffi_new(f"char[{lib.GxB_MAX_NAME_LEN}]") - info = lib.GxB_deserialize_type_name( - cname, - data_obj, - data.nbytes, - ) - if info != lib.GrB_SUCCESS: - raise _error_code_lookup[info]("Matrix deserialize failed to get the dtype name") - dtype_name = b"".join(itertools.takewhile(b"\x00".__ne__, cname)).decode() - if not dtype_name and hasattr(lib, "GxB_Serialized_get_String"): - # Handle UDTs. First get the size of name - dtype_size = ffi_new("size_t*") - info = lib.GxB_Serialized_get_SIZE(data_obj, dtype_size, lib.GrB_NAME, data.nbytes) - if info != lib.GrB_SUCCESS: - raise _error_code_lookup[info]( - "Matrix deserialize failed to get the size of name" - ) - # Then get the name - dtype_char = ffi_new(f"char[{dtype_size[0]}]") - info = lib.GxB_Serialized_get_String( - data_obj, dtype_char, lib.GrB_NAME, data.nbytes - ) - if info != lib.GrB_SUCCESS: - raise _error_code_lookup[info]("Matrix deserialize failed to get the name") - dtype_name = ffi.string(dtype_char).decode() - dtype = _string_to_dtype(dtype_name) + dtype = _resolve_serialized_dtype(data_obj, data.nbytes, "Matrix") else: dtype = lookup_dtype(dtype) desc = get_descriptor(**opts) @@ -4416,16 +4564,36 @@ def issorted(arr): # pragma: no cover (numba) @njit def indices_to_indptr(indices, size): # pragma: no cover (numba) - """Calculate the indptr for e.g. CSR from sorted COO rows.""" + """Calculate the indptr, e.g. for CSR, from uint64 COO rows, each below ``size - 1``. + + Each row's count goes at ``indptr[row + 1]``, then a running sum makes the offsets, + so every empty row, wherever it is, gets an empty range. + """ indptr = np.zeros(size, dtype=indices.dtype) - index = np.uint64(0) one = np.uint64(1) for i in range(indices.size): - row = indices[i] - if row != index: - indptr[index + one] = i - index = row - indptr[index + one] = indices.size + indptr[indices[i] + one] += one + for i in range(1, size): + indptr[i] += indptr[i - 1] + return indptr + + +def _coo_indptr(indices, size, name, size_name): + """Return the CSR (or CSC) indptr, of length ``size + 1``, for sorted COO rows (or columns).""" + # Read 2-D indices as a flat buffer, as other imports do. The kernel would index + # a 2-D array by its rows, past the end of its first axis. + indices = ints_to_numpy_buffer(indices, np.uint64, name=name).ravel() + # indices_to_indptr writes at each index + 1, with no bounds check under Numba + if indices.size > 0 and (biggest := indices.max()) >= size: + raise IndexOutOfBound(f"{name} must be less than {size_name}, {size}; got {biggest}") + if _has_numba: + return indices_to_indptr(indices, size + 1) + # Without Numba the kernel's loops run in Python, once per entry and once per row. + # Every index is below size, so it reads the same as int64; counts are never negative. + indptr = np.empty(size + 1, np.uint64) + indptr[0] = 0 + counts = np.bincount(indices.view(np.int64), minlength=size) + np.cumsum(counts.view(np.uint64), out=indptr[1:]) return indptr diff --git a/graphblas/core/ss/select.py b/graphblas/core/ss/select.py index 3ba135eee..323af9f40 100644 --- a/graphblas/core/ss/select.py +++ b/graphblas/core/ss/select.py @@ -38,7 +38,12 @@ def register_new(name, jit_c_definition, input_type, thunk_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. This will also create an IndexUnary operator under ``gb.indexunary.ss`` @@ -52,6 +57,8 @@ def register_new(name, jit_c_definition, input_type, thunk_type): The C definition as a string of the user-defined function. For example: ``"void woot (bool *z, const int32_t *x, GrB_Index i, GrB_Index j, int32_t *y) "`` ``"{ (*z) = ((*x) + i + j == (*y)) ; }"`` + A UDT made from a numpy dtype is named in C by its ``jit_c_name`` (usually ``gbudt_`` + and its name), not its Python name. input_type : dtype The dtype of the operand of the select operator. thunk_type : dtype diff --git a/graphblas/core/ss/unary.py b/graphblas/core/ss/unary.py index 0b7ced3c8..695dec5c1 100644 --- a/graphblas/core/ss/unary.py +++ b/graphblas/core/ss/unary.py @@ -5,6 +5,7 @@ from ..operator.base import TypedOpBase from ..operator.unary import TypedUserUnaryOp, UnaryOp from . import _IS_SSGB7 +from ._utils import _jit_c_source ffi_new = ffi.new @@ -38,7 +39,12 @@ def register_new(name, jit_c_definition, input_type, ret_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. Parameters ---------- @@ -49,6 +55,8 @@ def register_new(name, jit_c_definition, input_type, ret_type): jit_c_definition : str The C definition as a string of the user-defined function. For example: ``"void square (float *z, float *x) { (*z) = (*x) * (*x) ; } ;"`` + A UDT made from a numpy dtype is named in C by its ``jit_c_name`` (usually ``gbudt_`` + and its name), not its Python name. input_type : dtype The dtype of the operand of the unary operator. ret_type : dtype @@ -90,20 +98,25 @@ def register_new(name, jit_c_definition, input_type, ret_type): else: # We use `is_udt=True` to make dtype handling flexible and explicit. rv = UnaryOp(name, is_udt=True) + from .jit_config import _compiling_c_source + + c_name, c_definition = _jit_c_source(name, funcname, jit_c_definition) gb_obj = ffi_new("GrB_UnaryOp*") - check_status_carg( - lib.GxB_UnaryOp_new( - gb_obj, - NULL, - ret_type._carg, - input_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), - ), - "UnaryOp", - gb_obj[0], - ) + with _compiling_c_source(f"gb.unary.{name}"): + check_status_carg( + lib.GxB_UnaryOp_new( + gb_obj, + NULL, + ret_type._carg, + input_type._carg, + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), + ), + "UnaryOp", + gb_obj[0], + ) op = TypedJitUnaryOp(rv, funcname, input_type, ret_type, gb_obj[0], jit_c_definition) + op._jit_c_info = (c_name, c_definition) rv._add(op, is_jit=True) setattr(module, funcname, rv) return rv diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index fdde7eb92..4ad1c3eb9 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -1,26 +1,28 @@ import itertools import numpy as np -from suitesparse_graphblas.utils import claim_buffer, unclaim_buffer +from suitesparse_graphblas.utils import claim_buffer import graphblas as gb from ... import binary, monoid from ...dtypes import _INDEX, INT64, UINT64, lookup_dtype -from ...exceptions import _error_code_lookup, check_status, check_status_carg +from ...exceptions import InvalidValue, _error_code_lookup, check_status, check_status_carg from .. import NULL, ffi, lib from ..base import call -from ..dtypes import _string_to_dtype from ..operator import get_typed_op from ..scalar import Scalar, _as_scalar from ..utils import ( _CArray, + _distinct_buffers, _MatrixArray, + check_values_size, ints_to_numpy_buffer, normalize_chunks, values_to_numpy_buffer, wrapdoc, ) +from ._utils import _pack_securely, _resolve_serialized_dtype, _unclaim_taken from .config import BaseConfig from .descriptor import get_descriptor from .matrix import _concat_mn, njit @@ -120,9 +122,18 @@ def is_iso(self): @property def iso_value(self): + """The value every entry of an iso-valued Vector holds, as a Scalar. + + SuiteSparse can mark a Vector with no entries iso. It then holds no value, + and the Scalar is empty, as for an empty reduction. Raises ValueError if + the Vector is not iso-valued. + """ if self.is_iso: + parent = self._parent + if parent.nvals == 0: + return Scalar(parent.dtype, name="") # This may not be thread-safe if the parent is being modified in another thread - return Scalar.from_value(next(self.itervalues()), dtype=self._parent.dtype, name="") + return Scalar.from_value(next(self.itervalues()), dtype=parent.dtype, name="") raise ValueError("Vector is not iso-valued") @property @@ -262,7 +273,7 @@ def concat(self, tiles, **opts): graphblas.ss.concat """ - tiles, m, n, is_matrix = _concat_mn(tiles, is_matrix=False) + tiles, m, _n, _is_matrix = _concat_mn(tiles, is_matrix=False) self._concat(tiles, m, opts) def build_scalar(self, indices, value): @@ -504,7 +515,9 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m size = parent._size if method == "export": - vhandle = ffi_new("GrB_Vector*", parent._carg) + # GraphBLAS frees the Vector as it hands the arrays over, and sets this + # handle to NULL, so an error after that cannot leave it dangling + vhandle = parent.gb_obj type_ = ffi_new("GrB_Type*") size_ = ffi_new("GrB_Index*") args = (type_, size_) @@ -595,6 +608,12 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m if raw: rv["size"] = size elif format == "full": + if parent._nvals != size: + # SuiteSparse refuses with GrB_INVALID_VALUE and no message. + raise InvalidValue( + f"Vector with {parent._nvals} of {size} entries is not full, " + f"so it cannot {method} to format 'full'" + ) check_status( getattr(lib, f"GxB_Vector_{method}_Full")( vhandle, @@ -625,8 +644,6 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m format=format, values=values, ) - if method == "export": - parent.gb_obj[0] = NULL if parent.dtype._is_udt: rv["dtype"] = parent.dtype return rv @@ -864,12 +881,17 @@ def import_sparse( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (uint64 for indices) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Vector. If not specified, this will be inferred from ``values``. @@ -959,15 +981,27 @@ def _import_sparse( ): if format is not None and format.lower() != "sparse": raise ValueError(f"Invalid format: {format!r}. Must be None or 'sparse'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_sparse, + vector, + "sparse", + opts, + indices=indices, + values=values, + nvals=nvals, + is_iso=is_iso, + sorted_index=sorted_index, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import indices = ints_to_numpy_buffer(indices, np.uint64, copy=copy, ownable=True, name="indices") if method == "pack": dtype = vector.dtype values, dtype = values_to_numpy_buffer( values, dtype, copy=copy, ownable=True, subarray_after=1 ) - if indices is values: - values = np.copy(values) + indices, values = _distinct_buffers(indices, values) vi = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indices)) vx = ffi_new("void**", ffi.from_buffer("void*", values)) if nvals is None: @@ -983,6 +1017,7 @@ def _import_sparse( else: vhandle = vector._carg args = () + check_values_size(values, dtype, nvals, is_iso) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Vector_{method}_CSC")( vhandle, @@ -996,6 +1031,7 @@ def _import_sparse( not sorted_index, NULL if desc is None else desc._carg, ) + _unclaim_taken((indices, vi), (values, vx)) if method == "import": check_status_carg( status, @@ -1005,8 +1041,6 @@ def _import_sparse( vector = gb.Vector._from_obj(vhandle, dtype, size, name=name) else: check_status(status, vector) - unclaim_buffer(indices) - unclaim_buffer(values) return vector @classmethod @@ -1048,12 +1082,17 @@ def import_bitmap( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (bool for bitmap) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Vector. If not specified, this will be inferred from ``values``. @@ -1139,7 +1178,19 @@ def _import_bitmap( ): if format is not None and format.lower() != "bitmap": raise ValueError(f"Invalid format: {format!r}. Must be None or 'bitmap'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_bitmap, + vector, + "bitmap", + opts, + bitmap=bitmap, + values=values, + nvals=nvals, + is_iso=is_iso, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import bitmap = ints_to_numpy_buffer(bitmap, np.bool_, copy=copy, ownable=True, name="bitmap") if method == "pack": dtype = vector.dtype @@ -1147,8 +1198,7 @@ def _import_bitmap( values, dtype = values_to_numpy_buffer( values, dtype, copy=copy, ownable=True, subarray_after=1 ) - if bitmap is values: - values = np.copy(values) + bitmap, values = _distinct_buffers(bitmap, values) vhandle = ffi_new("GrB_Vector*") vb = ffi_new("int8_t**", ffi.from_buffer("int8_t*", bitmap)) vx = ffi_new("void**", ffi.from_buffer("void*", values)) @@ -1170,6 +1220,7 @@ def _import_bitmap( else: vhandle = vector._carg args = () + check_values_size(values, dtype, size, is_iso) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Vector_{method}_Bitmap")( vhandle, @@ -1182,6 +1233,7 @@ def _import_bitmap( nvals, NULL if desc is None else desc._carg, ) + _unclaim_taken((bitmap, vb), (values, vx)) if method == "import": check_status_carg( status, @@ -1191,8 +1243,6 @@ def _import_bitmap( vector = gb.Vector._from_obj(vhandle, dtype, size, name=name) else: check_status(status, vector) - unclaim_buffer(bitmap) - unclaim_buffer(values) return vector @classmethod @@ -1228,12 +1278,17 @@ def import_full( memory buffers to GraphBLAS, the arrays must: - be C contiguous - have the correct dtype (bool for bitmap) - - own its own data + - own its own data, allocated by GraphBLAS (arrays from ``export``, + ``unpack`` and ``suitesparse_graphblas.utils.empty``) or by NumPy's + default allocator (unless suitesparse-graphblas was initialized + with ``memory_manager="c"``) - be writeable - If all of these conditions are not met, then the data will be + If any of these conditions is not met, then the data will be copied and the original array will be unmodified. If zero copy to GraphBLAS is successful, then the array will be modified to be read-only and will no longer own the data. + With ``secure_import=True`` the data is always copied, so the + arrays stay valid if the check fails. dtype : dtype, optional dtype of the new Vector. If not specified, this will be inferred from ``values``. @@ -1311,7 +1366,17 @@ def _import_full( ): if format is not None and format.lower() != "full": raise ValueError(f"Invalid format: {format!r}. Must be None or 'full'.") - copy = not take_ownership + if method == "pack" and secure_import: + return _pack_securely( + cls._import_full, + vector, + "full", + opts, + values=values, + is_iso=is_iso, + take_ownership=take_ownership, + ) + copy = not take_ownership or secure_import if method == "pack": dtype = vector.dtype size = vector._size @@ -1331,6 +1396,7 @@ def _import_full( else: vhandle = vector._carg args = () + check_values_size(values, dtype, size, is_iso) desc = get_descriptor(secure_import=secure_import, **opts) status = getattr(lib, f"GxB_Vector_{method}_Full")( vhandle, @@ -1340,6 +1406,7 @@ def _import_full( is_iso, NULL if desc is None else desc._carg, ) + _unclaim_taken((values, vx)) if method == "import": check_status_carg( status, @@ -1349,7 +1416,6 @@ def _import_full( vector = gb.Vector._from_obj(vhandle, dtype, size, name=name) else: check_status(status, vector) - unclaim_buffer(values) return vector @wrapdoc(head) @@ -1630,10 +1696,10 @@ def serialize(self, compression="default", level=None, **opts): - "lz4hc": LZ4 compression that allows the compression level (1-9) to be set. Low compression level (1) is faster, high (9) is more compact. Default is 9. - "zstd": ZSTD compression, which allows compression level (1-19) to be set. - Low compression level (1) is faster, high (19) is more compact. Default is 19. + Low compression level (1) is faster, high (19) is more compact. Default is 1. - "none" or None: no compression - level : int [1-9], optional - The compression level, between 1 to 9, to use with "lz4hc" and "zstd" compression. + level : int, optional + The compression level to use with "lz4hc" (1-9) and "zstd" (1-19) compression. Level 1 is the fastest and largest, and is the default for "zstd" compression. Level 9 is the default when using "lz4hc" compression. nthreads : int, optional @@ -1641,7 +1707,7 @@ def serialize(self, compression="default", level=None, **opts): None, 0 or negative nthreads means to use the default number of threads. For best performance, this function returns a numpy array with uint8 dtype. - Use ``Vector.ss.deserialize(blob)`` to create a Vector from the result of serialization· + Use ``Vector.ss.deserialize(blob)`` to create a Vector from the result of serialization. This method is intended to support all serialization options from SuiteSparse:GraphBLAS. @@ -1659,13 +1725,22 @@ def serialize(self, compression="default", level=None, **opts): dtype_size = ffi_new("size_t*") status = lib.GrB_Type_get_SIZE(parent.dtype.gb_obj[0], dtype_size, lib.GrB_NAME) check_status_carg(status, "Type", parent.dtype.gb_obj[0]) - # Then get the name - dtype_char = ffi_new(f"char[{dtype_size[0]}]") - status = lib.GrB_Type_get_String(parent.dtype.gb_obj[0], dtype_char, lib.GrB_NAME) - check_status_carg(status, "Type", parent.dtype.gb_obj[0]) - # Then set the name - status = lib.GrB_Vector_set_String(parent._carg, dtype_char, lib.GrB_NAME) - check_status_carg(status, "Vector", parent._carg) + if dtype_size[0] >= lib.GxB_MAX_NAME_LEN: + # The dtype name is too long to safely store in the blob (GxB_Serialized_get_SIZE + # segfaults on names >= GxB_MAX_NAME_LEN). For named UDTs, use the short + # registered name instead; anonymous UDTs cannot round-trip without dtype=. + if not parent.dtype._is_anonymous: + val_obj = ffi.new("char[]", parent.dtype.name.encode()) + status = lib.GrB_Vector_set_String(parent._carg, val_obj, lib.GrB_NAME) + check_status_carg(status, "Vector", parent._carg) + else: + # Then get the name + dtype_char = ffi_new(f"char[{dtype_size[0]}]") + status = lib.GrB_Type_get_String(parent.dtype.gb_obj[0], dtype_char, lib.GrB_NAME) + check_status_carg(status, "Type", parent.dtype.gb_obj[0]) + # Then set the name + status = lib.GrB_Vector_set_String(parent._carg, dtype_char, lib.GrB_NAME) + check_status_carg(status, "Vector", parent._carg) check_status( lib.GxB_Vector_serialize( @@ -1709,33 +1784,7 @@ def deserialize(cls, data, dtype=None, *, name=None, **opts): data = np.frombuffer(data, np.uint8) data_obj = ffi.from_buffer("void*", data) if dtype is None: - # Get the dtype name first (for non-UDTs) - cname = ffi_new(f"char[{lib.GxB_MAX_NAME_LEN}]") - info = lib.GxB_deserialize_type_name( - cname, - data_obj, - data.nbytes, - ) - if info != lib.GrB_SUCCESS: - raise _error_code_lookup[info]("Vector deserialize failed to get the dtype name") - dtype_name = b"".join(itertools.takewhile(b"\x00".__ne__, cname)).decode() - if not dtype_name and hasattr(lib, "GxB_Serialized_get_String"): - # Handle UDTs. First get the size of name - dtype_size = ffi_new("size_t*") - info = lib.GxB_Serialized_get_SIZE(data_obj, dtype_size, lib.GrB_NAME, data.nbytes) - if info != lib.GrB_SUCCESS: - raise _error_code_lookup[info]( - "Vector deserialize failed to get the size of name" - ) - # Then get the name - dtype_char = ffi_new(f"char[{dtype_size[0]}]") - info = lib.GxB_Serialized_get_String( - data_obj, dtype_char, lib.GrB_NAME, data.nbytes - ) - if info != lib.GrB_SUCCESS: - raise _error_code_lookup[info]("Vector deserialize failed to get the name") - dtype_name = ffi.string(dtype_char).decode() - dtype = _string_to_dtype(dtype_name) + dtype = _resolve_serialized_dtype(data_obj, data.nbytes, "Vector") else: dtype = lookup_dtype(dtype) desc = get_descriptor(**opts) diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py new file mode 100644 index 000000000..27ebc7b90 --- /dev/null +++ b/graphblas/core/udt.py @@ -0,0 +1,936 @@ +"""The rules for user-defined types (UDTs), apart from generating code for them. + +GraphBLAS knows a UDT only by its size: it cannot cast one to or from any other +type, and it copies an element of one as raw bytes. The rules here fill that +gap, in this order: + +- which built-in ops lift to UDTs (``BUILTIN_UDT_BINARY_OPS``), and the layout + of an element (:func:`_iter_record_leaves`); +- which operands an op on UDTs takes and the type of its result + (:func:`_check_udt_pair`, :func:`_udt_op_types`); +- what may be stored in an object of another type, and how it is converted on + the way (:func:`_store_cast_op`); +- when ``ewise_add`` would need GraphBLAS to cast a UDT; +- how ``isequal``, and ``==`` without Numba, compare elements; +- how an element's value becomes bytes and back, for ``Scalar.value``. + +``operator.udt_codegen`` generates the operators that compute on UDTs, +``operator.udf`` wraps a user's function on them, and ``literals`` types a +literal operand beside one. + +``operator.udt_codegen`` imports this module while the operator package is still +loading, so a function here that needs an operator, ``base`` or ``scalar`` +imports it inside the function, once it knows a UDT is involved. +""" + +import copy +from functools import lru_cache, reduce +from operator import mul + +import numpy as np + +from ..exceptions import DomainMismatch +from . import _has_numba, ffi, utils +from .dtypes import _nats_in_leaf_units, _registry, _view_if_same_layout, lookup_dtype +from .utils import output_type + +# The built-in ops that lift to UDTs, each with the Python operator or function +# the generated code applies. +BUILTIN_UDT_BINARY_OPS = { + "plus": "+", + "minus": "-", + "times": "*", + "truediv": "/", + "floordiv": "//", + "min": "min", + "max": "max", +} +BUILTIN_UDT_UNARY_OPS = { + "ainv": "-", + "abs": "abs", +} + + +# Ops with no meaning on complex fields: no ordering for min and max, and no +# integer mod for floordiv. +_OPS_NOT_FOR_COMPLEX = frozenset({"min", "max", "floordiv"}) + + +# The layout of an element. + + +def _get_udt_info(dtype): + """Return ('record', field_names) or ('array', (base_dtype, flat_size)) or None.""" + np_type = dtype.np_type + if np_type.names is not None: + return "record", np_type.names + if np_type.subdtype is not None: + base, shape = np_type.subdtype + return "array", (base, reduce(mul, shape)) + return None + + +def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): + """Yield ``(python_access, c_access, leaf_dtype)`` for each leaf in a record. + + A non-nested record yields one entry per field, with paths like + ``"['a']"`` (Python) and ``"a"`` (C struct). A nested record recurses + into structured fields: a top-level field ``"outer"`` whose dtype is + itself a struct with field ``"inner_a"`` yields + ``("['outer']['inner_a']", "outer.inner_a", float64_dtype)``. + + Only record-in-record nesting is recognised. Array-typed fields are + yielded as a single leaf, which the Numba path handles: a binary op + computes it element by element, a unary op applies its expression to the + whole sub-array, and the wrapper slice-assigns the result back. The JIT + path never sees such a record, because a subarray dtype has no entry in + ``udt_codegen.NP_TO_C_TYPES`` and ``udt_codegen._udt_c_typedef`` bails out, + so the op keeps the cfunc for every call. + """ + for name in np_type.names: + field_dtype = np_type.fields[name][0] + py_path = f"{python_prefix}[{name!r}]" + c_path = f"{c_prefix}{name}" if not c_prefix else f"{c_prefix}.{name}" + if field_dtype.names is not None: + yield from _iter_record_leaves(field_dtype, py_path, c_path) + else: + yield py_path, c_path, field_dtype + + +def _strip_leading_ones(shape): + i = 0 + while i < len(shape) and shape[i] == 1: + i += 1 + return shape[i:] + + +# Which operands an op on UDTs takes, and the type of its result. + + +def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): + """Raise ``KeyError`` if two UDT operands disagree on shape. + + ``info_x`` and ``info_y`` are ``_get_udt_info`` results. When either is + ``None`` (a built-in scalar side broadcasts), no further check fires. A + UDT that is neither a record nor an array (bytes, say) holds one value, + which pairs only with its own type, on either side; for ``eq`` and ``ne`` + only, text pairs with text of the same kind at another width, and a date + or duration with one in another unit, as numpy and a record's fields + compare them. Without this gate the generated + wrapper code raises a cryptic Numba ``TypingError`` on the first field + access, or, for array fields that do not broadcast, fails inside the + cfunc without telling the caller. Element dtypes may differ; + :func:`_udt_op_types` promotes them. + """ + compared_as_numpy = ( + op_name in {"eq", "ne"} + and info_x is None + and info_y is None + and dtype._is_udt + and dtype2._is_udt + and dtype.np_type.kind == dtype2.np_type.kind + and dtype.np_type.kind in "SUmM" + ) + if ( + dtype != dtype2 + and not compared_as_numpy + and (info_x is None and dtype._is_udt or info_y is None and dtype2._is_udt) + ): + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): a UDT that is " + "neither a record nor an array pairs only with its own type." + ) + if info_x is None or info_y is None: + return + kind_x, detail_x = info_x + kind_y, detail_y = info_y + if kind_x != kind_y: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"cannot mix record and array UDTs in a single element-wise op." + ) + if kind_x == "record": + if detail_x != detail_y: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must share field names, in the same order; got " + f"{list(detail_x)} vs {list(detail_y)}." + ) + # Matching top-level names is not enough. The codegen reads both + # operands at the left one's leaf paths, so a field that is a + # sub-record on one side and a scalar on the other, or sub-records + # with different names, fails Numba's typing: a compile error + # reported for what is really the same shape disagreement the check + # above reports as a KeyError. The paths must also come in the same + # order, as the top-level names must, so that leaf ``i`` of one + # operand is leaf ``i`` of the other. + leaves_x = list(_iter_record_leaves(dtype.np_type)) + leaves_y = list(_iter_record_leaves(dtype2.np_type)) + if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_y]: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must nest the same way, with the same field names in the " + f"same order at every level; got {[c for _py, c, _d in leaves_x]} vs " + f"{[c for _py, c, _d in leaves_y]}." + ) + # Each field must also have the same shape in both (field dtypes may + # differ), as numpy requires before it promotes or compares two record + # dtypes. Shapes that do not broadcast would raise inside the cfunc, + # unseen, and leave whatever was in the output buffer. Shapes that do + # broadcast would work one way round only, since the arithmetic ops + # write into a field shaped like the left operand's: ``(3,)`` plus + # ``(1,)`` works, ``(1,)`` plus ``(3,)`` does not. + for (_py, c, leaf_x), (_py_y, _c_y, leaf_y) in zip(leaves_x, leaves_y, strict=True): + if leaf_x.shape != leaf_y.shape: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must have the same shape for each field; field {c} is " + f"{leaf_x.shape} vs {leaf_y.shape}." + ) + return + # Array UDTs broadcast as numpy arrays do, so ``(3, 1)`` with ``(1, 4)`` + # gives ``(3, 4)``; the generated code reads each operand through a flat + # index map (``udt_codegen._array_operand_source``). Shapes such as ``(2, 3)`` with + # ``(3, 2)`` do not broadcast, as numpy refuses them. + (base_x, shape_x), (base_y, shape_y) = dtype.np_type.subdtype, dtype2.np_type.subdtype + try: + np.broadcast_shapes(shape_x, shape_y) + except ValueError: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): array UDTs must " + f"have shapes that broadcast together; got {base_x}{list(shape_x)} vs " + f"{base_y}{list(shape_y)}." + ) from None + + +def _builtin_dtype(np_type): + """Return the built-in ``DataType`` for ``np_type``, or ``None`` if it is not one.""" + dt = _registry.get(np_type) + return None if dt is None or dt._is_udt else dt + + +def _complex_op_error(op_name, dtype, dtype2): + return KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"this op is not defined on complex fields. Use ``binary.plus``, " + f"``minus``, ``times``, or ``truediv`` for complex element-wise " + f"arithmetic, or register a custom binary op." + ) + + +def _element_types(op_name, dtype, dtype2, elem_x, elem_y): + """Return ``(in_x, in_y, out)``, the numpy dtypes ``binary.{op_name}`` uses on one element. + + They come from the op's typing on the built-in dtypes, so each element of an + array UDT, or each record leaf, is promoted as a Vector of that dtype would + be, and ``truediv`` on integers gives floats. ``in_x`` and ``in_y`` are the + types GraphBLAS casts the operands to before it applies the op. The + generated code casts to them too, because Numba promotes differently: it + adds uint64 and int64 in int64, where GraphBLAS uses float64. ``dtype`` and + ``dtype2`` are the operands, named in errors. Two elements that are not + both built-in dtypes (bytes, datetime, records inside an array field) have + no promotion rule, so they must be the same dtype, which the result keeps. + """ + # ``min``/``max``/``floordiv`` have no defined semantics on complex + # operands (no ordering, no integer mod). Reject early with a clear + # message; otherwise Numba's cfunc compile blows up several frames + # down with ``NotImplementedError: No definition for lowering lt``. + if op_name in _OPS_NOT_FOR_COMPLEX and "c" in (elem_x.kind, elem_y.kind): + raise _complex_op_error(op_name, dtype, dtype2) + builtin_x = _builtin_dtype(elem_x) + builtin_y = _builtin_dtype(elem_y) + if builtin_x is None or builtin_y is None: + if elem_x == elem_y: + return elem_x, elem_y, elem_x + # Builds without complex dtypes (the vanilla backend) still lift complex + # fields, so promote them as the built-in typing does where it has them. + if "c" in (elem_x.kind, elem_y.kind) and {elem_x.kind, elem_y.kind} <= set("biufc"): + promoted = np.promote_types(elem_x, elem_y) + return promoted, promoted, promoted + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): elements of " + f"{elem_x} and {elem_y} have no common type." + ) + from .. import binary + from .operator import get_typed_op + + typed = get_typed_op(getattr(binary, op_name), builtin_x, builtin_y) + return typed.type.np_type, typed.type2.np_type, typed.return_type.np_type + + +def _udt_op_types(op_name, dtype, dtype2): + """Return ``(ret_type, input_types)`` for ``binary.{op_name}`` on ``dtype`` and ``dtype2``. + + At least one operand is a record or array UDT, and ``_check_udt_pair`` has + accepted the pair; the other may be a plain scalar type. Each element, or + record leaf, follows :func:`_element_types`, so the result never narrows: + an int64 array UDT plus ``0.5`` is ``FP64[3]``. ``ret_type`` is an + operand's own type when that holds the result, so same-type arithmetic + keeps its UDT. Otherwise it is the promoted layout, which is the UDT a user + registered with that layout if there is one, else an anonymous one. + ``input_types`` holds the ``(in_x, in_y)`` pair of each record leaf, in + leaf order, or the single pair for an array UDT's elements. + """ + np_x, np_y = dtype.np_type, dtype2.np_type + if np_x.names is not None or np_y.names is not None: + leaves_x = list(_iter_record_leaves(np_x)) if np_x.names is not None else None + leaves_y = list(_iter_record_leaves(np_y)) if np_y.names is not None else None + input_types = [] + result_leaves = [] + for i, (_py, _c, leaf) in enumerate(leaves_x or leaves_y): + elem_x = np_x if leaves_x is None else leaves_x[i][2].base + elem_y = np_y if leaves_y is None else leaves_y[i][2].base + in_x, in_y, elem = _element_types(op_name, dtype, dtype2, elem_x, elem_y) + input_types.append((in_x, in_y)) + result_leaves.append(np.dtype((elem, leaf.shape)) if leaf.shape else elem) + for operand, leaves in [(dtype, leaves_x), (dtype2, leaves_y)]: + if leaves is not None and all( + result_leaf == leaf + for result_leaf, (_py, _c, leaf) in zip(result_leaves, leaves, strict=True) + ): + return operand, input_types + records = [np_type for np_type in (np_x, np_y) if np_type.names is not None] + return lookup_dtype(_promoted_record(records, iter(result_leaves))), input_types + base_x, shape_x = np_x.subdtype or (np_x, ()) + base_y, shape_y = np_y.subdtype or (np_y, ()) + in_x, in_y, elem = _element_types(op_name, dtype, dtype2, base_x, base_y) + result = np.dtype((elem, np.broadcast_shapes(shape_x, shape_y))) + for operand in (dtype, dtype2): + if operand.np_type == result: + return operand, [(in_x, in_y)] + return lookup_dtype(result), [(in_x, in_y)] + + +def _monoid_pair_type(binaryop_name, dtype, dtype2): + """Return the type a built-in Monoid indexed with a UDT pair takes, compiling nothing. + + It is the type the pair promotes to, as ``monoid.plus["INT8", "INT16"]`` is + ``monoid.plus["INT16"]``: the type lifted arithmetic gives the pair, leaf by + leaf. ``binary.any`` keeps the second operand's type, so ``monoid.any`` + promotes as ``plus`` does, the same in either order. A pair that does not + pair raises the ``KeyError`` compiling the op would. + """ + op_name = binaryop_name if binaryop_name in BUILTIN_UDT_BINARY_OPS else "plus" + _check_udt_pair(op_name, dtype, dtype2, _get_udt_info(dtype), _get_udt_info(dtype2)) + return _udt_op_types(op_name, dtype, dtype2)[0] + + +def _promoted_record(records, result_leaves): + """Build a record laid out like ``records`` with ``result_leaves`` for its leaves. + + ``records`` holds one or two record dtypes with the same field names and + nesting. The new record aligns its fields, at each level, if any of them + does there, so a C-compatible operand gives a C-compatible result. + """ + first = records[0] + fields = [] + for name in first.names: + field = first.fields[name][0] + if field.names is not None: + field = _promoted_record([r.fields[name][0] for r in records], result_leaves) + else: + field = next(result_leaves) + fields.append((name, field)) + return np.dtype(fields, align=any(r.isalignedstruct for r in records)) + + +def _value_kind(np_type): + """What a leaf of ``np_type`` holds, as ``eq`` compares it: a number, text, or its own kind.""" + if np_type.kind in "biufc": + return "number" + return "text" if np_type.kind in "SU" else np_type.kind + + +def _check_leaves_compare(op_name, dtype, dtype2): + """Raise ``KeyError`` unless each pair of leaves that ``eq`` compares holds one kind of value. + + Text, a date, a duration and a number compare only with their own kind, as + ``plus`` finds no common type for any other pair. ``dtype`` and ``dtype2`` + already pair (:func:`_check_udt_pair`); a side that is not a record stands + in every leaf of the other. + """ + + def leaves(np_type): + if np_type.names is not None: + return [leaf.base for _py, _c, leaf in _iter_record_leaves(np_type)] + return [np_type.base] + + xs, ys = leaves(dtype.np_type), leaves(dtype2.np_type) + if len(xs) == 1: + xs = xs * len(ys) + if len(ys) == 1: + ys = ys * len(xs) + for x, y in zip(xs, ys, strict=True): + if _value_kind(x) != _value_kind(y): + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): elements of " + f"{x} and {y} do not compare." + ) + + +# Storing a value of one type in an object of another. + + +def _corresponding_leaves(np_x, np_z): + """Return ``(reason, pairs)`` for putting an element of ``np_x`` into one of ``np_z``. + + ``reason`` says why the two layouts do not correspond, or is ``None``, and + then ``pairs`` holds the base dtypes of each pair of leaves. Records must + line up as the operands of a lifted op do (:func:`_check_udt_pair`): the + same field names, order and nesting and the same shape for each field. + Arrays must have the same shape apart from leading axes of length 1; + operands broadcast, but a stored element keeps every value in its place. + numpy casts records by position whatever their names, and repeats or + drops elements of an array field of another length, which would put + values in places nobody chose. + """ + if (np_x.names is None) != (np_z.names is None): + return "a record UDT and an array UDT do not correspond", None + if np_x.names is None: + if np_x.subdtype is None or np_z.subdtype is None: + return "only record and array UDTs cast", None + (base_x, shape_x), (base_z, shape_z) = np_x.subdtype, np_z.subdtype + if _strip_leading_ones(shape_x) != _strip_leading_ones(shape_z): + return "array UDTs must have the same shape, apart from leading axes of length 1", None + return None, [(base_x, base_z)] + leaves_x = list(_iter_record_leaves(np_x)) + leaves_z = list(_iter_record_leaves(np_z)) + if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_z]: + return "records must have the same field names, in the same order, at every level", None + pairs = [(x, z) for (_, _, x), (_, _, z) in zip(leaves_x, leaves_z, strict=True)] + if any(x.shape != z.shape for x, z in pairs): + return "records must have the same shape for each field", None + return None, [(x.base, z.base) for x, z in pairs] + + +def _udt_cast_error(dtype, to_dtype): + """Return why an element of UDT ``dtype`` cannot be stored as UDT ``to_dtype``, or ``None``. + + The layouts must correspond (:func:`_corresponding_leaves`), each pair of + elements must be built-in numbers, or the same type, and none float16, + which the Numba code that casts cannot compute on. + """ + reason, pairs = _corresponding_leaves(dtype.np_type, to_dtype.np_type) + for x, z in pairs or (): + if x != z and not (x.kind in "biufc" and z.kind in "biufc"): + return f"elements of {x} do not cast to {z}" + if reason is None: + for udt in (dtype, to_dtype): + if _has_half_leaf(udt.np_type): + return f"Numba cannot compute on float16, which {udt} holds" + return reason + + +def _leaf_dtypes(np_type): + """Yield the base dtype of each leaf of ``np_type``: a record's fields, an array's element.""" + if np_type.names is not None: + for _py, _c, leaf in _iter_record_leaves(np_type): + yield leaf.base + else: + yield np_type.base + + +def _casts_safely_into(literal_type, dtype): + """Whether each leaf of ``literal_type`` casts safely (``np.can_cast``) into UDT ``dtype``'s. + + A scalar type fills every leaf; otherwise the layouts must correspond + (:func:`_corresponding_leaves`). + """ + if literal_type == dtype: + return True + np_from, np_to = literal_type.np_type, dtype.np_type + if np_from.names is None and np_from.subdtype is None: + return all(np.can_cast(np_from, leaf, "safe") for leaf in _leaf_dtypes(np_to)) + reason, pairs = _corresponding_leaves(np_from, np_to) + return reason is None and all(np.can_cast(x, z, "safe") for x, z in pairs) + + +def _store_cast_op(result_dtype, output_dtype, *, is_value=False): + """Return the UnaryOp that converts a result for an object of ``output_dtype``, or ``None``. + + GraphBLAS casts between built-in types on store (C's casts, with a float + saturating into an integer type), but it cannot cast to or from a UDT; + left to GraphBLAS, that is a bare GrB_DOMAIN_MISMATCH or, for a single + element, bytes read as the wrong type. So when a UDT is involved, the + result is stored by applying the op returned here, which casts each + element or record leaf the way GraphBLAS casts built-in types, between + UDTs whose layouts correspond (:func:`_udt_cast_error`). Any other + pair raises ``DomainMismatch``. ``None`` means GraphBLAS stores the result + as it is. This is the one place that decides what may be stored where. + ``is_value`` says the error is for an object stored or accumulated as it + is, not a computed result. + """ + if result_dtype is output_dtype or not (result_dtype._is_udt or output_dtype._is_udt): + return None + reason = "UDTs do not cast" + if _has_numba and result_dtype._is_udt and output_dtype._is_udt: + if (reason := _udt_cast_error(result_dtype, output_dtype)) is None: + from .operator.unary import _udt_cast_op + + return _udt_cast_op(result_dtype, output_dtype) + elif _has_numba: + reason = "a UDT casts only to another UDT" + if is_value: + raise DomainMismatch( + f"cannot store a value of type {result_dtype} in an object of type " + f"{output_dtype}: {reason}." + ) + raise DomainMismatch( + f"cannot store a result of type {result_dtype} in an object of type " + f"{output_dtype}: {reason}. Use .new() for an object of the result's type, or " + "compute the result in the type you want to store." + ) + + +def _fused_store_expr(expr, output_dtype): + """Return ``expr`` computing straight into ``output_dtype``, or raise ``DomainMismatch``. + + A result of another UDT type is stored by the op that computes it, which + casts each element as it writes it (``udt_codegen.fused_store_op``), so no + temporary of the result's own type is made: data GraphBLAS works on is + often too large to hold twice. That works for an element-wise + computation by a built-in op that lifts to UDTs (``apply``, + ``ewise_mult``, ``ewise_union``). ``ewise_add`` copies an entry present in + only one input without the op, and a semiring or reduction would add in + the narrower type, so those and any other computation raise. + """ + from .operator.udt_codegen import fused_store_op # needs Numba + + op, cfunc_name = expr.op, expr.cfunc_name + if op is not None and op.opclass == "Monoid": + # Element-wise, a Monoid is its BinaryOp. + op, cfunc_name = op.binaryop, cfunc_name.replace("_Monoid", "_BinaryOp") + fused = None + if any(k in cfunc_name for k in ("_apply", "eWiseMult", "eWiseUnion")): + fused = fused_store_op(op, output_dtype) + if fused is None: + method = "extraction" if expr.method_name == "__getitem__" else expr.method_name + what = method if op is None else f"{method} with {op.name}" + raise DomainMismatch( + f"cannot store a result of type {expr.dtype} in an object of type " + f"{output_dtype}: GraphBLAS cannot cast a UDT, and only a built-in op that lifts " + "to UDTs, in apply, ewise_mult or ewise_union, can write another type as it " + f"computes ({what} cannot), so this would need a temporary copy of the whole " + "result. Make it explicitly by storing the computed result, as in " + "`obj << expr.new()`, which casts it in one more pass." + ) + rv = copy.copy(expr) + rv.op, rv.cfunc_name, rv.dtype = fused, cfunc_name, output_dtype + return rv + + +def _scalar_accum_operand(scalar, value): + """Return ``value`` to accumulate into ``scalar``, converted to its type if that is a UDT. + + ``__call__`` types the accumulator on the object's dtype, as GraphBLAS + needs, so for a built-in Scalar GraphBLAS converts the value to that type + before it accumulates: an INT8 Scalar's ``s += 0.5`` adds 0, as a Vector's + ``v(accum=binary.plus) << 0.5`` does. GraphBLAS cannot convert a value to a + UDT, so for a UDT Scalar it is converted here, as storing it would convert + it: a literal as ``Scalar.from_value`` converts it, and a Scalar or an + element by :func:`_cast_for_store`, which raises for a type that does not + cast. ``value`` is a literal, a Scalar, or an element to extract. + """ + if not scalar.dtype._is_udt: + return value + from .scalar import Scalar + + if output_type(value) is Scalar: + return _cast_for_store(value, scalar.dtype) + try: + return Scalar.from_value(value, scalar.dtype, is_cscalar=False, name="") + except TypeError: # not a scalar; ewise_add reports it + return value + + +def _cast_for_store(value, dtype): + """Return Scalar, Vector or Matrix ``value``, or its expression, to assign as ``dtype``. + + GraphBLAS's assign and setElement take a UDT value as raw bytes, so one of + another type would be read as the wrong type; :func:`_store_cast_op` + casts it or raises. Assign cannot cast as it writes, so a value of + another UDT type needs a converted copy: one element for a Scalar, which + is made here, but the whole value for a Vector or Matrix, which raises. + """ + from .base import BaseType + from .scalar import Scalar + + # An expression (``v[:] << w.apply(op)``) is a computed result, not a value. + is_value = isinstance(value, BaseType) + if (cast_op := _store_cast_op(value.dtype, dtype, is_value=is_value)) is None: + return value + + if output_type(value) is not Scalar: + raise DomainMismatch( + f"cannot assign a value of type {value.dtype} into an object of type {dtype}: " + "assign cannot cast a UDT as it writes, so this would need a converted copy of " + "the whole value. Make it explicitly with .dup(dtype=...) first." + ) + if not isinstance(value, BaseType): + value = value.new(name="") + return value.apply(cast_op).new(name=value.name) + + +# ewise_add, which copies an entry present in only one input without the op. + + +def _is_lifted_op(op): + """Whether ``op`` is a built-in BinaryOp that lifts to UDTs, or a Monoid of one.""" + from .operator.binary import BinaryOp + from .operator.monoid import Monoid + + if isinstance(op, Monoid): + op = op.binaryop + return type(op) is BinaryOp and not op._anonymous and op.name in BUILTIN_UDT_BINARY_OPS + + +def _ewise_add_needs_cast(op, left_dtype, right_dtype): + """Whether ``ewise_add`` with ``op`` would have GraphBLAS cast to or from a UDT. + + eWiseAdd copies an entry present in only one input into the result, cast + to the op's output type, and GraphBLAS checks that cast before it looks at + the data. A UDT casts only to itself, so such a pair is a bare + GrB_DOMAIN_MISMATCH even when every entry has a partner. + """ + out = op.return_type + if not (out._is_udt or left_dtype._is_udt or right_dtype._is_udt): + return False + return left_dtype != out or right_dtype != out + + +def _ewise_add_cast_error(op, left_dtype, right_dtype): + """Return the ``DomainMismatch`` for an ``ewise_add`` that would need a UDT cast.""" + return DomainMismatch( + f"ewise_add cannot use {op.parent!r} on {left_dtype} and {right_dtype}: an entry " + f"present in only one input is copied into the result as {op.return_type}, and " + "GraphBLAS cannot cast a UDT, so this would need converted copies of the operands. " + "Use ewise_union, which takes a default for each side (infix + does, with zeros), or " + "ewise_mult; or convert an operand first with .dup(dtype=...)." + ) + + +def _ewise_add_casts_udt(left, right, op): + """Whether ``left.ewise_add(right, op)`` would need GraphBLAS to cast a UDT.""" + try: + left_dtype, right_dtype = left.dtype, right.dtype + except TypeError: + # An infix operand that has no dtype here (``x & y`` of non-BOOL); + # ewise_add reports it as before. + return False + if not (left_dtype._is_udt or right_dtype._is_udt): + return False + from .operator import get_typed_op + + typed_op = get_typed_op(op, left_dtype, right_dtype, kind="binary") + return _ewise_add_needs_cast(typed_op, left_dtype, right_dtype) + + +def _scalar_ewise_add_converted(left, right, op): + """Return ``left``, ``right`` and ``op`` for an ``ewise_add`` of two Scalars that needs a cast. + + ``op`` is typed on the two, and :func:`_ewise_add_needs_cast` holds: + GraphBLAS would cast an unpaired entry to the result type, which a UDT + cannot be. For a built-in op that lifts to UDTs, both are converted to the + result type first, which changes no value, and the op applies in that one + type: an int8 UDT plus 0.5 is float64, and an empty Scalar still adds as 0. + A Scalar is one element, so the converted copies cost nothing; Vector and + Matrix raise instead, as any other op does here (:func:`_ewise_add_cast_error`). + """ + if not _is_lifted_op(op.parent): + raise _ewise_add_cast_error(op, left.dtype, right.dtype) + if left.dtype != op.return_type: + left = _converted_to(left, op.return_type, name="") + if right.dtype != op.return_type: + right = _converted_to(right, op.return_type, name="") + from .operator import get_typed_op + + return left, right, get_typed_op(op.parent, left.dtype, right.dtype, kind="binary") + + +def _plus_zero(dtype): + """Return a Scalar of ``dtype`` that ``plus`` adds to any value without changing it. + + Its float and complex leaves are ``-0.0``, because ``-0.0 + x`` is ``x`` + for every float, where ``0.0 + -0.0`` is ``0.0``. + """ + from .scalar import Scalar + + np_type = dtype.np_type + if np_type.subdtype is not None: + zero = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) + else: + zero = np.zeros(1, dtype=np_type) + for leaf in _leaf_arrays(zero): + if leaf.dtype.kind in "fc": + leaf[...] = complex(-0.0, -0.0) if leaf.dtype.kind == "c" else -0.0 + return Scalar.from_value(zero if np_type.subdtype is not None else zero[0], dtype, name="") + + +def _converted_to(operand, dtype, **new_kwargs): + """Return ``operand`` converted to UDT ``dtype``, which holds its elements, value for value. + + It adds ``dtype``'s zero on the left with the lifted ``plus``. With the + zero on the left, the result is ``dtype`` itself even when the operand has + the same elements in another layout (a packed and an aligned record), + since a left operand that holds the promoted elements is the result type. + """ + from .. import binary + + rv = operand.apply(binary.plus, left=_plus_zero(dtype)).new(**new_kwargs) + if rv.dtype != dtype: # pragma: no cover (safety) + raise DomainMismatch(f"cannot convert {operand.dtype} to {dtype}") + return rv + + +# isequal, and == without Numba. + + +def _element_shapes_differ(dtype1, dtype2): + """Whether elements of ``dtype1`` and ``dtype2`` differ in shape, for ``isequal``. + + ``==`` on array UDTs broadcasts, as numpy's does, but ``isequal`` asks + whether two values are the same array, as ``np.array_equal`` does: the + shapes must match, apart from leading axes of length 1 (the rule a store + uses), and an element of a built-in dtype has shape ``()``. Records pair + field by field instead (:func:`_check_udt_pair`). + """ + if dtype1 is dtype2 or not (dtype1._is_udt or dtype2._is_udt): + return False + shapes = [] + for np_type in (dtype1.np_type, dtype2.np_type): + if np_type.names is not None: + return False + shapes.append(_strip_leading_ones(np_type.subdtype[1]) if np_type.subdtype else ()) + return shapes[0] != shapes[1] + + +def _literal_shape_differs(dtype, value): + """Whether literal ``value`` differs in shape from an element of array UDT ``dtype``. + + As :func:`_element_shapes_differ`, for ``Scalar.isequal`` with a literal: a + number has shape ``()``, so it equals an ``FP64[1]`` element but not an + ``FP64[3]`` one, as in ``np.array_equal``. ``False`` for a record UDT. + """ + np_type = dtype.np_type + if np_type.subdtype is None or isinstance(value, dict): + return False + try: + shape = np.shape(value) + except ValueError: # ragged; Scalar.from_value reports it + return False + return _strip_leading_ones(shape) != _strip_leading_ones(np_type.subdtype[1]) + + +def _leaves_equal(x, x_type, y, y_type): + """Whether values ``x`` and ``y`` are equal in every leaf, as ``binary.eq`` compares them. + + For ``==`` on a UDT Scalar without Numba. Records pair their fields in + order (the caller checks, by ``_check_udt_pair``, that two records pair as + for ``binary.eq``), and a value that is not a record stands in every + field. Each pair of leaves compares in the type their dtypes promote to, + whatever the values (numpy 1 types a scalar by its value), and broadcasts. + """ + if x_type.names is not None or y_type.names is not None: + x_names = x_type.names or [None] * len(y_type.names) + y_names = y_type.names or [None] * len(x_type.names) + return all( + _leaves_equal( + x if x_name is None else x[x_name], + x_type if x_name is None else x_type.fields[x_name][0], + y if y_name is None else y[y_name], + y_type if y_name is None else y_type.fields[y_name][0], + ) + for x_name, y_name in zip(x_names, y_names, strict=True) + ) + x = np.asarray(x, dtype=x_type.base) + y = np.asarray(y, dtype=y_type.base) + common = np.result_type(x.dtype, y.dtype) + return bool(np.all(x.astype(common) == y.astype(common))) + + +def _udt_eq_without_numba(scalar, other): + """Return ``scalar == other`` as a bool, in numpy, for a UDT when Numba is not installed. + + ``binary.eq`` on a UDT needs Numba. ``scalar`` is nonempty, as is ``other`` + when it is a Scalar, and one of them is a UDT. A literal ``other`` is typed + as ``binary.eq`` types it, and the values compare leaf by leaf as + ``binary.eq`` compares them (:func:`_leaves_equal`). + """ + from .. import binary + from .literals import _literal_operand + from .scalar import Scalar + + if type(other) is not Scalar: + value, dtype = _literal_operand(scalar.dtype, other, binary.eq) + other = Scalar.from_value(value, dtype, is_cscalar=None, name="s_eq_other") + # Two UDTs pair as for binary.eq, or raise its KeyError (records with other + # field names, array shapes that do not broadcast, a record and an array). + info = _get_udt_info(scalar.dtype), _get_udt_info(other.dtype) + _check_udt_pair("eq", scalar.dtype, other.dtype, *info) + _check_leaves_compare("eq", scalar.dtype, other.dtype) + return _leaves_equal(scalar.value, scalar.dtype.np_type, other.value, other.dtype.np_type) + + +# An element's value as bytes and back, for Scalar.value. + + +def _element_bytes(dtype, value): + """Return ``value`` as the bytes of an element of UDT ``dtype``, padding zeroed. + + ``value`` is anything numpy assigns to one element: a tuple or dict for a + record, an array or sequence for an array UDT, or a number for every leaf. + """ + np_type = dtype.np_type + # ``_udt_bytes`` zeros a record's alignment padding, which numpy may leave + # as junk however it writes the record. + if np_type.subdtype is None: + arr = np.zeros(1, dtype=np_type) + if isinstance(value, dict) and np_type.names is not None: + value = _dict_to_record(np_type, value) + else: + arr = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) + value = _nats_in_leaf_units(_view_if_same_layout(value, arr.dtype), arr.dtype) + if _has_narrow_float_leaf(np_type) and np.geterr()["over"] == "warn": + # A float too large for a float32 leaf is infinity, and numpy warns; + # warn at the caller's line instead, as a built-in FP32 does, so the + # default filter shows each place it happens. + try: + with np.errstate(over="raise"): + arr[:] = value + except FloatingPointError: + utils._warn_from_caller( + f"overflow encountered in cast: {value!r} as {dtype}", RuntimeWarning + ) + with np.errstate(over="ignore"): + arr[:] = value + else: + arr[:] = value + return _udt_bytes(np_type, arr) + + +def _element_value(dtype, data): + """Return the element of UDT ``dtype`` that cffi ``data`` holds: a record or an array.""" + np_type = dtype.np_type + rv = np.array(ffi.buffer(data[0 : np_type.itemsize])) + if np_type.subdtype is None: + return rv.view(np_type)[0] + base, shape = np_type.subdtype + return rv.view(base).reshape(shape) + + +def _dict_to_record(np_type, d): + """Converts e.g. ``{"x": 1, "y": 2.3}`` to ``(1, 2.3)``.""" + rv = [] + for name, (dtype, _) in np_type.fields.items(): + val = d[name] + if dtype.names is not None and isinstance(val, dict): + rv.append(_dict_to_record(dtype, val)) + else: + rv.append(val) + return tuple(rv) + + +def _leaf_arrays(arr): + """Yield the array of each leaf of a structured array, in field order.""" + if arr.dtype.names is None: + yield arr + else: + for name in arr.dtype.names: + yield from _leaf_arrays(arr[name]) + + +@lru_cache +def _has_half_leaf(np_type): + """Whether ``np_type`` has a float16 leaf: a record field, an array's element, or itself.""" + if np_type.subdtype is not None: + return _has_half_leaf(np_type.subdtype[0]) + if np_type.names is not None: + return any(_has_half_leaf(np_type.fields[name][0]) for name in np_type.names) + return np_type.kind == "f" and np_type.itemsize == 2 + + +def _refuse_half(op, *dtypes): + """Raise ``KeyError`` if a UDT among ``dtypes`` has a float16 leaf, which ``op`` cannot take. + + Numba has a float16 type but no code for it: compiling an op on one raises + a bare ``NotImplementedError``. So such an op is refused up front, as an op + that does not take a type is. + """ + for dtype in dtypes: + if dtype is not None and dtype._is_udt and _has_half_leaf(dtype.np_type): + names = ", ".join(str(d) for d in dtypes if d is not None) + raise KeyError( + f"{op._modname}.{op.name} does not work with ({names}): Numba cannot " + f"compute on float16, which {dtype} holds." + ) + + +@lru_cache +def _has_narrow_float_leaf(np_type): + """Whether ``np_type`` has a float leaf narrower than float64, which a float may overflow.""" + if np_type.subdtype is not None: + return _has_narrow_float_leaf(np_type.subdtype[0]) + if np_type.names is not None: + return any(_has_narrow_float_leaf(np_type.fields[name][0]) for name in np_type.names) + kind, size = np_type.kind, np_type.itemsize + return (kind == "f" and size < 8) or (kind == "c" and size < 16) + + +@lru_cache +def _padding_ranges(np_type): + """The ``(start, stop)`` byte ranges of ``np_type`` that are alignment padding. + + Walks records, nested records, and array dtypes, and returns the gaps + between the field data. + """ + data = [] + _collect_data_ranges(np_type, 0, data) + data.sort() + ranges = [] + pos = 0 + for start, stop in data: + if start > pos: + ranges.append((pos, start)) + pos = max(pos, stop) + if pos < np_type.itemsize: + ranges.append((pos, np_type.itemsize)) + return tuple(ranges) + + +def _collect_data_ranges(np_type, offset, out): + """Append the ``(start, stop)`` byte ranges of ``np_type`` that hold field data.""" + if np_type.subdtype is not None: + base = np_type.subdtype[0] + for start in range(offset, offset + np_type.itemsize, base.itemsize): + _collect_data_ranges(base, start, out) + elif np_type.names is not None: + for field_type, field_offset in (np_type.fields[name][:2] for name in np_type.names): + _collect_data_ranges(field_type, offset + field_offset, out) + else: + out.append((offset, offset + np_type.itemsize)) + + +def _udt_bytes(np_type, arr): + """The bytes of a UDT buffer, with any alignment padding zeroed. + + ``tobytes()`` flattens any-rank numpy buffers (in particular, multi-dim + array UDTs like ``FP64[2, 3]``) to a 1-D byte buffer that cffi can copy + into GrB_Scalar storage. + + numpy promises nothing about the padding bytes it leaves behind, and which + way it goes has changed across releases: numpy 2.5 made ``arr[:] = record`` + a whole-item copy out of an uninitialized temporary, so even a buffer from + ``np.zeros`` comes back with junk between the fields. A caller's own array + can carry junk padding in too. Since the whole record is copied into the + GrB_Scalar and read back out of ``Scalar.value``, anything that compares or + hashes those raw bytes would otherwise see an equal-valued scalar as + unequal. + """ + raw = arr.tobytes() + ranges = _padding_ranges(np_type) + if not ranges: + return raw + buf = bytearray(raw) + for start, stop in ranges: + buf[start:stop] = bytes(stop - start) + return bytes(buf) diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index e9a29b3a9..2d3de41cb 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -4,6 +4,13 @@ from ..dtypes import _INDEX, lookup_dtype from . import ffi, lib +from .dtypes import _nats_in_leaf_units, _view_if_same_layout + +try: + from suitesparse_graphblas.utils import can_unclaim_buffer, empty +except ImportError: # pragma: no cover (import) + # MAINT 2026-10-18: suitesparse-graphblas < 10.5.1.1 has neither + can_unclaim_buffer = empty = None _NP2 = np.__version__.startswith("2.") @@ -33,6 +40,23 @@ def inner(func_wo_doc): return inner +def _module_attr_error(module_name, key, names): + """Build the AttributeError raised by an operator namespace's ``__getattr__``. + + ``names`` should be the module's ``__dir__()`` so lazily-registered operators + are offered as "did you mean" suggestions without forcing them to build. + """ + import difflib + + msg = f"module {module_name!r} has no attribute {key!r}" + candidates = [name for name in names if not name.startswith("_")] + matches = difflib.get_close_matches(key, candidates, n=3) + if matches: + hint = " or ".join(repr(match) for match in matches) + msg = f"{msg}. Did you mean {hint}?" + return AttributeError(msg) + + # Include most common types (even mistakes) _output_types = { int: int, @@ -48,6 +72,30 @@ def inner(func_wo_doc): _output_types.update((k, k) for k in set(np.sctypeDict.values())) +def _warn_from_caller(message, category): + """Warn at the first caller outside python-graphblas, so default filters can show it. + + Python shows a DeprecationWarning by default only when it is attributed to + ``__main__``, and a check that warns may sit at any depth below the public + method. python-graphblas's own tests count as callers. + """ + import sys + import warnings + from pathlib import Path + + package = Path(__file__).parent.parent + tests = package / "tests" + frame = sys._getframe(1) + stacklevel = 2 + while frame is not None: # pragma: no branch (safety) + filename = Path(frame.f_code.co_filename) + if not filename.is_relative_to(package) or filename.is_relative_to(tests): + break + frame = frame.f_back + stacklevel += 1 + warnings.warn(message, category, stacklevel=stacklevel) + + def output_type(val): try: return _output_types[type(val)] @@ -62,11 +110,65 @@ def ints_to_numpy_buffer(array, dtype, *, name="array", copy=False, ownable=Fals and not np.issubdtype(array.dtype, np.bool_) ): raise ValueError(f"{name} must be integers, not {array.dtype.name}") + if ownable: + return _ownable_array(array, dtype, copy=copy, order=order) # https://numpy.org/doc/stable/release/2.0.0-notes.html#new-copy-keyword-meaning-for-array-and-asarray-constructors - array = np.array(array, dtype, copy=copy or _NP2 and None, order=order) - if ownable and (not array.flags.owndata or not array.flags.writeable): - array = array.copy(order) - return array + return np.array(array, dtype, copy=copy or _NP2 and None, order=order) + + +def _ownable_array(array, dtype, *, copy, order): + """Return ``np.array(array, dtype, copy=copy, order=order)`` with data GraphBLAS can own. + + A NumPy array is copied (or converted) at most once, straight into memory + GraphBLAS can own. A copy NumPy made would have to be copied again whenever + GraphBLAS can't own NumPy's memory (``memory_manager="c"``, or a custom + NumPy memory handler). + """ + if not isinstance(array, np.ndarray): + # NumPy reads a list or a scalar into an array of its own + array = np.array(array, dtype, copy=copy or _NP2 and None, order=order) + return _ownable(array, order) + contiguous = array.flags.c_contiguous if order == "C" else array.flags.f_contiguous + if not copy and type(array) is np.ndarray and array.dtype == dtype and contiguous: + return _ownable(array, order) # np.array would return it as it is + return _ownable_copy(array, order, dtype) + + +def _ownable(array, order): + """Return ``array``, or a copy of it if GraphBLAS can't take ownership of its data. + + GraphBLAS frees the data it owns with its own allocator, and ``owndata`` can't + say which allocator made an array: one from a custom NumPy memory handler + (NEP 49) owns its data too, and the process aborts when GraphBLAS frees it. + """ + if can_unclaim_buffer is None: + # Older suitesparse-graphblas can't tell, so assume NumPy's default allocator + ownable = array.flags.owndata and array.flags.writeable + else: + ownable = can_unclaim_buffer(array) + return array if ownable else _ownable_copy(array, order) + + +def _ownable_copy(array, order, dtype=None): + """Copy ``array``, cast to ``dtype`` if given, into memory GraphBLAS can take ownership of.""" + if empty is None or can_unclaim_buffer(np.empty(0, np.uint8)): + # GraphBLAS can own what NumPy allocates, as it can by default + return np.array(array, dtype, copy=True, order=order) + # Not with libc's allocator (``memory_manager="c"``), nor a custom NumPy memory handler + buf = empty(array.shape, array.dtype if dtype is None else dtype, order) + np.copyto(buf, array, casting="unsafe") # the cast np.array(array, dtype) makes + return buf + + +def _distinct_buffers(*arrays): + """Copy each array given more than once, since GraphBLAS frees every buffer it takes.""" + rv = [] + for array in arrays: + if any(array is prev for prev in rv): + order = "F" if array.flags.f_contiguous and not array.flags.c_contiguous else "C" + array = _ownable_copy(array, order) + rv.append(array) + return rv def _get_subdtype(dtype): @@ -93,27 +195,81 @@ def values_to_numpy_buffer( """ if dtype is not None: dtype = lookup_dtype(dtype) + np_type = _get_subdtype(dtype.np_type) + array = _nats_in_leaf_units(_view_if_same_layout(array, np_type), np_type) + if ownable: + return _ownable_array(array, np_type, copy=copy, order=order), dtype # https://numpy.org/doc/stable/release/2.0.0-notes.html#new-copy-keyword-meaning-for-array-and-asarray-constructors - array = np.array( - array, _get_subdtype(dtype.np_type), copy=copy or _NP2 and None, order=order - ) + array = np.array(array, np_type, copy=copy or _NP2 and None, order=order) else: is_input_np = isinstance(array, np.ndarray) - # https://numpy.org/doc/stable/release/2.0.0-notes.html#new-copy-keyword-meaning-for-array-and-asarray-constructors - array = np.array(array, copy=copy or _NP2 and None, order=order) + if ownable and is_input_np and not array.dtype.hasobject: + # View nested subarrays flattened (as below) before copying, not after: + # a view of the copy would not own its data, so it would be copied again + array = _view_if_same_layout(array, lookup_dtype(array.dtype).np_type) + array = _ownable_array(array, array.dtype, copy=copy, order=order) + else: + # https://numpy.org/doc/stable/release/2.0.0-notes.html#new-copy-keyword-meaning-for-array-and-asarray-constructors + array = np.array(array, copy=copy or _NP2 and None, order=order) if array.dtype.hasobject: raise ValueError("object dtype for values is not allowed") if not is_input_np and array.dtype == np.int32: # pragma: no cover # fix for win64 numpy handling of ints array = array.astype(np.int64) dtype = lookup_dtype(array.dtype) + # Registration flattens nested subarrays, so values declared with them + # carry a dtype the DataType we just looked up no longer matches. The + # bytes are identical, so view them and return a pair that agrees. + array = _view_if_same_layout(array, dtype.np_type) if subarray_after is not None and array.ndim > subarray_after: dtype = lookup_dtype(np.dtype((dtype.np_type, array.shape[subarray_after:]))) - if ownable and (not array.flags.owndata or not array.flags.writeable): - array = array.copy(order) + if ownable and not is_input_np: + array = _ownable(array, order) # an array NumPy made from a list, say return array, dtype +def check_element_shape(values, dtype, ndim): + """Raise unless ``values`` is an ``ndim``-d array of whole ``dtype`` elements. + + ``values_to_numpy_buffer`` converts to the scalar type of a sub-array dtype, so each + element's shape must trail the ``ndim`` leading dimensions. GraphBLAS reads whole + elements, so a wrong shape reads the wrong numbers or reads past the end of the array. + """ + if values.shape[ndim:] != dtype.np_type.shape: + raise ValueError( + f"values must have shape {values.shape[:ndim] + dtype.np_type.shape} " + f"for dtype {dtype}, not {values.shape}" + ) + + +def check_values_size(values, dtype, count, is_iso, *, indptr=None, indices=None): + """Raise unless ``values`` holds ``count`` elements of ``dtype``, or one if ``is_iso``. + + With ``indptr``, the count is ``indptr[count]``, as GraphBLAS reads it, and must not + exceed the number of ``indices``. GraphBLAS imports before 10.5.2 compare the byte + size of the values with the number of entries, not with the bytes those entries need, + so a short array gets through and is read past its end. 10.5.2 raises InvalidValue + with no message instead, so this check stays. + """ + if indptr is not None: + # GraphBLAS rejects an indptr this short itself + count = int(indptr[count]) if count < indptr.size else 0 + if indices is not None and indices.size < count: + raise ValueError( + f"index pointers give {count} entries, more than the number of indices, " + f"{indices.size}" + ) + if is_iso: + count = 1 + if (have := values.nbytes // dtype.np_type.itemsize) < count: + if values.ndim == 0 and have == 0: + raise ValueError( + f"a scalar does not fill one element of dtype {dtype}, " + f"of shape {dtype.np_type.shape}" + ) + raise ValueError(f"values must have {count} or more elements of dtype {dtype}, not {have}") + + def normalize_values(self, values, dtype, shape=None, is_iso=False): """Expand and/or update dtype of values array.""" if dtype is not None: @@ -123,10 +279,8 @@ def normalize_values(self, values, dtype, shape=None, is_iso=False): else: dtype = self.dtype if is_iso: - if dtype.np_type.subdtype is None: - values = np.broadcast_to(values, shape) - else: - values = np.broadcast_to(values, shape + values.shape) + # The iso value of an array UDT is exported as shape (1, *subarray shape) + values = np.broadcast_to(values, shape + values.shape[1:]) return values @@ -144,7 +298,7 @@ def get_shape(nrows, ncols, dtype=None, **arrays): break else: raise ValueError( - "Either nrows and ncols must be provided, or one of the following arrays" + "Either nrows and ncols must be provided, or one of the following arrays " f'must be 2d (from which to get nrows and ncols): {", ".join(arrays)}' ) if nrows is None: @@ -361,6 +515,14 @@ def name(self): name = f"{'M' if c == 'M' else c.lower()}_temp" return f"&{name}" + @property + def _expr_name(self): + return getattr(self.val, "_expr_name", self.name) + + @property + def _expr_name_html(self): + return getattr(self.val, "_expr_name_html", self.name) + class _MatrixArray: __slots__ = "_carg", "_exc_arg", "name" diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 8bac4198e..82fb5270f 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -3,13 +3,15 @@ import numpy as np from .. import backend, binary, monoid, select, semiring, unary -from ..dtypes import _INDEX, FP64, INT64, lookup_dtype, unify +from ..dtypes import _INDEX, FP64, INT64, UINT64, lookup_dtype, unify from ..exceptions import DimensionMismatch, NoValue, check_status from . import _supports_udfs, automethods, ffi, lib, utils -from .base import BaseExpression, BaseType, _check_mask, call +from .base import BaseExpression, BaseType, call from .descriptor import lookup as descriptor_lookup +from .dtypes import _raise_dtype_or_arraylike from .expr import _ALL_INDICES, AmbiguousAssignOrExtract, IndexerResolver, InfixExprBase, Updater -from .mask import Mask, StructuralMask, ValueMask +from .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype +from .mask import Mask, StructuralMask, ValueMask, _check_mask from .operator import ( UNKNOWN_OPCLASS, _get_typed_op_from_exprs, @@ -18,6 +20,7 @@ get_typed_op, op_from_string, ) +from .operator.select import _value_select_in_own_type from .scalar import ( _COMPLETE, _MATERIALIZE, @@ -26,10 +29,18 @@ ScalarIndexExpr, _as_scalar, _scalar_index, + _union_default, +) +from .udt import ( + _cast_for_store, + _element_shapes_differ, + _ewise_add_cast_error, + _ewise_add_needs_cast, ) from .utils import ( _CArray, _Pointer, + check_element_shape, class_property, ints_to_numpy_buffer, normalize_values, @@ -99,8 +110,13 @@ def _select_mask(updater, obj, mask): updater << obj.dup(mask=mask) +def _abs_type(dtype): + return dtype if dtype.np_type.kind == "c" else FP64 + + def _isclose_recipe(self, other, rel_tol, abs_tol, **opts): - # x == y or abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol) + # x == y or abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol), except that + # an infinity is close only to itself isequal = self.ewise_mult(other, binary.eq).new(bool, name="isclose", **opts) if isequal._nvals != self._nvals: return False @@ -113,23 +129,38 @@ def _isclose_recipe(self, other, rel_tol, abs_tol, **opts): # So we can use structural mask below isequal(**opts) << select.value(isequal == True) # noqa: E712 - # abs(x) - x = self.apply(unary.abs).new(FP64, mask=~isequal.S, **opts) - # abs(y) - y = other.apply(unary.abs).new(FP64, mask=~isequal.S, **opts) + # abs(x) and abs(y), in the type x - y is in below, or in FP64 for a real type: + # abs wraps at an integer type's minimum + abs_type = _abs_type(unify(self.dtype, other.dtype)) + x = self.apply(unary.abs[abs_type]).new(FP64, mask=~isequal.S, **opts) + y = other.apply(unary.abs[abs_type]).new(FP64, mask=~isequal.S, **opts) # max(abs(x), abs(y)) x(**opts) << x.ewise_mult(y, binary.max) max_x_y = x + # As in math.isclose, an infinity is close only to itself, and those are equal + if max_x_y.ndim == 1: + biggest = max_x_y.reduce(monoid.max, allow_empty=False).new(**opts).value + else: + biggest = max_x_y.reduce_scalar(monoid.max, allow_empty=False).new(**opts).value + if biggest == np.inf: + return False # rel_tol * max(abs(x), abs(y)) max_x_y(**opts) << max_x_y.apply(binary.times, rel_tol) # max(rel_tol * max(abs(x), abs(y)), abs_tol) max_x_y(**opts) << max_x_y.apply(binary.max, abs_tol) - # x - y - y(~isequal.S, replace=True, **opts) << self.ewise_mult(other, binary.minus) - abs_x_y = y # abs(x - y) - abs_x_y(**opts) << abs_x_y.apply(unary.abs) + if unify(self.dtype, other.dtype).np_type.kind in "iu": + # x - y wraps in an integer type, and the difference may not fit it. It fits + # UINT64, where max(x, y) and min(x, y) convert modulo 2**64. + hi = self.ewise_mult(other, binary.max).new(UINT64, mask=~isequal.S, **opts) + lo = self.ewise_mult(other, binary.min).new(UINT64, mask=~isequal.S, **opts) + y(~isequal.S, replace=True, **opts) << hi.ewise_mult(lo, binary.minus) + else: + # abs before storing in y, whose FP64 would keep only a complex number's real part + diff = self.ewise_mult(other, binary.minus).new(mask=~isequal.S, **opts) + y(~isequal.S, replace=True, **opts) << diff.apply(unary.abs) + abs_x_y = y # abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol) isequal(**opts) << abs_x_y.ewise_mult(max_x_y, binary.le) @@ -152,21 +183,22 @@ class Vector(BaseType): """ - __slots__ = "_size", "_parent", "ss" + __slots__ = "_size", "_parent" ndim = 1 _name_counter = itertools.count() def __new__(cls, dtype=FP64, size=0, *, name=None): self = object.__new__(cls) - self.dtype = lookup_dtype(dtype) + try: + self.dtype = lookup_dtype(dtype) + except (ValueError, TypeError) as exc: + _raise_dtype_or_arraylike("Vector", dtype, exc) size = _as_scalar(size, _INDEX, is_cscalar=True) self.name = f"v_{next(Vector._name_counter)}" if name is None else name self.gb_obj = ffi_new("GrB_Vector*") call("GrB_Vector_new", [_Pointer(self), self.dtype, size]) self._size = size.value self._parent = None - if backend == "suitesparse": - self.ss = ss(self) return self @classmethod @@ -177,8 +209,6 @@ def _from_obj(cls, gb_obj, dtype, size, *, parent=None, name=None): self.dtype = dtype self._size = size self._parent = parent - if backend == "suitesparse": - self.ss = ss(self) return self def __del__(self): @@ -335,7 +365,7 @@ def __sizeof__(self): size = ffi_new("size_t*") check_status(lib.GxB_Vector_memoryUsage(size, self.gb_obj[0]), self) return size[0] + object.__sizeof__(self) - raise TypeError("Unable to get size of Vector with backend: {backend}") + raise TypeError(f"Unable to get size of Vector with backend: {backend}") def isequal(self, other, *, check_dtype=False, **opts): """Check for exact equality (same size, same structure). @@ -344,13 +374,18 @@ def isequal(self, other, *, check_dtype=False, **opts): ---------- other : Vector The vector to compare against - check_dtypes : bool, default=False + check_dtype : bool, default=False If True, also checks that dtypes match Returns ------- bool + Notes + ----- + Elements of array UDTs must have the same shape, apart from leading axes + of length 1, as in ``np.array_equal``; ``==`` broadcasts instead. + See Also -------- :meth:`isclose` : For equality check of floating point dtypes @@ -359,6 +394,8 @@ def isequal(self, other, *, check_dtype=False, **opts): other = self._expect_type(other, Vector, within="isequal", argname="other") if check_dtype and self.dtype != other.dtype: return False + if _element_shapes_differ(self.dtype, other.dtype): + return False if self._size != other._size: return False if self._nvals != other._nvals: @@ -381,7 +418,8 @@ def isequal(self, other, *, check_dtype=False, **opts): def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False, **opts): """Check for approximate equality (including same size and same structure). - Equivalent to: ``abs(a-b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)``. + Equivalent to: ``abs(a-b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)``, + except that, as in ``math.isclose``, an infinity is close only to itself. Parameters ---------- @@ -469,7 +507,7 @@ def to_coo(self, dtype=None, *, indices=True, values=True, sort=True): ---------- dtype : Requested dtype for the output values array. - indices :bool, default=True + indices : bool, default=True Whether to return indices; will return ``None`` for indices if ``False`` values : bool, default=True Whether to return values; will return ``None`` for values if ``False`` @@ -528,12 +566,11 @@ def build(self, indices, values, *, dup_op=None, clear=False, size=None): """ # TODO: accept `dtype` keyword to match the dtype of `values`? indices = ints_to_numpy_buffer(indices, np.uint64, name="indices") - values, dtype = values_to_numpy_buffer(values, self.dtype) + values, _dtype = values_to_numpy_buffer(values, self.dtype) n = values.shape[0] if indices.size != n: - raise ValueError( - f"`indices` and `values` lengths must match: {indices.size} != {values.size}" - ) + raise ValueError(f"`indices` and `values` lengths must match: {indices.size} != {n}") + check_element_shape(values, self.dtype, 1) if clear: self.clear() if size is not None: @@ -593,7 +630,11 @@ def dup(self, dtype=None, *, clear=False, mask=None, name=None, **opts): dtype = self.dtype rv = Vector(dtype, size=self._size, name=name) if not clear: - rv(mask=mask, **opts)[...] = self + if rv.dtype._is_udt or self.dtype._is_udt: + # ``<<``, not assign: a UDT of another type casts in one pass. + rv(mask=mask, **opts) << self + else: + rv(mask=mask, **opts)[...] = self else: if opts: # Ignore opts for now @@ -609,8 +650,6 @@ def diag(self, k=0, *, name=None): ---------- k : int Off-diagonal offset in the Matrix. - dtype : - Data type of the new Matrix. Normal typecasting rules apply. name : str, optional Name to give the new Matrix. @@ -682,6 +721,11 @@ def get(self, index, default=None): def from_coo(cls, indices, values=1.0, dtype=None, *, size=None, dup_op=None, name=None): """Create a new Vector from indices and values. + .. warning:: + When ``size`` is omitted, it is inferred from the largest index, so + trailing empty positions are dropped. Pass ``size`` explicitly to + pin the length. + Parameters ---------- indices : list or np.ndarray @@ -730,7 +774,8 @@ def from_coo(cls, indices, values=1.0, dtype=None, *, size=None, dup_op=None, na if backend == "suitesparse": w.ss.build_scalar(indices, values.tolist()) else: - w.build(indices, np.broadcast_to(values, indices.size), dup_op=binary.any) + values = np.broadcast_to(values, (indices.size, *dtype.np_type.shape)) + w.build(indices, values, dup_op=binary.any) else: # This needs to be the original data to get proper error messages w.build(indices, values, dup_op=dup_op) @@ -800,10 +845,8 @@ def from_scalar(cls, value, size, dtype=None, *, name=None, **opts): ---------- value : scalar Scalar value used to fill the Vector. - nrows : int - Number of rows. - ncols : int - Number of columns. + size : int + Size of the Vector. dtype : DataType, optional Data type of the Vector. If not provided, the scalar value will be inspected to choose an appropriate dtype. @@ -885,6 +928,7 @@ def from_dense(cls, values, missing_value=None, *, dtype=None, name=None, **opts raise ValueError("A >1d array is required to create a dense Vector with subdtype") if values.ndim > 1 and dtype.np_type.subdtype is None: raise ValueError(f"values array must be 1d to create dense Vector with dtype {dtype}") + check_element_shape(values, dtype, 1) if backend == "suitesparse": rv = cls.ss.import_full(values, dtype=dtype, name=name) else: @@ -931,14 +975,15 @@ def to_dense(self, fill_value=None, dtype=None, **opts): ) if backend == "suitesparse": info = self.ss.export("full") - return normalize_values(self, info["values"], dtype, self._size, info["is_iso"]) + return normalize_values(self, info["values"], dtype, (self._size,), info["is_iso"]) return self.to_coo(dtype, indices=False)[1] if dtype is None and not self.dtype._is_udt: # dtype of fill_value can upcast the dtype if type(fill_value) is not Scalar: try: - fill_value = Scalar.from_value(fill_value, is_cscalar=None, name="") + dtype = _weak_builtin_literal_dtype(self.dtype, fill_value) + fill_value = Scalar.from_value(fill_value, dtype, is_cscalar=None, name="") except TypeError: fill_value = self._expect_type( fill_value, @@ -1030,15 +1075,23 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): # Per the spec, op may be a semiring, but this is weird, so don't. self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") - if other.ndim == 2: - # Broadcast columnwise from the left - if other._nrows != self._size: - # Check this before we compute a possibly large matrix below + if other.ndim == 2 and other._nrows != self._size: + # Broadcast columnwise from the left. Check this before we compute a + # possibly large matrix below. + raise DimensionMismatch( + "Dimensions not compatible for broadcasting Vector from the left " + f"to columns of Matrix in {method_name}. Matrix.nrows (={other._nrows}) " + f"must equal Vector.size (={self._size})." + ) + if _ewise_add_needs_cast(op, self.dtype, other.dtype): + # GraphBLAS would report the UDT cast before the sizes. + if other.ndim == 1 and self._size != other._size: raise DimensionMismatch( - "Dimensions not compatible for broadcasting Vector from the left " - f"to columns of Matrix in {method_name}. Matrix.nrows (={other._nrows}) " - f"must equal Vector.size (={self._size})." + f"Vector sizes must match in {method_name}: {self._size} and {other._size}" ) + raise _ewise_add_cast_error(op, self.dtype, other.dtype) + + if other.ndim == 2: return MatrixExpression( method_name, None, @@ -1208,42 +1261,8 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): ) temp_op = get_typed_op(op, self.dtype, other.dtype, kind="binary") - left_dtype = temp_op.type - dtype = left_dtype if left_dtype._is_udt else None - if type(left_default) is not Scalar: - try: - left = Scalar.from_value( - left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - left = self._expect_type( - left_default, - Scalar, - within=method_name, - keyword_name="left_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar - right_dtype = temp_op.type2 - dtype = right_dtype if right_dtype._is_udt else None - if type(right_default) is not Scalar: - try: - right = Scalar.from_value( - right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - right = self._expect_type( - right_default, - Scalar, - within=method_name, - keyword_name="right_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar + left = _union_default(self, left_default, temp_op.type, op, "left_default") + right = _union_default(self, right_default, temp_op.type2, op, "right_default") if is_infix: op1 = _get_typed_op_from_exprs(op, self, right, kind="binary") @@ -1316,7 +1335,7 @@ def vxm(self, other, op=semiring.plus_times): Parameters ---------- - other: Matrix + other : Matrix The matrix on the right side in the computation op : :class:`~graphblas.core.operator.Semiring` Semiring used in the computation @@ -1417,13 +1436,19 @@ def apply(self, op, right=None, *, left=None): """ method_name = "apply" + right_literal = None extra_message = ( - "apply only accepts UnaryOp with no scalars or BinaryOp with `left` or `right` scalar" + "apply only accepts UnaryOp with no scalars or BinaryOp with `left` or `right` scalar " "or IndexUnaryOp with `right` thunk." ) if isinstance(op, str): op = op_from_string(op) op, opclass = find_opclass(op) + if opclass == "Monoid" and self.dtype._is_udt: + # Applying a Monoid applies its BinaryOp. Typing it as that op lets a + # literal promote as it would with the BinaryOp, where a Monoid takes + # one type and would need the literal converted into the UDT. + op, opclass = op.binaryop, "BinaryOp" if opclass in {"IndexUnaryOp", "SelectOp"}: # Provide default value for index unary if right is None: @@ -1445,7 +1470,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = None elif right is None: if type(left) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + left, dtype = _literal_operand(self.dtype, left, op, left=True) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: @@ -1481,7 +1506,8 @@ def apply(self, op, right=None, *, left=None): expr_repr = "{1.name}.apply({op}, left={0._expr_name})" elif left is None: if type(right) is not Scalar: - dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None + right_literal = right + right, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -1524,7 +1550,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = "{0.name}.apply({op}, right={1._expr_name})" else: raise TypeError("Cannot provide both `left` and `right` to apply") - return VectorExpression( + expr = VectorExpression( method_name, cfunc_name, args, @@ -1532,6 +1558,8 @@ def apply(self, op, right=None, *, left=None): expr_repr=expr_repr, size=self._size, ) + expr._right_literal = right_literal + return expr def select(self, op, thunk=None): """Create a new Vector by applying ``op`` to each element of the Vector and keeping @@ -1597,8 +1625,10 @@ def select(self, op, thunk=None): if thunk is None: thunk = False # most basic form of 0 when unifying dtypes + given = thunk if type(thunk) is not Scalar: - dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None + # A thunk beside a UDT becomes an element of it, so it must fit, as in apply. + thunk, dtype = _literal_operand(self.dtype, thunk, op) try: thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: @@ -1611,6 +1641,7 @@ def select(self, op, thunk=None): op=op, ) op = get_typed_op(op, self.dtype, thunk.dtype, is_right_scalar=True, kind="select") + op, thunk = _value_select_in_own_type(op, self.dtype, thunk, given) self._expect_op(op, ("SelectOp", "IndexUnaryOp"), within=method_name, argname="op") if thunk._is_cscalar: if thunk.dtype._is_udt: @@ -1847,6 +1878,15 @@ def _extract_element( dtype = self.dtype else: dtype = lookup_dtype(dtype) + if dtype is not self.dtype and (dtype._is_udt or self.dtype._is_udt): + # GraphBLAS copies a UDT element as raw bytes of this vector's type, + # past the end of a smaller one, so extract it in this type and cast + # it as a store does (or raise). + value = self._extract_element(resolved_indexes, None, opts, is_cscalar=is_cscalar) + if result is None: + result = Scalar(dtype, is_cscalar=is_cscalar, name=name) + result << value + return result idx = resolved_indexes.indices[0] if result is None: result = Scalar(dtype, is_cscalar=is_cscalar, name=name) @@ -1878,7 +1918,7 @@ def _prep_for_extract(self, resolved_indexes, mask=None, is_submask=False): def _assign_element(self, resolved_indexes, value): idx = resolved_indexes.indices[0] if type(value) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, value) try: value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except TypeError: @@ -1889,6 +1929,10 @@ def _assign_element(self, resolved_indexes, value): argname="value", extra_message="Literal scalars also accepted.", ) + if value.dtype._is_udt or self.dtype._is_udt: + # After the conversion: an autocomputed element of another UDT type + # arrives here as a Scalar too. + value = _cast_for_store(value, self.dtype) if value._is_cscalar: if value._empty: call("GrB_Vector_removeElement", [self, idx.index]) @@ -1905,6 +1949,8 @@ def _assign_element(self, resolved_indexes, value): def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, opts): method_name = "__setitem__" + if output_type(value) in {Scalar, Vector}: + value = _cast_for_store(value, self.dtype) idx = resolved_indexes.indices[0] size = idx.size cscalar = idx.cscalar @@ -1939,7 +1985,7 @@ def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, o expr_repr = "[[{2._expr_name} elements]] = {0.name}" else: if type(value) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, value) try: value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except (TypeError, ValueError): @@ -2082,10 +2128,7 @@ def from_dict(cls, d, dtype=None, *, size=None, name=None): else: # If we know the dtype, then using `np.fromiter` is much faster dtype = lookup_dtype(dtype) - if dtype.np_type.subdtype is not None and np.__version__[:5] in {"1.21.", "1.22."}: - values, dtype = values_to_numpy_buffer(list(d.values()), dtype) # FLAKY COVERAGE - else: - values = np.fromiter(d.values(), dtype.np_type) + values = np.fromiter(d.values(), dtype.np_type) if size is None and indices.size == 0: size = 0 return cls.from_coo(indices, values, dtype, size=size, name=name) @@ -2109,10 +2152,17 @@ def to_dict(self): if backend == "suitesparse": - Vector.ss = class_property(Vector.ss, ss) + # Built lazily per access, not stored, to avoid the ss/_parent reference + # cycle that would keep every Vector alive until the cyclic gc; see gh-559 + # and the matching note in matrix.py. + def _ss(self): + return ss(self) + + _ss.__name__ = _ss.__qualname__ = "ss" + Vector.ss = class_property(property(_ss), ss) else: Vector.ss = class_property( - Vector.ss, 'ss attribute is only available with "suitesparse" backend', exceptional=True + property(), 'ss attribute is only available with "suitesparse" backend', exceptional=True ) diff --git a/graphblas/exceptions.py b/graphblas/exceptions.py index 05cac988a..2c034c212 100644 --- a/graphblas/exceptions.py +++ b/graphblas/exceptions.py @@ -91,7 +91,32 @@ class JitError(GraphblasException): # Our errors class UdfParseError(GraphblasException): - """SuiteSparse:GraphBLAS unable to parse the user-defined function.""" + """Raised when a UDF can't be compiled for the requested operand types. + + Wraps Numba compilation errors (``TypingError``, ``LoweringError``, + ``UnsupportedError``) into a single graphblas-level exception with the + actionable diagnostic line surfaced from Numba's traceback. Common + causes: an operator that doesn't exist for the field type + (``binary.floordiv`` on a complex field), a UDF that returns a tuple + whose length doesn't match any input UDT, or use of a Python construct + Numba doesn't support in nopython mode. + """ + + +# Warnings +class NoJITWarning(UserWarning): + """Auto-lifted UDT op fell back to the Numba cfunc path. + + Emitted from :func:`graphblas.core.ss.jit_config._maybe_warn_no_jit` + when an op like ``binary.plus[udt]`` would otherwise have JIT-compiled, + but the UDT has no C form, so no kernel source was generated. Fires once + per ``(op, dtype)`` pair per process, so a user who registers several UDTs + gets one warning per pair. + + Inherits from :class:`UserWarning` so existing ``UserWarning`` filters + still match it; users can also filter by category for a tight scope: + ``warnings.filterwarnings("ignore", category=NoJITWarning)``. + """ _error_code_lookup = { @@ -137,11 +162,34 @@ def check_status(response_code, args): arg = arg._exc_arg if type(arg) is _Pointer: arg = arg.val - type_name = type(arg).__name__ + type_name = _error_type_name(type(arg)) carg = arg._carg return check_status_carg(response_code, type_name, carg) +def _error_type_name(cls): + """The type name whose ``GrB__error`` reports on objects of ``cls``. + + A subclass of Matrix or Vector has no error function of its own, so it uses that of + the nearest base class that has one. Only the classes of graphblas.core count, which + is where every such type is: another class, such as a user's mixin, may share a + GraphBLAS type's name. Any other class keeps its own name. + """ + # Every object a failing call can hand check_status is a Matrix, Vector, Scalar, + # Descriptor or Context, or a subclass of one, so the walk finds a match. Only + # BuiltinDescriptor and GlobalContext pass a class without an error function on the + # way, and their calls cannot fail. + for base in cls.__mro__: + if not base.__module__.startswith("graphblas.core."): + continue + try: + _libget(f"GrB_{base.__name__}_error") + except AttributeError: # pragma: no cover (defensive) + continue + return base.__name__ + return cls.__name__ # pragma: no cover (defensive) + + def check_status_carg(response_code, type_name, carg): if response_code == GrB_SUCCESS: return diff --git a/graphblas/graphblas.yaml b/graphblas/graphblas.yaml index 45aa47c01..93a6748ef 100644 --- a/graphblas/graphblas.yaml +++ b/graphblas/graphblas.yaml @@ -1,2 +1,8 @@ autocompute: True mapnumpy: True +# SuiteSparse:GraphBLAS C JIT control, set when graphblas initializes the "suitesparse" +# backend: "on", "load", "run", "pause" or "off". Environment variable: GRAPHBLAS_JIT_C_CONTROL. +# Unset (null) is "run", raised to "on" when UDT work needs a kernel compiled; a value set +# through this key is kept. +# See "The SuiteSparse:GraphBLAS JIT" in the user guide's initialization page. +jit_c_control: null diff --git a/graphblas/indexbinary/__init__.py b/graphblas/indexbinary/__init__.py new file mode 100644 index 000000000..60c73fa96 --- /dev/null +++ b/graphblas/indexbinary/__init__.py @@ -0,0 +1,33 @@ +# All items are dynamically added by classes in operator.py. +# This module acts as a container of IndexBinaryOp instances. +_delayed = {} + + +def __dir__(): + return globals().keys() | _delayed.keys() | {"ss"} + + +def __getattr__(key): + if key in _delayed: + func, kwargs = _delayed.pop(key) + rv = func(**kwargs) + globals()[key] = rv + return rv + if key == "ss": + from .. import backend + + if backend != "suitesparse": + raise AttributeError( + f'module {__name__!r} only has attribute "ss" when backend is "suitesparse"' + ) + from importlib import import_module + + ss = import_module(".ss", __name__) + globals()["ss"] = ss + return ss + raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + + +from ..core import operator # noqa: E402 isort:skip + +del operator diff --git a/graphblas/indexbinary/ss.py b/graphblas/indexbinary/ss.py new file mode 100644 index 000000000..6ffd91c01 --- /dev/null +++ b/graphblas/indexbinary/ss.py @@ -0,0 +1,6 @@ +from ..core import operator +from ..core.ss.indexbinary import register_new # noqa: F401 + +_delayed = {} + +del operator diff --git a/graphblas/indexunary/__init__.py b/graphblas/indexunary/__init__.py index a3cb06608..5fa548c5a 100644 --- a/graphblas/indexunary/__init__.py +++ b/graphblas/indexunary/__init__.py @@ -25,7 +25,9 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/io/_awkward.py b/graphblas/io/_awkward.py index b30984251..7b9f18906 100644 --- a/graphblas/io/_awkward.py +++ b/graphblas/io/_awkward.py @@ -14,7 +14,7 @@ def to_awkward(A, format=None): ---------- A : Matrix or Vector GraphBLAS object to be converted - format : str {'csr', 'csc', 'hypercsr', 'hypercsc', 'vec} + format : str {'csr', 'csc', 'hypercsr', 'hypercsc', 'vec'} Default format is csr for Matrix; vec for Vector The Awkward Array will have top-level attributes based on format: @@ -28,19 +28,10 @@ def to_awkward(A, format=None): awkward.Array """ - try: - # awkward version 1 - # MAINT: we can probably drop awkward v1 at the end of 2024 or 2025 - import awkward._v2 as ak - from awkward._v2.forms.listoffsetform import ListOffsetForm - from awkward._v2.forms.numpyform import NumpyForm - from awkward._v2.forms.recordform import RecordForm - except ImportError: - # awkward version 2 - import awkward as ak - from awkward.forms.listoffsetform import ListOffsetForm - from awkward.forms.numpyform import NumpyForm - from awkward.forms.recordform import RecordForm + import awkward as ak + from awkward.forms.listoffsetform import ListOffsetForm + from awkward.forms.numpyform import NumpyForm + from awkward.forms.recordform import RecordForm out_type = output_type(A) if format is None: diff --git a/graphblas/io/_matrixmarket.py b/graphblas/io/_matrixmarket.py index 8cf8738a3..4d7fcca22 100644 --- a/graphblas/io/_matrixmarket.py +++ b/graphblas/io/_matrixmarket.py @@ -1,3 +1,5 @@ +import warnings + from .. import backend from ..core.matrix import Matrix from ._scipy import to_scipy_sparse @@ -19,7 +21,7 @@ def mmread(source, engine="auto", *, dup_op=None, name=None, **kwargs): Parameters ---------- source : str or file - Filename (.mtx or .mtz.gz) or file-like object + Filename (.mtx or .mtx.gz) or file-like object engine : {"auto", "scipy", "fmm", "fast_matrix_market"}, default "auto" How to read the matrix market file. "scipy" uses ``scipy.io.mmread``, "fmm" and "fast_matrix_market" uses ``fast_matrix_market.mmread``, @@ -40,19 +42,50 @@ def mmread(source, engine="auto", *, dup_op=None, name=None, **kwargs): except ImportError: # pragma: no cover (import) raise ImportError("scipy is required to read Matrix Market files") from None engine = engine.lower() + using_scipy = True + if engine in {"fmm", "fast_matrix_market"}: + warnings.warn( + "fast_matrix_market is no longer maintained and will be removed in a future version. " + 'Use engine="scipy" instead.', + DeprecationWarning, + stacklevel=2, + ) if engine in {"auto", "fmm", "fast_matrix_market"}: try: from fast_matrix_market import mmread # noqa: F811 + + if engine == "auto": + warnings.warn( + "fast_matrix_market is installed but is no longer maintained and will be " + "removed in a future version. Uninstall it or use engine='scipy' to " + "silence this warning.", + DeprecationWarning, + stacklevel=2, + ) except ImportError: # pragma: no cover (import) if engine != "auto": raise ImportError( "fast_matrix_market is required to read Matrix Market files " f'using the "{engine}" engine' ) from None + else: + using_scipy = False elif engine != "scipy": raise ValueError( f'Bad engine value: {engine!r}. Must be "auto", "scipy", "fmm", or "fast_matrix_market"' ) + if using_scipy and "spmatrix" not in kwargs: + # scipy's `mmread` still returns a sparse matrix by default, but 1.18 + # deprecated that default and 1.20 flips it to a sparse array. We only + # read `.shape`, `.row`, `.col`, and `.data` below, which both types + # provide, so ask for the future default and skip the warning. + # MAINT: 2026-07-31 once we require scipy >= 1.20 this whole block goes + # away, since False becomes the default. + # Older scipy has no such argument; check rather than pin a version. + import inspect + + if "spmatrix" in inspect.signature(mmread).parameters: + kwargs["spmatrix"] = False array = mmread(source, **kwargs) if getattr(array, "format", None) == "coo": nrows, ncols = array.shape @@ -104,9 +137,25 @@ def mmwrite( except ImportError: # pragma: no cover (import) raise ImportError("scipy is required to write Matrix Market files") from None engine = engine.lower() + if engine in {"fmm", "fast_matrix_market"}: + warnings.warn( + "fast_matrix_market is no longer maintained and will be removed in a future version. " + 'Use engine="scipy" instead.', + DeprecationWarning, + stacklevel=2, + ) if engine in {"auto", "fmm", "fast_matrix_market"}: try: from fast_matrix_market import __version__, mmwrite # noqa: F811 + + if engine == "auto": + warnings.warn( + "fast_matrix_market is installed but is no longer maintained and will be " + "removed in a future version. Uninstall it or use engine='scipy' to " + "silence this warning.", + DeprecationWarning, + stacklevel=2, + ) except ImportError: # pragma: no cover (import) if engine != "auto": raise ImportError( diff --git a/graphblas/io/_networkx.py b/graphblas/io/_networkx.py index 8cf84e576..a804e187c 100644 --- a/graphblas/io/_networkx.py +++ b/graphblas/io/_networkx.py @@ -1,3 +1,5 @@ +import numpy as np + from ..dtypes import lookup_dtype from ._scipy import from_scipy_sparse @@ -31,10 +33,18 @@ def from_networkx(G, nodelist=None, dtype=None, weight="weight", name=None): return from_scipy_sparse(A, name=name) +# A DiGraph spends about 350 bytes on each node, so a graph with this many nodes would +# need hundreds of terabytes. Refusing up front beats a long climb to out-of-memory, +# and it never refuses a graph that could have been built. +_MAX_ISOLATES = 2**40 + + # TODO: add parameters to allow different networkx classes and attribute names -def to_networkx(m, edge_attribute="weight"): +def to_networkx(m, edge_attribute="weight", create_isolates=True): """Create a networkx DiGraph from a square adjacency Matrix. + The nodes are the indices of the Matrix, in order. + Parameters ---------- m : Matrix @@ -42,6 +52,11 @@ def to_networkx(m, edge_attribute="weight"): edge_attribute : str, optional Name of edge attribute from values of Matrix. If None, values will be skipped. Default is "weight". + create_isolates : bool, default=True + Include the nodes with no edges (index i, where row i and column i are both empty). + Every index then becomes a node, which costs memory for each one, so a Matrix + with huge dimensions (often hypersparse) needs False. More than 2**40 rows raise + ValueError unless this is False. Returns ------- @@ -50,10 +65,25 @@ def to_networkx(m, edge_attribute="weight"): """ import networkx as nx + if m.nrows != m.ncols: + # Row i and column i are both node i, so only a square Matrix is an adjacency matrix + raise ValueError(f"to_networkx needs a square Matrix, not shape {m.shape}") + if create_isolates and m.nrows > _MAX_ISOLATES: + raise ValueError( + f"Too many nodes to create: the Matrix has {m.nrows} rows. " + "Use create_isolates=False to create only the nodes that have edges." + ) rows, cols, vals = m.to_coo() + G = nx.DiGraph() + # Added first and in index order, so the nodes iterate in that order and not in + # the order the edges name them: from_networkx then gives back the same matrix + # (without its empty rows and columns if create_isolates is False). + if create_isolates: + G.add_nodes_from(range(m.nrows)) + else: + G.add_nodes_from(np.union1d(rows, cols).tolist()) rows = rows.tolist() cols = cols.tolist() - G = nx.DiGraph() if edge_attribute is None: G.add_edges_from(zip(rows, cols, strict=True)) else: diff --git a/graphblas/io/_scipy.py b/graphblas/io/_scipy.py index 228432eed..2003d3281 100644 --- a/graphblas/io/_scipy.py +++ b/graphblas/io/_scipy.py @@ -15,7 +15,8 @@ def from_scipy_sparse(A, *, dup_op=None, name=None): A : scipy.sparse Scipy sparse array or matrix dup_op : BinaryOp, optional - Aggregation function for formats that allow duplicate entries (e.g. coo) + Aggregation function for duplicate entries, which coo allows (as do csr and csc + not in canonical format). Duplicates without it raise ValueError. name : str, optional Name of resulting Matrix @@ -28,6 +29,11 @@ def from_scipy_sparse(A, *, dup_op=None, name=None): dtype = lookup_dtype(A.dtype) if A.nnz == 0: return Matrix(dtype, nrows=nrows, ncols=ncols, name=name) + if A.format in {"csr", "csc"} and not A.has_canonical_format: + # An index may repeat within a row (column), and an import takes the indices + # unchecked: a repeat would make an invalid Matrix. Build from coo instead, where + # dup_op combines repeats. + A = A.tocoo() if backend == "suitesparse" and A.format in {"csr", "csc"}: data = A.data is_iso = (data[[0]] == data).all() diff --git a/graphblas/io/_sparse.py b/graphblas/io/_sparse.py index c0d4beabb..e780ee96a 100644 --- a/graphblas/io/_sparse.py +++ b/graphblas/io/_sparse.py @@ -15,7 +15,9 @@ def from_pydata_sparse(s, *, dup_op=None, name=None): s : sparse PyData sparse array or matrix (see https://sparse.pydata.org) dup_op : BinaryOp, optional - Aggregation function for formats that allow duplicate entries (e.g. coo) + Combines the values of an index repeated within a row (or column) of a + two-dimensional gcxs array. By default they are summed, as pydata's todense does. + (A coo array sums repeated coordinates when it is built.) name : str, optional Name of resulting Matrix @@ -45,7 +47,14 @@ def from_pydata_sparse(s, *, dup_op=None, name=None): ) # handle two-dimensional arrays if isinstance(s, sparse.GCXS): - return from_scipy_sparse(s.to_scipy_sparse(), dup_op=dup_op, name=name) + A = s.to_scipy_sparse() + if dup_op is None and not A.has_canonical_format: + # pydata lets an index repeat within a row (column) and means the sum. scipy + # sums in any dtype it holds, which binary.plus may not take (a UDT). Copy + # first: A shares its arrays with s. + A = A.copy() + A.sum_duplicates() + return from_scipy_sparse(A, dup_op=dup_op, name=name) if isinstance(s, (sparse.DOK, sparse.COO)): _s = s.asformat("coo") return Matrix.from_coo( diff --git a/graphblas/monoid/__init__.py b/graphblas/monoid/__init__.py index 027fc0afe..30463e6a2 100644 --- a/graphblas/monoid/__init__.py +++ b/graphblas/monoid/__init__.py @@ -29,7 +29,9 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/monoid/numpy.py b/graphblas/monoid/numpy.py index b9ff2b502..f17495988 100644 --- a/graphblas/monoid/numpy.py +++ b/graphblas/monoid/numpy.py @@ -40,7 +40,7 @@ # Comparison functions "equal": {"BOOL": True}, "logical_and": {"BOOL": True}, - "logical_or": {"BOOL": True}, + "logical_or": {"BOOL": False}, "logical_xor": {"BOOL": False}, "maximum": { "BOOL": False, @@ -93,19 +93,21 @@ or _has_numba and not isinstance(_numba.njit(lambda x, y: _np.fmax(x, y))(1, 2), float) # pragma: no branch ): - # Incorrect behavior was introduced in numba 0.56.2 and numpy 1.23 + # Incorrect behavior was introduced in numba 0.56.2 and numpy 1.23. # See: https://github.com/numba/numba/issues/8478 - # MAINT: we may be able to remove the behavior-based check above in 2025 + # MAINT (2026-05): the behavioral check above still discriminates fixed + # vs. broken numba/numpy combinations. Once we drop everything older + # than the fix, this branch can become unconditional. _monoid_identities["fmax"].update( { "BOOL": False, "INT8": _np.iinfo(_np.int8).min, "UINT8": 0, - "INT16": _np.iinfo(_np.int8).min, + "INT16": _np.iinfo(_np.int16).min, "UINT16": 0, - "INT32": _np.iinfo(_np.int8).min, + "INT32": _np.iinfo(_np.int32).min, "UINT32": 0, - "INT64": _np.iinfo(_np.int8).min, + "INT64": _np.iinfo(_np.int64).min, "UINT64": 0, } ) diff --git a/graphblas/op/__init__.py b/graphblas/op/__init__.py index 1eb2b51d7..5a37b6121 100644 --- a/graphblas/op/__init__.py +++ b/graphblas/op/__init__.py @@ -47,7 +47,9 @@ def __getattr__(key): f"module {__name__!r} unable to compile UDF for {key!r}; " "install numba for UDF support" ) - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator, _supports_udfs # noqa: E402 isort:skip diff --git a/graphblas/select/__init__.py b/graphblas/select/__init__.py index b55766ff8..a2d94a0a0 100644 --- a/graphblas/select/__init__.py +++ b/graphblas/select/__init__.py @@ -3,6 +3,9 @@ # identical to IndexUnaryOp under the hood, except they must return # boolean and they default to `select` when called, rather than # `apply`). +import math +import numbers + _delayed = {} _binary_to_select = {} @@ -29,7 +32,9 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) def _resolve_expr(expr, callname, opname): @@ -42,21 +47,6 @@ def _resolve_expr(expr, callname, opname): ) tensor = expr.args[0] thunk = expr.args[1] - method = f"{opname}{expr.op.name}" - if method not in globals(): - # TODO: remove this once rowlt/rowge/collt/colge exist - # Convert thunk to Python int to avoid possible subtraction with uints - thunk = thunk.value - # Attempt to convert < into <= (rowlt is not part of official spec, but rowle is) - if expr.op.name == "lt": - method = f"{opname}le" - thunk -= 1 - # Attempt to convert >= into > (rowge is not part of official spec, but rowgt is) - elif expr.op.name == "ge": - method = f"{opname}gt" - thunk -= 1 - if method not in globals(): # pragma: no cover (sanity) - raise ValueError(f"Unknown or unregistered select method: {method}") if expr._is_scalar: # Handle ScalarExpressions that change their arguments to Vector if tensor._parent is not None: # e.g., suitesparse @@ -65,9 +55,55 @@ def _resolve_expr(expr, callname, opname): else: # e.g., suitesparse-vanilla tensor = tensor[0].new() thunk = thunk[0].new() + name = expr.op.name + if opname != "value": + if name in {"lt", "ge"} and f"{opname}{name}" in globals(): + # An op registered for it (a user's rowlt): the number as given + if expr._right_literal is not None: + thunk = expr._right_literal + else: + thunk, name = _index_bound(expr._right_literal, thunk, name) + method = f"{opname}{name}" + if method not in globals(): # pragma: no cover (sanity) + raise ValueError(f"Unknown or unregistered select method: {method}") return globals()[method](tensor, thunk) +def _index_bound(literal, thunk, name): + """Return the thunk and comparison that select an index ``name`` (lt, le, gt, ge) a bound. + + The bound is the number the caller gave (``literal``), else the thunk Scalar's value. + The comparison typed its thunk beside the values, which is wrong for an index: 300 + beside INT8 values is infinity, and 2**24 + 1 beside FP32 rounds. Only rowle and + rowgt exist (and colle, colgt, indexle, indexgt), and indices are whole numbers, so + ``i < t`` is ``i <= ceil(t) - 1``, ``i >= t`` is ``i > ceil(t) - 1``, and ``i <= t`` + and ``i > t`` take ``floor(t)``, where GraphBLAS would truncate a float toward zero + (``i > -0.5`` dropped index 0). A whole number beyond INT64, which holds every index, + is clamped to it. NaN compares False with every index, so it keeps none. An empty + Scalar goes to the op as it is, which raises EmptyObject, and so does a thunk that is + no Scalar (from ``A.apply(binary.le, left=3)``), which raises the op's TypeError. + """ + shift = name in {"lt", "ge"} + if shift: + name = "le" if name == "lt" else "gt" + # The caller's number, if it is one; a Scalar expression (``w[0]``), a 0-d array or + # ``np.True_`` made the thunk keep its own type, so the thunk's value is the bound. + if isinstance(literal, numbers.Real): + bound = literal + elif getattr(thunk, "_is_scalar", False): + bound = thunk.value + else: + return thunk, name + if not isinstance(bound, numbers.Real): + return thunk, name + if math.isnan(bound): + return (-1 if name == "le" else 2**63 - 1), name # no index is <= -1 or > INT64_MAX + if math.isfinite(bound): + bound = math.ceil(bound) - 1 if shift else math.floor(bound) + bound = min(max(bound, -(2**63)), 2**63 - 1) + return bound, name + + def _match_expr(parent, expr): """Match expressions to rewrite ``A.select(A < 5)`` into select expression. @@ -80,7 +116,7 @@ def _match_expr(parent, expr): and len(args) == 2 and op.parent in _binary_to_select and args[1]._is_scalar - and (args[0] is parent or getattr(args[0], "T", None) is parent) + and args[0] is parent and op.type in _binary_to_select[op.parent] ): selectop = _binary_to_select[op.parent][op.type] diff --git a/graphblas/semiring/__init__.py b/graphblas/semiring/__init__.py index 95a44261a..b0cf07bce 100644 --- a/graphblas/semiring/__init__.py +++ b/graphblas/semiring/__init__.py @@ -74,7 +74,9 @@ def __getattr__(key): f"module {__name__!r} unable to compile UDF for {key!r}; " "install numba for UDF support" ) - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/ss/__init__.py b/graphblas/ss/__init__.py index 1f059771b..87bc3e2cd 100644 --- a/graphblas/ss/__init__.py +++ b/graphblas/ss/__init__.py @@ -1,7 +1,24 @@ from suitesparse_graphblas import burble +from .. import backend from ._core import _IS_SSGB7, about, concat, config, diag if not _IS_SSGB7: - # Context was introduced in SuiteSparse:GraphBLAS 8.0 + # Context was introduced in SuiteSparse:GraphBLAS 8.0. from ..core.ss.context import Context, global_context + + # JIT compiler config helpers, introduced alongside JIT in SS 8. + from ..core.ss.jit_config import _auto_fix_jit_at_import as _auto_fix + from ..core.ss.jit_config import fix_jit_config, jit_compiler_is_usable + + # Repair the JIT compiler path on import. Conda-built psg bakes in a + # compiler path that almost never exists at runtime, so JIT silently falls + # back to the cfunc path (2-3x slowdown). This only rewrites the compiler + # name and flags; ``jit_c_control`` is left where the backend's init set + # it (``_set_jit_c_control_at_init``), so importing this submodule cannot + # change what a later operation computes. Compilation is enabled on demand + # by ``_enable_jit_for_udt``. Skipped on vanilla: the repair reads + # ``gb.ss.config``, which needs GxB callables vanilla strips. + if backend == "suitesparse": + _auto_fix() + del _auto_fix diff --git a/graphblas/ss/_core.py b/graphblas/ss/_core.py index b42ea72b4..9fd6eecb6 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -128,24 +128,32 @@ class GlobalConfig(BaseConfig): print_1based : bool gpu_control : str, {"always", "never"} Only available for SuiteSparse:GraphBLAS 7 - **GPU support is a work in progress--not recommended to use** + **GPU support is a work in progress; not recommended to use** gpu_chunk : double Only available for SuiteSparse:GraphBLAS 7 - **GPU support is a work in progress--not recommended to use** + **GPU support is a work in progress; not recommended to use** gpu_id : int Which GPU to use; default is -1, which means do not run on the GPU. Only available for SuiteSparse:GraphBLAS >=8 - **GPU support is a work in progress--not recommended to use** - jit_c_control : {"off", "pause", "run", "load", "on} + **GPU support is a work in progress; not recommended to use** + jit_c_control : {"off", "pause", "run", "load", "on"} Control the CPU JIT: "off" : do not use the JIT and free all JIT kernels if loaded "pause" : do not run JIT kernels, but keep any loaded "run" : run JIT kernels if already loaded, but don't load or compile "load" : able to load and run JIT kernels; may not compile "on" : full JIT: able to compile, load, and run - Only available for SuiteSparse:GraphBLAS >=8 + Only available for SuiteSparse:GraphBLAS >=8. + python-graphblas sets it at init from ``graphblas.config["jit_c_control"]``, + or to "run" when that is not set, and UDT work raises that default "run" to + "on" when it needs a kernel compiled. A "run" configured there is kept, and + so is one set here over another configured value; with nothing configured, + "run" set here is the default again, as None is in any case. The user + guide's initialization page covers the trade-offs, and when "on" is safe. jit_use_cmake : bool Whether to use cmake to compile the JIT kernels. + Always True for SuiteSparse built with MSVC (on Windows), where + python-graphblas does not turn the JIT on yet. Only available for SuiteSparse:GraphBLAS >=8 jit_c_compiler_name : str C compiler for JIT kernels. @@ -215,10 +223,12 @@ class GlobalConfig(BaseConfig): "jit_c_preface": (lib.GxB_JIT_C_PREFACE, "char*"), "jit_error_log": (lib.GxB_JIT_ERROR_LOG, "char*"), "jit_cache_path": (lib.GxB_JIT_CACHE_PATH, "char*"), - # CUDA GPU control - "gpu_id": (lib.GxB_GLOBAL_GPU_ID, "int"), } ) + # GxB_GLOBAL_GPU_ID existed in SS:GraphBLAS <=10.1 as a scalar int. + # It was renamed to GxB_GLOBAL_GPU_IDS in 10.2 and changed to a list type. + if hasattr(lib, "GxB_GLOBAL_GPU_ID"): + _options["gpu_id"] = (lib.GxB_GLOBAL_GPU_ID, "int") # Values to restore defaults _defaults = { "hyper_switch": lib.GxB_HYPER_DEFAULT, @@ -252,14 +262,37 @@ class GlobalConfig(BaseConfig): "on": lib.GxB_JIT_ON, } + def __setitem__(self, key, val): + key = key.lower() + if key != "jit_c_control": + super().__setitem__(key, val) + return + from ..core.ss import jit_config + + # UDT work keeps a 'run' the user set. One written here counts when it + # changes the configured control (graphblas.config, read at init) or repeats + # a configured 'run'. With nothing configured it is the default 'run' again, + # which UDT work raises, so saving and restoring the default changes nothing. + # None is that default whatever is configured. + super().__setitem__(key, "run" if val is None else val) + if isinstance(val, str): + val = val.lower() + elif val is not None: + val = self._reverse_enumerations[key].get(val, val) + jit_config._user_set_run = val == "run" and jit_config._configured_jit_c_control is not None + class About(Mapping): _modes = { lib.GrB_NONBLOCKING: "nonblocking", lib.GrB_BLOCKING: "blocking", - lib.GxB_NONBLOCKING_GPU: "nonblocking_gpu", - lib.GxB_BLOCKING_GPU: "blocking_gpu", } + # GxB_NONBLOCKING_GPU and GxB_BLOCKING_GPU were modes of their own in SS:GraphBLAS <=10.5.1. + # In 10.5.2 they became aliases of GrB_NONBLOCKING and GrB_BLOCKING, which python-suitesparse- + # graphblas leaves out; `setdefault` keeps "nonblocking" and "blocking" if they are aliases. + if hasattr(lib, "GxB_NONBLOCKING_GPU"): + _modes.setdefault(lib.GxB_NONBLOCKING_GPU, "nonblocking_gpu") + _modes.setdefault(lib.GxB_BLOCKING_GPU, "blocking_gpu") _mode_options = { "mode": lib.GxB_MODE, } @@ -314,6 +347,13 @@ def __getitem__(self, key): raise KeyError(key) raise _error_code_lookup[info](f"Failed to get info for {key}") # pragma: no cover (safety) + def __setattr__(self, key, value): + # About never sets instance attributes, so any attribute write is a + # mistake that would otherwise create a dead attribute shadowing reads. + if key.lower() in self: + raise AttributeError(f"About option {key.lower()!r} is read-only") + raise AttributeError(f"Unknown About option {key!r}; known options are {sorted(self)}.") + def __iter__(self): return iter( sorted( diff --git a/graphblas/tests/conftest.py b/graphblas/tests/conftest.py index 964325e0d..f68a597fd 100644 --- a/graphblas/tests/conftest.py +++ b/graphblas/tests/conftest.py @@ -2,6 +2,7 @@ import contextlib import functools import itertools +import os import platform import sys from pathlib import Path @@ -12,6 +13,52 @@ import graphblas as gb from graphblas.core import _supports_udfs as supports_udfs +try: + from numpy._core._multiarray_umath import _set_numpy_warn_if_no_mem_policy +except ImportError: # pragma: no cover (NumPy < 1.26) + try: + from numpy.core._multiarray_umath import _set_numpy_warn_if_no_mem_policy + except ImportError: # NumPy 1.24 lacks it, and it is private, so it may go away + _set_numpy_warn_if_no_mem_policy = None + +# Under ``pytest -n`` (xdist), give each worker its own ``GRAPHBLAS_CACHE_PATH`` +# so concurrent JIT compiles within a single session don't race on the same +# ``.c``/``.dylib`` paths. SuiteSparse reads this env var at first GraphBLAS +# init; it must be set before that init runs, or SS takes the unset +# (HOME-derived) default. ``import graphblas`` above does not trigger +# init (it's lazy), so setting it here is in time. +# +# The path is stable per worker id (``gw0``, ``gw1``, ...) and persistent +# across pytest invocations (not session-scoped). This way the JIT cache +# builds up over successive runs and most kernels are dlopen'd instead of +# recompiled. Total disk footprint is bounded by the number of unique kernels +# the suite touches across all worker assignments (roughly a few hundred +# ``.dylib`` files, low tens of MB). +# +# Caveat: two ``pytest -n`` invocations running simultaneously would both +# write to ``gw0/`` etc. and can race on shared kernels. We accept that +# occasional concurrent-run failure rather than reintroducing per-session +# isolation, which would defeat cache reuse. +if ( + _xdist_worker := os.environ.get("PYTEST_XDIST_WORKER") +) and "GRAPHBLAS_CACHE_PATH" not in os.environ: + _home = os.environ.get("HOME") or os.environ.get("LOCALAPPDATA") or str(Path.cwd()) + os.environ["GRAPHBLAS_CACHE_PATH"] = str(Path(_home) / ".SuiteSparse_xdist" / _xdist_worker) + +# Several ``test_formatting.py`` tests assert verbatim against a Matrix / +# Vector ``repr()``, and those reprs are produced via pandas. With pandas' +# default ``display.max_columns = 0`` (terminal mode), pandas auto-fits to +# ``shutil.get_terminal_size()``, which reads the ``COLUMNS`` env var. When +# pytest is invoked from a wide terminal the workers inherit ``COLUMNS=`` +# and pandas renders extra columns that the tests expected to be elided to +# ``...``. Drop ``COLUMNS`` here so terminal-size detection falls back to the +# 80x24 default the test expectations were written for. The pop has to run at +# conftest module load, before any worker imports pandas and before xdist +# forks. Setting ``display.width`` alone doesn't help; the terminal-size +# branch is taken whenever ``max_columns == 0``. +os.environ.pop("COLUMNS", None) + + orig_binaryops = set() orig_semirings = set() @@ -19,7 +66,22 @@ def pytest_configure(config): - rng = np.random.default_rng() + # Under ``pytest -n`` (xdist), every worker imports this conftest and + # calls ``pytest_configure`` independently. A fresh ``default_rng()`` in + # each worker would pick different backend / blocking / mapnumpy values, + # so the workers collect different tests (vanilla skips ``test_ssjit`` + # etc.) and xdist aborts with "Different tests were collected between + # gw0 and gw2." Seed once on the controller and propagate to workers via + # an env var; workers inherit it at spawn and derive identical choices. + # Set the env var unconditionally on the controller so reproducing a + # specific session is just ``GRAPHBLAS_TEST_SEED= pytest ...``. + seed_str = os.environ.get("GRAPHBLAS_TEST_SEED") + if seed_str is None: + seed = int.from_bytes(os.urandom(8), "big") + os.environ["GRAPHBLAS_TEST_SEED"] = str(seed) + else: + seed = int(seed_str) + rng = np.random.default_rng(seed) randomly = config.getoption("--randomly", None) if randomly is None: # pragma: no cover options_unavailable = True @@ -52,7 +114,7 @@ def pytest_configure(config): gb.init(backend, blocking=blocking) print( f"Running tests with {backend!r} backend, blocking={blocking}, " - f"record={record}, mapnumpy={mapnumpy}, runslow={runslow}" + f"record={record}, mapnumpy={mapnumpy}, runslow={runslow}, seed={seed}" ) if record: rec = gb.Recorder() @@ -108,6 +170,27 @@ def _reset_name_counters(): gb.Scalar._name_counter = itertools.count() +@pytest.fixture(scope="session", autouse=True) +def _warn_if_no_memory_handler(): + """Fail a test that frees an array that owns its data but has no memory handler. + + NumPy frees such data with libc ``free``, which is right only if libc's + ``malloc`` made it; GraphBLAS's allocator is NumPy's. NumPy warns about + this only when asked, and ``filterwarnings = error`` makes the warning fail + the test. suitesparse-graphblas before 10.5.1.1 (which lacks + ``can_unclaim_buffer``) exports arrays with no handler, so it would fail + every export. + """ + from suitesparse_graphblas import utils + + if _set_numpy_warn_if_no_mem_policy is None or not hasattr(utils, "can_unclaim_buffer"): + yield + return + prev = _set_numpy_warn_if_no_mem_policy(True) + yield + _set_numpy_warn_if_no_mem_policy(prev) + + @pytest.fixture(scope="session", autouse=True) def ic(): # pragma: no cover (debug) """Make `ic` available everywhere during testing for easier debugging.""" diff --git a/graphblas/tests/test_core.py b/graphblas/tests/test_core.py index 3586eb4a8..c324c7d7d 100644 --- a/graphblas/tests/test_core.py +++ b/graphblas/tests/test_core.py @@ -41,6 +41,27 @@ def test_bad_init(): gb.init(**params) +def test_jit_c_control_config(): + """``config["jit_c_control"]``, which init sets SuiteSparse's JIT control to.""" + import yaml + + path = pathlib.Path(gb.__file__).parent / "graphblas.yaml" + # Unset by default, which init takes as a 'run' that UDT work may raise + assert yaml.safe_load(path.read_text())["jit_c_control"] is None + # YAML 1.1 reads an unquoted on/off in a user's config file as a bool + cases = [("on", "on"), ("LOAD", "load"), ("run", "run"), ("pause", "pause"), ("off", "off")] + cases += [(True, "on"), (False, "off"), (None, None)] + for val, expected in cases: + with gb.config.set(jit_c_control=val): + assert gb._jit_c_control_config() == expected + for val in ["of", "", 1]: + with ( + gb.config.set(jit_c_control=val), + pytest.raises(ValueError, match="GRAPHBLAS_JIT_C_CONTROL"), + ): + gb._jit_c_control_config() + + def test_bad_libget(): with pytest.raises(AttributeError, match="GrB_bad_name"): gb.core.base.libget("GrB_bad_name") @@ -67,7 +88,7 @@ class bad: def test_version(): from packaging.version import parse - assert parse(gb.__version__) > parse("2022.11.0") + assert parse(gb.__version__) > parse("2024.2.0") @pytest.mark.skipif("not setuptools or not tomli or not gb.__file__") @@ -94,3 +115,48 @@ def test_packages(): def test_index_max(): assert gb.MAX_SIZE == 2**60 # True for all current backends + + +def test_config_attribute_assignment_raises(): + # Setting an option by plain attribute assignment used to silently no-op: + # donfig has no __setattr__ for options, so it created a dead instance + # attribute and left the real option alone. It should raise and name the + # canonical API instead. + orig = gb.config["autocompute"] + try: + with pytest.raises(AttributeError, match=r"config\.set\(autocompute"): + gb.config.autocompute = not orig + assert gb.config["autocompute"] == orig # the failed write changed nothing + # An unknown name is a mistaken write too, not a new attribute + with pytest.raises(AttributeError, match="not_a_real_option"): + gb.config.not_a_real_option = 5 + assert "not_a_real_option" not in vars(gb.config) + # Canonical idioms still work: read by item, write by set() + gb.config.set(autocompute=not orig) + assert gb.config["autocompute"] == (not orig) + with gb.config.set(autocompute=orig): + assert gb.config["autocompute"] == orig + assert gb.config["autocompute"] == (not orig) # restored on exit + finally: + gb.config.set(autocompute=orig) + + +def test_config_donfig_attrs_are_derived(monkeypatch): + # The allowlist of writable attributes must come from what donfig's + # __init__ actually set, not a hardcoded list. Simulate a future donfig + # that grows a field: the derived allowlist absorbs it, where a + # hardcoded list would start rejecting donfig's own writes. + import donfig + + real_init = donfig.Config.__init__ + + def init_with_extra_field(self, *args, **kwargs): + real_init(self, *args, **kwargs) + self.hypothetical_future_field = 1 + + monkeypatch.setattr(donfig.Config, "__init__", init_with_extra_field) + probe = type(gb.config)("test_derived_allowlist") + assert "hypothetical_future_field" in probe._donfig_attrs + probe.hypothetical_future_field = 2 # writable, not rejected + with pytest.raises(AttributeError, match="autocompute"): + probe.autocompute = False diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index ecbca707f..a7edc6d7b 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -2,14 +2,17 @@ import pickle import string import sys +import warnings import numpy as np import pytest import graphblas as gb from graphblas import core, dtypes +from graphblas.core import _supports_udfs as supports_udfs from graphblas.core import lib -from graphblas.core.utils import _NP2 +from graphblas.core.operator.udt_codegen import _has_jit_set +from graphblas.core.utils import _NP2, values_to_numpy_buffer from graphblas.dtypes import lookup_dtype suitesparse = gb.backend == "suitesparse" @@ -213,15 +216,408 @@ def test_default_names(): ) assert _default_name(np.dtype("(29,)uint8")) == "UINT8[29]" assert _default_name(np.dtype("(3,4)bool")) == "BOOL[3, 4]" - assert _default_name(np.dtype((np.dtype("(5,)float64"), (6,)))) == "FP64[5][6]" + assert _default_name(np.dtype((np.dtype([("x", np.int32)]), (6,)))) == "{'x': INT32}[6]" assert _default_name(np.dtype("S5")) == "dtype('S5')" +def test_nested_subarrays_register_flat(): + """An array-of-arrays dtype registers as the flat array UDT with its layout. + + numpy keeps ``np.dtype((np.dtype((np.float32, (7,))), (3,)))`` layered and + compares it unequal to ``np.dtype((np.float32, (3, 7)))``, but nothing else + tells the two apart: same bytes, same expanded ``(..., 3, 7)`` array, same + Numba type. As two DataTypes they had separate GrB_Types, which SuiteSparse + refused to mix, and the layered one broke code that reads ``subdtype`` as + ``(scalar, shape)``. + """ + flat = np.dtype((np.float32, (3, 7))) + layered = np.dtype((np.dtype((np.float32, (7,))), (3,))) + assert layered != flat + udt = lookup_dtype(layered) + assert udt.np_type == flat + assert udt.name == "FP32[3, 7]" + assert lookup_dtype(flat) is udt + assert dtypes.register_anonymous(layered) is udt + assert udt == layered + assert pickle.loads(pickle.dumps(udt)) is udt + if _has_jit_set: + assert udt.jit_c_definition is not None + + # INT16 keeps this under 128 bytes, the most SuiteSparse < 9 accepts for a + # UDT on builds without variable-length arrays (Windows). + deeper = np.dtype((np.dtype((np.dtype((np.int16, (7,))), (2,))), (3,))) + assert lookup_dtype(deeper).np_type == np.dtype((np.int16, (3, 2, 7))) + + +def test_nested_subarrays_register_flat_in_records(): + """Record fields flatten too, keeping offsets, itemsize, alignment, and titles.""" + point = np.dtype((np.float32, (5,))) + layered = np.dtype([("fl_id", np.int16), ("fl_tri", point, (3,))], align=True) + flat = np.dtype([("fl_id", np.int16), ("fl_tri", np.float32, (3, 5))], align=True) + udt = lookup_dtype(layered) + assert udt.np_type == flat + assert udt.np_type.isalignedstruct + assert lookup_dtype(flat) is udt + + packed = np.dtype([("fl_a", np.int8), ("fl_b", point, (2,))]) + assert lookup_dtype(packed).np_type == np.dtype( + [("fl_a", np.int8), ("fl_b", np.float32, (2, 5))] + ) + + titled = np.dtype({"names": ["fl_t"], "formats": [(point, (4,))], "titles": ["Tee"]}) + assert lookup_dtype(titled).np_type == np.dtype( + {"names": ["fl_t"], "formats": [(np.float32, (4, 5))], "titles": ["Tee"]} + ) + + # An array of records whose field nests subarrays. + assert lookup_dtype(np.dtype((layered, (2,)))).np_type == np.dtype((flat, (2,))) + + +def test_nested_subarray_record_values_load_intact(): + """Values whose record dtype nests subarrays load into the flattened UDT intact. + + numpy casts between a field of two ``FP64[3]`` elements and an + ``FP64[2, 3]`` field wrongly in both directions: ``astype``, + ``np.array(..., dtype=)``, and item assignment all scramble the values. The + bytes already match, so ingest views such values instead of casting them. + """ + point = np.dtype((np.float64, (3,))) + layered = np.dtype([("li_k", np.int64), ("li_pts", point, (2,))], align=True) + udt = lookup_dtype(layered) + vals = np.zeros(2, layered) + vals["li_k"] = [1, 2] + vals["li_pts"] = np.arange(12.0).reshape(2, 2, 3) + + # Inferring the dtype views the values too, so the pair agrees. + arr, dtype = values_to_numpy_buffer(vals) + assert dtype is udt + assert arr.dtype == udt.np_type + np.testing.assert_array_equal(arr["li_pts"], vals["li_pts"]) + + s = gb.Scalar.from_value(vals[0], dtype=udt) + np.testing.assert_array_equal(s.value["li_pts"], vals["li_pts"][0]) + + # Off SuiteSparse, ``build`` uses ``binary.any`` as its ``dup_op``, which a + # UDT gets only by compiling it with Numba. + if suitesparse or supports_udfs: + v = gb.Vector.from_coo([0, 1], vals, dtype=udt) + np.testing.assert_array_equal(v[1].new().value["li_pts"], vals["li_pts"][1]) + v = gb.Vector.from_coo([0, 1], vals) + assert v.dtype is udt + np.testing.assert_array_equal(v[1].new().value["li_pts"], vals["li_pts"][1]) + v[0] = vals[1] + np.testing.assert_array_equal(v[0].new().value["li_pts"], vals["li_pts"][1]) + + +@pytest.mark.skipif("not supports_udfs") +def test_lookup_numba_type_goes_through_numpy_dtype(): + """A UDT's Numba type resolves through its numpy dtype, not a registry key. + + Keyed directly, a layered and a flat array dtype, which share a Numba type, + resolved to whichever registered last, and a layout nobody had registered + yet did not resolve at all. + """ + import numba + + nb_type = numba.typeof(np.dtype((np.int32, (2, 9)))).dtype + udt = lookup_dtype(nb_type) + assert udt.np_type == np.dtype((np.int32, (2, 9))) + assert lookup_dtype(np.dtype((np.dtype((np.int32, (9,))), (2,)))) is udt + assert lookup_dtype(nb_type) is udt + + rec = dtypes.register_anonymous(np.dtype([("nb_q", np.int32), ("nb_r", np.float32)])) + assert lookup_dtype(numba.typeof(rec.np_type).dtype) is rec + + def test_record_dtype_from_dict(): dtype = dtypes.lookup_dtype({"x": int, "y": float}) assert dtype.name == "{'x': INT64, 'y': FP64}" +def test_register_anonymous_from_dataclass(): + """A ``@dataclass`` (class or instance) registers as a record UDT, fields and all.""" + from dataclasses import dataclass + + @dataclass + class DcEdge: + weight: float + count: int + + # Class form: defaults the UDT name to the class name. + udt = dtypes.register_anonymous(DcEdge) + assert udt.name == "DcEdge" + assert "weight" in udt.np_type.names + assert "count" in udt.np_type.names + + # Explicit name overrides the class name (reuses the same DataType object + # because the underlying np.dtype matches; rename behavior is intentional). + udt2 = dtypes.register_anonymous(DcEdge, "_DcEdgeAlias") + assert udt2 is udt + assert udt.name == "_DcEdgeAlias" + + # Instance form also works. + e = DcEdge(weight=1.5, count=10) + udt3 = dtypes.register_anonymous(e, "_DcEdgeFromInst") + assert udt3 is udt + + # Empty dataclasses are rejected; a zero-field record UDT isn't useful. + @dataclass + class _Empty: + pass + + with pytest.raises(ValueError, match="at least one field"): + dtypes.register_anonymous(_Empty) + + +def test_register_anonymous_zero_size(): + for dtype in [str, bytes, np.dtype("U"), np.dtype([]), np.dtype((np.int64, (0,)))]: + with pytest.raises(ValueError, match="size of 0"): + dtypes.register_anonymous(dtype) + with pytest.raises(ValueError, match="Unknown dtype"): + dtypes.lookup_dtype(str) + + +def test_register_anonymous_array_numba_cannot_type(): + # Numba cannot type an array of these, alone, in a record or in an array of records. + # Its own error was "AttributeError: 'Boolean' object has no attribute 'bitwidth'" + # and the like. + for np_type, element, instead in [ + (np.dtype((np.bool_, (2,))), "bool", "uint8"), + (np.dtype(([("_nbt_e", np.bool_, (2,))], (3,))), "bool", "uint8"), + (np.dtype(("M8[s]", (2, 3))), "datetime64", "int64"), + (np.dtype([("_nbt_a", np.int8), ("_nbt_b", "m8[ns]", (2,))]), "timedelta64", "int64"), + (np.dtype(("U3", (2,))), "str", "uint32"), + (np.dtype([("_nbt_c", [("_nbt_d", "S2", (2,))])]), "bytes", "uint8"), + ]: + if not supports_udfs: + # Nothing types it without Numba, and it registers + assert dtypes.register_anonymous(np_type).np_type == np_type + continue + match = f"cannot type an array of {element}.*use an array of {instead}" + with pytest.raises(ValueError, match=match) as exc_info: + dtypes.register_anonymous(np_type) + assert isinstance(exc_info.value.__cause__, AttributeError) + assert np_type not in core.dtypes._registry + + +def test_register_anonymous_from_dataclass_deferred_annotations(): + """Dataclass fields annotated as strings (PEP 563 / 649) resolve through ``lookup_dtype``.""" + from dataclasses import make_dataclass + + _DcDeferred = make_dataclass("_DcDeferred", [("deferred_x", "int"), ("deferred_y", "float")]) + udt = dtypes.register_anonymous(_DcDeferred) + assert udt.np_type.names == ("deferred_x", "deferred_y") + assert udt.np_type.fields["deferred_x"][0] == np.int64 + assert udt.np_type.fields["deferred_y"][0] == np.float64 + + +def test_register_new_from_dataclass(): + """``register_new(MyDataclass)`` infers the dtype name from the class name.""" + from dataclasses import dataclass + + @dataclass + class _DcRegNew: + a: int + b: float + + udt = dtypes.register_new(_DcRegNew) + assert udt.name == "_DcRegNew" + assert udt is dtypes._DcRegNew + # Explicit name + dtype still works. + + @dataclass + class _DcRegNew2: + x: int + + udt2 = dtypes.register_new("_DcRegNewAlias", _DcRegNew2) + assert udt2.name == "_DcRegNewAlias" + # Sole non-dataclass argument is a TypeError, not a cryptic AttributeError. + with pytest.raises(TypeError, match="requires both"): + dtypes.register_new(np.dtype([("x", np.int64)])) + + +def test_register_anonymous_from_dataclass_instance_default_name(): + """``register_anonymous(instance)`` without ``name=`` defaults to the class name.""" + from dataclasses import dataclass + + @dataclass + class _DcInstName: + v: float + + inst = _DcInstName(v=1.5) + udt = dtypes.register_anonymous(inst) + assert udt.name == "_DcInstName" + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not _has_jit_set") +def test_udt_with_macro_field_name_falls_back_to_cfunc(): + """A stdlib-macro field name (e.g., ``M_PI``) skips JIT and falls back to cfunc. + + Regression for the silent JIT-compile failure: the C preprocessor would + expand ``M_PI`` to a numeric literal inside the struct declarator + (``typedef struct { double M_PI ; ... }``), producing uncompilable C + that SuiteSparse swallowed without warning. The ``_C_RESERVED`` set + blocks macros so the warning fires and JIT is skipped cleanly. + """ + from graphblas.core.ss import jit_config + from graphblas.exceptions import NoJITWarning + + record_dtype = np.dtype([("M_PI", np.float64), ("other", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_MacroFieldUdt") + assert udt.jit_c_name is None + assert udt.jit_c_definition is None + + jit_config._warned_no_jit_for.discard(("plus", udt.name)) + with pytest.warns(NoJITWarning, match="without JIT"): + plus_udt = gb.binary.plus[udt] + assert plus_udt.jit_c_source is None + # Once per (op, dtype): the same pair again is silent + with warnings.catch_warnings(): + warnings.simplefilter("error") + jit_config._maybe_warn_no_jit(op_name="plus", dtype_name=udt.name) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not _has_jit_set") +def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): + """A C-reserved field name skips JIT silently; the op still runs via the cfunc path.""" + from graphblas.core.ss import jit_config + from graphblas.exceptions import NoJITWarning + + # ``class`` is a C++ reserved word and is in our ``_C_RESERVED`` set. + record_dtype = np.dtype([("class", np.int64), ("ok_field", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_CReservedUdt") + + # JIT info must be ``None``: no C struct to register. + assert udt.jit_c_name is None + assert udt.jit_c_definition is None + + # Auto-lift must succeed and emit a warning. Discard the keyed entry so + # this test sees the warning regardless of test ordering (the same + # ``(op, dtype)`` may have been warned about in an earlier test). + jit_config._warned_no_jit_for.discard(("plus", udt.name)) + with pytest.warns(NoJITWarning, match="without JIT") as record: + plus_udt = gb.binary.plus[udt] + # The cause has to name the UDT, on every host. Nothing armed a kernel for + # this dtype, so ``jit_c_control`` is still whatever init set, and blaming + # it would send the reader to a setting that cannot fix anything. + assert "not expressible as a C struct" in str(record[0].message) + assert plus_udt.jit_c_source is None # no JIT + # ``eq`` has a separate JIT codegen path from arithmetic (BOOL output, leaf + # comparisons), so pin its skip independently. + assert gb.binary.eq[udt].jit_c_source is None + v = gb.Vector(udt, size=2) + v[0] = (1, 10) + v[1] = (2, 20) + w = gb.Vector(udt, size=2) + w[0] = (100, 1000) + w[1] = (200, 2000) + result = v.ewise_mult(w, plus_udt).new() + assert tuple(result[0].new().value) == (101, 1010) + assert tuple(result[1].new().value) == (202, 2020) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not _has_jit_set") +def test_no_jit_warning_blames_the_udt_not_the_host_jit_state(): + """The cause named must not move with the host's JIT state. + + On a runner whose JIT probe fails, init leaves ``jit_c_control`` at + ``'run'`` and caches the probe's ``False``. Codegen + for this UDT read neither, so naming either as the cause points the + reader at a setting that cannot fix it. CI hit exactly that: on + ubuntu-latest the message blamed ``jit_c_control``. + """ + from graphblas.core.ss import jit_config + from graphblas.exceptions import NoJITWarning + + # ``switch`` is a C reserved word; the field names are unique to this test + # because ``register_anonymous`` caches one DataType per ``np.dtype``. + record_dtype = np.dtype([("switch", np.int64), ("plain", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_ProbeFailedUdt") + + prev_control = gb.ss.config["jit_c_control"] + prev_enabled = jit_config._jit_enabled_for_udt + try: + gb.ss.config["jit_c_control"] = "run" + jit_config._jit_enabled_for_udt = False + jit_config._warned_no_jit_for.discard(("plus", udt.name)) + with pytest.warns(NoJITWarning, match="not expressible as a C struct") as record: + gb.binary.plus[udt] + finally: + gb.ss.config["jit_c_control"] = prev_control + jit_config._jit_enabled_for_udt = prev_enabled + assert not any("jit_c_control" in str(w.message) for w in record) + + +@pytest.mark.skipif("not _has_jit_set") +def test_udt_jit_c_info_pinned_at_first_register(): + """SuiteSparse's ``GxB_JIT_C_DEFINITION`` is one-shot per ``GrB_Type``. + + Re-registering an existing ``np.dtype`` under a new Python-side name + updates ``rv.name`` but not the SS-side JIT identity. The introspection + properties ``jit_c_name`` and ``jit_c_definition`` must report what + SS actually has, not the renamed Python value, so callers can trust + the introspection to match what JIT-compiled kernels see. + """ + # Use a field/shape combo unique to this test to keep the C-side name + # stable. + record_dtype = np.dtype([("pin_a", np.int64), ("pin_b", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_PinUDT_first") + assert udt.jit_c_name == "gbudt__PinUDT_first" + + udt2 = dtypes.register_anonymous(record_dtype, "_PinUDT_renamed") + assert udt2 is udt + assert udt.name == "_PinUDT_renamed" + # SS-side name is frozen at first register; introspection must reflect that. + assert udt.jit_c_name == "gbudt__PinUDT_first" + assert "_PinUDT_first" in udt.jit_c_definition + + +@pytest.mark.skipif("not _has_jit_set") +def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): + """Two UDT layouts registered under one name get distinct JIT C names. + + SuiteSparse's JIT knows a type by its C name alone, so the second layout + ran kernels compiled for the first: making an iso Vector non-iso expanded + 28-byte elements with a kernel for 6-byte ones (wrong values), or 6-byte + elements with one for 28-byte ones (heap corruption, in either order of + the tests that did this). + """ + small = dtypes.register_anonymous(np.dtype((np.int16, (1, 3, 1))), "_SameJitName") + big = dtypes.register_anonymous(np.dtype((np.int16, (2, 1, 7))), "_SameJitName") + assert small.jit_c_name == "gbudt__SameJitName" + assert big.jit_c_name != small.jit_c_name + assert big.jit_c_name in big.jit_c_definition + # Nor may a typedef registered in C take a name a UDT already has. + with pytest.raises(ValueError, match="unavailable"): + dtypes.ss.register_new( + small.jit_c_name, f"typedef struct {{ int8_t q ; }} {small.jit_c_name} ;" + ) + # The rest compiles kernels for both layouts, which needs a C compiler the JIT can use. + # Asked under 'on': a 'run' the session set (GRAPHBLAS_JIT_C_CONTROL) would answer no. + from graphblas.core.ss.jit_config import _enable_jit_for_udt + + jit_c_control = gb.ss.config["jit_c_control"] + gb.ss.config["jit_c_control"] = "on" + try: + if not _enable_jit_for_udt(): + pytest.skip("the JIT cannot compile here (no usable C compiler)") + for udt in [small, big]: + shape = udt.np_type.shape + one = np.full(shape, 1, np.int16) + two = np.full(shape, 2, np.int16) + v = gb.Vector(udt, size=4) + v[0] = one + v[1] = one # iso + v[2] = two # no longer iso: SuiteSparse expands the iso value + values = [v[i].new().value.tolist() for i in range(3)] + assert values == [one.tolist(), one.tolist(), two.tolist()] + finally: + gb.ss.config["jit_c_control"] = jit_c_control + + def test_dtype_to_from_string(): types = [dtypes.BOOL, dtypes.FP64] for c in string.ascii_letters: @@ -248,6 +644,30 @@ def test_dtype_to_from_string(): assert dtype == dtype2 +# The 180-char repr exceeds GxB_MAX_NAME_LEN, which warns on SuiteSparse < 9. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_dtype_to_string_aligned_outer_packed_inner(): + """Regression: nested dtype with ``align=True`` outer and a packed inner + used to fail ``_dtype_to_string`` round-trip with ``ValueError: Unable + to reliably convert dtype to string and back``. Numpy's reconstruction + from ``str(np_type)`` enforces the outer's alignment on the inner, + rejecting the packed inner's offset. The explicit-offsets fallback + encodes the layout verbatim and round-trips cleanly. + """ + packed_inner = np.dtype([("inner_a", np.int32), ("inner_b", np.float64)]) + outer = np.dtype([("flag", np.int32), ("coord", packed_inner)], align=True) + s = core.dtypes._dtype_to_string(outer) + rt = core.dtypes._string_to_dtype(s) + # The DataType wraps an np.dtype; check the underlying layout matches. + rt_np = rt.np_type + assert rt_np == outer + assert outer.fields["flag"] == rt_np.fields["flag"] + assert outer.fields["coord"] == rt_np.fields["coord"] + # Registration through register_anonymous used to raise on this shape. + udt = dtypes.register_anonymous(outer, "_NestedAlignPacked") + assert udt.np_type == outer + + def test_has_complex(): """Only SuiteSparse has complex (with Windows support in Python after v7.4.3.1).""" if not suitesparse: diff --git a/graphblas/tests/test_formatting.py b/graphblas/tests/test_formatting.py index faadc983b..75393afa0 100644 --- a/graphblas/tests/test_formatting.py +++ b/graphblas/tests/test_formatting.py @@ -9,7 +9,6 @@ from graphblas import Matrix, Scalar, Vector # isort:skip (for dask-graphblas) - try: import pandas as pd except ImportError: # pragma: no cover (import) @@ -82,6 +81,17 @@ def html_printer(x, name="", indent=4): return _printer(repr_html(x), name, "repr_html", indent) +@pytest.fixture(autouse=True) +def _stable_repr_width(monkeypatch): + """Pin terminal width so pandas-rendered reprs are deterministic. + + Matrix/Vector reprs let pandas truncate columns to the terminal width, and + the expected strings assume 80. Under ``pytest -n`` workers inherit the + launching terminal's width, so a wide terminal would show extra columns. + """ + monkeypatch.setenv("COLUMNS", "80") + + @pytest.fixture def A(): return Matrix.from_coo([0, 0, 0], [0, 2, 4], [0, 1, 2], nrows=1, ncols=5, name="A_1") @@ -4000,6 +4010,64 @@ def test_infix_expr_repr(A, B, v): ) +@pytest.mark.skipif("not pd") +def test_scalar_infix_expr_repr_literal(): + # A literal shows its value and its own type, as it has no dtype until an op types it + s = Scalar.from_value(5, dtype=dtypes.INT8, name="s_1") + repr_printer(s & 1, "s & 1") + assert repr(s & 1) == ( + "gb.ScalarEwiseMultExpr left_dtype right_dtype\n" + "s_1 & 1 INT8 int\n" + "\n" + "Do op(expr) to create a ScalarExpression for ewise_mult.\n" + "For example: times(s_1 & 1)" + ) + repr_printer(0.5 | s, "0.5 | s") + assert repr(0.5 | s) == ( + "gb.ScalarEwiseAddExpr left_dtype right_dtype\n" + "0.5 | s_1 float INT8\n" + "\n" + "Do op(expr) to create a ScalarExpression for ewise_add.\n" + "For example: plus(0.5 | s_1)" + ) + html = repr_html(s & 1) + assert "
s1 & 1
" in html + assert "INT8\n int" in html + + +@autocompute +def test_scalar_infix_expr_repr_auto(): + if not pd: # pragma: no cover (import) + pytest.skip("needs pandas") + # As for Vector, only BOOL `x & y` and `x | y` compute without an op, to show the result + s = Scalar.from_value(5, dtype=dtypes.INT8, name="s_1") + repr_printer(s & s, "s & s") + assert repr(s & s) == ( + "gb.ScalarEwiseMultExpr left_dtype right_dtype\n" + "s_1 & s_1 INT8 INT8\n" + "\n" + "Do op(expr) to create a ScalarExpression for ewise_mult.\n" + "For example: times(s_1 & s_1)" + ) + assert "Result" not in repr(s | 1) + assert '
' not in repr_html(s & s) + t = Scalar.from_value(True, name="t") + repr_printer(t & True, "t & True") + assert repr(t & True) == ( + "gb.ScalarEwiseMultExpr left_dtype right_dtype\n" + "t & True BOOL bool\n" + "\n" + '"Result" value dtype\n' + "gb.Scalar True BOOL\n" + "\n" + "Do op(expr) to create a ScalarExpression for ewise_mult.\n" + "For example: times(t & True)" + ) + html = repr_html(False | t) + assert '
' in html + assert "True\n BOOL" in html + + @pytest.mark.skipif("not pd") def test_inner_outer_repr_html(v): html_printer(v.inner(v), "v.inner(v)") @@ -5022,7 +5090,12 @@ def test_empty(): @pytest.mark.skipif("not pd") def test_vector_as_matrix(): - v = Vector.from_coo([1], [2], name="v_A") + # Built iso outright: the repr asserted below names the format, and + # SuiteSparse 10.5.0 no longer infers iso from a build whose values happen + # to be equal, so ``from_coo`` would render "bitmapc" on 10.5.0 and + # "bitmapc (iso)" on everything older. + v = Vector(int, 2, name="v_A") + v.ss.build_scalar([1], 2) A = v._as_matrix() repr_printer(A, "A") assert repr(A) == ( diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py new file mode 100644 index 000000000..333c35fcc --- /dev/null +++ b/graphblas/tests/test_indexbinary.py @@ -0,0 +1,1004 @@ +import pickle + +import numpy as np +import pytest + +import graphblas as gb +from graphblas import Matrix, Scalar, Vector, dtypes, indexbinary +from graphblas.core import _supports_udfs as supports_udfs +from graphblas.core.operator.indexbinary import _has_idxbinop +from graphblas.exceptions import EmptyObject, UdfParseError + +pytestmark = [ + pytest.mark.skipif(not supports_udfs, reason="requires numba"), + pytest.mark.skipif(not _has_idxbinop, reason="requires SuiteSparse:GraphBLAS 9.4+"), +] + + +def test_register_anonymous(): + def add_with_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_with_theta) + assert "add_with_theta" in op.name + assert int in op.types or dtypes.INT64 in op.types + + +def test_register_new(): + def my_idxbin(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x * y + theta + + result = indexbinary.register_new("my_idxbin", my_idxbin) + assert hasattr(indexbinary, "my_idxbin") + assert result is indexbinary.my_idxbin + + A = Matrix.from_coo([0, 1], [1, 0], [3, 7]) + B = Matrix.from_coo([0, 1], [1, 0], [5, 2]) + binop = indexbinary.my_idxbin(100) + C = A.ewise_mult(B, binop).new() + assert list(C.to_coo()[2]) == [115, 114] # 3*5+100, 7*2+100 + + delattr(indexbinary, "my_idxbin") + + +def test_register_new_lazy(): + def lazy_op(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + + result = indexbinary.register_new("lazy_op", lazy_op, lazy=True) + assert result is None + assert "lazy_op" in dir(indexbinary) + + op = indexbinary.lazy_op + assert op.name == "lazy_op" + delattr(indexbinary, "lazy_op") + + +def test_typed_call(): + def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_theta) + typed = op[int] + binop = typed(10) + assert binop.opclass == "BinaryOp" + + A = Matrix.from_coo([0, 1], [1, 0], [3, 7]) + B = Matrix.from_coo([0, 1], [1, 0], [5, 2]) + C = A.ewise_mult(B, binop).new() + assert list(C.to_coo()[2]) == [18, 19] # 3+5+10, 7+2+10 + + +def test_untyped_call(): + def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_theta) + binop = op(10) + assert binop.opclass == "BinaryOp" + + A = Matrix.from_coo([0, 1], [1, 0], [3, 7]) + B = Matrix.from_coo([0, 1], [1, 0], [5, 2]) + C = A.ewise_mult(B, binop).new() + assert list(C.to_coo()[2]) == [18, 19] + + +def test_untyped_call_is_typed_when_it_runs(): + """``op(theta)`` takes its type from both inputs and theta each time it runs. + + Theta is typed as an IndexUnaryOp's thunk is: a Python int is INT64 and a + float FP64, a numpy scalar or a Scalar keeps its type, and the default, + False, changes no type. Before, the op had theta's type alone: ``op()`` + was BOOL, so ints became True; ``op(2)`` truncated floats to INT64; and + ``op(np.int8(2))`` computed in INT8, where 300.0 became 127 and wrapped. + """ + + def upper_from_x(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x if ix <= jx else y + + pick = indexbinary.register_anonymous(upper_from_x) + A = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [10, 20, 30, 40]) + B = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [-1, -2, -3, -4]) + expected = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [10, 20, -3, 40]) + for binop in [pick(), pick(0), pick(False)]: + assert A.ewise_mult(B, binop).new().isequal(expected, check_dtype=True) + C = A.ewise_mult(B, pick(0.0)).new() + assert C.isequal(expected.dup("FP64"), check_dtype=True) + + def weighted(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + theta * y + + wsum = indexbinary.register_anonymous(weighted) + u = Vector.from_coo([0, 1], [1.5, 300.0]) + w = Vector.from_coo([0, 1], [2.25, 1.0]) + expected = Vector.from_coo([0, 1], [6.0, 302.0]) + for theta in [2, 2.0, np.int8(2), Scalar.from_value(2, "INT8")]: + binop = wsum(theta) + assert u.ewise_mult(w, binop).new().isequal(expected, check_dtype=True) + assert binop(u & w).new().isequal(expected, check_dtype=True) + # One op runs at each type it meets + i8 = Vector.from_coo([0, 1], [1, 100], dtype="INT8") + binop = wsum(1) + assert u.ewise_mult(w, binop).new().isequal(Vector.from_coo([0, 1], [3.75, 301.0])) + result = i8.ewise_mult(i8, binop).new() + assert result.isequal(Vector.from_coo([0, 1], [2, 200], dtype="INT64"), check_dtype=True) + for binop in [wsum(), wsum(np.int8(1))]: + assert i8.ewise_mult(i8, binop).new().dtype == dtypes.INT8 + # A named type is the op's type, as before + for binop in [ + wsum["INT64"](2), + wsum(2, dtype="INT64"), + wsum(Scalar.from_value(2.0), dtype="INT64"), + wsum(2)["INT64"], + ]: + assert binop.type == dtypes.INT64 + result = u.ewise_mult(w, binop).new() + assert result.isequal(Vector.from_coo([0, 1], [5, 302]), check_dtype=True) + assert repr(wsum(2)) == "indexbinary.weighted(2)" + rec_udt = dtypes.register_anonymous(np.dtype([("theta_a", np.int64)]), "_UntypedThetaRec") + with pytest.raises(KeyError, match="weighted does not work with _UntypedThetaRec"): + wsum(Scalar.from_value((1,), rec_udt)) + # Binding when it runs uses theta's value from when it was given + theta = Scalar.from_value(2.0) + binop = wsum(theta) + theta.value = 100.0 + assert u.ewise_mult(w, binop).new().isequal(expected) + + +def test_index_aware(): + """The wrapper passes ``(ix, jx, iy, jy)`` through to the user function.""" + + def index_sum(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return ix + jx + iy + jy + theta + + op = indexbinary.register_anonymous(index_sum) + # For ewise_mult, both operands have the same indices: ix==iy and jx==jy + A = Matrix.from_coo([0, 1, 2], [0, 1, 2], [100, 200, 300]) + B = Matrix.from_coo([0, 1, 2], [0, 1, 2], [1, 1, 1]) + binop = op[int](0) + C = A.ewise_mult(B, binop).new() + # (0,0): 0+0+0+0+0=0, (1,1): 1+1+1+1+0=4, (2,2): 2+2+2+2+0=8 + assert list(C.to_coo()[2]) == [0, 4, 8] + + +def test_floating_point(): + def fp_add(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta * 0.5 + + op = indexbinary.register_anonymous(fp_add) + A = Matrix.from_coo([0], [0], [1.5]) + B = Matrix.from_coo([0], [0], [2.5]) + binop = op(4.0) + C = A.ewise_mult(B, binop).new() + assert abs(C.to_coo()[2][0] - 6.0) < 1e-10 # 1.5 + 2.5 + 4.0*0.5 = 6.0 + + +def test_vector_ewise(): + def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_theta) + v1 = Vector.from_coo([0, 1, 2], [10, 20, 30]) + v2 = Vector.from_coo([0, 1, 2], [1, 2, 3]) + binop = op(0) + v3 = v1.ewise_mult(v2, binop).new() + assert list(v3.to_coo()[1]) == [11, 22, 33] + + +def test_ewise_add(): + def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_theta) + A = Matrix.from_coo([0, 1], [0, 1], [3, 7]) + B = Matrix.from_coo([0, 1], [0, 1], [5, 2]) + binop = op(10) + C = A.ewise_add(B, binop).new() + assert list(C.to_coo()[2]) == [18, 19] + + +def test_default_theta(): + """``theta=0`` is a valid bind value and does not get treated as missing.""" + + def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_theta) + binop = op(0) # theta=0 as int + A = Matrix.from_coo([0], [0], [3]) + B = Matrix.from_coo([0], [0], [5]) + C = A.ewise_mult(B, binop).new() + assert C.to_coo()[2][0] == 8 # 3 + 5 + 0 + + +def test_bool_return(): + def is_close(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return abs(x - y) <= theta + + op = indexbinary.register_anonymous(is_close) + A = Matrix.from_coo([0, 1], [0, 1], [10, 20]) + B = Matrix.from_coo([0, 1], [0, 1], [11, 25]) + binop = op[int](2) + C = A.ewise_mult(B, binop).new() + assert list(C.to_coo()[2]) == [True, False] # |10-11|<=2, |20-25|>2 + + +def test_scalar_theta(): + """A graphblas ``Scalar`` is accepted as a theta bind value.""" + + def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_theta) + theta = Scalar.from_value(42) + binop = op[int](theta) + A = Matrix.from_coo([0], [0], [3]) + B = Matrix.from_coo([0], [0], [5]) + C = A.ewise_mult(B, binop).new() + assert C.to_coo()[2][0] == 50 # 3 + 5 + 42 + + +def test_parameterized(): + def make_op(scale): + def inner(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return (x + y) * scale + theta + + return inner + + op = indexbinary.register_anonymous(make_op, parameterized=True) + scaled_op = op(2) # scale=2 + binop = scaled_op(10) # theta=10 + A = Matrix.from_coo([0], [0], [3]) + B = Matrix.from_coo([0], [0], [5]) + C = A.ewise_mult(B, binop).new() + assert C.to_coo()[2][0] == 26 # (3+5)*2 + 10 = 26 + + +# Module-level so pickle can serialize the function reference. IBO +# ``__reduce__`` for user-registered ops returns a tuple containing the +# original function; pickling needs the function to be globally importable +# (i.e., not a closure / local function). +def _ibo_add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + +def _ibo_return_x(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + + +def test_pickle_registered(): + indexbinary.register_new("pickle_test_op", _ibo_add_theta) + try: + op = indexbinary.pickle_test_op + op2 = pickle.loads(pickle.dumps(op)) + assert op2.name == op.name + + typed = op[int] + typed2 = pickle.loads(pickle.dumps(typed)) + assert typed2.name == typed.name + + # Unpickled op must actually run, not just carry the right name. + bound = op2[int](7) + A = Matrix.from_coo([0, 1], [0, 1], [3, 11]) + B = Matrix.from_coo([0, 1], [0, 1], [5, 1]) + C = A.ewise_mult(B, bound).new() + # (3+5)+7 = 15; (11+1)+7 = 19 + assert C.isequal(Matrix.from_coo([0, 1], [0, 1], [15, 19])) + finally: + delattr(indexbinary, "pickle_test_op") + + +def test_pickle_bound(): + indexbinary.register_new("pickle_bound_test_op", _ibo_add_theta) + try: + bound = indexbinary.pickle_bound_test_op[int](7) + bound2 = pickle.loads(pickle.dumps(bound)) + assert bound2._theta == 7 + A = Matrix.from_coo([0, 1], [0, 1], [3, 11]) + B = Matrix.from_coo([0, 1], [0, 1], [5, 1]) + C1 = A.ewise_mult(B, bound).new() + C2 = A.ewise_mult(B, bound2).new() + assert C1.isequal(C2) + finally: + delattr(indexbinary, "pickle_bound_test_op") + + +def test_pickle_bound_array_udt_theta(): + """Bound IBO with an array-UDT theta value round-trips through pickle. + + Regression: ``_BoundIndexBinaryOp.__reduce__`` stored ``_theta`` as the + raw numpy array. On unpickle, ``_rebind_indexbinaryop`` passed it back to + ``TypedBuiltinIndexBinaryOp.__call__``, which called + ``Scalar.from_value(theta)`` with no dtype; that crashed inside the + multi-dim ndarray ``Scalar.value`` setter. The fix wraps the value as a + typed Scalar in ``_rebind_indexbinaryop`` when the thunk type is a UDT. + """ + from graphblas import dtypes + + arr_udt = dtypes.register_anonymous(np.dtype((np.int64, (3,))), "_BoundIboArr3") + indexbinary.register_new("pickle_bound_array_udt_op", _ibo_return_x, is_udt=True) + try: + op = indexbinary.pickle_bound_array_udt_op + theta_scalar = gb.Scalar(arr_udt) + theta_scalar.value = np.array([1, 2, 3], dtype=np.int64) + bound = op[arr_udt](theta_scalar) + bound2 = pickle.loads(pickle.dumps(bound)) + # ``_theta`` round-trips bit-identically. + assert np.array_equal(bound2._theta, np.array([1, 2, 3], dtype=np.int64)) + # The unpickled bound op is a fresh ``_BoundIndexBinaryOp`` wrapping a + # different ``GrB_BinaryOp`` pointing at the same kernel; pin the + # introspection state that previously crashed. + assert bound2.parent is bound.parent + assert bound2.type is arr_udt + finally: + delattr(indexbinary, "pickle_bound_array_udt_op") + + +def test_bind_raw_array_udt_theta(): + """A raw (non-Scalar) array-UDT theta is bound using the known thunk dtype. + + Regression: the ``not isinstance(theta, Scalar)`` branch of the IBO + ``__call__`` passed the raw value to ``Scalar.from_value`` with no dtype, + which inferred a scalar element type and crashed on the multi-dim ndarray. + The typed op now supplies its thunk dtype; the untyped op uses the explicit + ``dtype=`` argument. + """ + from graphblas import dtypes + + arr_udt = dtypes.register_anonymous(np.dtype((np.int64, (3,))), "_RawThetaArr3") + indexbinary.register_new("raw_array_udt_op", _ibo_return_x, is_udt=True) + try: + op = indexbinary.raw_array_udt_op + raw = np.array([1, 2, 3], dtype=np.int64) + # Typed op: ``op[arr_udt]`` knows the thunk dtype. + bound = op[arr_udt](raw) + assert bound.type is arr_udt + assert np.array_equal(bound._theta, raw) + # Untyped op: the explicit ``dtype=`` supplies the UDT. + bound2 = op(raw, dtype=arr_udt) + assert bound2.type is arr_udt + assert np.array_equal(bound2._theta, raw) + finally: + delattr(indexbinary, "raw_array_udt_op") + + +def test_array_udt_udf_shape_is_checked(): + """An IBO whose UDF builds a wrong-shape array is rejected when typed. + + The same guard the binary and unary paths get. Numba's ``Array`` type + carries ``ndim`` but not extents, so the mismatch is not a type error; left + to the wrapper it raises inside the cfunc, where Numba prints the traceback + and returns, leaving the element as SuiteSparse found it. + """ + # A shape of its own: ``register_anonymous`` caches by numpy dtype, so + # sharing one with another test would hand back that test's DataType. + arr_udt = dtypes.register_anonymous(np.dtype((np.float64, (12,))), "_IboShapeArr12") + + def _truncate(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return (x + y)[:2] + + op = indexbinary.register_anonymous(_truncate, is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op[arr_udt] + + def _combine(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + good = indexbinary.register_anonymous(_combine, is_udt=True) + assert good[arr_udt].return_type is arr_udt + + +def test_empty_scalar_theta(): + # SuiteSparse would read the empty Scalar's value: a crash, or a value it never held + op = indexbinary.register_anonymous(_ibo_add_theta) + opu = indexbinary.register_anonymous(_ibo_add_theta, is_udt=True) + cleared = Scalar.from_value(5) + cleared.clear() + for theta in [Scalar(int), Scalar(int, is_cscalar=True), cleared]: + for bind in [op, opu, op[int], opu[int], lambda t: op(t, dtype=int)]: + with pytest.raises(EmptyObject, match="theta is an empty Scalar"): + bind(theta) + + +def test_bind_raw_udt_theta_without_dtype(): + """A numpy array theta is typed as its array layout, as other array literals are. + + With ``is_udt=True``, theta keeps that type, as any theta does, so it is + never converted to y's type. A tuple has no type of its own: give it as a + Scalar of its UDT, or name the op's type. An op without ``is_udt`` does + not take a UDT theta. + """ + + def theta_only(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return theta + + arr_udt = dtypes.register_anonymous(np.dtype((np.int32, (3,))), "_RawThetaArr3i32") + indexbinary.register_new("raw_no_dtype_op", theta_only, is_udt=True) + try: + op = indexbinary.raw_no_dtype_op + v = Vector.from_coo([0], np.array([[1, 2, 3]], dtype=np.int32), dtype=arr_udt) + for theta in [np.array([10, 20, 30], np.int32), Scalar.from_value((10, 20, 30), arr_udt)]: + result = v.ewise_mult(v, op(theta)).new() + assert result.dtype == arr_udt + assert result[0].new().value.tolist() == [10, 20, 30] + # An int64 or float64 array keeps its type, where it would have become int32 + result = v.ewise_mult(v, op(np.array([10, 20, 2**40]))).new() + assert result.dtype.np_type == np.dtype((np.int64, (3,))) + assert result[0].new().value.tolist() == [10, 20, 2**40] + result = v.ewise_mult(v, op(np.array([10, 20.5, 30]))).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + assert result[0].new().value.tolist() == [10, 20.5, 30] + with pytest.raises(TypeError, match=r"Cannot infer a dtype for theta from \(10, 20, 30\)"): + op((10, 20, 30)) + with pytest.raises(TypeError, match="Cannot infer a dtype for theta"): + op(np.array([True, False])) + # Its array UDT is shared, under the name it got first (another test's) + with pytest.raises(KeyError, match="_ibo_add_theta does not work with "): + indexbinary.register_anonymous(_ibo_add_theta)(np.array([1, 2, 3])) + finally: + delattr(indexbinary, "raw_no_dtype_op") + + +def test_untyped_bind_theta_keeps_its_type_with_is_udt(): + """With ``is_udt=True``, ``op(theta)`` compiles for x, y and theta as they are typed. + + SuiteSparse gives theta a type of its own (``GxB_IndexBinaryOp_new`` takes + ``theta_type``), so theta is not converted: 0.5 stays 0.5 beside INT64 + inputs, and Numba types the result as it would for those three types. + Before, theta took y's type when the op ran. Naming two types, + ``op[X, Y](theta)``, still gives theta y's type and casts it, and + ``op[X, Y, T](theta)`` names all three. + """ + op = indexbinary.register_anonymous(_ibo_add_theta, is_udt=True) + v = Vector.from_coo([0, 1], [1, 2]) + expected = Vector.from_coo([0, 1], [4, 6]) + for theta in [2, np.int8(2), Scalar.from_value(2, "INT16")]: + assert v.ewise_mult(v, op(theta)).new().isequal(expected, check_dtype=True) + for theta in [2.0, np.float64(2), Scalar.from_value(2.0)]: + result = v.ewise_mult(v, op(theta)).new() + assert result.isequal(expected.dup(float), check_dtype=True) + result = v.ewise_mult(v, op(0.5)).new() + assert result.isequal(Vector.from_coo([0, 1], [2.5, 4.5]), check_dtype=True) + i8 = Vector.from_coo([0], [1], dtype="INT8") + result = i8.ewise_mult(i8, op(300)).new() + assert result.isequal(Vector.from_coo([0], [302]), check_dtype=True) + assert {key for key in op._udt_types if key[0] == dtypes.INT64} == { + (dtypes.INT64, dtypes.INT64, dtypes.INT64), + (dtypes.INT64, dtypes.INT64, dtypes.INT8), + (dtypes.INT64, dtypes.INT64, dtypes.INT16), + (dtypes.INT64, dtypes.INT64, dtypes.FP64), + } + # Named types + result = v.ewise_mult(v, op["INT64", "INT64"](0.5)).new() + assert result.isequal(Vector.from_coo([0, 1], [2, 4]), check_dtype=True) + typed = op["INT64", "INT64", "FP64"] + assert typed is op._udt_ops[dtypes.INT64, dtypes.INT64, dtypes.FP64] + assert typed.thunk_type == dtypes.FP64 + assert repr(typed) == "indexbinary._ibo_add_theta[INT64, INT64, FP64]" + assert repr(op["INT64", "INT64"]) == "indexbinary._ibo_add_theta[INT64, INT64]" + result = v.ewise_mult(v, typed(0.5)).new() + assert result.isequal(Vector.from_coo([0, 1], [2.5, 4.5]), check_dtype=True) + # Pickle keeps all three types, for the typed op and for the op it binds + typed2 = pickle.loads(pickle.dumps(typed)) # re-registers the anonymous op + assert repr(typed2) == repr(typed) + assert v.ewise_mult(v, typed2(0.5)).new().isequal(result, check_dtype=True) + bound = v.ewise_mult(v, op(0.5)).op + assert bound.type == bound._type2 == dtypes.INT64 + bound2 = pickle.loads(pickle.dumps(bound)) + assert bound2._theta_type == dtypes.FP64 + assert v.ewise_mult(v, bound2).new().isequal(result, check_dtype=True) + bound2 = pickle.loads(pickle.dumps(op["INT64", "FP64"](1))) # was op["INT64"](1) + assert (bound2.type, bound2._type2, bound2._theta_type) == ( + dtypes.INT64, + dtypes.FP64, + dtypes.FP64, + ) + # Without is_udt, the three types are one, as when the op runs + wsum = indexbinary.register_anonymous(_ibo_add_theta) + assert wsum["INT8", "INT8", "INT64"] is wsum["INT64"] + rec_udt = dtypes.register_anonymous(np.dtype([("theta_b", np.int64)]), "_UntypedThetaRec2") + with pytest.raises(KeyError, match="does not work with .INT64, INT64, _UntypedThetaRec2"): + wsum["INT64", "INT64", rec_udt] + + +def test_pickle_bound_record_udt_theta(): + """Bound IBO with a record-UDT theta value round-trips through pickle. + + Companion to :func:`test_pickle_bound_array_udt_theta` for the record + flavor of UDT; the array case was the one that crashed before the + ``_rebind_indexbinaryop`` Scalar-wrap fix, and the record case worked + only because ``_theta`` happened to be a 0-d ``np.void`` that + ``Scalar.value =`` accepted. Pin the contract for both shapes so a + regression in either is loud. + """ + from graphblas import dtypes + + rec_udt = dtypes.register_anonymous( + np.dtype([("a", np.int64), ("b", np.int64)], align=True), + "_BoundIboRec", + ) + indexbinary.register_new("pickle_bound_record_udt_op", _ibo_return_x, is_udt=True) + try: + op = indexbinary.pickle_bound_record_udt_op + theta_scalar = gb.Scalar(rec_udt) + theta_scalar.value = (5, 7) + bound = op[rec_udt](theta_scalar) + bound2 = pickle.loads(pickle.dumps(bound)) + assert tuple(bound2._theta) == (5, 7) + assert bound2.parent is bound.parent + assert bound2.type is rec_udt + finally: + delattr(indexbinary, "pickle_bound_record_udt_op") + + +@pytest.mark.slow +def test_bound_ibo_memory_bounded(): + """Regression: ``iop(theta)`` used to leak one ``GrB_BinaryOp`` per call. + + Bound IBOs are never cached, so a loop that re-binds with different + theta values leaks unboundedly without the ``TypedOpBase.__del__`` free. + 50k bound IBOs are well above the noise floor for the leak (each is + several hundred bytes of SS state plus a python-level wrapper); RSS + growth should stay below a few MB after GC drops them. + """ + import gc + import sys + + pytest.importorskip("resource") + import resource + + if hasattr(indexbinary, "leak_test_bound_ibo_op"): + delattr(indexbinary, "leak_test_bound_ibo_op") + indexbinary.register_new("leak_test_bound_ibo_op", _ibo_add_theta) + op = indexbinary.leak_test_bound_ibo_op + try: + # ``ru_maxrss`` is bytes on macOS, KB on Linux. + scale = 1024 if sys.platform.startswith("linux") else 1024 * 1024 + + def rss_mb(): + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / scale + + # Warm up to populate typed-op cache; first call also touches the + # numba sample-types path. + for _ in range(100): + _ = op[int](0) + gc.collect() + before = rss_mb() + for i in range(50_000): + _ = op[int](i) + gc.collect() + after = rss_mb() + # 50k leaked GrB_BinaryOps were on the order of tens of MB before the + # fix; under the fix we expect single-digit MB. + assert ( + after - before + ) < 25.0, f"bound-IBO leak: RSS grew {after - before:.1f} MB over 50k binds" + finally: + delattr(indexbinary, "leak_test_bound_ibo_op") + + +def test_bound_ibo_jit_introspection_returns_none(): + """Bound IBOs are user-defined and don't go through the built-in JIT C + codegen path, so the introspection properties must report ``None`` + without crashing. Regression for an earlier construction path that + used ``__new__`` directly and skipped the inherited ``_jit_c_info`` + slot's default-None initialization. Also checks the unpickled instance, + which goes through a different construction path (``_rebind_indexbinaryop`` + -> ``TypedBuiltinIndexBinaryOp.__call__``). + """ + indexbinary.register_new("_jit_introspect_test_op", _ibo_add_theta) + try: + op = indexbinary._jit_introspect_test_op + bound = op[int](3) + # User IBOs don't go through the built-in JIT C codegen path, so these + # are expected to be ``None``; the introspection properties must not + # crash either way. + assert bound.jit_c_name is None + assert bound.jit_c_source is None + bound2 = pickle.loads(pickle.dumps(bound)) + assert bound2.jit_c_name is None + assert bound2.jit_c_source is None + finally: + delattr(indexbinary, "_jit_introspect_test_op") + + +def test_bad_udf(): + with pytest.raises(UdfParseError, match="Unable to parse function using Numba"): + indexbinary.register_anonymous(lambda x, ix, jx, y, iy, jy, theta: result) # noqa: F821 + + +def test_bad_type(): + with pytest.raises(TypeError, match="UDF argument must be a function"): + indexbinary.register_anonymous(42) + + +def test_type_the_op_does_not_take(): + """``R in op`` is False and ``op[R]`` raises KeyError, as for an IndexUnaryOp. + + Both raised AttributeError: IndexBinaryOp had no ``is_positional``, which + ``op[R]`` reads for a type the op does not have. + """ + + def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_vals) + rec = dtypes.register_anonymous( + np.dtype([("ibt_a", np.float64), ("ibt_b", np.int64)]), "_IbtRec" + ) + assert dtypes.INT64 in op + assert rec not in op + for expr in [lambda: op[rec], lambda: op(Scalar.from_value((1.0, 2), rec), dtype=rec)]: + with pytest.raises(KeyError, match="does not work with _IbtRec"): + expr() + + +def test_with_mask(): + def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_vals) + A = Matrix.from_coo([0, 1, 2], [0, 1, 2], [3, 7, 11]) + B = Matrix.from_coo([0, 1, 2], [0, 1, 2], [5, 2, 1]) + mask = Matrix.from_coo([0, 2], [0, 2], [True, True], nrows=3, ncols=3) + C = Matrix(int, nrows=3, ncols=3) + binop = op(10) + C(mask=mask.S) << A.ewise_mult(B, binop) + rows, _, vals = C.to_coo() + assert list(rows) == [0, 2] + assert list(vals) == [18, 22] # 3+5+10, 11+1+10 + + +def test_with_accumulator(): + def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_vals) + A = Matrix.from_coo([0], [0], [3]) + B = Matrix.from_coo([0], [0], [5]) + C = Matrix.from_coo([0], [0], [100]) + binop = op(10) + C(accum=gb.binary.plus) << A.ewise_mult(B, binop) + assert C.to_coo()[2][0] == 118 # 100 + (3+5+10) + + +def test_reduce_refuses_a_bound_op(): + """A bound IndexBinaryOp gets the TypeError reduce gives a BinaryOp that is in no Monoid. + + It raised AttributeError: neither the bound op nor its IndexBinaryOp had + ``monoid``, which reduce and its error message read. + """ + + def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_vals) + v = Vector.from_coo([0, 1], [1, 2]) + for binop in [op(1), op[int](1)]: + with pytest.raises(TypeError, match="is not known to be part of a Monoid"): + v.reduce(binop) + + +def test_ewise_with_bound_binop(): + """Confirm bound IndexBinaryOp works in all ewise operations.""" + + def mul_plus(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x * y + theta + + op = indexbinary.register_anonymous(mul_plus) + binop = op[int](0) + A = Matrix.from_coo([0, 0], [0, 1], [2, 3]) + B = Matrix.from_coo([0, 0], [0, 1], [4, 5]) + C = A.ewise_mult(B, binop).new() + assert list(C.to_coo()[2]) == [8, 15] # 2*4+0, 3*5+0 + + +def test_find_opclass(): + from graphblas.core.operator import find_opclass + + def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + + op = indexbinary.register_anonymous(add_vals) + _, opclass = find_opclass(op) + assert opclass == "IndexBinaryOp" + + typed = op[int] + assert typed.opclass == "IndexBinaryOp" + + bound = typed(0) + assert bound.opclass == "BinaryOp" + assert find_opclass(op(0))[1] == "BinaryOp" + + +def test_udt_tuple_return(): + """An IndexBinaryOp UDF can return a tuple matching the output UDT's fields.""" + record_dtype = np.dtype([("a", np.int64), ("b", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype) + + v = Vector(udt, size=2) + v[0] = (1, 2.0) + v[1] = (3, 4.0) + w = Vector(udt, size=2) + w[0] = (10, 20.0) + w[1] = (30, 40.0) + + theta = Scalar.from_value(np.array((0, 0.0), dtype=record_dtype)) + + def _add_with_idx(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return (x["a"] + y["a"] + theta["a"], x["b"] + y["b"] + ix) + + op = indexbinary.register_anonymous(_add_with_idx, is_udt=True) + binop = op[udt](theta) + result = v.ewise_mult(w, binop).new() + assert result.dtype == udt + expected = Vector(udt, size=2) + expected[0] = (11, 22.0) # a: 1+10+0, b: 2+20+0 + expected[1] = (33, 45.0) # a: 3+30+0, b: 4+40+1 + assert result.isequal(expected) + + # UDT input, scalar output (no tuple unpacking needed) + def _sum_ab(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x["a"] + y["b"] + theta["a"] + + theta2 = Scalar.from_value(np.array((100, 0.0), dtype=record_dtype)) + op2 = indexbinary.register_anonymous(_sum_ab, is_udt=True) + binop2 = op2[udt](theta2) + result2 = v.ewise_mult(w, binop2).new() + assert result2[0].new() == 121.0 # 1 + 20.0 + 100 + assert result2[1].new() == 143.0 # 3 + 40.0 + 100 + + +def test_udt_matrix_ewise_with_bound_ibo(): + """End-to-end: register an IBO on a record UDT, bind a theta, and use it + as the binary op of ``Matrix.ewise_mult`` and ``Matrix.ewise_add``. + + The matrix path is more interesting than the vector path because the + IBO receives both row and column indices for each operand. This also + exercises the full ``is_udt`` codegen with two-field-tuple returns on + Matrix-shaped inputs. + """ + record_dtype = np.dtype([("a", np.int64), ("b", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "ibo_mxm_udt") + + def weighted_add(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + # Field a: x.a * theta.a + y.a; Field b: x.b + y.b + ix + jy + return (x["a"] * theta["a"] + y["a"], x["b"] + y["b"] + ix + jy) + + ibo = indexbinary.register_anonymous(weighted_add, is_udt=True) + theta = Scalar.from_value(np.array((10, 0), dtype=record_dtype)) + binop = ibo[udt](theta) + assert binop.opclass == "BinaryOp" + + A = Matrix(udt, nrows=3, ncols=3) + A[0, 0] = (1, 2) + A[1, 1] = (3, 4) + A[2, 0] = (5, 6) + B = Matrix(udt, nrows=3, ncols=3) + B[0, 0] = (1, 1) + B[1, 1] = (2, 2) + B[2, 0] = (3, 3) + + C = A.ewise_mult(B, binop).new() + assert C.dtype == udt + expected = Matrix(udt, nrows=3, ncols=3) + # (0,0): a = 1*10 + 1 = 11, b = 2 + 1 + 0 + 0 = 3 + # (1,1): a = 3*10 + 2 = 32, b = 4 + 2 + 1 + 1 = 8 + # (2,0): a = 5*10 + 3 = 53, b = 6 + 3 + 2 + 0 = 11 + expected[0, 0] = (11, 3) + expected[1, 1] = (32, 8) + expected[2, 0] = (53, 11) + assert C.isequal(expected) + + # ewise_add path with the same bound IBO. ewise_add applies the op only where + # both sides are present; pattern is the same as ewise_mult here. + D = A.ewise_add(B, binop).new() + assert D.isequal(expected) + + +def _semiring_with_bound_ibo_factory(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + # Module-level so pickle can find it; used by both the semiring and the + # cross-pickle test below. + return x * y + theta * (ix + jy) + + +def test_semiring_from_bound_ibo_mxm(): + """A bound IndexBinaryOp may serve as the multiplier of a Semiring. + + Per SS:GraphBLAS, ``GxB_BinaryOp_new_IndexOp`` turns an IBO + theta into + a regular ``GrB_BinaryOp``, which any Semiring can use as its multiplier. + Verify the end-to-end path: build a Semiring with ``monoid.plus`` plus a + bound IBO, and use it as ``A.mxm(B, sr)``. + """ + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + ibo = indexbinary.register_anonymous(_semiring_with_bound_ibo_factory) + bound = ibo[dtypes.INT64](theta=2) + sr = semiring_module.register_anonymous(monoid_module.plus, bound) + + # Diagonal-ish A and B so we can compute the result by hand. + A = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [1, 2, 3, 4], dtype=dtypes.INT64) + B = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [5, 6, 7, 8], dtype=dtypes.INT64) + C = A.mxm(B, sr).new() + + # mxm(A, B)[i, j] = sum_k f(A[i,k], i, k, B[k,j], k, j, theta=2) + # = sum_k A[i,k]*B[k,j] + 2*(i + j) + # For [[1, 2], [3, 4]] @ [[5, 6], [7, 8]] = [[19, 22], [43, 50]] + # plus 2*(i + j) per inner-product term (two terms per cell): + # C[0, 0] = 5 + 14 + 2*(0+0)*2 = 19 (theta term: 0 per k) + # C[0, 1] = 6 + 16 + 2*(0+1)*2 = 22 + 4 = 26 + # C[1, 0] = 15 + 28 + 2*(1+0)*2 = 43 + 4 = 47 + # C[1, 1] = 18 + 32 + 2*(1+1)*2 = 50 + 8 = 58 + assert C.dtype == dtypes.INT64 + expected = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [19, 26, 47, 58], dtype=dtypes.INT64) + assert C.isequal(expected) + + +def test_semiring_from_bound_ibo_pickle(): + """A Semiring built from a bound IBO survives pickle round-trip. + + The bound IBO already pickles (via its module-level parent op); the + Semiring inherits the same path. + """ + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + indexbinary.register_new("semiring_pickle_ibo", _semiring_with_bound_ibo_factory) + bound = indexbinary.semiring_pickle_ibo[dtypes.INT64](theta=2) + sr = semiring_module.register_anonymous(monoid_module.plus, bound) + sr2 = pickle.loads(pickle.dumps(sr)) + + A = Matrix.from_coo([0, 1], [0, 1], [1, 3], dtype=dtypes.INT64) + B = Matrix.from_coo([0, 1], [0, 1], [5, 7], dtype=dtypes.INT64) + assert A.mxm(B, sr).new().isequal(A.mxm(B, sr2).new()) + + +def test_semiring_from_bound_ibo_type_mismatch_errors(): + """Monoid must accept the bound IBO's return type.""" + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + ibo = indexbinary.register_anonymous(_semiring_with_bound_ibo_factory) + bound = ibo[dtypes.INT64](theta=0) + with pytest.raises(TypeError, match="band.*INT64"): + # ``band`` is bitwise-AND on unsigned ints; INT64 is not in its types. + semiring_module.register_anonymous(monoid_module.band, bound) + + +def test_semiring_register_anonymous_rejects_bad_binaryop(): + """The error message points users at the bound-IBO syntax.""" + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + with pytest.raises(TypeError, match=r"bound IndexBinaryOp.*ibo\[dtype\]\(theta\)"): + semiring_module.register_anonymous(monoid_module.plus, 42) + + +def test_semiring_from_bound_ibo_over_udt_mxm(): + """A bound IBO over a UDT also composes into a Semiring via the + auto-lifted ``monoid.plus[udt]`` path. + """ + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + record_dtype = np.dtype([("a", np.int64), ("b", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "ibo_mxm_sr_udt") + + def weighted(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return (x["a"] * y["a"] + theta["a"], x["b"] + y["b"]) + + ibo = indexbinary.register_anonymous(weighted, is_udt=True) + theta = Scalar.from_value(np.array((1, 0), dtype=record_dtype)) + bound = ibo[udt](theta) + + # ``monoid.plus[udt]`` is auto-lifted (field-wise add); the semiring + # composes the bound IBO multiplier with that monoid. + sr = semiring_module.register_anonymous(monoid_module.plus, bound) + + A = Matrix(udt, nrows=2, ncols=2) + A[0, 0] = (2, 3) + A[0, 1] = (1, 1) + A[1, 1] = (4, 5) + B = Matrix(udt, nrows=2, ncols=2) + B[0, 0] = (5, 1) + B[1, 1] = (6, 2) + + C = A.mxm(B, sr).new() + # mxm(A, B)[i, j] = reduce_plus_k(weighted(A[i, k], i, k, B[k, j], k, j, theta=(1, 0))) + # weighted has fields a = A.a * B.a + theta.a (=1), b = A.b + B.b + # C[0, 0]: only k=0 contributes -> (2*5 + 1, 3 + 1) = (11, 4) + # C[0, 1]: only k=1 contributes -> (1*6 + 1, 1 + 2) = (7, 3) + # C[1, 1]: only k=1 contributes -> (4*6 + 1, 5 + 2) = (25, 7) + assert C.dtype == udt + assert C[0, 0].new() == (11, 4) + assert C[0, 1].new() == (7, 3) + assert C[1, 1].new() == (25, 7) + + +def _ibo_decay(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + d = ix - jx if ix >= jx else jx - ix + return x * y * theta**d + + +def test_semiring_from_untyped_bound_ibo(): + """A Semiring from ``ibo(theta)`` is typed when it runs, by the inputs and theta. + + Before, ``ibo(0.5)`` was an FP64 op, so its semiring refused INT64 + matrices with KeyError, and ``ibo(1)`` was INT64 and refused FP64 ones. + """ + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + decay = indexbinary.register_anonymous(_ibo_decay) + A = Matrix.from_coo([0, 0, 1, 2], [0, 2, 1, 0], [1, 1, 1, 1], nrows=3, ncols=3) + v = Vector.from_coo([0, 1, 2], [4, 4, 4]) + half = semiring_module.register_anonymous(monoid_module.plus, decay(0.5)) + one = semiring_module.register_anonymous(monoid_module.plus, decay(1)) + expected = Vector.from_coo([0, 1, 2], [5.0, 4.0, 1.0]) + assert A.mxv(v, half).new().isequal(expected, check_dtype=True) + assert A.dup("FP32").mxv(v.dup("FP32"), half).new().isequal(expected, check_dtype=True) + assert A.mxv(v, one).new().isequal(Vector.from_coo([0, 1, 2], [8, 4, 4]), check_dtype=True) + result = A.dup(float).mxv(v.dup(float), one).new() + assert result.isequal(Vector.from_coo([0, 1, 2], [8.0, 4.0, 4.0]), check_dtype=True) + sr2 = pickle.loads(pickle.dumps(half)) + assert A.mxv(v, sr2).new().isequal(expected, check_dtype=True) + # Theta keeps its type through pickle: FP32 data stays FP32 with an FP32 theta + bound2 = pickle.loads(pickle.dumps(decay(np.float32(0.5)))) + sr2 = semiring_module.register_anonymous(monoid_module.plus, bound2) + result = A.dup("FP32").mxv(v.dup("FP32"), sr2).new() + assert result.isequal(expected.dup("FP32"), check_dtype=True) + + +def test_argmax_semiring_from_untyped_bound_ibo(): + """The SuiteSparse user guide's row argmax, with theta unused, bound with ``iop()``. + + The multiplier turns ``A(i, k)`` into the pair ``(k, A(i, k))`` and a max + monoid on pairs keeps the larger value (the smaller column on a tie). The + guide binds a BOOL theta because theta cannot be NULL. Here the matrix + holds pairs, so the UDF's tuple return takes their type. Before, ``iop()`` + and ``iop(0)`` were ops on BOOL and INT64 values, not the matrix's pairs. + Theta keeps its own type, so each bind compiles the op for it. + """ + from graphblas import binary + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + pair = dtypes.register_anonymous(np.dtype([("k", np.int64), ("v", np.float64)]), "IboPair") + + def make_pair(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return (np.int64(jx), x["v"]) + + def max_pair(a, b): # pragma: no cover (numba) + if a[1] > b[1] or (a[1] == b[1] and a[0] < b[0]): + return a + return b + + iop = indexbinary.register_anonymous(make_pair, is_udt=True) + max_monoid = monoid_module.register_anonymous( + binary.register_anonymous(max_pair, is_udt=True), (np.iinfo(np.int64).max, -np.inf) + ) + values = np.array([(0, 3.0), (0, 7.0), (0, 7.0), (0, 5.0), (0, -1.0)], dtype=pair.np_type) + A = Matrix.from_coo([0, 0, 0, 1, 1], [0, 1, 2, 0, 2], values, dtype=pair) + ones = Vector.from_scalar(True, A.ncols) + for binop in [iop(), iop(False), iop(0), iop(0.0)]: + argmax = semiring_module.register_anonymous(max_monoid, binop) + c = A.mxv(ones, argmax).new() + assert c.dtype == pair + assert c[0].new() == (1, 7.0) + assert c[1].new() == (0, 5.0) + assert set(iop._udt_types) == { + (pair, dtypes.BOOL, dtypes.BOOL), + (pair, dtypes.BOOL, dtypes.INT64), + (pair, dtypes.BOOL, dtypes.FP64), + } + + +def test_dir_and_module(): + assert "register_new" in dir(indexbinary) + assert "register_anonymous" in dir(indexbinary) + assert "ss" in dir(indexbinary) + # Actually access the ss module to verify it exists (not just in __dir__) + ss = indexbinary.ss + assert hasattr(ss, "_delayed") + assert hasattr(ss, "register_new") diff --git a/graphblas/tests/test_infix.py b/graphblas/tests/test_infix.py index 601f282a7..75150d7d0 100644 --- a/graphblas/tests/test_infix.py +++ b/graphblas/tests/test_infix.py @@ -1,6 +1,11 @@ +import operator + +import numpy as np import pytest -from graphblas import binary, monoid, op +from graphblas import binary, dtypes, monoid, op, unary +from graphblas.core import _supports_udfs as supports_udfs +from graphblas.core.utils import output_type from graphblas.exceptions import DimensionMismatch from .conftest import autocompute @@ -119,6 +124,11 @@ def test_bad_ewise(s1, v1, A1, A2): (s1, A1.T), (A1, 1), (1, A1), + # A Scalar takes a literal, but not just any object, as `s1 + object()` does not + (s1, object()), + (object(), s1), + (s1, "x"), + (None, s1), ]: with pytest.raises(TypeError, match="Bad type for argument"): left | right @@ -163,6 +173,13 @@ def test_bad_ewise(s1, v1, A1, A2): A1.T | A1 with pytest.raises(DimensionMismatch): A1.T & A1 + # Beside an infix expression, on either side + for left, right in [(v1 | v1, w), (w, v1 | v1), (A2 | A2, A1), (A1, A2 | A2)]: + with pytest.raises(DimensionMismatch): + left | right + for left, right in [(v1 & v1, w), (w, v1 & v1), (A2 & A2, A1), (A1, A2 & A2)]: + with pytest.raises(DimensionMismatch): + left & right # These are okay now # with pytest.raises(TypeError): @@ -220,6 +237,15 @@ def test_bad_matmul(s1, v1, A1, A2): A1.T @ A1.T with pytest.raises(DimensionMismatch): A1 @= A1 + # Beside an infix expression, on either side + for left, right in [ + (A1 @ A2, v1), + (v1, A1 @ A2), + (A2 @ A1, A2.T), + (A2, A2 @ A1), + ]: + with pytest.raises(DimensionMismatch): + left @ right with pytest.raises(TypeError): s1 @ 1 with pytest.raises(TypeError): @@ -454,34 +480,34 @@ def test_multi_infix_vector(): result = op.plus_plus(D0 @ (D0 @ v1)).new() assert result.isequal(v1) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | v3 - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2).__ror__(v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1 | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1.__ror__(v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1 & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1.__rand__(v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & v3 - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2).__rand__(v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & (v2 & v3) # We differentiate between infix and methods @@ -497,6 +523,8 @@ def test_multi_infix_vector(): v1.ewise_mult(v2 | v3) with pytest.raises(TypeError, match="Automatic computation"): (v1 | v2).ewise_mult(v3) + with pytest.raises(TypeError, match="Automatic computation"): + (v1 | v2) + v3 @autocompute @@ -564,34 +592,34 @@ def test_multi_infix_matrix(): assert op.plus_plus((v1.T @ D0) @ v1).new()[0, 0].new().value == 6 assert op.plus_plus(D0 @ D0 @ D0 @ D0 @ D0).new().isequal(D0) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | v3 - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2).__ror__(v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1 | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1.__ror__(v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1 & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1.__rand__(v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & v3 - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2).__rand__(v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & (v2 & v3) # We differentiate between infix and methods @@ -716,34 +744,34 @@ def test_multi_infix_scalar(): result = binary.plus(v1 | (2 | v3), left_default=10, right_default=10).new() assert result.isequal(expected) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | v3 - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2).__ror__(v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) | (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1 | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1.__ror__(v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) | (v2 & v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1 & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): v1.__rand__(v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 & v2) & (v2 | v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & v3 - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2).__rand__(v3) - with pytest.raises(TypeError, match="XXX"): # TODO + with pytest.raises(TypeError, match="Cannot mix"): (v1 | v2) & (v2 & v3) # We differentiate between infix and methods @@ -761,6 +789,183 @@ def test_multi_infix_scalar(): (v1 | v2).ewise_mult(v3) +def test_scalar_infix_literal_keeps_its_side(): + # `op(2 & s)` applies op to 2 and s in that order, as `op(2, s)` does + s = Scalar.from_value(5, dtype=dtypes.INT8) + expr = 2 & s + assert expr.left == 2 + assert expr.right is s + for binop, expected, dtype in [ + (binary.minus, -3, dtypes.INT8), + (binary.rminus, 3, dtypes.INT8), + (binary.cdiv, 0, dtypes.INT8), + (binary.rdiv, 2, dtypes.INT8), + (binary.truediv, 0.4, dtypes.FP64), + (binary.first, 2, dtypes.INT8), + (binary.second, 5, dtypes.INT8), + (binary.pow, 32, dtypes.INT8), + (binary.lt, True, dtypes.BOOL), + (binary.gt, False, dtypes.BOOL), + (binary.isgt, 0, dtypes.INT8), + ]: + for expr in [2 & s, 2 | s]: + result = binop(expr).new() + assert result.value == expected + assert result.dtype == dtype + assert result.isequal(binop(2, s).new(), check_dtype=True) + # Each default goes with its side + empty = Scalar(dtypes.INT8) + assert binary.minus(2 | empty, left_default=0, right_default=10).new().value == -8 + assert binary.minus(empty | 2, left_default=0, right_default=10).new().value == -2 + assert binary.minus(2 | s, left_default=0, right_default=10).new().value == -3 + # And so does each literal in a chain + assert binary.minus(1 & (2 & s)).new().value == 4 + assert binary.minus((2 & s) & 1).new().value == -4 + assert binary.minus(1 | (2 | s)).new().value == 4 + + # A Python literal is typed beside the Scalar as by `-` (weak), once the op is known + f32 = Scalar.from_value(2.5, dtype=dtypes.FP32) + t = Scalar.from_value(True) + for lit, x in [(2, s), (0.5, s), (0.5, f32), (2, f32), (2, t)]: + for expr in [lit & x, lit | x]: + assert binary.minus(expr).new().isequal((lit - x).new(), check_dtype=True) + assert binary.minus(0.5 & s).new().dtype == dtypes.FP64 + assert binary.minus(0.5 & f32).new().dtype == dtypes.FP32 + assert binary.first(False & t).new().dtype == dtypes.BOOL + assert binary.lt(False | t).new().value is True + # 300 does not fit INT8, which `-` reports, and `<` compares it as it is + with pytest.raises(OverflowError): + binary.minus(300 & s) + assert binary.lt(300 & s).new().value is False + assert binary.gt(300 | s).new().value is True + + # BOOL `&` and `|` with a literal compute themselves, by LAND and LOR + for expr, expected in [ + (t & True, True), + (False & t, False), + (t | False, True), + (False | t, True), + ]: + result = expr.new() + assert result.dtype == dtypes.BOOL + assert result.value is expected + with pytest.raises(TypeError, match="The argument dtypes are INT64 and BOOL"): + (1 & t).new() + with pytest.raises(TypeError, match="only valid for BOOL"): + (s | 2).new() + + +@pytest.mark.skipif("not supports_udfs") +def test_scalar_infix_literal_keeps_its_side_udt(): + # The op types a literal beside a UDT, on the left as on the right: a lifted op as numpy + # would, and `first` as an element of the UDT, which 0.5 is not + udt = dtypes.register_anonymous(np.dtype((np.int8, (2,))), "_InfixLiteralI8") + s = Scalar.from_value([5, 7], udt) + result = binary.minus((1, 2) & s).new() + assert result.dtype == udt + np.testing.assert_array_equal(result.value, [-4, -5]) + result = binary.minus(0.5 | s).new() + assert result.dtype.np_type == np.dtype((np.float64, (2,))) + np.testing.assert_array_equal(result.value, [-4.5, -6.5]) + np.testing.assert_array_equal(binary.first([1, 2] & s).new().value, [1, 2]) + with pytest.raises(ValueError, match="does not fit"): + binary.first(0.5 & s) + with pytest.raises(TypeError, match="Bad type for argument `left`"): + object() & s + # 1e300 overflows a float32 leaf: the op warns once, as `+` does, and `&` does not + f32s = Scalar.from_value( + [1, 2], dtypes.register_anonymous(np.dtype((np.float32, (2,))), "_InfixLiteralF32") + ) + expr = f32s & (1e300, 1.0) # the tests turn warnings into errors + with pytest.warns(RuntimeWarning, match="overflow") as record: + result = binary.plus(expr).new() + assert len(record) == 1 + np.testing.assert_array_equal(result.value, [np.inf, 3]) + + +@autocompute +def test_numpy_value_on_the_left(): + # numpy's operators defer to ours (__array_priority__), so a numpy scalar or 0-d array on + # the left is a literal operand, as on the right, and numpy does not compute the result + s = Scalar.from_value(5, dtype=dtypes.INT8) + v = Vector.from_coo([0, 2], [5, 6], dtype=dtypes.INT8, size=3) + A = Matrix.from_coo([0, 1], [1, 0], [5, 6], dtype=dtypes.INT8) + ops = [ + (operator.add, binary.plus), + (operator.sub, binary.minus), + (operator.mul, binary.times), + (operator.truediv, binary.truediv), + (operator.pow, binary.pow), + (operator.lt, binary.lt), + (operator.le, binary.le), + (operator.gt, binary.gt), + (operator.ge, binary.ge), + ] + if supports_udfs: # integer floordiv and numpy.mod are Numba functions + ops += [(operator.floordiv, binary.floordiv), (operator.mod, binary.numpy.mod)] + operands = [ + (s, s), + (s.apply(unary.identity), s), + (v[0], s), + (v @ v, (v @ v).new()), + (v, v), + (v.apply(unary.identity), v), + (v[[0, 1, 2]], v), + (A, A), + (A.T, A.T.new()), + (A.apply(unary.identity), A), + (A[:, :], A), + (A @ A, (A @ A).new()), + ] + for x, computed in operands: + is_scalar = output_type(x) is Scalar + cases = ops if is_scalar else [*ops, (operator.eq, binary.eq), (operator.ne, binary.ne)] + for lit in [np.int16(2), np.float32(0.5), np.array(3)]: + for infix, binop in cases: + expected = binop(lit, computed).new() + assert infix(lit, x).new().isequal(expected, check_dtype=True) + if is_scalar: + # A Scalar's `==` is a bool, not numpy's np.bool_ + assert (np.int16(computed.value) == x) is True + assert (np.int16(computed.value) != x) is False + with pytest.raises(TypeError, match="unsupported operand"): + np.int16(2) @ x + else: + with pytest.raises(TypeError, match="Bad type for argument `left`"): + np.int16(2) @ x + with pytest.raises(TypeError, match="Bad type for argument `left`"): + np.ones(x.shape[0]) @ x + # Before, numpy computed these from s.__array__() and raised for v and A + with pytest.raises(TypeError, match="Bad type"): + np.array([1, 2]) - x + with pytest.raises(TypeError, match="unsupported operand"): + np.int16(2) << x + + # `&`, `|` and `^` as with a Python literal + t = Scalar.from_value(True) + bv = Vector.from_coo([0, 2], [True, False], size=3) + for infix in [operator.and_, operator.or_, operator.xor]: + assert infix(np.False_, t).new().isequal(infix(False, t).new(), check_dtype=True) + assert (np.True_ ^ bv).new().isequal((True ^ bv).new(), check_dtype=True) + for infix in [operator.and_, operator.or_]: + with pytest.raises(TypeError, match="Bad type for argument `left`"): + infix(np.True_, bv) + with pytest.raises(TypeError, match="only valid for BOOL"): + (np.int16(2) & s).new() + + # In place, numpy no longer computes an array minus a Scalar + arr = np.array([1, 2]) + with pytest.raises(TypeError, match="Bad type"): + arr -= s + np.testing.assert_array_equal(arr, [1, 2]) + # Ufuncs still convert a Scalar with __array__ + assert np.sin(s) == np.sin(np.int8(5)) + np.testing.assert_array_equal(np.asarray(s), np.array(5, np.int8)) + result = np.subtract(np.int16(2), s) + assert isinstance(result, np.int16) + assert result == -3 + + @autocompute def test_multi_infix_scalar_auto(): v1 = Scalar.from_value(1) diff --git a/graphblas/tests/test_io.py b/graphblas/tests/test_io.py index 7e786f0da..1536fa4c7 100644 --- a/graphblas/tests/test_io.py +++ b/graphblas/tests/test_io.py @@ -1,3 +1,4 @@ +import functools from io import BytesIO, StringIO import numpy as np @@ -159,12 +160,66 @@ def test_matrix_to_from_networkx(): assert M.shape == (1, 1) +@pytest.mark.skipif("not nx") +def test_to_networkx_keeps_rows_without_entries(): + # Nodes 1 and 4 have no edges, and the entries name the others in the order + # 0, 3, 2: every row is still a node, in index order, so the round trip gives + # back the same matrix. + M = gb.Matrix.from_coo([2, 0, 3], [0, 3, 3], [5, 6, 7], nrows=5, ncols=5) + G = gb.io.to_networkx(M) + assert list(G) == [0, 1, 2, 3, 4] + # The dtype is given so the round trip does not depend on the one inferred + assert gb.io.from_networkx(G, dtype=M.dtype).isequal(M, check_dtype=True) + G = gb.io.to_networkx(M, edge_attribute=None) + assert list(G) == [0, 1, 2, 3, 4] + assert sorted(G.edges) == [(0, 3), (2, 0), (3, 3)] + assert len(gb.io.to_networkx(Matrix(int, 3, 3))) == 3 + + +@pytest.mark.skipif("not nx") +def test_to_networkx_without_isolates(): + # The nodes are the indices with entries, still in index order, so the round + # trip gives back the matrix without its empty rows and columns. + M = gb.Matrix.from_coo([2, 0, 3], [0, 3, 3], [5, 6, 7], nrows=5, ncols=5) + G = gb.io.to_networkx(M, create_isolates=False) + assert list(G) == [0, 2, 3] + expected = M[[0, 2, 3], [0, 2, 3]].new() + assert gb.io.from_networkx(G, dtype=M.dtype).isequal(expected, check_dtype=True) + G = gb.io.to_networkx(M, edge_attribute=None, create_isolates=False) + assert list(G) == [0, 2, 3] + assert sorted(G.edges) == [(0, 3), (2, 0), (3, 3)] + assert len(gb.io.to_networkx(Matrix(int, 3, 3), create_isolates=False)) == 0 + # Too many rows to make each a node, as in a hypersparse matrix + n = 2**60 + M = gb.Matrix.from_coo([5, n - 1], [n - 1, 5], [1, 2], nrows=n, ncols=n) + with pytest.raises(ValueError, match="create_isolates=False"): + gb.io.to_networkx(M) + G = gb.io.to_networkx(M, create_isolates=False) + assert list(G) == [5, n - 1] + assert dict(G.edges) == {(5, n - 1): {"weight": 1}, (n - 1, 5): {"weight": 2}} + + +@pytest.mark.skipif("not nx") +def test_to_networkx_limits(monkeypatch): + for create_isolates in [True, False]: + with pytest.raises(ValueError, match="square"): + gb.io.to_networkx(Matrix(int, 2, 3), create_isolates=create_isolates) + monkeypatch.setattr(gb.io._networkx, "_MAX_ISOLATES", 4) + assert len(gb.io.to_networkx(Matrix(int, 4, 4))) == 4 + with pytest.raises(ValueError, match="create_isolates=False"): + gb.io.to_networkx(Matrix(int, 5, 5)) + assert len(gb.io.to_networkx(Matrix(int, 5, 5), create_isolates=False)) == 0 + + @pytest.mark.skipif("not ss") @pytest.mark.parametrize("engine", ["auto", "scipy", "fmm"]) def test_mmread_mmwrite(engine): if engine == "fmm" and fmm is None: # pragma: no cover (import) pytest.skip("needs fast_matrix_market") try: + # Importing this module evaluates ``pytest.mark.thread_unsafe``, which + # fails unless the marker is registered; see the ``markers`` entry in + # pyproject.toml. from scipy.io.tests import test_mmio except ImportError: # Test files are mysteriously missing from some conda-forge builds @@ -272,7 +327,7 @@ def test_mmread_mmwrite(engine): example == "_empty_lines_example" and engine in {"fmm", "auto"} and fmm is not None - and fmm.__version__ in {"1.4.5"} + and fmm.__version__ == "1.4.5" ): # `fast_matrix_market` __version__ v1.4.5 does not handle this, but v1.5.0 does continue @@ -304,6 +359,42 @@ def test_mmread_mmwrite(engine): assert success > 0 +@pytest.mark.skipif("not ss") +def test_mmread_asks_scipy_for_a_sparse_array(monkeypatch): + """We pass ``spmatrix=False`` so scipy's changing default never reaches us. + + Deliberately does not use scipy's own example files: they are absent from + some conda-forge builds, which skips ``test_mmread_mmwrite`` entirely. + """ + import inspect + + import scipy.io as sio + + if "spmatrix" not in inspect.signature(sio.mmread).parameters: # pragma: no cover (old scipy) + pytest.skip("scipy too old to accept spmatrix=") + + seen = [] + real = sio.mmread + + @functools.wraps(real) # keep the signature the caller inspects + def spy(source, **kwargs): + seen.append(kwargs) + return real(source, **kwargs) + + monkeypatch.setattr(sio, "mmread", spy) + mm = StringIO("%%MatrixMarket matrix coordinate real general\n2 2 2\n1 1 3.0\n2 2 4.0\n") + M = gb.io.mmread(mm) + + assert seen == [{"spmatrix": False}] + assert M.isequal(Matrix.from_coo([0, 1], [0, 1], [3.0, 4.0])) + + # An explicit value wins: the guard only fills in a default. + seen.clear() + mm.seek(0) + gb.io.mmread(mm, spmatrix=True) + assert seen == [{"spmatrix": True}] + + @pytest.mark.skipif("not ss") def test_from_scipy_sparse_duplicates(): a = ss.coo_array(([1, 2, 3, 4], ([0, 1, 2, 2], [2, 1, 0, 0]))) @@ -315,6 +406,31 @@ def test_from_scipy_sparse_duplicates(): assert a2.isequal(expected) +@pytest.mark.skipif("not ss") +@pytest.mark.parametrize("format", ["csr", "csc"]) +def test_from_scipy_sparse_compressed_duplicates(format): + # Outside canonical format an index may repeat within a row (column); importing the + # arrays as they are made an invalid Matrix (nvals 4, two entries at one position). + cls, shape = (ss.csr_array, (2, 3)) if format == "csr" else (ss.csc_array, (3, 2)) + a = cls(([1, 2, 3, 4], [1, 1, 0, 2], [0, 3, 4]), shape=shape) + assert not a.has_canonical_format + with pytest.raises(ValueError, match="Duplicate indices found"): + gb.io.from_scipy_sparse(a) + A = gb.io.from_scipy_sparse(a, dup_op=gb.binary.plus) + assert A.nvals == 3 + np.testing.assert_array_equal(A.to_dense(fill_value=0), a.toarray()) + A = gb.io.from_scipy_sparse(a, dup_op=gb.binary.min) + expected = np.array([[3, 1, 0], [0, 0, 4]]) + if format == "csc": + expected = expected.T + np.testing.assert_array_equal(A.to_dense(fill_value=0), expected) + # Unsorted but distinct indices need no dup_op + b = cls(([1, 2, 3], [2, 1, 0], [0, 3, 3]), shape=shape) + assert not b.has_canonical_format + B = gb.io.from_scipy_sparse(b) + np.testing.assert_array_equal(B.to_dense(fill_value=0), b.toarray()) + + @pytest.mark.skipif("not ss") @pytest.mark.parametrize("engine", ["auto", "scipy", "fast_matrix_market"]) def test_matrix_market_sparse_duplicates(engine): @@ -348,6 +464,24 @@ def test_matrix_market_bad_engine(): gb.io.mmread(mm_out, engine="bad_engine") +@pytest.mark.skipif("not ss") +@pytest.mark.skipif("fmm is not None") +def test_matrix_market_fmm_engine_unavailable(): + # Naming the deprecated engine warns before anything else, so the caller hears + # about the deprecation even when the install that would satisfy it is missing. + A = gb.Matrix.from_coo([0, 0, 3, 5], [1, 4, 0, 2], [1, 0, 2, -1], nrows=7, ncols=6) + with ( + pytest.warns(DeprecationWarning, match="fast_matrix_market is no longer maintained"), + pytest.raises(ImportError, match="required to write Matrix Market files"), + ): + gb.io.mmwrite(BytesIO(), A, engine="fmm") + with ( + pytest.warns(DeprecationWarning, match="fast_matrix_market is no longer maintained"), + pytest.raises(ImportError, match="required to read Matrix Market files"), + ): + gb.io.mmread(BytesIO(), engine="fast_matrix_market") + + @pytest.mark.skipif("not ss") def test_scipy_sparse(): a = np.arange(12).reshape(3, 4) @@ -390,18 +524,28 @@ def test_awkward_roundtrip(): @pytest.mark.skipif("not ak") @pytest.mark.xfail(np.__version__[:5] in {"1.25.", "1.26."}, reason="awkward bug with numpy >=1.25") def test_awkward_iso_roundtrip(): + # Build iso outright rather than relying on all-equal values: SuiteSparse + # 10.5.0 no longer checks after a build whether every value is the same, + # so ``from_coo`` alone would leave these non-iso and the test would stop + # exercising the iso path it is named for. # Vector - v = gb.Vector.from_coo([1, 3, 5], [20, 20, 20], size=22) if suitesparse: + v = gb.Vector(int, 22) + v.ss.build_scalar([1, 3, 5], 20) assert v.ss.is_iso + else: + v = gb.Vector.from_coo([1, 3, 5], [20, 20, 20], size=22) kv = gb.io.to_awkward(v) assert isinstance(kv, ak.Array) v2 = gb.io.from_awkward(kv) assert v2.isequal(v) # Matrix - m = gb.Matrix.from_coo([0, 0, 3, 5], [1, 4, 0, 2], [1, 1, 1, 1], nrows=7, ncols=6) if suitesparse: + m = gb.Matrix(int, 7, 6) + m.ss.build_scalar([0, 0, 3, 5], [1, 4, 0, 2], 1) assert m.ss.is_iso + else: + m = gb.Matrix.from_coo([0, 0, 3, 5], [1, 4, 0, 2], [1, 1, 1, 1], nrows=7, ncols=6) for format in ["csr", "csc", "hypercsr", "hypercsc"]: km = gb.io.to_awkward(m, format=format) assert isinstance(km, ak.Array) @@ -481,3 +625,40 @@ def test_matrix_to_from_pydata_sparse(): gb.io.to_pydata_sparse(w, format="badformat") with pytest.raises(TypeError, match="sparse.pydata"): gb.io.from_pydata_sparse(w) + + +@pytest.mark.skipif("not sparse") +@pytest.mark.slow +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("dtype", ["int64", "float64", "bool"]) +def test_from_pydata_sparse_gcxs_duplicates(axis, dtype): + # pydata checks neither the order nor the distinctness of the indices in a GCXS row + # (column), and todense sums a repeat. Importing the arrays as they are made an invalid + # Matrix (nvals 4, two entries at one position). + data = np.array([1, 2, 3, 4], dtype=dtype) + indptr = np.array([0, 3, 4]) + shape = (2, 3) if axis == 0 else (3, 2) + s = sparse.GCXS((data, np.array([1, 1, 0, 2]), indptr), shape=shape, compressed_axes=[axis]) + A = gb.io.from_pydata_sparse(s) + assert A.nvals == 3 + assert A.dtype == dtype + np.testing.assert_array_equal(A.to_dense(fill_value=0), s.todense()) + A = gb.io.from_pydata_sparse(s, dup_op=gb.binary.min) + expected = np.array([[3, 1, 0], [0, 0, 4]], dtype=dtype) + if axis == 1: + expected = expected.T + np.testing.assert_array_equal(A.to_dense(fill_value=0), expected) + # Unsorted but distinct indices need no combining + s = sparse.GCXS((data, np.array([2, 1, 0, 2]), indptr), shape=shape, compressed_axes=[axis]) + A = gb.io.from_pydata_sparse(s) + assert A.nvals == 4 + np.testing.assert_array_equal(A.to_dense(fill_value=0), s.todense()) + # The default sum needs no binary op for the dtype (complex is a UDT on vanilla), + # and it leaves the array alone + cdata = data + 1j + s = sparse.GCXS((cdata, np.array([1, 1, 0, 2]), indptr), shape=shape, compressed_axes=[axis]) + indices = s.indices.copy() + A = gb.io.from_pydata_sparse(s) + assert A.nvals == 3 + np.testing.assert_array_equal(A.to_dense(fill_value=0), s.todense()) + np.testing.assert_array_equal(s.indices, indices) diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 24f0e73d7..611e1a13e 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -1,5 +1,6 @@ import inspect import itertools +import math import pickle import sys import types @@ -21,6 +22,7 @@ InvalidObject, InvalidValue, NotImplementedException, + NullPointer, OutputNotEmpty, ) @@ -177,6 +179,40 @@ def test_from_coo_scalar(): C.ss.iso_value +@pytest.mark.skipif("not suitesparse") +def test_ss_iso_value_of_an_empty_iso_matrix(): + # SuiteSparse marks a Matrix built from no tuples iso, and may mark other empty + # results iso too; iso_value used to leak StopIteration on them (RuntimeError + # inside a generator). There is no value, so the Scalar is empty. A new empty + # Matrix is not iso, and which way D goes depends on the SuiteSparse version. + C = Matrix.from_coo([], [], 7, nrows=3, ncols=2) + D = Matrix.from_coo([0], [1], [1.0], nrows=3, ncols=2).select("==", 2).new() + E = Matrix(int, 3, 2) + assert C.ss.is_iso + assert not E.ss.is_iso + for M in [C, D, E]: + if M.ss.is_iso: + value = M.ss.iso_value + assert isinstance(value, Scalar) + assert value.dtype == M.dtype + assert value.is_empty + else: + with pytest.raises(ValueError, match="not iso"): + M.ss.iso_value + + +@pytest.mark.skipif("not suitesparse") +@pytest.mark.parametrize("method", ["export", "unpack"]) +@pytest.mark.parametrize("fmt", ["fullr", "fullc"]) +def test_ss_export_full_of_a_matrix_that_is_not_full(method, fmt): + # SuiteSparse refuses with GrB_INVALID_VALUE and sets no message, so the + # error used to be blank. + A = Matrix.from_coo([0], [1], [2.0], nrows=2, ncols=2) + with pytest.raises(InvalidValue, match="1 of 4 entries is not full"): + getattr(A.ss, method)(fmt) + assert A.isequal(Matrix.from_coo([0], [1], [2.0], nrows=2, ncols=2)) + + def test_from_coo_default_values(): B = Matrix.from_coo([1, 3], [2, 4]) C = Matrix.from_coo([1, 3], [2, 4], [1.0, 1.0]) @@ -217,6 +253,34 @@ def test_nvals(A): assert A.nvals == 12 +@pytest.mark.skipif("not suitesparse") +def test_ss_export_failing_after_graphblas_frees_the_object(monkeypatch): + # GraphBLAS frees the object as it hands its arrays over. An error after that, + # here in claiming the values, left the handle pointing at the freed object, so a + # later call read freed memory and the object was freed again when collected. + import gc + + import graphblas.core.ss.matrix as ss_matrix + import graphblas.core.ss.vector as ss_vector + + def fail(*args): + raise RuntimeError("claim failed") + + monkeypatch.setattr(ss_matrix, "claim_buffer_2d", fail) + monkeypatch.setattr(ss_vector, "claim_buffer", fail) + A = Matrix.from_dense([[1.5, 2.5], [3.5, 4.5]]) + v = Vector.from_dense([1.5, 2.5]) + for obj, format, nvals in [(A, "fullr", 4), (v, "full", 2)]: + with pytest.raises(RuntimeError, match="claim failed"): + obj.ss.export(format) # exports a copy + gc.collect() # frees the copy + assert obj.nvals == nvals + with pytest.raises(RuntimeError, match="claim failed"): + obj.ss.export(format, give_ownership=True) + with pytest.raises(NullPointer): + obj._nvals + + def test_build(A): assert A.nvals == 12 A.clear() @@ -1235,6 +1299,24 @@ def test_apply_indexunary(A): assert pickle.loads(pickle.dumps(indexunary.tril[int])) is indexunary.tril[int] +@autocompute +def test_select_by_a_transposed_comparison(): + """``A.select(A.T == x)`` keeps the entries where ``A.T == x``, as a computed mask does. + + Rewriting it to ``select.valueeq(A, x)`` would drop the transpose and, for a matrix + that is not square, the DimensionMismatch. + """ + S = Matrix.from_coo([0, 1], [1, 0], [1, 5], nrows=2, ncols=2) + expected = Matrix.from_coo([1], [0], [5], nrows=2, ncols=2) + assert S.select(S.T == 1).new().isequal(expected) + assert S.select((S.T == 1).new()).new().isequal(expected) + # The rewrite still serves the comparison of the object itself + assert S.select(S == 1).new().isequal(Matrix.from_coo([0], [1], [1], nrows=2, ncols=2)) + B = Matrix.from_coo([0, 1], [1, 2], [1, 1], nrows=2, ncols=3) + with pytest.raises(DimensionMismatch): + B.select(B.T == 1).new() + + def test_select(A): A3 = Matrix.from_coo([0, 3, 3, 6], [3, 0, 2, 4], [3, 3, 3, 3], nrows=7, ncols=7) w1 = A.select(select.valueeq, 3).new() @@ -1273,6 +1355,57 @@ def test_select(A): assert pickle.loads(pickle.dumps(select.tril[bool])) is select.tril[bool] +def test_select_row_and_column_compare_indices_with_the_bound_given(): + # The comparison types its literal beside the values, where 300 beside INT8 or UINT8 + # is infinity and 2**24 + 1 beside FP32 rounds; the helpers used that for indices + # (row(B < 300) kept all 400 rows, row(B > 299) none; main was right). + n = 400 + for dtype in ["INT8", "UINT8"]: + B = Matrix.from_coo(np.arange(n), np.zeros(n, int), 1, dtype=dtype, nrows=n, ncols=1) + assert select.row(B < 300).new().nvals == 300 + assert select.row(B > 299).new().nvals == 100 + assert select.row(B <= 299).new().nvals == 300 + assert select.row(B >= 300).new().nvals == 100 + assert select.column(B.T <= 299).new().nvals == 300 + assert select.column(B.T > 299).new().nvals == 100 + assert select.row(B < 2**70).new().nvals == select.row(B > -(2**70)).new().nvals == n + rows = [2**24 - 1, 2**24, 2**24 + 1] + F = Matrix.from_coo(rows, [0, 0, 0], [1.0, 2.0, 3.0], dtype="FP32", nrows=2**24 + 3, ncols=1) + assert select.row(F < 2**24 + 1).new().to_coo()[0].tolist() == rows[:2] + # A fractional bound: i < 2.5 keeps 2, i > -0.5 keeps 0 (a float thunk was truncated) + A = Matrix.from_coo(np.arange(5), np.zeros(5, int), 1.0, nrows=5, ncols=1) + for expr, expected in [ + (A < A.nrows / 2, [0, 1, 2]), + (A <= 2.5, [0, 1, 2]), + (A >= 2.5, [3, 4]), + (A > 2.5, [3, 4]), + (A > -0.5, [0, 1, 2, 3, 4]), + (A < -0.5, []), + (A < float("inf"), [0, 1, 2, 3, 4]), + (A < Scalar.from_value(2.5), [0, 1, 2]), + # A bound that is not a Python or numpy number: the thunk keeps its own type + (A < np.array(3), [0, 1, 2]), + (A >= np.array(3), [3, 4]), + (A < np.True_, [0]), + ]: + assert select.row(expr).new().to_coo()[0].tolist() == expected, expr + assert select.column(A.T < 2.5).new().to_coo()[1].tolist() == [0, 1, 2] + w = Vector.from_coo([0], [3]) + with gb.config.set(autocompute=True): + assert select.row(A < w[0]).new().to_coo()[0].tolist() == [0, 1, 2] + assert select.row(A >= w[0]).new().to_coo()[0].tolist() == [3, 4] + # A comparison with the scalar on the left gives the op's TypeError, as on main; it + # was AttributeError ('Matrix' object has no attribute 'value') + with pytest.raises(TypeError, match="Bad type for keyword argument `thunk=`"): + select.row(A.apply(binary.le, left=3)) + # NaN compares False with every index, so it keeps none; NaN cast to the INT64 thunk + # was 0, so < and <= kept row 0, and > and >= every other row (main too) + nan = float("nan") + for expr in [A < nan, A <= nan, A > nan, A >= nan, A < Scalar.from_value(nan)]: + assert select.row(expr).new().nvals == 0, expr + assert select.column(A.T >= np.float32(nan)).new().nvals == 0 + + @autocompute def test_select_bools_and_masks(A): A3 = Matrix.from_coo([0, 3, 3, 6], [3, 0, 2, 4], [3, 3, 3, 3], nrows=7, ncols=7) @@ -1850,6 +1983,38 @@ def test_isclose(A, v): assert C6.isclose(A, rel_tol=1e-3) +def test_isclose_infinity(): + # As in math.isclose, an infinity is close only to itself. It was close to every + # number (abs(inf - y) is inf, and so is rel_tol * inf), with Numba and without. + from graphblas.core.vector import _isclose_recipe + + values = [math.inf, -math.inf, 1.0, 0.0, math.nan] + for x, y in itertools.product(values, repeat=2): + left = Matrix.from_coo([0, 1], [1, 0], [x, 5.0]) + right = Matrix.from_coo([0, 1], [1, 0], [y, 5.0]) + for abs_tol in [0.0, math.inf]: + expected = math.isclose(x, y, rel_tol=1e-7, abs_tol=abs_tol) + if supports_udfs: + assert left.isclose(right, abs_tol=abs_tol) is expected, (x, y) + assert _isclose_recipe(left, right, 1e-7, abs_tol) is expected, (x, y) + + +def test_isclose_integers_do_not_wrap(): + # x - y or abs(x) wrapped in an integer type; see test_vector.ISCLOSE_WRAPPING_CASES + from graphblas.core.vector import _isclose_recipe + + from .test_vector import ISCLOSE_WRAPPING_CASES + + for xtype, ytype, x, y, rel_tol in ISCLOSE_WRAPPING_CASES: + expected = math.isclose(x, y, rel_tol=rel_tol) + for (lt, lv), (rt, rv) in [((xtype, x), (ytype, y)), ((ytype, y), (xtype, x))]: + left = Matrix.from_coo([0, 1], [1, 0], [lv, 5], dtype=lt) + right = Matrix.from_coo([1, 0], [0, 1], [rv, 5], dtype=rt).T + if supports_udfs: + assert left.isclose(right, rel_tol=rel_tol) is expected, (lt, lv, rt, rv) + assert _isclose_recipe(left, right, rel_tol, 0.0) is expected, (lt, lv, rt, rv) + + @pytest.mark.slow def test_transpose_equals(A): data = [ @@ -1886,7 +2051,7 @@ def test_transpose_exceptional(): with pytest.raises(TypeError, match="not callable"): B.T()[1, 0] << 10 # with pytest.raises(AttributeError): - # should use new instead--Now okay. + # should use new instead. Now okay. assert B.T.dup().isequal(B.T.new()) # Not exceptional, but while we're here... C = B.T.new(mask=A.V) @@ -2084,12 +2249,16 @@ def test_ss_import_export(A, do_iso, methods): assert A5.isequal(A) assert A5.ss.is_iso is do_iso d["bitmap"] = np.concatenate([d["bitmap"], d["bitmap"]], axis=0) + if not do_iso: + # Values for the whole bitmap; GraphBLAS reads them all + d["values"] = np.concatenate([d["values"], d["values"]], axis=0) B5b = Matrix.ss.import_any(**d) if in_method == "import": # B5b == [A, A] (i.e, get 2d shape from bitmap first if possible) assert B5b.nvals == 2 * A.nvals assert B5b.nrows == 2 * A.nrows assert B5b.ncols == A.ncols + assert B5b[7:, :].new().isequal(A) else: A5.ss.pack_any(**d) assert A5.isequal(A) @@ -2794,7 +2963,7 @@ def test_ss_concat(A, v): expected[:, A.ncols] = v assert B5.isequal(expected) - with pytest.raises(TypeError, match=""): + with pytest.raises(TypeError): gb.ss.concat([v, [v]]) with pytest.raises(TypeError): gb.ss.concat([[v], v]) @@ -2837,7 +3006,7 @@ def test_auto(A, v): # print(type(expr).__name__, method) val1 = getattr(expected, method)(expected).new() if method in {"__or__", "__ror__"} and type(expr) is MatrixEwiseMultExpr: - # Doing e.g. `plus(A & B | C)` isn't allowed--make user be explicit + # Doing e.g. `plus(A & B | C)` isn't allowed; make user be explicit with pytest.raises(TypeError): val2 = getattr(expected, method)(expr) with pytest.raises(TypeError): @@ -2929,6 +3098,7 @@ def test_expr_is_like_matrix(A): "__call__", "__del__", "__delitem__", + "__getattr__", "__lshift__", "__setitem__", "_assign_element", @@ -2995,6 +3165,7 @@ def test_index_expr_is_like_matrix(A): expected = { "__del__", "__delitem__", + "__getattr__", "__setitem__", "_assign_element", "_delete_element", @@ -3062,6 +3233,25 @@ def test_dup_expr(A): assert result.isequal(B.dup(clear=True)) +@pytest.mark.skipif("not suitesparse") +@pytest.mark.parametrize( + "fmt", ["csr", "csc", "hypercsr", "hypercsc", "bitmapr", "bitmapc", "fullr", "fullc"] +) +@pytest.mark.parametrize("order", ["rowwise", "columnwise"]) +def test_ss_flatten_size(fmt, order): + # SuiteSparse holds a Vector by column. A reshape of a Matrix held by row was + # cast to a Vector as it stood, and SuiteSparse then saw a Vector of size 1: + # .size was 1 and every index past 0 raised InvalidIndex, though isequal and + # to_coo, which read the Python-side size, looked right. + A = Matrix.from_coo([0, 0, 0, 1, 1, 1], [0, 1, 2, 0, 1, 2], [1, 2, 3, 4, 5, 6]) + B = Matrix.ss.import_any(**A.ss.export(format=fmt)) + v = B.ss.flatten(order=order) + expected = A.to_dense().ravel(order="C" if order == "rowwise" else "F").tolist() + assert v.size == 6 + assert [v.get(i) for i in range(6)] == expected + assert v.dup().size == 6 + + @pytest.mark.skipif("not suitesparse") def test_ss_flatten(A): data = [ @@ -3429,7 +3619,7 @@ def test_ss_lastk(A): @pytest.mark.slow def test_ss_compactify(A, do_iso): if do_iso: - r, c, v = A.to_coo() + r, c, _v = A.to_coo() A = Matrix.from_coo(r, c, 1) rows = [0, 0, 1, 1, 2, 3, 3, 4, 5, 6, 6, 6] new_cols = [0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 2] @@ -3693,11 +3883,12 @@ def test_udt(): assert s == t t = A.reduce_scalar(monoid.any, allow_empty=False).new() assert s == t - A << binary.first(1, A) + # A literal fills every field by type; an int does not fit the bool field, True does. + A << binary.first(True, A) assert A.isequal(ones) A << binary.first(s, A) assert A.isequal(zeros) - A << binary.second(A, 1) + A << binary.second(A, True) assert A.isequal(ones) A << binary.second(A, s) assert A.isequal(zeros) @@ -3848,11 +4039,11 @@ def test_to_coo_sort(): rng.shuffle(r) rng.shuffle(c) A = Matrix.from_coo(r, c, r, nrows=N, ncols=N) - rows, cols, values = A.to_coo(sort=False) + rows, cols, _values = A.to_coo(sort=False) A = Matrix.from_coo(r, c, r, nrows=N, ncols=N) - rows, cols, values = A.to_coo(sort=True) + rows, cols, _values = A.to_coo(sort=True) assert_array_equal(rows, expected_rows) - rows, cols, values = A.T.to_coo(sort=True) + rows, cols, _values = A.T.to_coo(sort=True) assert_array_equal(cols, expected_rows) @@ -4104,6 +4295,17 @@ def test_to_dicts_from_dicts(A): assert D.to_dicts() == d +@pytest.mark.parametrize("shape", [(3, 2), (0, 0)]) +@pytest.mark.parametrize("order", ["rowwise", "columnwise"]) +def test_to_dicts_of_a_matrix_with_no_entries(shape, order): + # The compressed indptr of a Matrix with no entries is [0], which + # sliding_window_view(indptr, 2) refused with a ValueError. + A = Matrix(int, *shape) + assert A.to_dicts(order) == {} + assert A.T.to_dicts(order) == {} + assert Matrix.from_dicts({}, nrows=shape[0], ncols=shape[1]).isequal(A) + + def test_from_list_of_dicts(): list_of_dicts = [{1: 1}, {}, {0: 10, 2: 3}, {}] A1 = Matrix.from_dicts(list_of_dicts) @@ -4337,6 +4539,19 @@ def test_wait_chains(A): assert result == 47 +def test_subarray_dtypes_iso(): + # An iso array UDT exports its one value with a leading axis of length 1 + A = Matrix("INT64[4]", nrows=2, ncols=3) + A[:, :] = [0, 1, 2, 3] + if suitesparse: + assert A.ss.is_iso + expected = np.broadcast_to([0, 1, 2, 3], (6, 4)) + np.testing.assert_array_equal(A.to_dense(), expected.reshape(2, 3, 4)) + for method in ["coo", "csr", "csc", "dcsr", "dcsc"]: + np.testing.assert_array_equal(getattr(A, f"to_{method}")()[-1], expected) + assert A.to_dicts() == {i: {j: [0, 1, 2, 3] for j in range(3)} for i in range(2)} + + @pytest.mark.skipif("not supports_udfs") def test_subarray_dtypes(): a = np.arange(3 * 4, dtype=np.int64).reshape(3, 4) @@ -4376,6 +4591,146 @@ def test_subarray_dtypes(): assert Full1.isequal(Full2, check_dtype=True) +# Without SuiteSparse, building a UDT uses dup_op=binary.any, which needs Numba. +@pytest.mark.skipif("not suitesparse and not supports_udfs") +def test_subarray_dtype_values_shape(): + # Each of these has one entry per index but not two numbers per entry. Building + # from too few numbers read past the end of the array. (The dtype's name is + # whichever another test gave this anonymous dtype.) + match = r"values must have shape \(2, 2\) for dtype \S+, not \(2," + for values in [np.ones(2), [1.0, 2.0], np.ones((2, 1)), np.arange(6.0).reshape(2, 3)]: + with pytest.raises(ValueError, match=match): + Matrix.from_coo([0, 1], [0, 1], values, "FP64[2]") + A = Matrix("FP64[2]", 2, 2) + with pytest.raises(ValueError, match=match): + A.build([0, 1], [0, 1], values) + with pytest.raises(ValueError, match=match): + Matrix.from_csr([0, 1, 2], [0, 1], values, "FP64[2]") + with pytest.raises(ValueError, match=match): + Matrix.from_csc([0, 1, 2], [0, 1], values, "FP64[2]") + with pytest.raises(ValueError, match=r"shape \(2, 2, 2\) for dtype \S+, not \(2, 2, 1\)"): + Matrix.from_dense(np.ones((2, 2, 1)), dtype="FP64[2]") + # A scalar fills whole elements + expected = [[1.5, 1.5]] * 2 + A = Matrix.from_coo([0, 1], [0, 1], 1.5, "FP64[2]") + assert A.to_coo()[2].tolist() == expected + A = Matrix.from_csr([0, 1, 2], [0, 1], 1.5, "FP64[2]") + assert A.to_coo()[2].tolist() == expected + # on both backends (SuiteSparse raised "values must have 1 or more elements ... not 0") + A = Matrix.from_dcsr([0, 1], [0, 1, 2], [0, 1], 1.5, "FP64[2]") + assert A.to_coo()[2].tolist() == expected + A = Matrix.from_dcsc([0, 1], [0, 1, 2], [0, 1], 1.5, "FP64[2]") + assert A.to_coo()[2].tolist() == expected + # The lengths compared are the number of elements, not of numbers (it said 6) + with pytest.raises(ValueError, match="lengths must match: 2, 2, 3$"): + Matrix.from_coo([0, 1], [0, 1], np.ones((3, 2)), "FP64[2]") + + +@pytest.mark.skipif("not suitesparse") +def test_ss_import_coo_with_empty_rows_and_columns(): + # indptr was set only at the rows that appear, so every empty row after the first + # left a 0 in it: an empty row in the middle moved entries to another row (and + # isequal still said True), and an empty last row lost every entry + def entries(M): + return sorted(zip(*(x.tolist() for x in M.to_coo()), strict=True)) + + cases = [ + ([0, 2], [1, 0], 3, 2), # row 1, in the middle + ([0, 0], [0, 1], 3, 2), # rows 1 and 2, at the end + ([2, 2], [0, 1], 3, 2), # rows 0 and 1, at the start (this one worked) + ([1, 3], [2, 2], 5, 4), # rows and columns empty at the start, middle and end + ([], [], 3, 2), # no entries + ] + for rows, cols, nrows, ncols in cases: + A = Matrix.from_coo(rows, cols, np.arange(1.0, len(rows) + 1), nrows=nrows, ncols=ncols) + for M in [A, A.T.new()]: + for info in [M.ss.export("coor"), M.ss.export("cooc"), M.ss.export("coo", sort=True)]: + B = Matrix.ss.import_any(**info) + assert B.shape == M.shape + assert entries(B) == entries(M), (rows, cols, info["format"]) + # Lists work, and an index beyond the shape is refused before indptr is written + B = Matrix.ss.import_coor(rows=[0, 2], cols=[1, 0], values=[1.0, 2.0], nrows=3, ncols=2) + assert B.to_coo()[0].tolist() == [0, 2] + with pytest.raises(IndexOutOfBound, match="rows must be less than nrows, 3; got 3"): + Matrix.ss.import_coor(rows=[0, 3], cols=[0, 0], values=[1.0, 2.0], nrows=3, ncols=1) + with pytest.raises(IndexOutOfBound, match="cols must be less than ncols, 2; got 5"): + Matrix.ss.import_cooc(rows=[0, 0], cols=[1, 5], values=[1.0, 2.0], nrows=1, ncols=2) + # 2-D indices are a flat buffer, as in from_coo; they read past the end (or crashed) + rows = np.array([[0, 1], [1, 2]], np.uint64) + B = Matrix.ss.import_coor(rows=rows, cols=[0, 0, 1, 1], values=[1, 2, 3, 4], nrows=3, ncols=2) + assert entries(B) == entries(Matrix.from_coo(rows.ravel(), [0, 0, 1, 1], [1, 2, 3, 4])) + cols = np.array([[0], [0], [1]], np.uint64) + B = Matrix.ss.import_cooc(rows=[0, 2, 1], cols=cols, values=[1, 2, 3], nrows=3, ncols=2) + assert entries(B) == entries(Matrix.from_coo([0, 2, 1], cols.ravel(), [1, 2, 3])) + + +@pytest.mark.skipif("not suitesparse") +def test_ss_import_values_size(): + # GraphBLAS before 10.5.2 compares the byte size of imported values with the number + # of entries, not the bytes they need, so it read past the end of each of these + match = "have 3 or more elements of dtype FP64, not 1" + with pytest.raises(ValueError, match=match): + Matrix.ss.import_csr( + nrows=2, ncols=2, indptr=[0, 1, 3], col_indices=[0, 0, 1], values=np.ones(1) + ) + A = Matrix(float, 2, 2) + with pytest.raises(ValueError, match=match): + A.ss.pack_csc(indptr=[0, 1, 3], row_indices=[0, 0, 1], values=np.ones(1)) + with pytest.raises(ValueError, match=match): + Matrix.ss.import_hypercsr( + nrows=3, ncols=3, rows=[0, 2], indptr=[0, 1, 3], col_indices=[0, 0, 1], values=[1.0] + ) + with pytest.raises(ValueError, match="have 4 or more elements of dtype FP64, not 3"): + Matrix.ss.import_fullr(values=np.ones(3), nrows=2, ncols=2) + with pytest.raises(ValueError, match="have 4 or more elements of dtype FP64, not 2"): + Matrix.ss.import_bitmapc(bitmap=np.ones((2, 2), bool), values=np.ones(2), nrows=2, ncols=2) + # The dtype's name is whichever another test gave this anonymous dtype + match = r"have 2 or more elements of dtype \S+, not 1" + with pytest.raises(ValueError, match=match): + Matrix.ss.import_hypercsc( + nrows=3, + ncols=3, + cols=[0, 2], + indptr=[0, 1, 2], + row_indices=[0, 1], + values=np.ones(2), + dtype="FP64[2]", + ) + with pytest.raises(ValueError, match=match): + Matrix.from_dcsr([0, 1], [0, 1, 2], [0, 1], np.ones(2), "FP64[2]") + with pytest.raises(ValueError, match=r"have 4 or more elements of dtype \S+, not 2"): + Matrix.ss.import_fullc(values=np.ones((2, 2, 1)), dtype="FP64[2]") + # A scalar is not a whole element of a sub-array dtype (it said "not 0") + with pytest.raises(ValueError, match=r"a scalar does not fill one element of dtype \S+, of"): + Matrix.ss.import_fullr(values=1.5, nrows=2, ncols=2, dtype="FP64[2]", is_iso=True) + # A scalar that is one whole element, where more are needed, is just too few values + with pytest.raises(ValueError, match=r"have 4 or more elements of dtype FP64, not 1"): + Matrix.ss.import_fullr(values=1.5, nrows=2, ncols=2) + # More entries than indices is an indptr problem; it was reported as too few values + match = "index pointers give 4611686018427387904 entries, more than the number of indices, 1" + with pytest.raises(ValueError, match=match): + Matrix.ss.import_csr(nrows=1, ncols=3, indptr=[0, 2**62], col_indices=[0], values=[1.0]) + with pytest.raises(ValueError, match="index pointers give 3 entries, more than the number"): + Matrix.ss.import_hypercsc( + nrows=3, ncols=3, cols=[0, 2], indptr=[0, 1, 3], row_indices=[0, 1], values=[1.0] + ) + # One element is enough for iso + A = Matrix.ss.import_fullr(values=[7.0], nrows=2, ncols=2, is_iso=True) + assert A.ss.is_iso + assert A.to_coo()[2].tolist() == [7.0] * 4 + A = Matrix.ss.import_csr( + nrows=2, + ncols=2, + indptr=[0, 1, 2], + col_indices=[0, 1], + values=[1.0, 2.0], + dtype="FP64[2]", + is_iso=True, + ) + assert A.ss.is_iso + assert A.to_coo()[2].tolist() == [[1.0, 2.0]] * 2 + + def test_power(A): expected = A.dup() for i in range(1, 50): @@ -4545,3 +4900,224 @@ def test_setdiag(): A.setdiag(30, mask=v.S) expected[0, 0] = 30 assert A.isequal(expected) + + +def test_constructor_rejects_arraylike_first_arg(): + # The first positional arg is the dtype; passing data instead used to raise a + # confusing "Unknown dtype" ValueError. It should now raise TypeError pointing + # at the from_* constructors. + with pytest.raises(TypeError, match="Matrix.*dtype.*from_coo.*from_dense"): + Matrix([[1, 2], [3, 4]]) + with pytest.raises(TypeError, match="Matrix.*expects a dtype"): + Matrix(np.zeros((2, 2))) + # Valid dtype-first signatures still work, including list/tuple dtype specs + assert Matrix(int, 2, 2).dtype == dtypes.INT64 + assert Matrix("INT64", nrows=2, ncols=2).dtype == dtypes.INT64 + assert Matrix([("x", "i8"), ("y", "f8")], nrows=2, ncols=2).dtype._is_udt + assert Matrix((np.int32, (2, 2)), nrows=2, ncols=2).dtype._is_udt + # A non-array-like bad dtype keeps the original ValueError + with pytest.raises(ValueError, match="Unknown dtype"): + Matrix("not_a_dtype", nrows=2, ncols=2) + + +def test_wrong_kind_mask_on_matrix_raises(): + # A Vector mask on a Matrix output used to leak a raw cffi/C-signature + # error ("initializer for ctype 'struct GB_Matrix_opaque'"). It should + # now raise a clear TypeError, mirroring the Vector-output guard. + A = Matrix(int, 3, 3) + A[0, 0] = 1 + v = Vector(int, 3) + v[0] = 1 + # update path (full-matrix op with a Vector mask) + C = Matrix(int, 3, 3) + with pytest.raises(TypeError, match="Mask object must be type Matrix"): + C(mask=v.S) << A.ewise_mult(A) + # extract path + with pytest.raises(TypeError, match="Mask object must be type Matrix"): + A[:, :].new(mask=v.S) + # Valid Matrix mask on Matrix output still works + C = Matrix(int, 3, 3) + C(mask=A.S) << A.ewise_mult(A) + assert C.nvals == A.nvals + # Valid Vector input_mask broadcast on a Matrix extract still works + m = Vector(bool, 3) + m[0] = True + m[2] = True + assert A[0, [0, 1, 2]].new(input_mask=m.S) is not None + + +def test_wrong_kind_mask_on_dup_and_full_assign(): + # dup and whole-object update assign with `C(mask)[...] = A`, so a + # Vector mask used to reach GrB_Matrix_assign and leak the same raw + # cffi error as full-matrix operations. + A = Matrix(int, 3, 3) + A[0, 0] = 1 + v = Vector(bool, 3) + v[0] = True + v[1] = True + err = "Unable to use Vector mask on Matrix assignment to a Matrix" + with pytest.raises(TypeError, match=err): + A.dup(mask=v.S) + C = Matrix(int, 3, 3) + with pytest.raises(TypeError, match=err): + C(v.S) << A + with pytest.raises(TypeError, match=err): + C(v.V) << A + with pytest.raises(TypeError, match=err): + C(~v.S) << A + with pytest.raises(TypeError, match=err): + C(v.S)[:, :] << A + with pytest.raises(TypeError, match=err): + C[:, :](v.S) << A + # A Matrix mask on these paths still works + assert A.dup(mask=A.S).isequal(A) + C(A.S) << A + assert C.isequal(A) + # A Vector mask on a Matrix row/column assignment is still valid + w = Vector(int, 3) + w[0] = 10 + w[2] = 30 + C = Matrix(int, 3, 3) + C(v.S)[0, :] << w + assert C.isequal(Matrix.from_coo([0], [0], [10], nrows=3, ncols=3)) + C.clear() + C(v.S)[:, 0] << w + assert C.isequal(Matrix.from_coo([0], [0], [10], nrows=3, ncols=3)) + C.clear() + C[0, :](v.S) << w + assert C.isequal(Matrix.from_coo([0], [0], [10], nrows=3, ncols=3)) + C.clear() + C[:, 0](v.S) << w + assert C.isequal(Matrix.from_coo([0], [0], [10], nrows=3, ncols=3)) + + # The guard keys on dimension, not exact type, so a Vector subclass's + # mask is caught too instead of leaking the raw cffi error + class _VecSub(Vector): + pass + + mv = _VecSub(bool, 3) + mv[0] = True + with pytest.raises(TypeError, match="Unable to use Vector mask"): + C(mv.S) << A + + +def test_new_constructor_misuse_hint(): + # `.new()` resolves expressions; it is not a constructor or an instance + # method on concrete objects. Both misuses should hint the right API. + A = Matrix(int, 3, 3) + v = Vector(int, 3) + s = Scalar.from_value(5) + # Class-level access stays a plain AttributeError: hinting it would need + # a metaclass, and both metaclass choices break something (see + # test_abc_mixin_subclass). + for cls in (Matrix, Vector, Scalar): + with pytest.raises(AttributeError, match="has no attribute 'new'"): + cls.new + # Instance-level misuse: A.new() -> hint .dup() + for obj in (A, v, s): + with pytest.raises(AttributeError, match=r"has no attribute 'new'.*\.dup\(\)"): + obj.new + # A genuinely-missing attribute keeps the plain AttributeError (no hint) + with pytest.raises(AttributeError, match="has no attribute 'frobnicate'"): + A.frobnicate + assert not hasattr(A, "new") + assert not hasattr(Matrix, "new") + + +def test_abc_mixin_subclass(): + # BaseType must stay metaclass-free: a plain `type` hint metaclass broke + # abc-based mixins with "metaclass conflict", and an ABCMeta-derived one + # leaks metaclass attributes into dir(Matrix), which the + # expression-surface guards report as drift. + import collections.abc + + class SizedMatrix(Matrix, collections.abc.Sized): + def __len__(self): + return self.nvals + + assert issubclass(SizedMatrix, collections.abc.Sized) + A = SizedMatrix(int, 2, 3) + assert isinstance(A, collections.abc.Sized) + assert len(A) == 0 + A[0, 1] = 5 + A[1, 2] = 6 + assert len(A) == 2 + + +def test_matrix_subclass_repr(): + # The repr is a Matrix's but for the type's name. It used to take the + # TransposedMatrix branch for any type that is not exactly Matrix. + class SubMatrix(Matrix): + pass + + A = SubMatrix(int, 2, 3, name="A") + B = Matrix(int, 2, 3, name="A") + A[0, 1] = B[0, 1] = 5 + + def words(M): + return [word for word in repr(M).split() if set(word) != {"-"}] + + for a, b in [(A, B), (A.S, B.S), (~A.V, ~B.V)]: + assert words(a) == [word.replace("gb.Matrix", "gb.SubMatrix") for word in words(b)] + assert a._repr_html_() == b._repr_html_().replace("gb.Matrix", "gb.SubMatrix") + assert type(A.T) is gb.core.matrix.TransposedMatrix + assert repr(A.T) == repr(B.T) + assert A.T._repr_html_() == B.T._repr_html_() + + +@pytest.mark.skipif("not suitesparse") +def test_subclass_error_text(): + # An error's text comes from GrB__error, and a subclass has none of its own: + # it said "Unable to get the error string for type SubMatrix", asking for a bug report. + class SubMatrix(Matrix): + pass + + class SubVector(Vector): + pass + + class Context: # a user's class that shares a GraphBLAS type's name + pass + + class ContextVector(Context, Vector): + pass + + for base, sub, shape in [ + (Matrix, SubMatrix, (2, 3)), + (Vector, SubVector, (3,)), + (Vector, ContextVector, (3,)), + ]: + texts = [] + for cls in [base, sub]: + x = cls(int, *shape) + x.ss.export(give_ownership=True) # leaves no GraphBLAS object, so clear fails + with pytest.raises(gb.exceptions.NullPointer) as exc_info: + x.clear() + texts.append(str(exc_info.value)) + assert texts[0] == texts[1], texts + + +@pytest.mark.skipif("not suitesparse") +def test_transposed_error_text(): + # GraphBLAS has no transposed type, so an error on A.T is reported against A. It said + # "Unable to get the error string for type TransposedMatrix", asking for a bug report. + A = Matrix(int, 2, 3) + A.ss.export(give_ownership=True) # leaves no GraphBLAS object, so these calls fail + with pytest.raises(gb.exceptions.NullPointer) as expected: + A.wait() + for func in [A.T.wait, lambda: repr(A.T)]: + with pytest.raises(gb.exceptions.NullPointer) as exc_info: + func() + assert str(exc_info.value) == str(expected.value) + + +def test_transpose_hint(): + # A.transpose() should point at the .T property rather than fail with a + # bare "no attribute" message; transpose is not a method here. + A = Matrix(int, 3, 3) + v = Vector(int, 3) + with pytest.raises(AttributeError, match=r"has no attribute 'transpose'.*\.T"): + A.transpose + with pytest.raises(AttributeError, match=r"has no attribute 'transpose'.*\.T"): + v.transpose + # .T still works + assert A.T.shape == (3, 3) diff --git a/graphblas/tests/test_numpyops.py b/graphblas/tests/test_numpyops.py index 999c6d5e0..9e9a8fdd9 100644 --- a/graphblas/tests/test_numpyops.py +++ b/graphblas/tests/test_numpyops.py @@ -10,6 +10,7 @@ import graphblas as gb import graphblas.binary.numpy as npbinary import graphblas.monoid.numpy as npmonoid +import graphblas.op.numpy as npop import graphblas.semiring.numpy as npsemiring import graphblas.unary.numpy as npunary from graphblas import Vector, backend, config @@ -40,7 +41,7 @@ def test_bool_doesnt_get_too_large(): z = a.ewise_mult(b, gb.monoid.numpy.add).new() else: z = a.ewise_mult(b, gb.monoid.numpy.add).new() - x, y = z.to_coo() + _x, y = z.to_coo() np.testing.assert_array_equal(y, (True, True, True, False)) def func(x): # pragma: no cover (numba) @@ -48,10 +49,142 @@ def func(x): # pragma: no cover (numba) op = gb.core.operator.UnaryOp.register_anonymous(func) z = a.apply(op).new() - x, y = z.to_coo() + _x, y = z.to_coo() np.testing.assert_array_equal(y, (True, False, True, False)) +@pytest.mark.skipif("not supports_udfs") +# On a broken numba x numpy combo (numba #8478, still within our supported floor) +# the module sets _fmin_is_float=True and does not populate the integer fmax/fmin +# identities, so skip rather than KeyError on that config. +@pytest.mark.skipif("npmonoid._fmin_is_float") +@pytest.mark.parametrize("dtype", ["INT16", "INT32", "INT64"]) +def test_numpy_fmax_fmin_identity_outside_int8(dtype): + # Regression: the fmax identity for INT16/INT32/INT64 was set to int8's min + # (-128) instead of the dtype's own min, so a fmax-reduce over values all + # below -128 wrongly returned -128 rather than the true maximum. fmin always + # used the correct per-dtype max; it is checked here as a control. Only the + # mapnumpy=False path (UDF monoids) reads these identities; mapnumpy=True + # resolves fmax/fmin to the builtin max/min, which were already correct. + np_dtype = getattr(np, dtype.lower()) + lo = int(np.iinfo(np_dtype).min) + hi = int(np.iinfo(np_dtype).max) + # Pin the identity table directly (deterministic, independent of config). + assert npmonoid._monoid_identities["fmax"][dtype] == lo + assert npmonoid._monoid_identities["fmin"][dtype] == hi + + # Exercise the UDF monoids end-to-end. Force a fresh mapnumpy=False build so + # this hits the affected path regardless of the session's random mapnumpy + # (the resolved monoid is cached on the module after first access). + orig = config["mapnumpy"] + # The monoid is built from whatever binary.numpy has cached (monoid/numpy.py + # does getattr(_binary.numpy, name)), so the binary cache must be cleared too: + # under mapnumpy=True it holds the *builtin* binary.max/min, whose typed ops + # are TypedBuiltinBinaryOp, which has no `_monoid` slot for Monoid to write + # its back-reference to. Leaving it cached makes this test raise AttributeError + # whenever an earlier test in the session already resolved binary.numpy.fmax. + # op.numpy is cleared as well because registering the UDF binary op populates + # it as a side effect. Save every entry first so the exact prior state, cached + # or absent, is restored for later tests. + names = ("fmax", "fmin") + modules = (npmonoid, npbinary, npop) + saved = [(module, name, module.__dict__.get(name)) for module in modules for name in names] + config.set(mapnumpy=False) + for module, name, _ in saved: + module.__dict__.pop(name, None) + try: + below = [-1000, -2000, -3000] # all below int8's min, valid for INT16+ + above = [1000, 2000, 3000] # all above int8's max + v_below = Vector.from_coo([0, 1, 2], below, dtype=dtype) + v_above = Vector.from_coo([0, 1, 2], above, dtype=dtype) + assert v_below.reduce(npmonoid.fmax).new().value == max(below) # -1000 + assert v_above.reduce(npmonoid.fmin).new().value == min(above) # 1000 + assert npmonoid.fmax[dtype].identity == lo + assert npmonoid.fmin[dtype].identity == hi + finally: + config.set(mapnumpy=orig) + for module, name, obj in saved: + if obj is None: + module.__dict__.pop(name, None) + else: + module.__dict__[name] = obj + + +def test_numpy_monoid_identity_matches_builtin(): + # Invariant: mapnumpy must not change a numpy monoid's identity. For every op + # in the numpy->builtin mapping, the numpy identity (read straight from + # _monoid_identities, which register_new uses verbatim on the mapnumpy=False + # path) must equal the builtin monoid's identity for every shared dtype, + # compared cast-to-dtype so benign encodings agree (bitwise_and's -1 and + # band's 255 are both all-ones); only a genuinely wrong value (e.g. logical_or + # True vs lor False) is flagged. This guards the whole table and caught both + # fmax/fmin and logical_or. Reading the table directly keeps it deterministic + # and mutates no operator caches, so it cannot desync monoid/binary/op.numpy. + identities = npmonoid._monoid_identities + mismatches = [] + for np_name, gb_name in npmonoid._numpy_to_graphblas.items(): + table = identities[np_name] + gb_monoid = getattr(gb.monoid, gb_name) + for dtype in gb_monoid.types: + if isinstance(table, dict): + if dtype.name not in table: + continue # numpy op does not define this dtype + np_id = table[dtype.name] + else: + np_id = table # scalar identity applies to every dtype + gb_id = gb_monoid[dtype].identity + np_cast = np.asarray(np_id).astype(dtype.np_type) + gb_cast = np.asarray(gb_id).astype(dtype.np_type) + if np_cast != gb_cast: + mismatches.append(f"{np_name}[{dtype.name}]={np_id!r} vs {gb_name}={gb_id!r}") + assert not mismatches, mismatches + + +def test_numpy_monoid_unmapped_identity_consistency(): + # The six numpy monoids with no builtin counterpart (gcd, hypot, logaddexp, + # logaddexp2, maximum, minimum) are absent from _numpy_to_graphblas, so the + # sibling test above never reads their identities. Guard the two that have a + # byte-identical builtin-backed twin: numpy's maximum/minimum share fmax's/ + # fmin's identity table exactly (nan vs. non-nan changes only propagation of + # a present value, not the neutral element). fmax/fmin are themselves pinned + # against builtin max/min by the sibling test, so asserting maximum == fmax + # and minimum == fmin transitively validates maximum/minimum. Read the table + # directly and compare cast-to-dtype, matching the sibling test: no config + # toggle, no operator-cache mutation, no monoid resolution. + identities = npmonoid._monoid_identities + mismatches = [] + for np_name, twin in (("maximum", "fmax"), ("minimum", "fmin")): + table = identities[np_name] + twin_table = identities[twin] + # fmax/fmin drop their integer keys on a broken numba/numpy combo + # (_fmin_is_float), so intersect rather than assume every dtype is present. + shared = table.keys() & twin_table.keys() + assert shared, f"{np_name}/{twin} share no dtypes" + for name in sorted(shared): + np_type = gb.dtypes.lookup_dtype(name).np_type + np_cast = np.asarray(table[name]).astype(np_type) + twin_cast = np.asarray(twin_table[name]).astype(np_type) + if np_cast != twin_cast: + mismatches.append( + f"{np_name}[{name}]={table[name]!r} vs {twin}[{name}]={twin_table[name]!r}" + ) + assert not mismatches, mismatches + + # Light self-check for the log-add monoids: their neutral element is -inf + # (logaddexp(-inf, x) == x). gcd and hypot are deliberately left out of any + # functional identity check: their conventional identity is 0, but numpy's + # gcd/hypot are not sign-preserving (gcd(0, -5) == 5 != -5), so + # f(identity, x) == x is not a valid invariant for them even though 0 is the + # correct identity. So only their presence, not a value, is asserted here. + for np_name in ("logaddexp", "logaddexp2"): + table = identities[np_name] + assert table, f"{np_name} identity table is empty" + for name, value in table.items(): + assert value == -np.inf, f"{np_name}[{name}]={value!r} != -inf" + for np_name in ("gcd", "hypot"): + assert identities[np_name], f"{np_name} identity table is empty" + + @pytest.mark.slow def test_npunary(): L = list(range(5)) @@ -116,7 +249,7 @@ def test_npunary(): import numba if ( - unary_name in {"sign"} + unary_name == "sign" and np.__version__.startswith("2.") and parse(numba.__version__) < parse("0.61.0") ): @@ -201,7 +334,7 @@ def test_npbinary(): compare_op = isclose else: np_result = getattr(np, binary_name)(np_left, np_right) - if binary_name in {"arctan2"}: + if binary_name == "arctan2": compare_op = isclose else: compare_op = equal diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 41fae80ae..64bc1880c 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1,4 +1,8 @@ import itertools +import os +import subprocess +import sys +from pathlib import Path import numpy as np import pytest @@ -28,6 +32,7 @@ UnaryOp, get_semiring, ) +from graphblas.core.ss import version_major as ss_version_major # noqa: F401 (used in skipif) from graphblas.dtypes import ( BOOL, FP32, @@ -41,7 +46,7 @@ UINT32, UINT64, ) -from graphblas.exceptions import DomainMismatch, UdfParseError +from graphblas.exceptions import DimensionMismatch, DomainMismatch, InvalidValue, UdfParseError from .conftest import shouldhave @@ -420,6 +425,42 @@ def bad_identity(x=0): assert monoid.is_idempotent +def test_monoid_rejects_builtin_binaryop(): + """A built-in binaryop cannot back a user monoid. + + Built-ins already own their monoid association (binary.max resolves to + monoid.max), and SuiteSparse silently ignores the identity passed to + GrB_Monoid_new for them, so accepting one would produce a monoid whose + Python-side identity disagrees with what GraphBLAS computes with. + """ + with pytest.raises(TypeError, match="must be a user-defined BinaryOp"): + Monoid.register_new("_bad_builtin_monoid", binary.max, 0) + with pytest.raises(TypeError, match="must be a user-defined BinaryOp"): + Monoid.register_anonymous(binary.plus, 1) + # The rejection happens before any object is created, so the built-in's + # own monoid association is untouched. + assert binary.max.monoid is monoid.max + assert not hasattr(gb.monoid, "_bad_builtin_monoid") + + +def test_monoid_rejects_builtin_binaryop_with_udt_fallback(): + """binary.any, first, pair, and second are built-ins that carry a Python + func as their UDT fallback, so a func-based probe alone misses them. They + must be rejected like the other built-ins: their typed ops have no _monoid + slot, and a failed registration used to leave binary.._monoid + pointing at the half-built monoid. + """ + for binop in [binary.any, binary.first, binary.pair, binary.second]: + with pytest.raises(TypeError, match="must be a user-defined BinaryOp"): + Monoid.register_anonymous(binop, 0) + # State is untouched: any's built-in monoid association survives, and the + # ops without a built-in monoid still have none. + assert binary.any.monoid is monoid.any + assert binary.first.monoid is None + assert binary.pair.monoid is None + assert binary.second.monoid is None + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_semiring_parameterized(): @@ -869,13 +910,13 @@ def test_monoid_attributes(): assert monoid.plus[int].binaryop is binary.plus[int] assert monoid.plus[int].identity == 0 assert monoid.plus.binaryop is binary.plus - assert monoid.plus.identities == {typ: 0 for typ in monoid.plus.types} + assert monoid.plus.identities == dict.fromkeys(monoid.plus.types, 0) if shouldhave(monoid.numpy, "add"): assert monoid.numpy.add[int].binaryop is binary.numpy.add[int] assert monoid.numpy.add[int].identity == 0 assert monoid.numpy.add.binaryop is binary.numpy.add - assert monoid.numpy.add.identities == {typ: 0 for typ in monoid.numpy.add.types} + assert monoid.numpy.add.identities == dict.fromkeys(monoid.numpy.add.types, 0) def plus(x, y): # pragma: no cover (numba) return x + y @@ -1168,6 +1209,9 @@ def test_lazy_op(): Monoid.register_new("lazy2", binary.lazy, 0, lazy=True) assert isinstance(op.lazy2, Monoid) assert isinstance(monoid.lazy2, Monoid) + BinaryOp.register_new("lazy_idem", _isidem_pick_first, lazy=True) + Monoid.register_new("lazy_idem", "lazy_idem", 0, is_idempotent=True, lazy=True) + assert monoid.lazy_idem.is_idempotent Semiring.register_new("lazy", "lazy", "lazy", lazy=True) assert isinstance(semiring.lazy, Semiring) Semiring.register_new("lazy_lazy", monoid.lazy, binary.lazy, lazy=True) @@ -1262,9 +1306,12 @@ def _udt_index(val, idx, _, thunk): # pragma: no cover (numba) _udt_index = IndexUnaryOp.register_anonymous(_udt_index, "_udt_index", is_udt=True) assert udt in _udt_index - result = v.apply(_udt_index, 3).new() + result = v.apply(_udt_index, (False, 3)).new() expected = Vector.from_coo([0, 1, 2], [3, -3, -3]) assert result.isequal(expected) + # A number fills every field, and an int is not a bool: 3 does not fit. + with pytest.raises(ValueError, match="3 does not fit"): + v.apply(_udt_index, 3) def _udt_first(x, y): return x # pragma: no cover (numba) @@ -1273,7 +1320,11 @@ def _udt_first(x, y): assert udt in udt_first assert operator.get_typed_op(udt_first, udt) is udt_first[udt] assert udt_first(v & w).new().isequal(v) - assert udt_first(v, 1).new().isequal(v) + # A literal fills every field, by type: an int does not fit the bool field, + # True does. 1 still converts, being exact, but that is deprecated. + assert udt_first(v, True).new().isequal(v) + with pytest.warns(DeprecationWarning, match="1 does not fit TestUDT"): + assert udt_first(v, 1).new().isequal(v) assert udt_first[udt, dtypes.INT64].return_type == udt assert udt_first[dtypes.INT64, udt].return_type == dtypes.INT64 assert udt_first[udt, dtypes.BOOL].return_type == udt @@ -1350,130 +1401,5653 @@ def _this_or_that(val, idx, _, thunk): # pragma: no cover (numba) sel = SelectOp.register_anonymous(_this_or_that, is_udt=True) sel[udt] assert udt in sel - result = v.select(sel, 0).new() + # The thunk becomes an element of the record, whose bool field takes False, not 0. + result = v.select(sel, False).new() assert result.nvals == 0 assert result.dtype == v.dtype - result = w.select(sel, 0).new() + result = w.select(sel, False).new() assert result.nvals == 3 assert result.isequal(w) -def test_dir(): - for mod in [unary, binary, monoid, semiring, op]: - assert not set(mod._delayed) - set(dir(mod)) +@pytest.mark.skipif("not supports_udfs") +def test_udf_division_by_zero_follows_numpy(): + """Dividing by zero in a UDF returns numpy's answer instead of losing the element. + + Under Numba's default error model the division raises ZeroDivisionError + inside the cfunc, where Numba prints the traceback and returns, so + GraphBLAS keeps whatever was in the output element (in practice the + previous element's value). ``error_model="numpy"`` fixes that, but only if + it is set on the ``njit`` Dispatcher: a Dispatcher holds one compilation + per signature, and ``_build`` calls ``.compile(sig)`` before the wrapper + exists, so setting it on the ``cfunc`` alone comes too late to matter. + """ + + def _idiv(x, y): # pragma: no cover (numba) + return x // y + + op = BinaryOp.register_anonymous(_idiv, "_udf_zero_idiv") + v = Vector.from_coo([0, 1], [10, 20], dtype=dtypes.INT64) + w = Vector.from_coo([0, 1], [2, 0], dtype=dtypes.INT64) + assert op(v & w).new().to_coo()[1].tolist() == [5, 0] + + def _tdiv(x, y): # pragma: no cover (numba) + return x / y + + op = BinaryOp.register_anonymous(_tdiv, "_udf_zero_tdiv") + v = Vector.from_coo([0, 1], [1.0, 1.0], dtype=dtypes.FP64) + w = Vector.from_coo([0, 1], [2.0, 0.0], dtype=dtypes.FP64) + assert op(v & w).new().to_coo()[1].tolist() == [0.5, float("inf")] + + # Same guarantee for a UDT UDF, which reaches the cfunc by another route. + udt = dtypes.register_anonymous( + np.dtype([("dz_a", np.int64), ("dz_b", np.int64)], align=True), "_UdfDivZeroRec" + ) + def _rec_idiv(x, y): # pragma: no cover (numba) + return (x["dz_a"] // y["dz_a"], x["dz_b"]) -def test_semiring_commute_exists(): - from .conftest import orig_semirings + op = BinaryOp.register_anonymous(_rec_idiv, "_udf_zero_rec", is_udt=True) + v = Vector(udt, size=1) + v[0] = (10, 5) + w = Vector(udt, size=1) + w[0] = (0, 1) + got = v.ewise_mult(w, op).new()[0].new().value + assert (got["dz_a"], got["dz_b"]) == (0, 5) - vals = { - semiring._deprecated[key] if key in semiring._deprecated else getattr(semiring, key) - for key in orig_semirings - } - missing = set() - for key in orig_semirings: - val = semiring._deprecated[key] if key in semiring._deprecated else getattr(semiring, key) - commutes_to = val.commutes_to - if commutes_to is not None and commutes_to not in vals: # pragma: no cover (debug) - missing.add(commutes_to.name) - if missing: - raise AssertionError("Missing semirings: " + ", ".join(sorted(missing))) +@pytest.mark.skipif("not supports_udfs") +def test_select_op_outlives_source_indexunary(): + """A SelectOp keeps alive the IndexUnaryOp whose GraphBLAS handle it borrows. + + ``SelectOp._from_indexunary`` reuses the IndexUnaryOp's ``gb_obj`` rather + than allocating a second one, and ``register_anonymous`` drops that + IndexUnaryOp on the way out. Without an explicit reference the handle is + freed as soon as it is collected, and every use of the SelectOp raises + ``UninitializedObject``. + """ + import gc + + def _ne_thunk(x, i, j, thunk): # pragma: no cover (numba) + return x != thunk + + sel = SelectOp.register_anonymous(_ne_thunk) + gc.collect() + v = Vector.from_coo([0, 1, 2], [1, 5, 9]) + assert v.select(sel, 5).new().isequal(Vector.from_coo([0, 2], [1, 9], size=3)) + + +_jit_can_compile_cache = [] + + +def _jit_can_compile(): + """True when SuiteSparse has a C compiler it can actually use. + + Without one it falls back to the Numba cfunc and says nothing, so the + ``jit`` parameter would run the cfunc and report a pass. + + ``jit_compiler_is_usable`` alone is not enough: it only checks that the + configured compiler path exists on disk, and a runner can have the file + yet fail every compile (broken toolchain, missing headers). + ``_enable_jit_for_udt`` is the same call the UDT auto-lift path makes, + and it does a real compile once per process, so its answer is the + honest one here; ``test_ssjit`` keys its skips on the same signal. The + fixture calls this before it mutates the control, and the cache keeps + later per-test mutations from flipping it. + """ + if not _jit_can_compile_cache: + from graphblas.core.ss.jit_config import _enable_jit_for_udt + + _jit_can_compile_cache.append(gb.ss.jit_compiler_is_usable() and _enable_jit_for_udt()) + return _jit_can_compile_cache[0] + + +@pytest.fixture(params=["jit", "cfunc"]) +def udt_op_path(request): + """Pin SuiteSparse to one execution path for built-in UDT operators. + + Each auto-lifted UDT op carries both a C JIT definition and a Numba + cfunc, and SuiteSparse chooses between them per call depending on + whether a C compiler is available. A machine with one and a machine + without therefore run different code, so results have to hold on both. + """ + path = request.param + if backend != "suitesparse" or "jit_c_control" not in gb.ss.config: + if path == "jit": + pytest.skip("no SuiteSparse C JIT on this backend") + yield path + return + previous = gb.ss.config["jit_c_control"] + if path == "jit": + if not _jit_can_compile(): + pytest.skip("C JIT compilation not available (probe failed or compiler missing)") + # Set it rather than assume it. SuiteSparse demotes ``on`` to ``load`` + # after a failed compile, and a demoted control routes to the cfunc + # silently, so this parameter would pass while running the other path. + gb.ss.config["jit_c_control"] = "on" + else: + gb.ss.config["jit_c_control"] = "off" + try: + yield path + finally: + # Read before restoring: a demotion during the test is the signal that + # the kernel never compiled, which no assertion in the test can see. + demoted = path == "jit" and gb.ss.config["jit_c_control"] != "on" + gb.ss.config["jit_c_control"] = previous + if demoted: + pytest.fail( # pragma: no cover (only on failure) + "SuiteSparse demoted jit_c_control; the JIT path did not run" + ) -def test_binaryop_commute_exists(): - from .conftest import orig_binaryops - vals = { - binary._deprecated[key] if key in binary._deprecated else getattr(binary, key) - for key in orig_binaryops - } - missing = set() - for key in orig_binaryops: - val = binary._deprecated[key] if key in binary._deprecated else getattr(binary, key) - commutes_to = val.commutes_to - if commutes_to is not None and commutes_to not in vals: # pragma: no cover (debug) - missing.add(commutes_to.name) - if missing: - raise AssertionError("Missing binaryops: " + ", ".join(sorted(missing))) +def _udt_vectors(udt, xs, ys=None): + """Build one or two dense UDT vectors whose leaves all hold the given values. + + Values that repeat down the whole vector make it iso-valued, and + SuiteSparse answers those from a single element without reaching for a + C JIT kernel, so callers pass varied data. + """ + names = udt.np_type.names + out = [] + for vals in (xs, ys): + if vals is None: + continue + v = Vector(udt, size=len(vals)) + for i, val in enumerate(vals): + v[i] = tuple(val for _ in names) if names else np.full(udt.np_type.subdtype[1], val) + out.append(v) + return out @pytest.mark.skipif("not supports_udfs") -def test_binom(): - v = Vector.from_coo([0, 1, 2], [3, 4, 5]) - result = v.apply(binary.binom, 2).new() - expected = Vector.from_coo([0, 1, 2], [3, 6, 10]) +@pytest.mark.slow +def test_udt_mixed_record_dtypes_use_each_operands_own_dtype(udt_op_path): + """Two records sharing field names but not field types promote to the wider one. + + ``_check_udt_pair`` matches record operands on field names only, so their + leaf dtypes can differ. Offering the return-type resolver just the left + operand left it no choice but that record, so an int record over a float + record came back as the int record: ``7 / 2.0`` landed as 3 and ``6 / 0.0`` + as INT64_MAX. Swapping the operands changed the answer for the same pair. + """ + int_udt = dtypes.register_anonymous(np.dtype([("mxd_a", np.int64)], align=True), "_MixedRecInt") + float_udt = dtypes.register_anonymous( + np.dtype([("mxd_a", np.float64)], align=True), "_MixedRecFloat" + ) + v = Vector(int_udt, size=2) + v[0] = (6,) + v[1] = (7,) + w = Vector(float_udt, size=2) + w[0] = (0.0,) + w[1] = (2.0,) + + result = v.ewise_mult(w, binary.truediv).new() + assert result.dtype == float_udt, "result should promote to the float record" + assert result[0].new().value["mxd_a"] == float("inf") + assert result[1].new().value["mxd_a"] == 3.5 + + # The same pair the other way round must agree, which it did not when the + # resolver only ever saw the left operand. + swapped = w.ewise_mult(v, binary.truediv).new() + assert swapped.dtype == float_udt + assert swapped[1].new().value["mxd_a"] == 2.0 / 7.0 + + +def _bitwise_eq(got, want): + """Compare two floats by bit pattern, treating any two NaNs as equal. + + Bit patterns rather than ``==`` because ``-0.0 == 0.0``, and the sign of + a zero is exactly what a min/max tie-break decides. NaNs are exempted + because ``fmin`` may hand back either operand's NaN payload. + """ + if np.isnan(got) and np.isnan(want): + return True + return got.tobytes() == want.tobytes() + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +@pytest.mark.parametrize("np_dtype", [np.float64, np.float32]) +def test_udt_min_max_answer_what_the_builtin_dtype_answers(udt_op_path, np_dtype): + """``binary.min[udt]`` must give what ``binary.min[FP64]`` gives, bit for bit. + + An operator that means one thing on FP64 and another on a record of + FP64 is not one operator. SuiteSparse's ``GrB_MIN_FP64`` is C99 ``fmin``, + so that is what the UDT kernels have to be, and this compares them + directly against the built-in rather than against a convention chosen on + the Python side. The grid is every ordered pair drawn from NaN, both + infinities, both zeros and two ordinary values, so it covers a NaN on + either side, two NaNs, and a signed-zero tie either way round. + + The signed-zero tie itself is compared by value only. C99 leaves + ``fmin(-0.0, 0.0)`` unspecified and the built-in answers differently per + platform (left operand on macOS x86, right operand on Linux x86, IEEE + minNum on arm64 and Windows), so bit-for-bit agreement on that one pair + is not something any implementation can promise. Everything else, + including which zero a mixed zero/nonzero pair keeps, stays bit-exact. + + What this catches, in the two spellings it replaces: Python's builtin + ``min``, which the generated code reached through the exec namespace, + ordered NaN by position, and ``np.fmin`` under Numba gets the NaN rule + right but keeps the left operand on a signed-zero tie, so it drifts from + the C JIT kernel on ``min(0.0, -0.0)``. Both execution paths are checked + because SuiteSparse picks between them without telling anyone. + """ + nan, inf = float("nan"), float("inf") + values = [nan, inf, -inf, -0.0, 0.0, 1.5, -2.5] + xs = [x for x in values for _ in values] + ys = list(values) * len(values) + + udt = dtypes.register_anonymous( + np.dtype([("mmb_a", np_dtype)], align=True), f"_MinMaxBuiltin{np.dtype(np_dtype).name}" + ) + v, w = _udt_vectors(udt, xs, ys) + ref_v = Vector.from_dense(np.array(xs, dtype=np_dtype)) + ref_w = Vector.from_dense(np.array(ys, dtype=np_dtype)) + + for gb_op in (binary.min, binary.max): + expected = gb_op(ref_v & ref_w).new().to_dense() + result = gb_op(v & w).new() + for i, (x, y) in enumerate(zip(xs, ys, strict=True)): + got = result[i].new().value[0] + if x == 0 and y == 0 and np.signbit(x) != np.signbit(y): + # The one unspecified cell of the grid: either signed zero is + # a correct answer from either implementation, so only agree + # that both produced a zero. + msg = ( + f"{udt_op_path} {gb_op.name}({x}, {y}) on {udt.name}: " + f"got {got!r}, built-in {np.dtype(np_dtype).name} gives {expected[i]!r}" + ) + assert got == 0, msg + assert expected[i] == 0, msg + continue + assert _bitwise_eq(got, expected[i]), ( + f"{udt_op_path} {gb_op.name}({x}, {y}) on {udt.name}: " + f"got {got!r}, built-in {np.dtype(np_dtype).name} gives {expected[i]!r}" + ) + + # A NaN anywhere in the input must not change where a reduce lands. Under + # the Python-builtin semantics this same multiset reduced to 1.0 or to nan + # depending on which index the NaN sat at. + for data in ([1.0, 2.0, 3.0, nan], [nan, 1.0, 2.0, 3.0], [1.0, nan, 3.0, 2.0]): + (u,) = _udt_vectors(udt, data) + assert u.reduce(monoid.min[udt]).new().value[0] == 1.0, f"{udt_op_path} {data}" + assert u.reduce(monoid.max[udt]).new().value[0] == 3.0, f"{udt_op_path} {data}" + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_truediv_divides_in_floating_point(udt_op_path): + """``binary.truediv`` on integer fields divides in float64 and keeps the quotient. + + Regression: the C JIT kernel emitted C ``/``, which is integer division. + ``10**18 / 3`` came out as 333333333333333333 under the C JIT and + 333333333333333312 (float64, like numpy) through the cfunc, so the same + program gave different answers depending on whether a C compiler was + installed. The result fields are now float64, as ``truediv`` on two INT64 + vectors is FP64, so the quotient is not truncated either. + """ + udt = dtypes.register_anonymous( + np.dtype([("tdv_i", np.int64), ("tdv_j", np.int64)], align=True), "_TrueDivIntUDT" + ) + xs = [10**18, 10**18 + 1, 7, 22] + ys = [3, 3, 2, 7] + expected = np.array(xs, np.int64) / np.array(ys, np.int64) + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + assert [result.dtype.np_type[name] for name in ("tdv_i", "tdv_j")] == [np.dtype(np.float64)] * 2 + got = [result[i].new().value[0] for i in range(len(xs))] + assert got == list(expected), udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_floordiv_matches_numpy_on_floats(udt_op_path): + """``binary.floordiv`` on float fields is not ``floor(a / b)``. + + Regression: the C JIT kernel computed ``floor(a / b)``, which rounds + differently from the remainder-based algorithm numpy and CPython use and + treats infinities as ordinary values. ``1.0 // 0.1`` came out as 10.0 + instead of 9.0, and ``inf // 2.0`` as ``inf`` instead of NaN, while the + cfunc agreed with numpy all along. + """ + nan = float("nan") + inf = float("inf") + # ``floor(a / b)`` disagrees with numpy on the first four pairs: inf and + # -inf where numpy gives NaN, 10.0 rather than 9.0 for 1.0 // 0.1, and + # -0.0 rather than -1.0 for -2.0 // inf. + xs = [inf, -inf, 1.0, -2.0, 2.0, -2.0, 0.0, nan, -7.0, 7.0, -0.0, 7.5] + ys = [2.0, 2.0, 0.1, inf, 0.0, 0.0, 0.0, 2.0, 2.0, -2.0, 4.0, 2.5] + for np_dtype, name in ((np.float64, "_FloorDivF64UDT"), (np.float32, "_FloorDivF32UDT")): + udt = dtypes.register_anonymous( + np.dtype([("fdv_a", np_dtype), ("fdv_b", np_dtype)], align=True), name + ) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.floor_divide(np.array(xs, np_dtype), np.array(ys, np_dtype)) + v, w = _udt_vectors(udt, xs, ys) + result = binary.floordiv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))], np_dtype) + np.testing.assert_array_equal(got, expected, err_msg=f"{udt_op_path} {np_dtype.__name__}") + # ``assert_array_equal`` reads -0.0 and 0.0 as equal, so the sign of a + # zero quotient needs its own assertion. It is the whole job of the + # ``copysign`` branch the C JIT kernel emits for an exact-zero result. + np.testing.assert_array_equal( + np.signbit(got), + np.signbit(expected), + err_msg=f"{udt_op_path} {np_dtype.__name__} sign of zero", + ) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_integer_division_by_zero_is_defined(udt_op_path): + """Integer division by zero must return a value rather than trap. + + The C JIT kernel divided in integers, where a zero divisor is undefined + behaviour: on x86-64 ``idiv`` raises #DE, which is SIGFPE and process + death rather than an exception. AArch64's ``sdiv`` returns 0 and does not + trap, so this cannot be exhibited on an arm64 machine. The same trap + fires on ``INT_MIN / -1``, whose quotient is not representable. + + ``truediv`` now gives float64 fields, as on built-in integer vectors, so a + zero divisor gives an infinity there. ``floordiv`` stays in integers, where + the values are a choice: a zero divisor gives 0, as ``np.floor_divide`` + does, and ``INT_MIN // -1`` wraps to ``INT_MIN``, as numpy does. + """ + signed = dtypes.register_anonymous( + np.dtype([("dvz_a", np.int32), ("dvz_b", np.int8)], align=True), "_DivZeroSignedUDT" + ) + unsigned = dtypes.register_anonymous( + np.dtype([("dvz_c", np.uint32), ("dvz_d", np.uint64)], align=True), "_DivZeroUnsignedUDT" + ) + inf = float("inf") + v, w = _udt_vectors(signed, [7, -7, 100, -128], [0, 0, 3, -1]) + result = binary.truediv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [inf, -inf, 100 / 3, 128.0], udt_op_path + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got[:3] == [0, 0, 33], udt_op_path + # ``-128 // -1`` is the second trapping case; numpy wraps it to INT8_MIN. + result = binary.floordiv(v & w).new() + assert result[3].new().value[1] == np.iinfo(np.int8).min, udt_op_path + + v, w = _udt_vectors(unsigned, [7, 9, 100, 5], [0, 0, 3, 2]) + result = binary.truediv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [inf, inf, 100 / 3, 2.5], udt_op_path + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [0, 0, 33, 2], udt_op_path + + # Floor division still floors for signed operands of mixed sign. + v, w = _udt_vectors(signed, [-7, 7, -9, 11], [2, -2, 2, 3]) + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [-4, -4, -5, 3], udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not dtypes._supports_complex") +def test_udt_complex_truediv_by_zero(udt_op_path): + """``binary.truediv`` on a complex field survives a zero divisor. + + Numba's complex division raises ``ZeroDivisionError`` unconditionally, + outside the error model's control, so the cfunc left the element unwritten + while the C JIT kernel returned numpy's infinities. Reading back the + abandoned element gave uninitialized memory, or the previous element's + answer, either of which looks like a plausible value. + """ + udt = dtypes.register_anonymous( + np.dtype([("cxz_a", np.complex128)], align=True), "_ComplexDivZeroUDT" + ) + xs = [3 + 4j, 0j, 1 + 1j, 2 - 2j] + ys = [0j, 0j, 2 + 0j, 0j] + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.array(xs) / np.array(ys) + np.testing.assert_array_equal(got, expected, err_msg=udt_op_path) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not dtypes._supports_complex") +def test_udt_complex_times_truediv_with_infinities_as_numpy(udt_op_path): + """Complex ``*`` and ``/`` on a UDT field answer as numpy does at infinities. + + From Python 3.14, CPython recovers infinities that the textbook formulas + leave as ``nan`` (C99 Annex G), so ``1j * (inf+infj)`` is ``-inf+infj`` + there and ``nan+nanj`` in numpy. Numba 0.68 follows CPython, so a cfunc + built on Numba's operators answered as Python and the C JIT kernel as + numpy. Both now spell numpy's formulas out, on any Python and Numba. + """ + inf, nan = np.inf, np.nan + # The last pair is finite, with an exact product and quotient: numpy + # divides by multiplying with ``1 / denominator``, which can round the last + # bit differently from CPython's division, which the kernels use. + xs = [1j, -2.5j, complex(nan, inf), complex(inf, nan), 1 + 1j, complex(inf, 0.0), -1 + 7j] + ys = [complex(inf, inf), complex(inf, -inf), 2 + 1j, 1 + 1j, complex(inf, inf), 1j, 1 + 1j] + rec = dtypes.register_anonymous( + np.dtype([("cxi_s", np.complex64), ("cxi_d", np.complex128)], align=True), "_CxInfRec" + ) + arr = dtypes.register_anonymous(np.dtype((np.complex64, (3,))), "_CxInfArr3") + for udt, leaves in [(rec, [0, 1]), (arr, [0, 2])]: + v, w = _udt_vectors(udt, xs, ys) + for gb_op, reference in [(binary.times, np.multiply), (binary.truediv, np.true_divide)]: + result = gb_op(v & w).new() + for leaf in leaves: + got = np.array([result[i].new().value[leaf] for i in range(len(xs))]) + with np.errstate(invalid="ignore"): + expected = reference(np.array(xs, got.dtype), np.array(ys, got.dtype)) + msg = f"{gb_op.name} {got.dtype} {udt_op_path}" + np.testing.assert_array_equal(got, expected, err_msg=msg) + + +@pytest.mark.skipif("not supports_udfs") +# 136-byte UDT, which SS < 9 rejects; see test_udt_large_array. +@pytest.mark.skipif( + "ss_version_major < 9", + reason="SuiteSparse < 9 rejects a 136-byte UDT on builds without VLA support", +) +def test_udt_float_truediv_by_zero_is_infinite(udt_op_path): + """A zero divisor on a float field gives numpy's infinity, not a lost element. + + Unlike the integer case, nothing here is guarded: the generated code + divides and lets IEEE produce the infinity. That only holds because the + generated wrapper is compiled under Numba's numpy error model, which + nothing else in the suite pins down. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (17,))), "_FloatDivZeroArr17") + xs = [1.0, -1.0, 0.0, 6.0] + ys = [0.0, 0.0, 0.0, 3.0] + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.array(xs) / np.array(ys) + np.testing.assert_array_equal(got, expected, err_msg=udt_op_path) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_ops_match_record_ops(udt_op_path): + """The array-UDT codegen carries the same division and NaN fixes as records. + + Records and flat arrays go through separate branches in both the Numba + and the C JIT generators, so each fix has to land in both. + """ + nan = float("nan") + inf = float("inf") + float_udt = dtypes.register_anonymous(np.dtype((np.float64, (13,))), "_ArrOpsF64") + xs = [inf, 1.0, -2.0, nan, -7.0, 2.0] + ys = [2.0, 0.1, inf, 2.0, 2.0, 1.0] + v, w = _udt_vectors(float_udt, xs, ys) + # The reference computations touch inf and nan, and numpy raises the FP + # invalid flag for them on some platforms (Linux and Windows, via fmod) + # but not others; pyproject promotes the RuntimeWarning to an error. + with np.errstate(divide="ignore", invalid="ignore"): + expected_floordiv = np.floor_divide(np.array(xs), np.array(ys)) + expected_min = np.fmin(np.array(xs), np.array(ys)) + np.testing.assert_array_equal( + [binary.floordiv(v & w).new()[i].new().value[0] for i in range(len(xs))], + expected_floordiv, + err_msg=udt_op_path, + ) + # ``np.fmin``, not ``np.minimum``: ``binary.min`` is SuiteSparse's + # ``GrB_MIN_FP64``, which ignores a NaN operand rather than propagating it. + np.testing.assert_array_equal( + [binary.min(v & w).new()[i].new().value[0] for i in range(len(xs))], + expected_min, + err_msg=udt_op_path, + ) + + int_udt = dtypes.register_anonymous(np.dtype((np.int64, (6,))), "_ArrOpsI64") + ixs = [10**18, 7, -7, 100, -9, 5] + iys = [3, 0, 0, 3, 2, 2] + v, w = _udt_vectors(int_udt, ixs, iys) + result = binary.truediv(v & w).new() + got = [result[i].new().value[0] for i in range(len(ixs))] + with np.errstate(divide="ignore"): + assert got == list(np.array(ixs) / np.array(iys)), udt_op_path + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(len(ixs))] + assert got == [333333333333333333, 0, 0, 33, -5, 2], udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_multidim_array_ops_match_numpy(udt_op_path): + """Built-in ops on a multi-dimensional array UDT agree with numpy on both paths. + + The C JIT typedef flattens any rank to ``double v [N]`` and the Numba + wrapper walks the same flat run, so a 2-D UDT covers codegen that the 1-D + cases reach only by accident of both being contiguous. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (3, 2))), "_ArrOps2D") + xs = [1.0, -7.0, float("inf"), 2.0] + ys = [0.1, 2.0, 2.0, 0.0] + v, w = _udt_vectors(udt, xs, ys) + for gb_op, reference in ( + (binary.floordiv, np.floor_divide), + (binary.truediv, np.true_divide), + # ``fmin`` rather than ``minimum``: these inputs carry no NaN, so the + # two agree here, but ``binary.min`` is the NaN-ignoring one. + (binary.min, np.fmin), + ): + result = gb_op(v & w).new() + element = result[0].new().value + assert element.shape == (3, 2) + got = np.array([result[i].new().value[1, 1] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = reference(np.array(xs), np.array(ys)) + np.testing.assert_array_equal(got, expected, err_msg=f"{udt_op_path} {gb_op.name}") + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_tuple_return_binaryop(record_udt): + v, w = _record_pair(record_udt) + + def _add_udt(x, y): + return (x["a"] + y["a"], x["b"] + y["b"]) # pragma: no cover (numba) + + add_udt = BinaryOp.register_anonymous(_add_udt, "test_add_udt_b", is_udt=True) + result = add_udt(v & w).new() + assert result.dtype == record_udt + expected = _record_expected(record_udt, [(11, 22.0), (33, 44.0), (55, 66.0)]) assert result.isequal(expected) - assert op.binom is binary.binom -def test_builtins(): - v1 = Vector.from_coo([0, 1, 2], [1, 2, 3]) - v2 = Vector.from_coo([0, 1, 2], [3, 2, 1]) - result = v1.ewise_mult(v2, min).new() - expected = Vector.from_coo([0, 1, 2], [1, 2, 1]) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_tuple_return_unaryop_vector(record_udt): + v, _ = _record_pair(record_udt) + + def _double_udt(val): + return (val["a"] * 2, val["b"] * 2.0) # pragma: no cover (numba) + + double_udt = UnaryOp.register_anonymous(_double_udt, "test_double_udt_v", is_udt=True) + result = v.apply(double_udt).new() + assert result.dtype == record_udt + expected = _record_expected(record_udt, [(2, 4.0), (6, 8.0), (10, 12.0)]) assert result.isequal(expected) - v1(max) << v2 - expected = Vector.from_coo([0, 1, 2], [3, 2, 3]) - assert v1.isequal(expected) -def test_op_ss(): - if suitesparse: - gb.unary.ss.positioni - gb.binary.ss.firsti - gb.semiring.ss.max_secondj - gb.op.ss.positionj - gb.agg.ss.argmin - else: - with pytest.raises(AttributeError, match="suitesparse"): - gb.unary.ss - with pytest.raises(AttributeError, match="suitesparse"): - gb.binary.ss - with pytest.raises(AttributeError, match="suitesparse"): - gb.semiring.ss - with pytest.raises(AttributeError, match="suitesparse"): - gb.op.ss - with pytest.raises(AttributeError, match="suitesparse"): - gb.agg.ss +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_tuple_return_unaryop_matrix(record_udt): + def _double_udt(val): + return (val["a"] * 2, val["b"] * 2.0) # pragma: no cover (numba) -def test_deprecated(): - with pytest.warns(DeprecationWarning, match="please use"): - gb.unary.erf - with pytest.warns(DeprecationWarning, match="please use `gb.indexunary.rowindex`"): - gb.unary.positioni - with pytest.warns(DeprecationWarning, match="please use"): - gb.binary.firsti - with pytest.warns(DeprecationWarning, match="please use"): - gb.semiring.min_firsti - with pytest.warns(DeprecationWarning, match="please use"): - gb.op.secondj - with pytest.warns(DeprecationWarning, match="please use"): - gb.agg.argmin + double_udt = UnaryOp.register_anonymous(_double_udt, "test_double_udt_m", is_udt=True) + M = Matrix(record_udt, nrows=2, ncols=2) + M[0, 0] = (1, 2.0) + M[0, 1] = (3, 4.0) + M[1, 0] = (5, 6.0) + M[1, 1] = (7, 8.0) + result = M.apply(double_udt).new() + assert result[0, 0].new() == (2, 4.0) + assert result[1, 1].new() == (14, 16.0) +@pytest.mark.skipif("not supports_udfs") @pytest.mark.slow -def test_is_idempotent(): - assert monoid.min.is_idempotent - assert monoid.max[int].is_idempotent - assert monoid.lor.is_idempotent - assert monoid.band.is_idempotent - if shouldhave(monoid.numpy, "gcd"): - assert monoid.numpy.gcd.is_idempotent - assert not monoid.plus.is_idempotent - assert not monoid.times[float].is_idempotent - if config["mapnumpy"] or shouldhave(monoid.numpy, "equal"): - assert not monoid.numpy.equal.is_idempotent - with pytest.raises(AttributeError): - binary.min.is_idempotent +def test_udt_tuple_return_monoid(record_udt): + """A Monoid built from a tuple-returning BinaryOp reduces field-by-field via ewise_add.""" + v, w = _record_pair(record_udt) + def _add_udt(x, y): + return (x["a"] + y["a"], x["b"] + y["b"]) # pragma: no cover (numba) -def test_ops_have_ss(): - modules = [unary, binary, monoid, semiring, indexunary, select, op] - if suitesparse: - for mod in modules: - assert mod.ss is not None + add_udt = BinaryOp.register_anonymous(_add_udt, "test_add_udt_mon", is_udt=True) + add_monoid = Monoid.register_anonymous(add_udt, (0, 0.0)) + result = add_monoid(v | w).new() + expected = _record_expected(record_udt, [(11, 22.0), (33, 44.0), (55, 66.0)]) + assert result.isequal(expected) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_monoid_op_must_return_its_input_type(record_udt): + """A UDT monoid whose op returns another type fails up front, with a message. + + It used to reach ``GrB_Monoid_new_UDT``, which raised DomainMismatch with an + empty message, also from ``udt in monoid``. + """ + + def _sum_a(x, y): + return x["a"] + y["a"] # pragma: no cover (numba) + + sum_a = BinaryOp.register_anonymous(_sum_a, "test_sum_a_mon", is_udt=True) + sum_a_monoid = Monoid.register_anonymous(sum_a, 0) + # Not "_BuiltinOpsRec": registering this record dtype again under another name, as + # test_udt_op_pickle does, renames the one DataType + with pytest.raises(TypeError, match=f"must return its input type, {record_udt}, not INT64"): + sum_a_monoid[record_udt] + assert record_udt not in sum_a_monoid + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_tuple_return_semiring(record_udt): + """A Semiring built from tuple-returning ops drives mxm correctly.""" + v, _ = _record_pair(record_udt) + + def _add_udt(x, y): + return (x["a"] + y["a"], x["b"] + y["b"]) # pragma: no cover (numba) + + def _first_udt(x, y): + return x # pragma: no cover (numba) + + add_udt = BinaryOp.register_anonymous(_add_udt, "test_add_udt_sr", is_udt=True) + add_monoid = Monoid.register_anonymous(add_udt, (0, 0.0)) + first_udt = BinaryOp.register_anonymous(_first_udt, "test_first_udt_sr", is_udt=True) + sr = Semiring.register_anonymous(add_monoid, first_udt) + assert sr(v @ v).new() == (9, 12.0) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_tuple_return_indexunary(record_udt): + v, _ = _record_pair(record_udt) + + def _idx_udt(val, idx, _col, thunk): + return (val["a"] * (idx + 1), val["b"] + thunk["b"]) # pragma: no cover (numba) + + idx_op = IndexUnaryOp.register_anonymous(_idx_udt, "test_idx_udt", is_udt=True) + result = v.apply(idx_op, (0, 100.0)).new() + assert result.dtype == record_udt + expected = _record_expected(record_udt, [(1, 102.0), (6, 104.0), (15, 106.0)]) + assert result.isequal(expected) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_tuple_return_3field(): + """Tuple-return works for records with more than two fields.""" + dtype3 = np.dtype([("x", np.int32), ("y", np.float64), ("z", np.int64)], align=True) + udt3 = dtypes.register_anonymous(dtype3) + v3 = Vector(udt3, size=2) + v3[0] = (1, 2.0, 3) + v3[1] = (4, 5.0, 6) + w3 = Vector(udt3, size=2) + w3[0] = (10, 20.0, 30) + w3[1] = (40, 50.0, 60) + + def _add3(x, y): + return (x["x"] + y["x"], x["y"] + y["y"], x["z"] + y["z"]) # pragma: no cover (numba) + + add3 = BinaryOp.register_anonymous(_add3, "test_add3", is_udt=True) + result = add3(v3 & w3).new() + expected3 = Vector(udt3, size=2) + expected3[0] = (11, 22.0, 33) + expected3[1] = (44, 55.0, 66) + assert result.isequal(expected3) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_input_with_scalar_return(record_udt): + """A UDF that reads a UDT but returns a scalar still works (no tuple unpacking).""" + v, w = _record_pair(record_udt) + + def _sum_fields(x, y): + return x["a"] + y["b"] # pragma: no cover (numba) + + sum_op = BinaryOp.register_anonymous(_sum_fields, "test_sum_fields", is_udt=True) + result = sum_op(v & w).new() + assert result[0].new() == 21.0 # 1 + 20.0 + assert result[1].new() == 43.0 # 3 + 40.0 + assert result[2].new() == 65.0 # 5 + 60.0 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_return_type_errors(): + """Friendly error when a UDT UDF returns a shape that doesn't match the input.""" + record_dtype = np.dtype([("a", np.int64), ("b", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_RetErrUDT") + + # Wrong-arity tuple return: UDT has 2 fields, UDF returns 3. + def _three(x, y): # pragma: no cover (numba; raises before execution) + return (x["a"] + y["a"], x["b"] + y["b"], 0) + + op_three = BinaryOp.register_anonymous(_three, is_udt=True) + with pytest.raises(UdfParseError, match="tuple of length 3.*expected 2"): + op_three[udt] + + # Neither a scalar, nor a tuple, nor an array that fits an array UDT input. + def _nothing(x, y): # pragma: no cover (numba; raises before execution) + return None + + def _array(x, y): # pragma: no cover (numba; raises before execution) + return np.zeros(2) + + for func in [_nothing, _array]: + with pytest.raises(UdfParseError, match="returned an unsupported type"): + BinaryOp.register_anonymous(func, is_udt=True)[udt] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_return_type_errors_array_udt(): + """Tuple return against an array UDT input should suggest a numpy array, + not "record UDT's fields" (which would be misleading). + """ + arr_dtype = np.dtype((np.float64, (4,))) + audt = dtypes.register_anonymous(arr_dtype, "_RetErrArrUDT") + + def _bad_tuple(x, y): # pragma: no cover (numba; raises before execution) + return (x[0] + y[0], x[1] + y[1], x[2] + y[2]) + + op = BinaryOp.register_anonymous(_bad_tuple, is_udt=True) + with pytest.raises(UdfParseError, match="array UDTs of shape.*numpy array"): + op[audt] + + # A UDT that is neither a record nor an array takes no tuple either, and no + # built-in op lifts to it. + raw = dtypes.register_anonymous(np.dtype("S8"), "_RetErrBytesUDT") + + def _pair(x, y): # pragma: no cover (numba; raises before execution) + return (1, 2) + + with pytest.raises(UdfParseError, match="returned an unsupported type"): + BinaryOp.register_anonymous(_pair, is_udt=True)[raw] + # An operand returned as it is keeps its type. + v = Vector(raw, size=2) + v[0] = b"abc" + first = BinaryOp.register_anonymous(lambda x, y: x, "_ret_err_first", is_udt=True) + assert v.apply(first, right=b"zz").new().to_coo()[1].tolist() == [b"abc"] + for lifted in [binary.plus, unary.ainv]: + with pytest.raises(KeyError, match="require a record dtype .* or an array dtype"): + lifted[raw] + + +@pytest.fixture(scope="module") +def record_udt(): + return dtypes.register_anonymous( + np.dtype([("a", np.int64), ("b", np.float64)], align=True), + "_BuiltinOpsRec", + ) + + +@pytest.fixture(scope="module") +def array_udt(): + return dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_BuiltinOpsArr") + + +def _record_pair(udt): + """Return ``(v, w)`` with overlapping entries used by the record-UDT ops tests.""" + v = Vector(udt, size=3) + v[0] = (1, 2.0) + v[1] = (3, 4.0) + v[2] = (5, 6.0) + w = Vector(udt, size=3) + w[0] = (10, 20.0) + w[1] = (30, 40.0) + w[2] = (50, 60.0) + return v, w + + +def _record_expected(udt, rows): + out = Vector(udt, size=len(rows)) + for i, row in enumerate(rows): + out[i] = row + return out + + +@pytest.mark.parametrize( + ("op_name", "expected_rows"), + [ + ("plus", [(11, 22.0), (33, 44.0), (55, 66.0)]), + ("minus", [(-9, -18.0), (-27, -36.0), (-45, -54.0)]), + ("times", [(10, 40.0), (90, 160.0), (250, 360.0)]), + ("truediv", [(0.1, 0.1), (0.1, 0.1), (0.1, 0.1)]), + ], +) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_builtin_binary_record(record_udt, op_name, expected_rows): + """Per-field arithmetic on a record UDT matches the scalar definition.""" + v, w = _record_pair(record_udt) + result = getattr(binary, op_name)(v & w).new() + if op_name == "truediv": + # The int64 field divides to float64, as INT64 does, so the result is + # the record with both fields float64. + expected_type = dtypes.lookup_dtype( + np.dtype([("a", np.float64), ("b", np.float64)], align=True) + ) else: - for mod in modules: - with pytest.raises(AttributeError): - mod.ss + expected_type = record_udt + assert result.dtype == expected_type + assert result.isequal(_record_expected(expected_type, expected_rows)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_builtin_floordiv_record(record_udt): + """``binary.floordiv`` on a record UDT applies ``//`` per field.""" + v, w = _record_pair(record_udt) + result = binary.floordiv(w & v).new() + assert result.isequal(_record_expected(record_udt, [(10, 10), (10, 10), (10, 10)])) + + +@pytest.mark.parametrize( + ("op_name", "expected_rows"), + [ + ("min", [(3, 1.0), (2, 6.0)]), + ("max", [(5, 4.0), (7, 8.0)]), + ], +) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_builtin_minmax_record(record_udt, op_name, expected_rows): + """``binary.min`` / ``binary.max`` pick winners per field independently.""" + v = Vector(record_udt, size=2) + v[0] = (5, 1.0) + v[1] = (2, 8.0) + w = Vector(record_udt, size=2) + w[0] = (3, 4.0) + w[1] = (7, 6.0) + result = getattr(binary, op_name)(v & w).new() + assert result.isequal(_record_expected(record_udt, expected_rows)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_unary_ainv_abs_record(record_udt): + v = Vector(record_udt, size=3) + v[0] = (1, 2.0) + v[1] = (3, 4.0) + v[2] = (5, 6.0) + neg = v.apply(unary.ainv).new() + assert neg.isequal(_record_expected(record_udt, [(-1, -2.0), (-3, -4.0), (-5, -6.0)])) + + mixed = Vector(record_udt, size=3) + mixed[0] = (-1, -2.0) + mixed[1] = (3, -4.0) + mixed[2] = (-5, 6.0) + assert ( + mixed.apply(unary.abs) + .new() + .isequal(_record_expected(record_udt, [(1, 2.0), (3, 4.0), (5, 6.0)])) + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_matrix_apply_unary(record_udt): + M = Matrix(record_udt, nrows=2, ncols=2) + M[0, 0] = (1, 2.0) + M[0, 1] = (3, 4.0) + M[1, 0] = (5, 6.0) + M[1, 1] = (7, 8.0) + result = M.apply(unary.ainv).new() + assert result[0, 0].new() == (-1, -2.0) + assert result[0, 1].new() == (-3, -4.0) + assert result[1, 0].new() == (-5, -6.0) + assert result[1, 1].new() == (-7, -8.0) + + +def _array_pair(udt): + a = Vector(udt, size=2) + a[0] = [1.0, 2.0, 3.0] + a[1] = [4.0, 5.0, 6.0] + b = Vector(udt, size=2) + b[0] = [10.0, 20.0, 30.0] + b[1] = [40.0, 50.0, 60.0] + return a, b + + +@pytest.mark.parametrize( + ("op_name", "expected"), + [ + ("plus", [[11.0, 22.0, 33.0], [44.0, 55.0, 66.0]]), + ("times", [[10.0, 40.0, 90.0], [160.0, 250.0, 360.0]]), + ("minus", [[-9.0, -18.0, -27.0], [-36.0, -45.0, -54.0]]), + ], +) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_builtin_binary_array(array_udt, op_name, expected): + """Per-element arithmetic on a fixed-shape array UDT matches the scalar definition.""" + a, b = _array_pair(array_udt) + result = getattr(binary, op_name)(a & b).new() + for i, row in enumerate(expected): + np.testing.assert_array_equal(result[i].new().value, row) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_unary_ainv_abs_array(array_udt): + a, _b = _array_pair(array_udt) + np.testing.assert_array_equal(a.apply(unary.ainv).new()[0].new().value, [-1.0, -2.0, -3.0]) + c = Vector(array_udt, size=2) + c[0] = [-1.0, 2.0, -3.0] + c[1] = [4.0, -5.0, 6.0] + abs_c = c.apply(unary.abs).new() + np.testing.assert_array_equal(abs_c[0].new().value, [1.0, 2.0, 3.0]) + np.testing.assert_array_equal(abs_c[1].new().value, [4.0, 5.0, 6.0]) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_jit_typedef(): + """Registering a UDT sets GxB_JIT_C_NAME and GxB_JIT_C_DEFINITION.""" + from graphblas.core import ffi, lib + from graphblas.core.operator.udt_codegen import _has_jit_set + + if not _has_jit_set: + pytest.skip("JIT not available") + + # Record UDT with valid identifier name. Use unique field names to avoid + # collisions with UDTs registered by other tests (the registry caches by dtype). + record_dtype = np.dtype([("jx", np.int64), ("jy", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "JitTypeTest") + buf = ffi.new("char[512]") + lib.GrB_Type_get_String(udt._carg, buf, lib.GxB_JIT_C_DEFINITION) + defn = ffi.string(buf).decode() + assert "int64_t jx" in defn + assert "double jy" in defn + assert "JitTypeTest" in defn + + # Array UDT, of a length no other test uses: a literal can register float64 + # arrays of common lengths under another name first, in any test order. + arr_dtype = np.dtype((np.float64, (31,))) + arr_udt = dtypes.register_anonymous(arr_dtype, "Vec31") + lib.GrB_Type_get_String(arr_udt._carg, buf, lib.GxB_JIT_C_DEFINITION) + defn = ffi.string(buf).decode() + assert "double v [31]" in defn + assert "Vec31" in defn + + # 2D array UDT + mat_dtype = np.dtype((np.int32, (5, 5))) + mat_udt = dtypes.register_anonymous(mat_dtype, "Mat5x5") + lib.GrB_Type_get_String(mat_udt._carg, buf, lib.GxB_JIT_C_DEFINITION) + defn = ffi.string(buf).decode() + assert "int32_t" in defn + assert "Mat5x5" in defn + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_jit_op_definitions(): + """Auto-compiled UDT ops carry the JIT C name and source.""" + from graphblas.core import ffi, lib + from graphblas.core.operator.udt_codegen import _has_jit_set + + if not _has_jit_set: + pytest.skip("JIT not available") + + record_dtype = np.dtype([("jp", np.int64), ("jq", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "JitOpTest") + buf = ffi.new("char[1024]") + + # Binary op JIT definitions + for op_name, expected_c_op in [("plus", "+"), ("minus", "-"), ("times", "*")]: + typed = getattr(binary, op_name)[udt] + lib.GrB_BinaryOp_get_String(typed.gb_obj, buf, lib.GxB_JIT_C_DEFINITION) + defn = ffi.string(buf).decode() + assert f"{op_name}_gbudt_JitOpTest" in defn + assert "jp" in defn + assert "jq" in defn + assert expected_c_op in defn + + # Unary op JIT definitions + typed = unary.ainv[udt] + lib.GrB_UnaryOp_get_String(typed.gb_obj, buf, lib.GxB_JIT_C_DEFINITION) + defn = ffi.string(buf).decode() + assert "ainv_gbudt_JitOpTest" in defn + assert "jp" in defn + + # Array UDT JIT definitions, of a length no other test uses (see test_udt_jit_typedef) + arr_dtype = np.dtype((np.float64, (37,))) + arr_udt = dtypes.register_anonymous(arr_dtype, "Vec37Jit") + typed = binary.plus[arr_udt] + lib.GrB_BinaryOp_get_String(typed.gb_obj, buf, lib.GxB_JIT_C_DEFINITION) + defn = ffi.string(buf).decode() + assert "plus_gbudt_Vec37Jit" in defn + assert "z->v[i] = (x->v[i]) + (y->v[i])" in defn + assert "i < 37" in defn + + typed = unary.ainv[arr_udt] + lib.GrB_UnaryOp_get_String(typed.gb_obj, buf, lib.GxB_JIT_C_DEFINITION) + defn = ffi.string(buf).decode() + assert "ainv_gbudt_Vec37Jit" in defn + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_auto_monoid(): + """Built-in monoids auto-lift to UDTs with the right identity per field.""" + record_dtype = np.dtype([("p", np.int64), ("q", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype) + + v = Vector(udt, size=3) + v[0] = (1, 2.0) + v[1] = (3, 4.0) + v[2] = (5, 6.0) + w = Vector(udt, size=3) + w[0] = (10, 20.0) + w[1] = (30, 40.0) + w[2] = (50, 60.0) + + # monoid.plus: reduce and ewise_add + result = v.reduce(monoid.plus).new() + assert result == (9, 12.0) + result = monoid.plus(v | w).new() + expected = Vector(udt, size=3) + expected[0] = (11, 22.0) + expected[1] = (33, 44.0) + expected[2] = (55, 66.0) + assert result.isequal(expected) + + # monoid.times: reduce + result = v.reduce(monoid.times).new() + assert result == (15, 48.0) + + # monoid.min: reduce + result = v.reduce(monoid.min).new() + assert result == (1, 2.0) + + # monoid.max: reduce + result = v.reduce(monoid.max).new() + assert result == (5, 6.0) + + # Identity correctness: reduce of single element returns element + single = Vector(udt, size=1) + single[0] = (42, 99.5) + for mon in [monoid.plus, monoid.times, monoid.min, monoid.max]: + assert single.reduce(mon).new() == (42, 99.5) + + # __contains__ + assert udt in monoid.plus + assert udt in monoid.times + assert udt in monoid.min + assert udt in monoid.max + + # ---- Array UDT ---- + arr_dtype = np.dtype((np.float64, (4,))) + arr_udt = dtypes.register_anonymous(arr_dtype) + + a = Vector(arr_udt, size=3) + a[0] = [1.0, 2.0, 3.0, 4.0] + a[1] = [5.0, 6.0, 7.0, 8.0] + a[2] = [9.0, 10.0, 11.0, 12.0] + + result = a.reduce(monoid.plus).new() + np.testing.assert_array_equal(result.value, [15.0, 18.0, 21.0, 24.0]) + + result = a.reduce(monoid.min).new() + np.testing.assert_array_equal(result.value, [1.0, 2.0, 3.0, 4.0]) + + result = a.reduce(monoid.max).new() + np.testing.assert_array_equal(result.value, [9.0, 10.0, 11.0, 12.0]) + + # monoid.any on UDTs must return an actual input value, never the identity. + # Regression: previously ``binary.any._numba_func`` used ``_first`` semantics, + # so the UDT-reduce fold ``acc = first(acc, v_i) = acc`` always left the + # accumulator at the (zero) identity. Now ``_second`` semantics, so the fold + # captures an actual value. + arr_any = a.reduce(monoid.any).new() + np.testing.assert_array_equal(arr_any.value, [9.0, 10.0, 11.0, 12.0]) + + rec_any = Vector(udt, size=2) + rec_any[0] = (7, 8.0) + rec_any[1] = (11, 12.0) + any_res = rec_any.reduce(monoid.any).new() + # Result must be one of the input tuples; reject the (0, 0.0) identity. + # Compare via Scalar.__eq__ over a tuple of candidates (a set would require + # Scalar to be hashable, which it isn't). + assert any_res in ((7, 8.0), (11, 12.0)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_auto_monoid_time_fields(): + """Monoid min and max generate no identity for a datetime64 or timedelta64 field. + + numpy files timedelta64 under ``np.integer``, so a timedelta64 field used to + reach ``np.iinfo`` and raise ``ValueError: Invalid integer data type 'm'``, + from ``udt in agg.min`` too, where a datetime64 field gets a KeyError. + """ + for code in ["M8[s]", "m8[s]"]: + udt = dtypes.register_anonymous(np.dtype([("tf_n", np.int64), ("tf_t", code)])) + for mon in [monoid.min, monoid.max]: + with pytest.raises(KeyError, match="field 'tf_t' has unsupported dtype"): + mon[udt] + assert udt not in mon + assert udt not in agg.min + assert udt not in agg.max + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_monoid_named_like_builtin_gets_no_identity(): + """Only the built-in monoids generate a UDT identity, not a user monoid of the same name. + + A user monoid named plus, times, min or max (after its op's function, or by + ``name=``) used to get that built-in's generated identity. The op here takes + the field-wise max, so the generated ``min`` identity, INT64_MAX in every + field, absorbed every value: the reduce below gave that instead of (5, 2). + """ + udt = dtypes.register_anonymous(np.dtype([("ni_a", np.int64), ("ni_b", np.int64)])) + v = Vector(udt, size=2) + v[0] = (1, 2) + v[1] = (5, -3) + + def _fieldwise_max(x, y): + return (max(x["ni_a"], y["ni_a"]), max(x["ni_b"], y["ni_b"])) # pragma: no cover (numba) + + fieldwise_max = BinaryOp.register_anonymous(_fieldwise_max, "min", is_udt=True) + named_by_op = Monoid.register_anonymous(fieldwise_max, None) + assert named_by_op.name == "min" + for mon in [named_by_op] + [ + Monoid.register_anonymous(fieldwise_max, None, name) for name in ["plus", "times", "max"] + ]: + with pytest.raises(KeyError, match="no identity value"): + mon[udt] + assert udt not in mon + lowest = np.iinfo(np.int64).min + assert v.reduce(Monoid.register_anonymous(fieldwise_max, (lowest, lowest))).new() == (5, 2) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize("shape", [(4,), (3, 2)]) +def test_udt_array_wrapper_stays_within_element(shape): + """The array-UDT cfunc wrapper writes one element's bytes and nothing more. + + The wrapper used to load and store each operand as a ``NestedArray`` + *value*, which Numba models as its full array descriptor (meminfo, parent, + nitems, itemsize, data, shape, strides): 56 bytes on 64-bit for a 1-D + element and 72 for a 2-D one, whatever the payload size. Both elements here + are smaller, and the old wrapper overran each by 24 bytes. SuiteSparse's + generic reduce keeps a UDT accumulator in a stack array sized to the + element, so ``monoid.any`` overflowed it on every fold; depending on the + build that clobbered a spilled pointer (segfault or SIGBUS) or silently + produced a wrong answer. + + ``binary.any`` compiles through the generic ``_numba_func`` branch of + ``BinaryOp._compile_udt``, the same one a user's ``is_udt=True`` op takes, + so its wrapper stands in for both. Driving it over heap buffers with slack + makes a regression trip an assert instead of corrupting a stack frame. The + two sentinels must differ: the old load/store copied the source element + plus its trailing bytes, so if ``y``'s slack held ``z``'s guard value, the + overrun rewrote the guard with identical bytes and went unseen. + """ + import ctypes + + import numba + + from graphblas.core.operator.udf import _get_udt_wrapper + + udt = dtypes.register_anonymous(np.dtype((np.float64, shape)), f"_ArrWrapPin{len(shape)}D") + + # Mirror the generic ``_numba_func`` branch of ``BinaryOp._compile_udt``. + numba_func = binary.any._numba_func + sig = (udt.numba_type, udt.numba_type) + numba_func.compile(sig) + numba_ret_type = numba_func.overloads[sig].signature.return_type + wrapper, wrapper_sig = _get_udt_wrapper( + numba_func, udt, udt, udt, numba_ret_type=numba_ret_type + ) + cfunc = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) + call = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p)(cfunc.address) + + itemsize = udt.np_type.itemsize + slack = 128 # the old wrapper overran by 24 bytes; leave generous headroom + z = np.full(itemsize + slack, 0xAB, dtype=np.uint8) + x = np.full(itemsize + slack, 0xCD, dtype=np.uint8) + y = np.full(itemsize + slack, 0xCD, dtype=np.uint8) + xvals = np.arange(1.0, itemsize // 8 + 1).reshape(shape) + yvals = 10 * xvals + x[:itemsize] = xvals.ravel().view(np.uint8) + y[:itemsize] = yvals.ravel().view(np.uint8) + call(z.ctypes.data, x.ctypes.data, y.ctypes.data) + + # ``any`` uses ``_second`` semantics, so the payload must be ``y``'s. + np.testing.assert_array_equal(z[:itemsize].view(np.float64).reshape(shape), yvals) + overrun = np.flatnonzero(z[itemsize:] != 0xAB) + assert overrun.size == 0, f"wrote {overrun.size} bytes past the element at offsets {overrun}" + + # Public-path smoke: the reduce whose stack accumulator the old wrapper + # overflowed. Kept after the byte-level checks so a regression fails the + # assert above instead of reaching code that may crash the process. + v = Vector(udt, size=3) + rows = [xvals, yvals, xvals + yvals] + for i, row in enumerate(rows): + v[i] = row + res = v.reduce(monoid.any).new() + assert any(np.array_equal(res.value, row) for row in rows) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_returns_new_array(): + """An array-UDT UDF may build its result instead of returning an operand. + + The wrapper hands the UDF a numpy view of each element in the UDT's + declared shape, so ordinary array expressions work and the result is copied + back element-wise. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (8,))), "_ArrRetUDT") + v = Vector(udt, size=2) + v[0] = np.arange(8.0) + v[1] = np.arange(8.0, 16.0) + w = Vector(udt, size=2) + w[0] = np.full(8, 100.0) + w[1] = np.full(8, 200.0) + + def _add(x, y): + return x + y # pragma: no cover (numba) + + add_op = BinaryOp.register_anonymous(_add, "_arr_ret_add", is_udt=True) + result = add_op(v & w).new() + np.testing.assert_array_equal(result[0].new().value, np.arange(8.0) + 100.0) + np.testing.assert_array_equal(result[1].new().value, np.arange(8.0, 16.0) + 200.0) + + def _double(x): + return x * 2 # pragma: no cover (numba) + + double_op = UnaryOp.register_anonymous(_double, "_arr_ret_double", is_udt=True) + np.testing.assert_array_equal(double_op(v).new()[0].new().value, np.arange(8.0) * 2) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_multidim_array_keeps_shape_in_udf(): + """A multi-dimensional array UDT reaches the UDF in its declared shape. + + The wrapper builds the operand with ``numba.carray(ptr, shape)``, so 2-D + indexing and ``.shape`` work. Addressing it as a flat run of elements + would still be memory-safe but would silently drop that metadata. + """ + udt2d = dtypes.register_anonymous(np.dtype((np.float64, (2, 3))), "_ArrRet2D") + m = Vector(udt2d, size=1) + m[0] = np.arange(6.0).reshape(2, 3) + n = Vector(udt2d, size=1) + n[0] = np.full((2, 3), 10.0) + + def _add_corner(x, y): + # Fails to compile unless `x` really is 2-D with shape metadata. + return x + y[0, 0] + x.shape[1] # pragma: no cover (numba) + + add_2d = BinaryOp.register_anonymous(_add_corner, "_arr_ret_add_2d", is_udt=True) + np.testing.assert_array_equal( + add_2d(m & n).new()[0].new().value, np.arange(6.0).reshape(2, 3) + 10.0 + 3 + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_shape_errors(): + """An array-UDT UDF whose result can't fill the element is rejected when typed. + + Numba's ``Array`` type records ``ndim`` but not extents, so this is not a + type error. It used to reach the cfunc, where the shape mismatch raises in + a context that swallows the exception, handing the caller an uninitialized + element and no error. + """ + # Shapes unique to this test: ``register_anonymous`` caches by dtype and + # freezes the JIT C name at first registration, so sharing a shape with + # another test makes both order-dependent. + udt9 = dtypes.register_anonymous(np.dtype((np.float64, (9,))), "_ShapeErr9") + + def _truncate(x): # pragma: no cover (numba) + return x[:2] + + op = UnaryOp.register_anonymous(_truncate, "_shape_err_trunc", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op[udt9] + + # IndexUnaryOp passes the row and column between x and the thunk, and so + # must the probe: with the thunk's array in an index slot, ``i + j`` would + # not type as a slice bound. The probe's indices are 1, so this is (2,). + def _head(x, i, j, thunk): # pragma: no cover (numba) + return (x + thunk)[: i + j] + + op = IndexUnaryOp.register_anonymous(_head, "_shape_err_head", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op[udt9] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_array_leaf_shape_errors(): + """A record UDF that under-fills an array leaf, or can return None, is rejected. + + The wrapper slice-assigns array leaves, so a short return raises inside + the cfunc and abandons the write part-way: leaves after it keep whatever + SuiteSparse had in the buffer, scalar leaves included. A ``None`` leaf + raises there too. + """ + spec = np.dtype([("rl_vec", np.float64, (3,)), ("rl_tag", np.int64)], align=True) + udt = dtypes.register_anonymous(spec, "_RecLeafShape") + + def _short(x, y): # pragma: no cover (numba) + return (x["rl_vec"][:2], x["rl_tag"]) + + op = BinaryOp.register_anonymous(_short, "_rec_leaf_short", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) for field .* holds \(3,\)"): + op[udt] + + # ``v if cond else None`` types as Optional. Numba unwraps it where the + # wrapper assigns, and a ``None`` raises there like a short array does, so + # the type alone rejects it: no probe, and no dependence on which branch + # the sample values take. A scalar leaf fails the same way. + def _maybe_none(x, y): # pragma: no cover (numba) + return (x["rl_vec"] + y["rl_vec"] if x["rl_tag"] > 0 else None, x["rl_tag"]) + + op = BinaryOp.register_anonymous(_maybe_none, "_rec_leaf_maybe_none", is_udt=True) + with pytest.raises(UdfParseError, match=r"can return None for field \['rl_vec'\]"): + op[udt] + + def _maybe_no_tag(x, y): # pragma: no cover (numba) + return (x["rl_vec"], x["rl_tag"] if x["rl_tag"] > 0 else None) + + op = BinaryOp.register_anonymous(_maybe_no_tag, "_rec_leaf_maybe_no_tag", is_udt=True) + with pytest.raises(UdfParseError, match=r"can return None for field \['rl_tag'\]"): + op[udt] + + # A leaf that is always ``None`` would otherwise fail the wrapper compile + # with Numba's full traceback; it gets the same one-line diagnostic. + def _no_tag(x, y): # pragma: no cover (numba) + return (x["rl_vec"], None) + + op = BinaryOp.register_anonymous(_no_tag, "_rec_leaf_no_tag", is_udt=True) + with pytest.raises(UdfParseError, match=r"can return None for field \['rl_tag'\]"): + op[udt] + + # An array for a scalar leaf would fail the wrapper compile the same way + # (Numba has no record setitem for it), whether it is an operand's field + # or built. + def _arr_for_tag(x, y): # pragma: no cover (numba) + return (x["rl_vec"], x["rl_vec"]) + + op = BinaryOp.register_anonymous(_arr_for_tag, "_rec_leaf_arr_for_tag", is_udt=True) + with pytest.raises(UdfParseError, match=r"returned an array for field \['rl_tag'\]"): + op[udt] + + def _built_for_tag(x, y): # pragma: no cover (numba) + return (x["rl_vec"], x["rl_vec"] + y["rl_vec"]) + + op = BinaryOp.register_anonymous(_built_for_tag, "_rec_leaf_built_for_tag", is_udt=True) + with pytest.raises(UdfParseError, match=r"returned an array for field \['rl_tag'\]"): + op[udt] + + def _full(x, y): # pragma: no cover (numba) + return (x["rl_vec"] + y["rl_vec"], x["rl_tag"] + y["rl_tag"]) + + op = BinaryOp.register_anonymous(_full, "_rec_leaf_full", is_udt=True) + v = Vector(udt, size=1) + v[0] = ([1.0, 2.0, 3.0], 7) + w = Vector(udt, size=1) + w[0] = ([4.0, 5.0, 6.0], 8) + got = v.ewise_mult(w, op).new()[0].new().value + np.testing.assert_array_equal(got["rl_vec"], [5.0, 7.0, 9.0]) + assert got["rl_tag"] == 15 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_broadcast_return(): + """A return that broadcasts to the element fills it, and is not rejected. + + The wrapper slice-assigns and numpy broadcasts on assignment, so a ``(1,)`` + return legitimately fills every slot of a ``(6,)`` element. Requiring an + exact shape would refuse this, which works. + """ + udt6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_BCast6") + + def _fill(x): # pragma: no cover (numba) + return x[:1] + 10.0 + + op1 = UnaryOp.register_anonymous(_fill, "_bcast_fill", is_udt=True) + assert op1[udt6].return_type is udt6 + v = Vector(udt6, size=1) + v[0] = np.arange(1.0, 7.0) + np.testing.assert_array_equal(v.apply(op1).new()[0].new().value, [11.0] * 6) + + # A row broadcast across a 2-D element: the same rule one rank up. + udt42 = dtypes.register_anonymous(np.dtype((np.float64, (4, 2))), "_BCast42") + + def _fill_rows(x): # pragma: no cover (numba) + return x[:1, :] + 100.0 + + op2 = UnaryOp.register_anonymous(_fill_rows, "_bcast_fill_rows", is_udt=True) + assert op2[udt42].return_type is udt42 + v2 = Vector(udt42, size=1) + v2[0] = np.arange(8.0).reshape(4, 2) + np.testing.assert_array_equal( + v2.apply(op2).new()[0].new().value, np.tile([100.0, 101.0], (4, 1)) + ) + + # The other side of the boundary: (2,) does not broadcast to (6,), Numba's + # slice-assign raises on it, and it stays rejected. + def _short(x): # pragma: no cover (numba) + return x[:2] + 10.0 + + op3 = UnaryOp.register_anonymous(_short, "_bcast_short", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op3[udt6] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_leaf_broadcast_return(): + """A broadcastable array leaf fills its field, and later leaves still land. + + Same boundary as the array case, and + ``test_udt_record_array_leaf_shape_errors`` holds the rejecting side. The + scalar leaf is worth asserting because a leaf that raises in the cfunc + abandons the write, leaving every leaf after it as SuiteSparse had it. + """ + spec = np.dtype([("bc_vec", np.float64, (11,)), ("bc_tag", np.int64)], align=True) + udt = dtypes.register_anonymous(spec, "_RecLeafBCast") + + def _fill_leaf(x, y): # pragma: no cover (numba) + return (x["bc_vec"][:1] + y["bc_vec"][:1], x["bc_tag"] + y["bc_tag"]) + + op1 = BinaryOp.register_anonymous(_fill_leaf, "_rec_leaf_bcast", is_udt=True) + v = Vector(udt, size=1) + v[0] = (np.arange(11.0), 7) + w = Vector(udt, size=1) + w[0] = (np.arange(11.0) + 1.0, 8) + got = v.ewise_mult(w, op1).new()[0].new().value + np.testing.assert_array_equal(got["bc_vec"], [1.0] * 11) + assert got["bc_tag"] == 15 + + # A scalar fills the leaf the same way. Its type says it is no array, so + # the check lets it through without running the UDF. + def _fill_leaf_scalar(x, y): # pragma: no cover (numba) + return (x["bc_vec"][0] + y["bc_vec"][0], x["bc_tag"] + y["bc_tag"]) + + op2 = BinaryOp.register_anonymous(_fill_leaf_scalar, "_rec_leaf_bcast_scalar", is_udt=True) + got = v.ewise_mult(w, op2).new()[0].new().value + np.testing.assert_array_equal(got["bc_vec"], [1.0] * 11) + assert got["bc_tag"] == 15 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_leaf_sequence_return(): + """A tuple or list fills an array leaf only at its exact length. + + Numba copies a sequence into the field element by element, and unlike an + array it does not broadcast one. A wrong length raises inside the cfunc, + and the element comes back holding whatever was in the buffer. A tuple's + length is in its type; a list's is learned by running the UDF. + """ + spec = np.dtype([("sq_v", np.float64, (3,)), ("sq_t", np.int64)], align=True) + udt = dtypes.register_anonymous(spec, "_RecLeafSeq") + v = Vector(udt, size=1) + v[0] = ([1.0, 2.0, 3.0], 7) + + def _tuple3(x, y): # pragma: no cover (numba) + return ((x["sq_v"][0], x["sq_v"][1], 9.0), x["sq_t"]) + + def _list3(x, y): # pragma: no cover (numba) + return ([x["sq_v"][0], x["sq_v"][1], 9.0], x["sq_t"]) + + for func in [_tuple3, _list3]: + op = BinaryOp.register_anonymous(func, f"_rec_leaf_seq{func.__name__}", is_udt=True) + got = v.ewise_mult(v, op).new()[0].new().value + np.testing.assert_array_equal(got["sq_v"], [1.0, 2.0, 9.0]) + assert got["sq_t"] == 7 + + def _tuple2(x, y): # pragma: no cover (numba) + return ((x["sq_v"][0], 9.0), x["sq_t"]) + + op = BinaryOp.register_anonymous(_tuple2, "_rec_leaf_seq_tuple2", is_udt=True) + with pytest.raises(UdfParseError, match=r"tuple of length 2 for field \['sq_v'\] of _RecL"): + op[udt] + + # As a (1,) array this would broadcast; as a tuple it does not. + def _tuple1(x, y): # pragma: no cover (numba) + return ((9.0,), x["sq_t"]) + + op = BinaryOp.register_anonymous(_tuple1, "_rec_leaf_seq_tuple1", is_udt=True) + with pytest.raises(UdfParseError, match=r"tuple of length 1 .* holds \(3,\)\. Return 3 values"): + op[udt] + + def _list2(x, y): # pragma: no cover (numba) + return ([x["sq_v"][0], 9.0], x["sq_t"]) + + op = BinaryOp.register_anonymous(_list2, "_rec_leaf_seq_list2", is_udt=True) + with pytest.raises(UdfParseError, match=r"list of length 2 .* when run on sample values"): + op[udt] + + # A sequence never fills a field of higher rank; Numba will not compile it. + spec2 = np.dtype([("sq_m", np.float64, (2, 2)), ("sq_n", np.int64)], align=True) + udt2 = dtypes.register_anonymous(spec2, "_RecLeafSeq2D") + + def _flat4(x, y): # pragma: no cover (numba) + return ((1.0, 2.0, 3.0, 4.0), x["sq_n"]) + + op = BinaryOp.register_anonymous(_flat4, "_rec_leaf_seq_flat4", is_udt=True) + with pytest.raises(UdfParseError, match=r"only fills a one-dimensional field"): + op[udt2] + + # Nor a scalar field, like an array (test_udt_record_array_leaf_shape_errors). + def _tuple_for_t(x, y): # pragma: no cover (numba) + return (x["sq_v"], (x["sq_t"], 1)) + + op = BinaryOp.register_anonymous(_tuple_for_t, "_rec_leaf_seq_tuple_for_t", is_udt=True) + with pytest.raises(UdfParseError, match=r"returned a tuple for field \['sq_t'\] .* a scalar"): + op[udt] + + def _list_for_t(x, y): # pragma: no cover (numba) + return (x["sq_v"], [x["sq_t"], 1]) + + op = BinaryOp.register_anonymous(_list_for_t, "_rec_leaf_seq_list_for_t", is_udt=True) + with pytest.raises(UdfParseError, match=r"returned a list for field \['sq_t'\] .* a scalar"): + op[udt] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_broadcast_matches_numba_slice_assign(): + """The shape check accepts exactly what the wrapper's slice-assign accepts. + + The check turns a silent cfunc failure into a registration error, so a + shape it rejects that Numba would have assigned is a false rejection, and + one it accepts that Numba raises on is the failure it exists to catch. Pin + both directions against Numba itself, including the two ranks where + broadcasting alone gives the wrong answer: ``(1, 6)`` fills a ``(6,)`` + destination because assignment drops leading ones, ``(6, 1)`` does not. + """ + import numba + + from graphblas.core.operator.udf import _fits_by_broadcast + + @numba.njit + def _assign(z, src): # pragma: no cover (numba) + z[:] = src + + for dst, src in [ + ((6,), ()), + ((6,), (1,)), + ((6,), (6,)), + ((6,), (2,)), + ((6,), (12,)), + ((6,), (1, 6)), + ((6,), (6, 1)), + ((2, 3), (1, 3)), + ((2, 3), (2, 1)), + ((2, 3), (1, 1)), + ((2, 3), (3,)), + ((2, 3), (2, 3)), + ((2, 3), (6,)), + ((2, 3), (3, 2)), + # Every leading one is dropped, not just the first. + ((6,), (1, 1, 6)), + ((2, 3), (1, 1, 3)), + ((2, 3), (2, 1, 3)), + ]: + try: + _assign(np.zeros(dst), np.ones(src)) + except ValueError: + numba_assigns = False + else: + numba_assigns = True + assert _fits_by_broadcast(src, dst) is numba_assigns, (src, dst) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +# _ProbeWhenRec's numpy repr is 142 chars; see test_udt_eq_nested_record_with_nan_leaf. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_udf_shape_check_runs_udf_only_for_built_arrays(monkeypatch): + """The shape check runs the UDF only when Numba's type lacks the extents. + + An operand, or operand field, returned as-is keeps its extents in Numba's + ``NestedArray`` type, so it is checked without running user code. That + includes rejecting one returned into a field of another shape. A tuple is + checked the same way, since its length is in its type. Only an array the + UDF builds, typed as a plain ``Array``, needs the probe; so does a list, + whose type has no length (``test_udt_record_leaf_sequence_return``). + """ + from graphblas.core.operator import udf as _udf + + probed = [] + real_probe = _udf._run_udf_probe + + def spy(numba_func, arg_dtypes): + probed.append(numba_func.py_func.__name__) + return real_probe(numba_func, arg_dtypes) + + monkeypatch.setattr(_udf, "_run_udf_probe", spy) + udt = dtypes.register_anonymous(np.dtype((np.float32, (18,))), "_ProbeWhen18") + rec = dtypes.register_anonymous( + np.dtype( + [("pw_a", np.float64, (3,)), ("pw_b", np.float64, (2,)), ("pw_n", np.int64)], + align=True, + ), + "_ProbeWhenRec", + ) + + def _second(x, y): # pragma: no cover (numba) + return y + + op = BinaryOp.register_anonymous(_second, "_probe_when_second", is_udt=True) + assert op[udt].return_type is udt + + def _fields(x, y): # pragma: no cover (numba) + return (x["pw_a"], y["pw_b"], x["pw_n"]) + + op = BinaryOp.register_anonymous(_fields, "_probe_when_fields", is_udt=True) + assert op[rec].return_type is rec + + def _swapped(x, y): # pragma: no cover (numba) + return (x["pw_b"], x["pw_a"], x["pw_n"]) + + op = BinaryOp.register_anonymous(_swapped, "_probe_when_swapped", is_udt=True) + with pytest.raises( + UdfParseError, match=r"shape \(2,\) for field \['pw_a'\] of _ProbeWhenRec, but" + ): + op[rec] + + # A scalar for an array leaf broadcasts into it; its type says it is no + # array, so there is nothing to probe. + def _scalars(x, y): # pragma: no cover (numba) + return (x["pw_n"], y["pw_n"], x["pw_n"] + y["pw_n"]) + + op = BinaryOp.register_anonymous(_scalars, "_probe_when_scalars", is_udt=True) + assert op[rec].return_type is rec + + # A tuple carries its length in its type, so it is checked without a run. + def _tuple_leaf(x, y): # pragma: no cover (numba) + return ((x["pw_a"][0], x["pw_a"][1], 9.0), y["pw_b"], x["pw_n"]) + + op = BinaryOp.register_anonymous(_tuple_leaf, "_probe_when_tuple_leaf", is_udt=True) + assert op[rec].return_type is rec + assert probed == [] + + def _sum(x, y): # pragma: no cover (numba) + return x + y + + op = BinaryOp.register_anonymous(_sum, "_probe_when_sum", is_udt=True) + assert op[udt].return_type is udt + assert probed == ["_sum"] + + # A UDF that raises on the probe's values is not checked, and still types. + def _no_ones(x, y): # pragma: no cover (numba) + if x[0] == 1: + raise ValueError("the probe's values") + return x + y + + op = BinaryOp.register_anonymous(_no_ones, "_probe_when_raises", is_udt=True) + assert op[udt].return_type is udt + assert probed == ["_sum", "_no_ones"] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_array_field_roundtrip(): + """A record UDT with an array field writes exactly that field's extent. + + Numba's record-field setitem copies the *destination* extent whatever the + source's length, so a short return used to read past the end of the source + array. The wrapper slice-assigns array leaves to make that a shape error. + """ + spec = np.dtype([("count", np.int64), ("vec", np.float64, (3,))], align=True) + udt = dtypes.register_anonymous(spec, "_RecArrField") + v = Vector(udt, size=2) + v[0] = (1, [1.0, 2.0, 3.0]) + v[1] = (2, [4.0, 5.0, 6.0]) + + def _combine(x, y): # pragma: no cover (numba) + return (x["count"] + y["count"], x["vec"] + y["vec"]) + + op = BinaryOp.register_anonymous(_combine, "_rec_arr_field", is_udt=True) + result = v.ewise_mult(v, op).new() + assert result[0].new().value["count"] == 2 + np.testing.assert_array_equal(result[0].new().value["vec"], [2.0, 4.0, 6.0]) + np.testing.assert_array_equal(result[1].new().value["vec"], [8.0, 10.0, 12.0]) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_returning_operand_keeps_its_type(): + """A UDF that returns an array-UDT operand as-is produces that operand's UDT. + + A layered ``FP64[5][2]`` and the flat ``FP64[2, 5]`` share the Numba type + ``nestedarray(float64, (2, 5))``, and ``lookup_dtype`` used to key UDTs on + it, so ``return y`` named whichever of the two had registered last. + SuiteSparse rejects the other one as a domain mismatch in a monoid or on + assignment into the operand's vector. The two now register as one UDT, and + a Numba type resolves through its numpy dtype, not registration order. + """ + flat = dtypes.register_anonymous(np.dtype((np.float64, (2, 5))), "_RetOperandFlat") + assert dtypes.register_anonymous(np.dtype((np.dtype((np.float64, (5,))), (2,)))) is flat + + def _second(x, y): # pragma: no cover (numba) + return y + + op = BinaryOp.register_anonymous(_second, "_ret_operand_second", is_udt=True) + assert op[flat].return_type is flat + + v = Vector(flat, size=2) + v[0] = np.arange(10.0).reshape(2, 5) + v[1] = np.arange(10.0, 20.0).reshape(2, 5) + w = Vector(flat, size=2) + w << op(v & v) + np.testing.assert_array_equal(w[1].new().value, np.arange(10.0, 20.0).reshape(2, 5)) + # ``binary.any`` resolves its return type the same way. + res = v.reduce(monoid.any).new().value + assert any(np.array_equal(res, v[i].new().value) for i in range(2)) + + # A record's array field returned as-is carries its extents too, so it + # resolves to the array UDT of that shape, whether or not one existed. + # FP32 keeps each UDT under 128 bytes, the most SuiteSparse < 9 accepts on + # builds without variable-length arrays (Windows). + rec = dtypes.register_anonymous( + np.dtype([("ro_n", np.int64), ("ro_vec", np.float32, (19,))], align=True), + "_RetOperandRec", + ) + udt14 = dtypes.register_anonymous(np.dtype((np.float32, (14,))), "_RetOperand14") + + def _field(x, y): # pragma: no cover (numba) + return x["ro_vec"] + + op2 = BinaryOp.register_anonymous(_field, "_ret_operand_field", is_udt=True) + assert op2[rec, udt14].return_type is dtypes.lookup_dtype(np.dtype((np.float32, (19,)))) + r = Vector(rec, size=1) + r[0] = (1, np.arange(19.0)) + u = Vector(udt14, size=1) + u[0] = np.zeros(14) + np.testing.assert_array_equal(r.ewise_mult(u, op2).new()[0].new().value, np.arange(19.0)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_building_array_takes_operand_shape(): + """An array the UDF builds takes the type of the one operand it can match. + + Numba types ``x + 1.0`` as a plain ``Array`` with an element type and a + rank but no extents, so only an operand can say how long the result is. + Two same-rank candidates, or none, is an error rather than a guess. + """ + udt15 = dtypes.register_anonymous(np.dtype((np.float64, (15,))), "_BuildArr15") + udt16 = dtypes.register_anonymous(np.dtype((np.float64, (16,))), "_BuildArr16") + + def _shift(x, y): # pragma: no cover (numba) + return x + 1.0 + + op = BinaryOp.register_anonymous(_shift, "_build_arr_shift", is_udt=True) + assert op[udt15].return_type is udt15 + with pytest.raises(UdfParseError, match="matches more than one input array UDT"): + op[udt15, udt16] + + def _fold(x, y): # pragma: no cover (numba) + return x.reshape(3, 5) + + op2 = BinaryOp.register_anonymous(_fold, "_build_arr_fold", is_udt=True) + with pytest.raises(UdfParseError, match="matches no input array UDT"): + op2[udt15] + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_udf_known_extents_never_match_by_rank(monkeypatch): + """A return that kept its extents is not matched to a same-rank operand. + + ``lookup_dtype`` names a UDT for known extents, registering one if no UDT + has that layout yet, so the rank matcher normally sees only an ``Array`` + the UDF built, which carries no extents. Registration can fail, though: + SuiteSparse builds without variable-length arrays reject a UDT over 128 + bytes before 9.0 and over 1024 from then on. Matching known extents to a + same-rank operand would then hand SuiteSparse an element of the wrong + size, so refuse instead. Simulated here, since this build registers the + type. + """ + import numba + + from graphblas.core.operator import udf as _udf + + # FP32 keeps this UDT under that 128-byte cap, so it registers everywhere. + udt21 = dtypes.register_anonymous(np.dtype((np.float32, (21,))), "_RankGuard21") + ret = numba.typeof(np.dtype((np.float32, (23,)))).dtype + # Same element type and rank, different length: what the rank matcher sees. + assert ret.dtype == udt21.numba_type.dtype + assert ret.ndim == udt21.numba_type.ndim + + real_lookup = _udf.lookup_dtype + + def refuse_nested(key, value=None): + if isinstance(key, numba.core.types.NestedArray): + raise ValueError("simulated: SuiteSparse refused to register this UDT") + return real_lookup(key, value) + + monkeypatch.setattr(_udf, "lookup_dtype", refuse_nested) + with pytest.raises(UdfParseError, match="matches no input array UDT"): + _udf._resolve_udt_return_type(ret, udt21) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_declared_with_nested_subarrays_gets_builtin_ops(): + """Built-in ops work on UDTs declared with nested subarray dtypes. + + numpy reports such a dtype's ``subdtype`` as ``(inner array dtype, outer + shape)``, which the auto-lift codegen took for ``(scalar, shape)``: + ``plus``, ``eq``, and ``ainv`` failed in Numba lowering, ``monoid.plus`` and + ``agg.sum`` could not build an identity, and a UDF returning a record + operand could come back as its flat-field twin. Registration flattens these + dtypes now, so they take the ordinary path. + """ + udt = dtypes.register_anonymous(np.dtype((np.dtype((np.int64, (3,))), (4,))), "_NestedDeclArr") + assert udt.np_type == np.dtype((np.int64, (4, 3))) + vals = np.arange(12).reshape(4, 3) + v = Vector(udt, size=2) + v[0] = vals + v[1] = vals * 10 + np.testing.assert_array_equal(binary.plus(v & v).new()[1].new().value, vals * 20) + assert binary.eq(v & v).new().to_coo()[1].all() + np.testing.assert_array_equal(unary.ainv(v).new()[0].new().value, -vals) + np.testing.assert_array_equal(v.reduce(monoid.plus).new().value, vals * 11) + np.testing.assert_array_equal(v.reduce(agg.sum).new().value, vals * 11) + + point = np.dtype((np.float64, (3,))) + rec = dtypes.register_anonymous( + np.dtype([("nd_n", np.int64), ("nd_tri", point, (2,))], align=True), "_NestedDeclRec" + ) + flat_twin = np.dtype([("nd_n", np.int64), ("nd_tri", np.float64, (2, 3))], align=True) + assert dtypes.register_anonymous(flat_twin) is rec + r = Vector(rec, size=2) + r[0] = (1, np.ones((2, 3))) + r[1] = (2, np.full((2, 3), 5.0)) + res = r.reduce(monoid.plus).new().value + assert res["nd_n"] == 3 + np.testing.assert_array_equal(res["nd_tri"], np.full((2, 3), 6.0)) + + def _first(x, y): # pragma: no cover (numba) + return x + + op = BinaryOp.register_anonymous(_first, "_nested_decl_first", is_udt=True) + assert op[rec].return_type is rec + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_auto_semiring(): + """Built-in semirings auto-lift to UDTs and drive ``mxm``/``mxv``/``vxm``.""" + record_dtype = np.dtype([("r", np.float64), ("s", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype) + + # Matrix-vector multiply with plus_times + A = Matrix(udt, nrows=2, ncols=2) + A[0, 0] = (1.0, 2.0) + A[0, 1] = (3.0, 4.0) + A[1, 0] = (5.0, 6.0) + A[1, 1] = (7.0, 8.0) + x = Vector(udt, size=2) + x[0] = (1.0, 1.0) + x[1] = (1.0, 1.0) + + result = semiring.plus_times(A @ x).new() + # [0] = (1*1 + 3*1, 2*1 + 4*1) = (4, 6) + # [1] = (5*1 + 7*1, 6*1 + 8*1) = (12, 14) + assert result[0].new() == (4.0, 6.0) + assert result[1].new() == (12.0, 14.0) + + # __contains__ + assert udt in semiring.plus_times + + # vxm + result = semiring.plus_times(x @ A).new() + # [0] = (1*1 + 1*5, 1*2 + 1*6) = (6, 8) + # [1] = (1*3 + 1*7, 1*4 + 1*8) = (10, 12) + assert result[0].new() == (6.0, 8.0) + assert result[1].new() == (10.0, 12.0) + + # mxm + eye = Matrix(udt, nrows=2, ncols=2) + eye[0, 0] = (1.0, 1.0) + eye[1, 1] = (1.0, 1.0) + result = semiring.plus_times(A @ eye).new() + assert result.isequal(A) + + # Array UDT semiring (the use case from GH discussion #298) + arr_dtype = np.dtype((np.float64, (3,))) + arr_udt = dtypes.register_anonymous(arr_dtype) + + M = Matrix(arr_udt, nrows=2, ncols=2) + M[0, 0] = [1.0, 2.0, 3.0] + M[0, 1] = [4.0, 5.0, 6.0] + M[1, 0] = [7.0, 8.0, 9.0] + M[1, 1] = [10.0, 11.0, 12.0] + ones = Vector(arr_udt, size=2) + ones[0] = [1.0, 1.0, 1.0] + ones[1] = [1.0, 1.0, 1.0] + + result = semiring.plus_times(M @ ones).new() + np.testing.assert_array_equal(result[0].new().value, [5.0, 7.0, 9.0]) + np.testing.assert_array_equal(result[1].new().value, [17.0, 19.0, 21.0]) + + # min_plus semiring on array UDT + result = semiring.min_plus(M @ ones).new() + np.testing.assert_array_equal(result[0].new().value, [2.0, 3.0, 4.0]) + np.testing.assert_array_equal(result[1].new().value, [8.0, 9.0, 10.0]) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_single_field_record(): + single_dtype = np.dtype([("val", np.float64)], align=True) + single_udt = dtypes.register_anonymous(single_dtype) + v = Vector(single_udt, 2) + v[0] = (3.0,) + v[1] = (7.0,) + w = Vector(single_udt, 2) + w[0] = (10.0,) + w[1] = (20.0,) + result = binary.plus(v & w).new() + assert result[0].new() == (13.0,) + assert result[1].new() == (27.0,) + assert v.reduce(monoid.plus).new() == (10.0,) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_bool_field_record(): + """A record UDT with a bool field works for plus (bool + bool yields int in Numba).""" + bool_dtype = np.dtype([("flag", np.bool_), ("count", np.int64)], align=True) + bool_udt = dtypes.register_anonymous(bool_dtype) + bv = Vector(bool_udt, 2) + bv[0] = (True, 1) + bv[1] = (False, 2) + bw = Vector(bool_udt, 2) + bw[0] = (True, 10) + bw[1] = (True, 20) + result = binary.plus(bv & bw).new() + # count field sums normally; flag is ``bool(a + b)``, so it's False only + # when both inputs are False. + assert result[0].new().value.tolist() == (True, 11) + assert result[1].new().value.tolist() == (True, 22) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_op_compilation_is_lazy(): + """Registering a UDT does not compile any ops for it; the first use does.""" + lazy_dtype = np.dtype([("lazy_a", np.int64), ("lazy_b", np.float64)], align=True) + lazy_udt = dtypes.register_anonymous(lazy_dtype, "LazyCheck") + assert (lazy_udt, lazy_udt) not in binary.plus._udt_ops + assert lazy_udt not in monoid.plus._udt_ops + binary.plus[lazy_udt] + assert (lazy_udt, lazy_udt) in binary.plus._udt_ops + # The monoid is independent of the binary op cache. + assert lazy_udt not in monoid.plus._udt_ops + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +# A SuiteSparse compiled without variable-length arrays (MSVC, so the Windows +# builds) rejects a user-defined type larger than GB_VLA_MAXSIZE with +# GrB_INVALID_VALUE. That ceiling was 128 bytes through SS 8.x and is 1024 as +# of 9.0, so skipping on SS < 9 covers the affected builds and costs one test +# on the rest. +@pytest.mark.skipif( + "ss_version_major < 9", + reason="SuiteSparse < 9 rejects an 800-byte UDT on builds without VLA support", +) +def test_udt_large_array(): + big_dtype = np.dtype((np.float64, (100,))) + big_udt = dtypes.register_anonymous(big_dtype) + a = Vector(big_udt, 2) + a[0] = list(range(100)) + a[1] = list(range(100, 200)) + b = Vector(big_udt, 2) + b[0] = [1.0] * 100 + b[1] = [1.0] * 100 + result = binary.plus(a & b).new() + np.testing.assert_array_equal(result[0].new().value, np.arange(1.0, 101.0)) + np.testing.assert_array_equal(result[1].new().value, np.arange(101.0, 201.0)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_int_array(): + int_arr_dtype = np.dtype((np.int32, (4,))) + int_arr_udt = dtypes.register_anonymous(int_arr_dtype) + iv = Vector(int_arr_udt, 2) + iv[0] = [1, 2, 3, 4] + iv[1] = [5, 6, 7, 8] + iw = Vector(int_arr_udt, 2) + iw[0] = [10, 20, 30, 40] + iw[1] = [50, 60, 70, 80] + result = binary.times(iv & iw).new() + np.testing.assert_array_equal(result[0].new().value, [10, 40, 90, 160]) + np.testing.assert_array_equal(result[1].new().value, [250, 360, 490, 640]) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_expr_repr_does_not_crash(): + """``repr`` of UDT expressions returns a non-empty string mentioning the UDT. + + Regression pin: the expression types used to lack ``_expr_name`` for UDT + pointer return types, so ``repr`` raised. Pin both that the call returns + a non-empty string and that the UDT's dtype name appears in it, so a + future regression returning ``""`` or a generic placeholder still fails. + """ + record_dtype2 = np.dtype([("rx", np.float64), ("ry", np.float64)], align=True) + repr_udt = dtypes.register_anonymous(record_dtype2, "_ReprPinUdt") + rv = Vector(repr_udt, 2) + rv[:] = (1.0, 2.0) + rw = Vector(repr_udt, 2) + rw[:] = (10.0, 20.0) + M = Matrix(repr_udt, 2, 2) + M[:, :] = (1.0, 2.0) + for expr in (rv + 1, 1 + rv, rv * 2, -rv, rv + rw, M + 1): + text = repr(expr) + assert text, f"repr returned empty string for {expr!r}" + assert ( + repr_udt.name in text + ), f"repr did not mention the UDT name {repr_udt.name!r}: {text!r}" + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_ne_nan_simple_record(): + """Simple float-field record: NaN-bearing entries compare unequal under eq. + + Regression: the original implementation byte-compared records (with a + padding-byte mask) so two records whose float fields both held NaN + compared *equal*. The cfunc now reads each leaf and applies scalar + ``==`` / ``!=``, matching ``binary.eq[FP64](nan, nan) == False``. + """ + spec = np.dtype([("eq_a", np.float64), ("eq_b", np.float64)], align=True) + udt = dtypes.register_anonymous(spec, "_NaNEqSimple") + v1 = Vector(udt, size=2) + v2 = Vector(udt, size=2) + v1[0] = (1.0, np.nan) + v2[0] = (1.0, np.nan) + v1[1] = (1.0, 2.0) + v2[1] = (1.0, 2.0) + eq = v1.ewise_mult(v2, binary.eq[udt]).new() + ne = v1.ewise_mult(v2, binary.ne[udt]).new() + assert eq[0].new().value is False + assert eq[1].new().value is True + assert ne[0].new().value is True + assert ne[1].new().value is False + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_packed_mixed_width(): + """Packed (non-aligned) records compare by leaf, so padding bytes don't matter.""" + spec_packed = np.dtype([("pk_a", np.int32), ("pk_b", np.float64)]) + assert spec_packed.itemsize == 12 # packed: int32 + float64, no padding + udt_pk = dtypes.register_anonymous(spec_packed, "_NaNEqPacked") + vp1 = Vector(udt_pk, size=1) + vp2 = Vector(udt_pk, size=1) + vp1[0] = (1, 2.5) + vp2[0] = (1, 2.5) + assert vp1.ewise_mult(vp2, binary.eq[udt_pk]).new()[0].new().value is True + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +# SS < 9 has no GrB_NAME setter, so registration falls back to storing the +# numpy repr in the type name and warns when it does not fit in 128 chars. +# This dtype's repr is 133; how it serializes is not what the test is about. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_eq_nested_record_with_nan_leaf(): + """A NaN in a nested-record leaf still makes the outer record compare unequal.""" + nested = np.dtype( + [("n_id", np.int32), ("n_pt", [("n_x", np.float64), ("n_y", np.float64)])], + align=True, + ) + udt_n = dtypes.register_anonymous(nested, "_NaNEqNested") + vn1 = Vector(udt_n, size=2) + vn2 = Vector(udt_n, size=2) + vn1[0] = (1, (np.nan, 2.0)) + vn2[0] = (1, (np.nan, 2.0)) + vn1[1] = (1, (3.0, 4.0)) + vn2[1] = (1, (3.0, 4.0)) + eq_n = vn1.ewise_mult(vn2, binary.eq[udt_n]).new() + assert eq_n[0].new().value is False + assert eq_n[1].new().value is True + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_ne_array_with_nan_element(): + """Array UDT with a NaN element compares unequal under eq, equal under ne.""" + arr = np.dtype((np.float64, (3,))) + udt_a = dtypes.register_anonymous(arr, "_NaNEqArr") + va1 = Vector(udt_a, size=2) + va2 = Vector(udt_a, size=2) + va1[0] = [1.0, np.nan, 3.0] + va2[0] = [1.0, np.nan, 3.0] + va1[1] = [1.0, 2.0, 3.0] + va2[1] = [1.0, 2.0, 3.0] + eq_a = va1.ewise_mult(va2, binary.eq[udt_a]).new() + ne_a = va1.ewise_mult(va2, binary.ne[udt_a]).new() + assert eq_a[0].new().value is False + assert eq_a[1].new().value is True + assert ne_a[0].new().value is True + assert ne_a[1].new().value is False + + +@pytest.fixture(scope="module") +def broadcast_record_udt(): + return dtypes.register_anonymous( + np.dtype([("u", np.float64), ("v", np.float64)], align=True), + "_BroadcastRecUdt", + ) + + +@pytest.fixture(scope="module") +def broadcast_array_udt(): + return dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_BroadcastArrUdt") + + +def _broadcast_record_vec(udt): + v = Vector(udt, size=3) + v[0] = (1.0, 2.0) + v[1] = (3.0, 4.0) + v[2] = (5.0, 6.0) + return v + + +@pytest.mark.parametrize( + ("op_name", "scalar_dtype", "scalar_values", "expected_rows"), + [ + # commutative ops applied with UDT on the left + ("plus", "int", [10, 20, 30], [(11.0, 12.0), (23.0, 24.0), (35.0, 36.0)]), + ("times", "float", [2.0, 0.5, 10.0], [(2.0, 4.0), (1.5, 2.0), (50.0, 60.0)]), + ("min", "int", [10, 20, 30], [(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)]), + ("max", "int", [10, 20, 30], [(10.0, 10.0), (20.0, 20.0), (30.0, 30.0)]), + ], +) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_scalar_broadcast_udt_lhs( + broadcast_record_udt, op_name, scalar_dtype, scalar_values, expected_rows +): + """Scalar broadcasts to every field of a record UDT (UDT on the left).""" + udt = broadcast_record_udt + vec_udt = _broadcast_record_vec(udt) + vec_s = Vector.from_coo([0, 1, 2], scalar_values, dtype=scalar_dtype) + result = getattr(binary, op_name)(vec_udt & vec_s).new() + expected = _record_expected(udt, expected_rows) + assert result.isequal(expected) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_scalar_broadcast_commutativity(broadcast_record_udt): + """``plus`` is commutative across UDT/scalar broadcast; ``minus`` is not.""" + udt = broadcast_record_udt + vec_udt = _broadcast_record_vec(udt) + vec_int = Vector.from_coo([0, 1, 2], [10, 20, 30]) + + expected_plus = _record_expected(udt, [(11.0, 12.0), (23.0, 24.0), (35.0, 36.0)]) + assert binary.plus(vec_udt & vec_int).new().isequal(expected_plus) + assert binary.plus(vec_int & vec_udt).new().isequal(expected_plus) + + expected_minus_udt_lhs = _record_expected(udt, [(-9.0, -8.0), (-17.0, -16.0), (-25.0, -24.0)]) + expected_minus_int_lhs = _record_expected(udt, [(9.0, 8.0), (17.0, 16.0), (25.0, 24.0)]) + assert binary.minus(vec_udt & vec_int).new().isequal(expected_minus_udt_lhs) + assert binary.minus(vec_int & vec_udt).new().isequal(expected_minus_int_lhs) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_scalar_broadcast_plus_times(broadcast_array_udt): + """Scalar broadcasts to every element of an array UDT for commutative ops.""" + arr_udt = broadcast_array_udt + vec_arr = Vector(arr_udt, size=2) + vec_arr[0] = [1.0, 2.0, 3.0] + vec_arr[1] = [4.0, 5.0, 6.0] + vec_s = Vector.from_coo([0, 1], [10.0, 100.0]) + + # plus is commutative; both directions yield the same per-element broadcast. + res_lhs = binary.plus(vec_arr & vec_s).new() + res_rhs = binary.plus(vec_s & vec_arr).new() + np.testing.assert_array_equal(res_lhs[0].new().value, [11.0, 12.0, 13.0]) + np.testing.assert_array_equal(res_lhs[1].new().value, [104.0, 105.0, 106.0]) + np.testing.assert_array_equal(res_rhs[0].new().value, [11.0, 12.0, 13.0]) + np.testing.assert_array_equal(res_rhs[1].new().value, [104.0, 105.0, 106.0]) + + res_times = binary.times(vec_arr & vec_s).new() + np.testing.assert_array_equal(res_times[0].new().value, [10.0, 20.0, 30.0]) + np.testing.assert_array_equal(res_times[1].new().value, [400.0, 500.0, 600.0]) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_scalar_broadcast_minus_direction(broadcast_array_udt): + """For non-commutative ops on array UDTs, operand order is respected.""" + arr_udt = broadcast_array_udt + vec_arr = Vector(arr_udt, size=2) + vec_arr[0] = [1.0, 2.0, 3.0] + vec_arr[1] = [4.0, 5.0, 6.0] + vec_s = Vector.from_coo([0, 1], [10.0, 100.0]) + + res_scalar_lhs = binary.minus(vec_s & vec_arr).new() + np.testing.assert_array_equal(res_scalar_lhs[0].new().value, [9.0, 8.0, 7.0]) + + res_udt_lhs = binary.minus(vec_arr & vec_s).new() + np.testing.assert_array_equal(res_udt_lhs[0].new().value, [-9.0, -8.0, -7.0]) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_matrix_scalar_broadcast(broadcast_record_udt): + """Scalar/UDT broadcast also works on matrix-shaped operands.""" + udt = broadcast_record_udt + mat = Matrix(udt, nrows=2, ncols=2) + mat[:, :] = (1.0, 2.0) + mat_int = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [10, 20, 30, 40], nrows=2, ncols=2) + result = binary.plus(mat & mat_int).new() + assert result[0, 0].new() == (11.0, 12.0) + assert result[1, 1].new() == (41.0, 42.0) + + +# eq/ne broadcasting between a UDT and a scalar type. +# +# Before this fix, ``binary.eq(udt_vec & int_vec)`` silently reinterpreted +# the int cell as a UDT struct (reading past the cell) and produced +# byte-comparison nonsense that happened to look like plausible False/True. +# Now the scalar broadcasts to every leaf, so ``eq`` is true only when all +# leaves equal the scalar. + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_eq_ne_scalar_broadcast_record(): + record = dtypes.register_anonymous( + np.dtype([("u", np.float64), ("v", np.float64)], align=True), + name="_EqBcastUV", + ) + v_udt = Vector(record, size=3) + v_udt[0] = (10.0, 10.0) + v_udt[1] = (10.0, 20.0) # partial match + v_udt[2] = (1.0, 2.0) # no match + v_int = Vector.from_coo([0, 1, 2], [10, 10, 10]) + + eq_result = binary.eq(v_udt & v_int).new() + assert eq_result.dtype == dtypes.BOOL + expected_eq = Vector.from_coo([0, 1, 2], [True, False, False]) + assert eq_result.isequal(expected_eq) + + ne_result = binary.ne(v_udt & v_int).new() + assert ne_result.isequal(Vector.from_coo([0, 1, 2], [False, True, True])) + + # Reverse direction (scalar on left). + eq_rev = binary.eq(v_int & v_udt).new() + assert eq_rev.isequal(expected_eq) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_eq_ne_scalar_broadcast_nan_propagates(): + """A NaN leaf never equals anything, even another NaN.""" + # Distinct field names from the sibling record test: ``register_anonymous`` + # keys on ``np.dtype``, so reusing ``("u", "v")`` would alias the cached + # DataType across tests (see CLAUDE.md). + record = dtypes.register_anonymous( + np.dtype([("nan_u", np.float64), ("nan_v", np.float64)], align=True), + name="_EqBcastUVNan", + ) + v_nan = Vector(record, size=3) + v_nan[0] = (np.nan, 5.0) + v_nan[1] = (5.0, 5.0) + v_nan[2] = (np.nan, np.nan) + v_five = Vector.from_coo([0, 1, 2], [5.0, 5.0, 5.0]) + eq_nan = binary.eq(v_nan & v_five).new() + assert eq_nan.isequal(Vector.from_coo([0, 1, 2], [False, True, False])) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_ne_type_a_literal_as_builtin_comparisons_do(): + """``eq`` and ``ne`` beside a UDT type a Python number as a built-in comparison does. + + The number is weak, leaf by leaf, so ``fp32_udt == 0.1`` compares in + float32 and is True, as ``fp32_vec == 0.1`` and numpy are; it was strong, + and compared float32 0.1 with float64 0.1. A comparison keeps no result + type, so a literal with an int out of a leaf's range equals no element, + even past int64, where it raised. Beside an int64 leaf, it is not rounded + to equal one (in float64, the int64 ``2**63 - 1`` and ``2**63`` are equal). + numpy scalars stay strong. + """ + f32 = dtypes.register_anonymous(np.dtype((np.float32, (7,))), "_EqWeakF32") + i8 = dtypes.register_anonymous(np.dtype((np.int8, (7,))), "_EqWeakI8") + v = Vector(f32, size=1) + v[0] = np.full(7, 0.1) + assert (v == 0.1).new()[0].new().value + assert not (v != 0.1).new()[0].new().value + assert (v == (0.1,) * 7).new()[0].new().value + assert not (v == np.float64(0.1)).new()[0].new().value + builtin = Vector.from_coo([0], [0.1], dtype=dtypes.FP32) + assert (builtin == 0.1).new()[0].new().value + w = Vector(i8, size=1) + w[0] = np.arange(7) + for number in [300, -(2**63) - 1, 2**63, 2**64]: + assert not (w == number).new()[0].new().value + assert (w != number).new()[0].new().value + assert (w == tuple(range(7))).new()[0].new().value + assert not (w == (0.5, *range(1, 7))).new()[0].new().value + i64 = dtypes.register_anonymous(np.dtype((np.int64, (2,))), "_EqWeakI64") + x = Vector(i64, size=2) + x[0] = [2**63 - 1, 2**63 - 1] + x[1] = [-(2**63), 0] + for literal in [2**63, (2**63, 2**63), [2**63 - 1, 2**63], -(2**63) - 1, 10**400]: + assert (x == literal).new().to_coo()[1].tolist() == [False, False] + assert (x != literal).new().to_coo()[1].tolist() == [True, True] + assert (x == (2**63 - 1, 2**63 - 1)).new().to_coo()[1].tolist() == [True, False] + rec = dtypes.register_anonymous( + np.dtype([("a", np.int64), ("b", np.uint64)], align=True), "_EqWeakRec64" + ) + r = Vector(rec, size=1) + r[0] = (2**63 - 1, 2**64 - 1) + for literal in [(2**63, 2**64 - 1), {"a": 2**63 - 1, "b": 2**64}, (2**63 - 1, -1)]: + assert not (r == literal).new()[0].new().value + assert (r != literal).new()[0].new().value + assert (r == {"a": 2**63 - 1, "b": 2**64 - 1}).new()[0].new().value + + +def test_udt_float32_overflow_warns_at_the_callers_line(): + """A float too large for a float32 leaf is infinity, with a warning at the caller's line. + + numpy warned from inside python-graphblas, so the default filter showed it + once per process; a built-in FP32 warns at the caller's line, as numpy does. + """ + f32 = dtypes.register_anonymous(np.dtype((np.float32, (2,))), "_OverflowF32") + v = Vector(f32, size=1) + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + v[0] = [1e300, 1] + assert [w.filename for w in record] == [__file__] + assert v[0].new().value.tolist() == [np.inf, 1] + if supports_udfs: # UDT eq and plus need Numba where there is no C JIT + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + (v == 1e300).new() + assert [w.filename for w in record] == [__file__] + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + (v + 1e300).new() + assert [w.filename for w in record] == [__file__] + with np.errstate(over="ignore"): + v[0] = [1e300, 1] # no warning, as numpy + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_eq_ne_scalar_broadcast_array_1d(): + arr1d = dtypes.register_anonymous(np.dtype((np.float64, (3,))), name="_EqBcastA3") + v_a = Vector(arr1d, size=2) + v_a[0] = [1.0, 1.0, 1.0] + v_a[1] = [1.0, 2.0, 1.0] + v_one = Vector.from_coo([0, 1], [1.0, 1.0]) + eq_a = binary.eq(v_a & v_one).new() + assert eq_a.isequal(Vector.from_coo([0, 1], [True, False])) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_eq_ne_scalar_broadcast_array_2d(): + arr2d = dtypes.register_anonymous(np.dtype((np.float64, (2, 2))), name="_EqBcastA22") + v_a22 = Vector(arr2d, size=2) + v_a22[0] = [[3.0, 3.0], [3.0, 3.0]] + v_a22[1] = [[3.0, 3.0], [4.0, 3.0]] + v_three = Vector.from_coo([0, 1], [3.0, 3.0]) + eq_a22 = binary.eq(v_a22 & v_three).new() + assert eq_a22.isequal(Vector.from_coo([0, 1], [True, False])) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +# 131-char numpy repr; see test_udt_eq_nested_record_with_nan_leaf. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_eq_ne_scalar_broadcast_nested_record(): + inner = np.dtype([("a", np.float64), ("b", np.float64)], align=True) + nested = dtypes.register_anonymous( + np.dtype([("outer", np.float64), ("inner", inner)], align=True), + name="_EqBcastNest", + ) + v_n = Vector(nested, size=3) + v_n[0] = (5.0, (5.0, 5.0)) + v_n[1] = (5.0, (5.0, 6.0)) + v_n[2] = (1.0, (1.0, 1.0)) + v_5 = Vector.from_coo([0, 1, 2], [5.0, 5.0, 5.0]) + eq_n = binary.eq(v_n & v_5).new() + assert eq_n.isequal(Vector.from_coo([0, 1, 2], [True, False, False])) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_eq_ne_scalar_broadcast_record_with_array_subfield(): + rec_arr = dtypes.register_anonymous( + np.dtype([("a", np.float64), ("v", np.float64, (3,))], align=True), + name="_EqBcastRecArr", + ) + v_ra = Vector(rec_arr, size=2) + v_ra[0] = (7.0, [7.0, 7.0, 7.0]) + v_ra[1] = (7.0, [7.0, 8.0, 7.0]) + v_7 = Vector.from_coo([0, 1], [7.0, 7.0]) + eq_ra = binary.eq(v_ra & v_7).new() + assert eq_ra.isequal(Vector.from_coo([0, 1], [True, False])) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_eq_ne_rejects_incompatible_pairs(): + """eq/ne between two UDTs must reject mismatched structure rather than + silently byte-compare. + + The old code only consulted the first dtype when generating the leaf + chain, so two records with different field names (but identical + byte layout) compared as equal, and a record-vs-array pair compared + by reinterpreting one side as the other. + """ + uv = dtypes.register_anonymous( + np.dtype([("u", np.float64), ("v", np.float64)], align=True), + name="_EqRejUV", + ) + uw = dtypes.register_anonymous( + np.dtype([("u", np.float64), ("w", np.float64)], align=True), + name="_EqRejUW", + ) + arr = dtypes.register_anonymous(np.dtype((np.float64, (2,))), name="_EqRejA") + + v_uv = Vector(uv, size=1) + v_uv[0] = (1.0, 2.0) + v_uw = Vector(uw, size=1) + v_uw[0] = (1.0, 2.0) + v_arr = Vector(arr, size=1) + v_arr[0] = [1.0, 2.0] + + with pytest.raises(KeyError, match="record UDTs must share field names"): + binary.eq(v_uv & v_uw).new() + with pytest.raises(KeyError, match="record UDTs must share field names"): + binary.ne(v_uv & v_uw).new() + with pytest.raises(KeyError, match="cannot mix record and array UDTs"): + binary.eq(v_uv & v_arr).new() + + +@pytest.mark.skipif("not supports_udfs") +# SS < 9 has no GrB_NAME setter, so registration falls back to storing the +# numpy repr in the type name and warns when it does not fit in 128 chars. +# Each nested record's repr here is 141 or 142; how they serialize is not what +# the test is about. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +@pytest.mark.parametrize("op", ["plus", "eq"]) +@pytest.mark.parametrize("case", ["leaf_count", "nesting", "inner_names", "inner_order"]) +def test_udt_record_nesting_mismatch_is_a_keyerror(case, op): + """Records sharing top-level field names but nesting differently are a KeyError. + + ``_check_udt_pair`` matched on top-level names only, but the codegen reads + both operands at the left one's leaf paths. A field that is a sub-record on + one side and a scalar on the other, or sub-records with different names, + reached Numba, whose typing failure was a ``UdfParseError`` for ``plus`` + and a raw ``TypingError`` for ``eq``: compile errors for what is really the + shape disagreement its sibling checks raise ``KeyError`` for. Sub-record + fields in a different order are rejected as top-level ones are, so that + leaf ``i`` of one operand is leaf ``i`` of the other. + """ + f8 = np.float64 + sub = [("nst_n1", f8), ("nst_n2", f8)] + left, right = { + # A scalar field on the left is a sub-record on the right. + "leaf_count": ([("nst_a", f8), ("nst_b", f8)], [("nst_a", sub), ("nst_b", f8)]), + # Three leaves each, but nested under different fields. + "nesting": ([("nst_a", sub), ("nst_b", f8)], [("nst_a", f8), ("nst_b", sub)]), + "inner_names": ( + [("nst_a", sub), ("nst_b", f8)], + [("nst_a", [("nst_m1", f8), ("nst_m2", f8)]), ("nst_b", f8)], + ), + "inner_order": ([("nst_a", sub), ("nst_b", f8)], [("nst_a", sub[::-1]), ("nst_b", f8)]), + }[case] + v = Vector(dtypes.register_anonymous(np.dtype(left, align=True), "_NestLeft"), size=1) + w = Vector(dtypes.register_anonymous(np.dtype(right, align=True), "_NestRight"), size=1) + with pytest.raises(KeyError, match="must nest the same way, with the same field names"): + v.ewise_mult(w, getattr(binary, op)).new() + + +def _record_with_field_shape(shape, name, b_dtype=np.float64): + """Register a record UDT whose ``fsh_a`` field has ``shape`` (a scalar for ``()``).""" + field = ("fsh_a", np.float64, shape) if shape else ("fsh_a", np.float64) + return dtypes.register_anonymous(np.dtype([field, ("fsh_b", b_dtype)]), name) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize("op", ["plus", "eq"]) +@pytest.mark.parametrize( + ("left", "right"), + [((3,), (4,)), ((2, 3), (3, 2)), ((3,), ()), ((3,), (1,)), ((3,), (1, 3))], +) +def test_udt_record_field_shape_mismatch_is_a_keyerror(left, right, op): + """Record fields whose shapes differ are a KeyError, whichever operand is on the left. + + The codegen combines array fields as numpy arrays, so a ``(3,)`` field with + a ``(4,)`` one raised inside the cfunc, where the caller never saw it, and + the result element kept whatever was in the buffer (``eq`` read it as its + bool). Shapes that broadcast worked only one way round: ``plus`` writes + into a field shaped like the left operand's, so a ``(3,)`` field took a + ``(1,)`` one on its right but failed with it on its left, and a scalar + field on the left of an array one never compiled. numpy will not promote + or compare two record dtypes whose fields differ in shape either. + """ + X = _record_with_field_shape(left, "_FshLeft") + Y = _record_with_field_shape(right, "_FshRight") + for a, b in [(X, Y), (Y, X)]: + v = Vector(a, size=1) + w = Vector(b, size=1) + with pytest.raises(KeyError, match="must have the same shape for each field"): + v.ewise_mult(w, getattr(binary, op)).new() + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_fields_of_equal_shape_combine(): + """Two record UDTs with the same field shapes combine, though a field's dtype differs. + + ``eq`` and ``ne`` used to fail on this pair with GrB_DOMAIN_MISMATCH: + ``get_typed_op`` unified the two records to one of them, since numpy + promotes their field dtypes, so GraphBLAS saw an operand of the other + type. They now compile the pair, as the arithmetic ops do. + """ + X = _record_with_field_shape((2, 3), "_FshLeft") + v = Vector(X, size=2) + v[0] = (np.arange(6.0).reshape(2, 3), 1.0) + v[1] = (np.ones((2, 3)), 2.0) + w = Vector(_record_with_field_shape((2, 3), "_FshRight", b_dtype=np.float32), size=2) + w[0] = (np.full((2, 3), 10.0), 0.5) + w[1] = (np.ones((2, 3)), 2.0) + total = v.ewise_mult(w, binary.plus).new() + assert total.dtype == X + value = total[0].new().value + np.testing.assert_array_equal(value["fsh_a"], np.arange(6.0).reshape(2, 3) + 10.0) + assert value["fsh_b"] == 1.5 + assert binary.eq(v & w).new().isequal(Vector.from_coo([0, 1], [False, True])) + assert binary.ne(v & w).new().isequal(Vector.from_coo([0, 1], [True, False])) + # A semiring dispatches through its BinaryOp, so ``lor_eq`` takes the pair too. + assert v.inner(w, semiring.lor_eq).new().value + assert not v.inner(w, semiring.land_eq).new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_positional_ops_take_two_different_udts(): + """``first``, ``second``, ``any`` and ``pair`` take any two UDTs; they never read a field. + + ``any`` was the odd one out: it had no ``_custom_dtype``, so ``get_typed_op`` + unified two records of the same layout to one of them and GraphBLAS raised + GrB_DOMAIN_MISMATCH on the other operand. ``get_typed_op`` now never unifies + two different UDTs. + """ + # Own field names: the helper's scalar-field record is ``_FshRight``'s dtype, + # and ``register_anonymous`` returns one cached DataType per dtype. + X = dtypes.register_anonymous( + np.dtype([("pos_a", np.float64), ("pos_b", np.float64)]), "_PosLeft" + ) + Y = dtypes.register_anonymous( + np.dtype([("pos_a", np.float64), ("pos_b", np.int64)]), "_PosRight" + ) + v = Vector(X, size=1) + v[0] = (1.0, 3.0) + w = Vector(Y, size=1) + w[0] = (4.0, 6) + assert binary.first(v & w).new().isequal(v) + assert binary.second(v & w).new().isequal(w) + assert binary.any(v & w).new().isequal(w) + assert binary.pair(v & w).new().isequal(Vector.from_coo([0], [1])) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_pair_promotes_each_field(): + """Two records whose field dtypes differ give each field the dtype the built-in op gives. + + The result was always the left record, so ``x + y`` truncated an int64 field + that ``y + x`` kept as float64, and a complex field on the right of ``min`` + slipped past the check that rejects complex ``min`` and failed in Numba. + When one operand's record holds every promoted field, the result is that + record, on either side; otherwise it is the record with the promoted + fields, which is the UDT registered with that layout if there is one. Each + record also has an array field, which no Numba return type can be matched + against, so the result type cannot be recovered from what the generated + function returns. + """ + + def record(a_dtype, b_dtype, name): + fields = [("wid_a", a_dtype), ("wid_b", b_dtype), ("wid_v", np.float64, (2,))] + return dtypes.register_anonymous(np.dtype(fields), name) + + f8 = np.float64 + narrow = record(np.int64, f8, "_WidN") + wide = record(f8, f8, "_WidW") + crossed = record(f8, np.int64, "_WidX") + cplx = record(np.complex128, f8, "_WidC") + v = Vector(narrow, size=1) + v[0] = (1, 2.0, [1.0, 2.0]) + w = Vector(wide, size=1) + w[0] = (0.5, 0.25, [10.0, 20.0]) + for a, b in [(v, w), (w, v)]: + result = binary.plus(a & b).new() + assert result.dtype == wide + value = result[0].new().value + assert (value["wid_a"], value["wid_b"]) == (1.5, 2.25) + np.testing.assert_array_equal(value["wid_v"], [11.0, 22.0]) + # Each of ``narrow`` and ``crossed`` is wider in one field, so neither holds + # both. The promoted record has ``wide``'s layout, so it is ``wide``. + x = Vector(crossed, size=1) + x[0] = (1.0, 2, [1.0, 2.0]) + for a, b in [(v, x), (x, v)]: + result = binary.times(a & b).new() + assert result.dtype == wide + value = result[0].new().value + assert (value["wid_a"], value["wid_b"]) == (1.0, 4.0) + np.testing.assert_array_equal(value["wid_v"], [1.0, 4.0]) + assert binary.eq(v & x).new().isequal(Vector.from_coo([0], [True])) + c = Vector(cplx, size=1) + c[0] = (1 + 1j, 2.0, [0.0, 0.0]) + for a, b in [(w, c), (c, w)]: + assert binary.plus(a & b).new().dtype == cplx + with pytest.raises(KeyError, match="not defined on complex fields"): + binary.min(a & b).new() + # The result is ``wide`` on either side. Stored in ``narrow``, its float + # field is cast to int64 as a built-in store would cast it. + v << binary.plus(v & w) + assert v.dtype == narrow + value = v[0].new().value + assert (value["wid_a"], value["wid_b"]) == (1, 2.25) + + +@pytest.mark.skipif("not supports_udfs") +# SS < 9 names a UDT by its numpy repr and warns when that is over 128 chars, +# as these nested records' reprs are. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_record_pair_promotes_to_a_third_record(): + """Records that are each wider in some field give a third record, as numpy would. + + Each field is promoted as the built-in op promotes two Vectors, so int16 + with float32 is float32, and the result keeps the operands' names, nesting + and alignment. This pair used to be a KeyError, since the result had to be + one of the two operands. + """ + i2, i8, f4, f8 = np.int16, np.int64, np.float32, np.float64 + X = dtypes.register_anonymous( + np.dtype([("pro_a", [("pro_i", i2), ("pro_j", f8)]), ("pro_b", f4)], align=True), "_ProX" + ) + Y = dtypes.register_anonymous( + np.dtype([("pro_a", [("pro_i", f4), ("pro_j", i8)]), ("pro_b", i2)], align=True), "_ProY" + ) + expected = np.dtype([("pro_a", [("pro_i", f4), ("pro_j", f8)]), ("pro_b", f4)], align=True) + v = Vector(X, size=1) + v[0] = ((3, 0.5), 1.5) + w = Vector(Y, size=1) + w[0] = ((0.25, 4), 2) + for a, b in [(v, w), (w, v)]: + result = binary.plus(a & b).new() + assert result.dtype.np_type == expected + assert result[0].new().value.tolist() == ((3.25, 4.5), 3.5) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_scalar_promotes_each_field(): + """A number beside a record UDT keeps its dtype, and each field promotes as a Vector's would. + + The number used to be converted to the record first, so ``int_record + 0.5`` + added 0 to each field. + """ + ints = dtypes.register_anonymous( + np.dtype([("rsp_x", np.int64), ("rsp_y", np.int64)]), "_RspInts" + ) + floats = dtypes.register_anonymous( + np.dtype([("rsp_f", np.float64), ("rsp_g", np.float64)]), "_RspFloats" + ) + promoted = np.dtype([("rsp_x", np.float64), ("rsp_y", np.float64)]) + v = Vector(ints, size=1) + v[0] = (2, 3) + for expr, expected in [ + (v + 0.5, (2.5, 3.5)), + (1.5 * v, (3.0, 4.5)), + (v.apply(binary.truediv, right=4), (0.5, 0.75)), + ]: + result = expr.new() + assert result.dtype.np_type == promoted + assert result[0].new().value.tolist() == expected + assert (v + 1).new().dtype == ints + w = Vector(floats, size=1) + w[0] = (0.5, 1.5) + result = (w + 1).new() + assert result.dtype == floats + assert result[0].new().value.tolist() == (1.5, 2.5) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_packed_record_with_aligned_takes_an_operand_type(): + """A packed and an aligned record with the same leaves give the left operand's type. + + Their fields are the same dtypes at different offsets, so neither holds the + other's layout, but each holds every result field. Taking the left one lets + ``x << x + y`` update ``x`` in place, whichever of the two ``x`` is. + """ + fields = [("pak_a", np.int8), ("pak_b", np.float64)] + aligned = dtypes.register_anonymous(np.dtype(fields, align=True), "_PakAligned") + packed = dtypes.register_anonymous(np.dtype(fields), "_PakPacked") + assert aligned.np_type.itemsize != packed.np_type.itemsize + x = Vector(aligned, size=1) + x[0] = (1, 0.5) + y = Vector(packed, size=1) + y[0] = (2, 0.25) + for a, b in [(x, y), (y, x)]: + assert binary.plus(a & b).new().dtype == a.dtype + a << binary.plus(a & b) + assert x[0].new().value.tolist() == (3, 0.75) + assert y[0].new().value.tolist() == (5, 1.0) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize("op", ["plus", "eq"]) +@pytest.mark.parametrize(("left", "right"), [((2, 3), (3, 2)), ((6,), (2, 3)), ((3,), (4,))]) +def test_udt_array_pair_shape_mismatch_is_a_keyerror(left, right, op): + """Array UDTs whose shapes do not broadcast are a KeyError, as numpy refuses them. + + Two of these with the same size used to combine over their flat buffers, + so ``(2, 3)`` plus ``(3, 2)`` added elements numpy never pairs. The base + dtypes differ, which is allowed, so only the shapes are at fault. + """ + # int16 and int32: other tests register float64 arrays of these shapes, and + # ``register_anonymous`` returns one cached DataType per dtype. + X = dtypes.register_anonymous(np.dtype((np.int16, left)), "_PairArrL") + Y = dtypes.register_anonymous(np.dtype((np.int32, right)), "_PairArrR") + for a, b in [(X, Y), (Y, X)]: + with pytest.raises(KeyError, match="shapes that broadcast together"): + getattr(binary, op)(Vector(a, size=1) & Vector(b, size=1)).new() + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize( + ("x_type", "y_type"), + [ + ((np.uint16, (3, 1)), (np.int32, (1, 4))), + ((np.uint16, (2, 1, 3)), (np.float32, (4, 1))), + ((np.uint16, (1,)), (np.int16, (2, 2))), + ((np.int16, (3, 1)), (np.int16, (2, 1, 4))), + ], +) +def test_udt_array_pairs_broadcast_as_numpy(x_type, y_type): + """Two array UDTs broadcast as numpy arrays do, in either order. + + The result has numpy's broadcast shape and the promoted element dtype. It + is an operand's type when one has both, else the structural array type. + ``eq`` is True when every pair of elements numpy compares is equal. + """ + X = dtypes.register_anonymous(np.dtype(x_type), "_BcastX") + Y = dtypes.register_anonymous(np.dtype(y_type), "_BcastY") + xv = np.arange(1, 1 + np.prod(x_type[1])).reshape(x_type[1]).astype(x_type[0]) + yv = 10 * np.arange(1, 1 + np.prod(y_type[1])).reshape(y_type[1]).astype(y_type[0]) + v = Vector(X, size=3) + v[0] = xv + v[1] = xv + w = Vector(Y, size=3) + w[0] = yv + w[2] = yv + M = Matrix(X, nrows=1, ncols=1) + M[0, 0] = xv + N = Matrix(Y, nrows=1, ncols=1) + N[0, 0] = yv + for a, b, A, B, av, bv in [(v, w, M, N, xv, yv), (w, v, N, M, yv, xv)]: + expected = np.subtract(av, bv) + result_type = dtypes.lookup_dtype(np.dtype((expected.dtype, expected.shape))) + for result in [binary.minus(a & b).new(), binary.minus(A & B).new()]: + assert result.dtype == result_type + assert result.nvals == 1 + value = result[0].new().value if result.ndim == 1 else result[0, 0].new().value + np.testing.assert_array_equal(value, expected) + assert not binary.eq(a & b).new()[0].new().value + assert binary.ne(a & b).new()[0].new().value + # ewise_add would copy an entry without a partner as the result type, a + # cast GraphBLAS cannot make for a UDT; ewise_union computes it instead, + # so it broadcasts to the result's shape. + with pytest.raises(DomainMismatch, match="ewise_add cannot use binary.minus"): + a.ewise_add(b, binary.minus) + added = a.ewise_union(b, binary.minus, 0, 0).new() + assert added.dtype == result_type + np.testing.assert_array_equal(added[0].new().value, expected) + lone = 1 if a is v else 2 + np.testing.assert_array_equal(added[lone].new().value, np.broadcast_to(av, expected.shape)) + # Equal values: every compared pair is equal, as ``(xv == yv).all()``. + w[0] = np.broadcast_to(xv[(0,) * xv.ndim], y_type[1]) + v[0] = np.broadcast_to(xv[(0,) * xv.ndim], x_type[1]) + assert binary.eq(v & w).new()[0].new().value + assert binary.eq(w & v).new()[0].new().value + assert not binary.ne(v & w).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_broadcast_semiring_and_reduce(): + """A semiring on two array UDTs that broadcast sums the broadcast products.""" + X = dtypes.register_anonymous(np.dtype((np.int16, (3, 1))), "_BcastSemiX") + Y = dtypes.register_anonymous(np.dtype((np.int16, (1, 5))), "_BcastSemiY") + a = [np.arange(3).reshape(3, 1) + k for k in range(2)] + b = [10 * np.arange(5).reshape(1, 5) - k for k in range(2)] + A = Matrix(X, nrows=1, ncols=2) + A[0, 0] = a[0] + A[0, 1] = a[1] + B = Matrix(Y, nrows=2, ncols=1) + B[0, 0] = b[0] + B[1, 0] = b[1] + C = A.mxm(B, semiring.plus_times).new() + assert C.dtype.np_type == np.dtype((np.int16, (3, 5))) + expected = a[0] * b[0] + a[1] * b[1] + np.testing.assert_array_equal(C[0, 0].new().value, expected) + s = A.reduce_scalar(monoid.plus).new() + np.testing.assert_array_equal(s.value, a[0] + a[1]) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_2d_field_compares_with_a_number(): + """``eq`` and ``ne`` of a record with a 2-D array field and a number compare every element. + + The generated code read the field at flat positions, which in a 2-D field + are rows, so it failed to compile once a number was compared as it is + rather than converted into the record first. + """ + T = dtypes.register_anonymous( + np.dtype([("cf_grid", np.float64, (2, 3)), ("cf_n", np.int64)], align=True), "_CmpField2D" + ) + v = Vector(T, size=2) + v[0] = (np.full((2, 3), 2.0), 2) + v[1] = (np.array([[2.0, 2, 2], [2, 2, np.nan]]), 2) + assert binary.eq(v, 2.0).new().isequal(Vector.from_coo([0, 1], [True, False])) + assert binary.ne(2.0, v).new().isequal(Vector.from_coo([0, 1], [False, True])) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_ops_on_large_elements(): + """Lifted ops on a 64 by 64 array compile in a loop, not one statement per element. + + Unrolled, Numba took about two minutes to type ``minus`` on this element. + """ + try: + T = dtypes.register_anonymous(np.dtype((np.float64, (64, 64))), "_LoopBig") + except InvalidValue: + if sys.platform != "win32": + raise # pragma: no cover (only on failure) + # SuiteSparse:GraphBLAS built by MSVC (no variable-length arrays) refuses a + # UDT larger than 1024 bytes (GB_VLA_MAXSIZE); this one is 32768. + pytest.skip("SuiteSparse:GraphBLAS limits a UDT to 1024 bytes on Windows") + S = dtypes.register_anonymous(np.dtype((np.float64, (64, 1))), "_LoopBigCol") + x = np.arange(64 * 64, dtype=np.float64).reshape(64, 64) + col = np.arange(64, dtype=np.float64).reshape(64, 1) + v = Vector(T, size=1) + v[0] = x + u = Vector(S, size=1) + u[0] = col + np.testing.assert_array_equal((v - u).new()[0].new().value, x - col) + np.testing.assert_array_equal(unary.ainv(v).new()[0].new().value, -x) + assert binary.eq(v & v).new()[0].new().value + assert not binary.ne(v & v).new()[0].new().value + w = v.dup() + w[0] = np.where(x == 4095, np.nan, x) + assert not binary.eq(w & w).new()[0].new().value + assert binary.ne(w & w).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_pair_leading_ones(): + """Array UDTs that differ only in leading axes of length 1 combine, as numpy broadcasts them. + + This is what a ``Matrix.to_csr`` round trip produces: values of shape + ``(1, 3)`` for a ``(3,)`` UDT. The result takes the operand with more axes, + numpy's broadcast shape, whichever side it is on. + """ + short = dtypes.register_anonymous(np.dtype((np.int16, (3,))), "_PairArrShort") + long = dtypes.register_anonymous(np.dtype((np.int16, (1, 3))), "_PairArrLong") + v = Vector(short, size=1) + v[0] = [1, 2, 3] + w = Vector(long, size=1) + w[0] = [[10, 20, 30]] + for a, b in [(v, w), (w, v)]: + result = binary.plus(a & b).new() + assert result.dtype == long + np.testing.assert_array_equal(result[0].new().value, [[11, 22, 33]]) + assert not binary.eq(a & b).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_pair_promotes_the_base(): + """Array UDTs with different base dtypes combine, each element promoted as the built-in op does. + + The pair used to be a KeyError. The result is the operand whose type holds + the promoted elements, on either side, so ``int64`` with ``float64`` gives + the ``float64`` one. ``eq`` and ``ne`` must read each operand as its own + dtype: reading both as the left one's compares the bits of a float64 as an + int64. + """ + ints = dtypes.register_anonymous(np.dtype((np.int64, (2, 2, 2))), "_PromI222") + floats = dtypes.register_anonymous(np.dtype((np.float64, (2, 2, 2))), "_PromF222") + base = np.arange(8).reshape(2, 2, 2) + v = Vector(ints, size=2) + v[0] = base + v[1] = base + w = Vector(floats, size=2) + w[0] = base + 0.5 + w[1] = base + for a, b in [(v, w), (w, v)]: + result = binary.plus(a & b).new() + assert result.dtype == floats + np.testing.assert_array_equal(result[0].new().value, 2 * base + 0.5) + assert binary.eq(a & b).new().isequal(Vector.from_coo([0, 1], [False, True])) + assert binary.ne(a & b).new().isequal(Vector.from_coo([0, 1], [True, False])) + # A semiring's monoid is typed by its multiplier's result. + assert v.inner(w, semiring.plus_times).new().dtype == floats + # GraphBLAS would cast an entry of ``v`` alone to the result type, which a + # UDT cannot be; ewise_union computes that entry with the op instead, and + # infix ``+`` uses it, with 0 for the missing value. + with pytest.raises(DomainMismatch, match="ewise_add cannot use binary.plus"): + v.ewise_add(w) + assert v.ewise_union(w, binary.plus, 0, 0.0).new().dtype == floats + assert (v + w).new().isequal(v.ewise_union(w, binary.plus, 0, 0.0).new()) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_scalar_promotes_like_builtin(): + """A number beside an array UDT keeps its dtype, and each element promotes as the op does. + + ``apply`` converted a Python number to the UDT first (``v + 0.5`` is an + ``apply``), and the result took the UDT's base whatever the other operand + was, so ``int_udt + 0.5`` added 0 and ``int_udt / 4`` gave integers. Each + element now has the dtype the op gives on the built-in dtypes, and the + result is the operand's UDT only when that holds it. + """ + ints = dtypes.register_anonymous(np.dtype((np.int64, (2, 4))), "_UpcI24") + int8s = dtypes.register_anonymous(np.dtype((np.int8, (2, 4))), "_UpcB24") + fp32s = dtypes.register_anonymous(np.dtype((np.float32, (2, 4))), "_UpcF24") + fp64s = np.dtype((np.float64, (2, 4))) + base = np.arange(8).reshape(2, 4) + v = Vector(ints, size=1) + v[0] = base + M = Matrix(ints, nrows=1, ncols=1) + M[0, 0] = base + for expr, expected in [ + (v + 0.5, base + 0.5), + (0.5 + v, base + 0.5), + (v.apply(binary.times, right=2.5), base * 2.5), + (v.apply(binary.truediv, right=4), base / 4), + (binary.minus(10.5, v), 10.5 - base), + (v.apply(binary.plus, right=gb.Scalar.from_value(0.5)), base + 0.5), + (v.ewise_mult(Vector.from_coo([0], [0.5]), binary.times), base * 0.5), + (M * 0.5, base * 0.5), + (0.5 * M, base * 0.5), + # A 0-d numpy array is a number too, not a whole element. + (v + np.array(0.5, np.float32), base + 0.5), + (v / np.array(4, np.int8), base / 4), + ]: + result = expr.new() + assert result.dtype.np_type == fp64s + value = result[0].new().value if result.ndim == 1 else result[0, 0].new().value + np.testing.assert_array_equal(value, expected) + # Same element type: the result keeps the UDT, as does a literal that is a + # whole element rather than a number. + for expr in [v + 1, v + [[1] * 4] * 2]: + result = expr.new() + assert result.dtype == ints + np.testing.assert_array_equal(result[0].new().value, base + 1) + # A Python number is weak, as in numpy 2: it takes the elements' dtype when + # their kind holds it, so INT8 elements plus 1 stay INT8 and FP32 elements + # plus 0.5 stay FP32, and an in-place update can store the result. A typed + # number is strong, and an int out of the elements' range is an error. + w = Vector(int8s, size=1) + w[0] = base + assert (w + 1).new().dtype == int8s + w += 1 + np.testing.assert_array_equal(w[0].new().value, base + 1) + assert (w + np.int64(1)).new().dtype == ints # the UDT registered with that layout + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + w + 300 + # truediv gives float elements whatever the int, as numpy divides in float64. + result = (w / 300).new() + assert result.dtype.np_type.base == np.float64 + np.testing.assert_array_equal(result[0].new().value, (base + 1) / 300) + x = Vector(fp32s, size=1) + x[0] = base + result = (x + 0.5).new() + assert result.dtype == fp32s + np.testing.assert_array_equal(result[0].new().value, base + 0.5) + x *= 2 + np.testing.assert_array_equal(x[0].new().value, base * 2) + assert (x + np.float64(0.5)).new().dtype.np_type == fp64s + if dtypes._supports_complex: + for expr in [lambda: binary.min(x, 1j), lambda: binary.min(1j, x)]: + with pytest.raises(KeyError, match="not defined on complex fields"): + expr() + # eq and ne compare with the number itself; converted to the UDT, 0.5 was 0. + zeros = Vector(ints, size=1) + zeros[0] = np.zeros((2, 4)) + assert not (zeros == 0.5).new()[0].new().value + assert (zeros != 0.5).new()[0].new().value + assert (zeros == 0).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_uint64_with_signed_ints_computes_like_builtin(): + """UINT64 elements with signed ones are computed in FP64, as the built-in ops compute them. + + The result type was right, but the generated code let Numba do the + arithmetic, and Numba computes uint64 with int64 in int64: ``2**63 + 1`` + came out as ``-9.22e18`` and ``2**63 * 2`` as 0. Each operand is now + converted to the type the built-in op computes in first. + """ + big = 2**63 + uints = dtypes.register_anonymous(np.dtype((np.uint64, (2, 2))), "_MixU22") + ints = dtypes.register_anonymous(np.dtype((np.int64, (2, 2))), "_MixI22") + u = Vector(uints, size=1) + u[0] = [[big, 7], [0, 1]] + i = Vector(ints, size=1) + i[0] = [[1, -1], [0, 0]] + for expr, expected in [ + (binary.plus(u & i), [[big + 1.0, 6.0], [0.0, 1.0]]), + (binary.plus(i & u), [[big + 1.0, 6.0], [0.0, 1.0]]), + (binary.floordiv(i & u), [[0.0, -1.0], [np.nan, 0.0]]), + # A numpy int64 is strong, so it is a signed operand like ``i``. + (u + np.int64(1), [[big + 1.0, 8.0], [1.0, 2.0]]), + (binary.plus(np.int64(1), u), [[big + 1.0, 8.0], [1.0, 2.0]]), + (u * np.int64(2), [[2.0 * big, 14.0], [0.0, 2.0]]), + ]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (2, 2))) + np.testing.assert_array_equal(result[0].new().value, expected) + # A Python int is weak, so it takes the uint64 elements' type, exactly. + result = (u + 1).new() + assert result.dtype == uints + np.testing.assert_array_equal(result[0].new().value, [[big + 1, 8], [1, 2]]) + # BOOL with UINT64 is UINT64, so True // 2**63 is 0, not -1 through int64. + result = binary.floordiv(True, u).new() + assert result.dtype == uints + np.testing.assert_array_equal(result[0].new().value[0], [0, 0]) + record = dtypes.register_anonymous( + np.dtype([("mxs_u", np.uint64), ("mxs_f", np.float64)]), "_MixRec" + ) + other = dtypes.register_anonymous( + np.dtype([("mxs_u", np.int64), ("mxs_f", np.float64)]), "_MixRecI" + ) + r = Vector(record, size=1) + r[0] = (big, 1.5) + s = Vector(other, size=1) + s[0] = (-1, 0.5) + assert (r * np.int64(2)).new()[0].new().value.tolist() == (2.0 * big, 3.0) + assert binary.times(np.int64(2), r).new()[0].new().value.tolist() == (2.0 * big, 3.0) + # A weak 2 keeps the uint64 field, which wraps as numpy 2's uint64 does. + assert (r * 2).new()[0].new().value.tolist() == (0, 3.0) + assert binary.plus(r & s).new()[0].new().value.tolist() == (big - 1.0, 2.0) + assert binary.floordiv(s & r).new()[0].new().value.tolist() == (-1.0, 0.0) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_user_op_named_like_a_lifted_op(): + """A UDF registered under a lifted op's name runs its own function, and literals become UDTs. + + ``_compile_udt`` matched built-in ops by name only, so a UDF registered as + ``"plus"`` or ``"max"`` was compiled as the built-in, its function ignored, + and on a built-in dtype it raised KeyError. A literal beside a UDT is + converted to the UDT for it, since nothing says how a UDF combines a scalar + with an element, but only when its type fits: ``0.5`` became an element of + zeros. + """ + udt = dtypes.register_anonymous(np.dtype((np.int64, (2, 3))), "_NamedLikeLifted") + + def scaled_plus(x, y): # pragma: no cover (numba) + return x * 10 + y + + def difference(x, y): # pragma: no cover (numba) + return x - y + + plus_like = BinaryOp.register_anonymous(scaled_plus, "plus", is_udt=True) + max_like = BinaryOp.register_anonymous(difference, "max", is_udt=True) + v = Vector(udt, size=1) + v[0] = np.ones((2, 3), dtype=np.int64) + for udf_op, right, expected in [ + (plus_like, 3, 13), + (plus_like, np.int8(2), 12), + (max_like, 3, -2), + ]: + result = v.apply(udf_op, right=right).new() + assert result.dtype == udt + np.testing.assert_array_equal(result[0].new().value, np.full((2, 3), expected)) + assert udf_op[udt].jit_c_source is None + # A float is not an int, whatever its value, as numpy refuses ``ints += 2.0``; + # 2.0 still converts, being exact, but that is deprecated. + with pytest.raises(ValueError, match="0.5 does not fit"): + v.apply(plus_like, right=0.5) + with pytest.warns(DeprecationWarning, match="2.0 does not fit"): + v.apply(plus_like, right=2.0) + assert plus_like(Vector.from_coo([0], [5]), 2).new().isequal(Vector.from_coo([0], [52])) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_python_number_is_weak_per_leaf(): + """A Python number beside a UDT is typed leaf by leaf, as numpy 2 types it. + + Each leaf keeps its dtype when its kind holds the number (an int beside + int8, a float beside float32) and takes the number's default dtype when it + does not (a float beside int8 is float64). So an in-place update with a + number of the leaf's kind keeps the UDT, where typing every Python int as + INT64 and every float as FP64 made ``int8_udt += 1`` and ``f32_udt *= 2`` + unstorable. numpy scalars are strong, and so are typed Scalars. + """ + + # Aligned, so the same-type ops get a JIT kernel and warn nothing. + def record(i_dtype, f_dtype): + return np.dtype([("wk_i", i_dtype), ("wk_f", f_dtype)], align=True) + + rec = dtypes.register_anonymous(record(np.int8, np.float32), "_WeakRec") + v = Vector(rec, size=1) + v[0] = (1, 2.0) + v += 1 + v *= 2 + assert v.dtype == rec + assert v[0].new().value.tolist() == (4, 6.0) + result = (v + 0.5).new() + assert result.dtype.np_type == record(np.float64, np.float32) + assert result[0].new().value.tolist() == (4.5, 6.5) + assert (v + np.float32(0.5)).new().dtype.np_type == record(np.float32, np.float32) + # The result has a float64 leaf, which is cast back to int8 to be stored. + v += 0.5 + assert v.dtype == rec + assert v[0].new().value.tolist() == (4, 6.5) + with pytest.raises(OverflowError, match="200 out of bounds for int8"): + v + 200 + # A UDT Scalar follows the same rule, through infix and apply alike. + arr = dtypes.register_anonymous(np.dtype((np.int16, (5,))), "_WeakArr") + s = gb.Scalar(arr) + s.value = [1, 2, 3, 4, 5] + for expr in [s * 2.5, 2.5 * s, s.apply(binary.times, right=2.5)]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (5,))) + np.testing.assert_array_equal(result.value, np.arange(1, 6) * 2.5) + assert (s + 1).new().dtype == arr + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ewise_add_needs_operands_of_the_result_type(): + """``ewise_add`` on UDT operands of another type than the result's: Scalars only. + + GraphBLAS copies an entry present in only one input into the result, cast + to the result type, which a UDT cannot be. A Scalar is one element, so its + operands are converted to the result type first (which changes no value, + not even ``-0.0``). Vectors and Matrices would need converted copies of + whole operands, so they raise, and ``ewise_union`` computes the same sum + with the op instead. + """ + int8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_UniPlusI8") + fp32s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_UniPlusF32") + f64s = dtypes.lookup_dtype(np.dtype((np.float64, (3,)))) + s = gb.Scalar.from_value([1, 2, 3], dtype=int8s) + for expr in [s + 0.5, 0.5 + s]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + np.testing.assert_array_equal(result.value, [1.5, 2.5, 3.5]) + result = (gb.Scalar(int8s) + 0.5).new() # an empty Scalar adds as 0 + np.testing.assert_array_equal(result.value, [0.5, 0.5, 0.5]) + neg = gb.Scalar.from_value([-0.0, 0.0, -1.5], dtype=fp32s) + result = (neg + gb.Scalar(f64s)).new() + assert result.dtype == f64s + np.testing.assert_array_equal(np.signbit(result.value), [True, False, True]) + v = Vector(int8s, size=3) + v[0] = [1, 2, 3] + v[1] = [4, 5, 6] + w = Vector(fp32s, size=3) + w[1] = [0.5, 0.5, 0.5] + w[2] = [7, 8, 9] + A = Matrix(int8s, nrows=1, ncols=2) + A[0, 0] = [1, 2, 3] + B = Matrix(fp32s, nrows=1, ncols=2) + B[0, 1] = [0.5, 0.5, 0.5] + for expr in [ + lambda: v.ewise_add(w), + lambda: w.ewise_add(v), + lambda: binary.plus(v | w), + lambda: A.ewise_add(B), + lambda: A.ewise_add(w[:2].new()), + ]: + with pytest.raises(DomainMismatch, match="ewise_add cannot use .* converted copies"): + expr() + with pytest.raises(DimensionMismatch, match="Matrix shapes must match"): + A.ewise_add(Matrix(fp32s, nrows=2, ncols=2)) + # A Scalar converts only for an op that lifts to UDTs, so ``first`` raises too. + s2 = gb.Scalar.from_value([0.5, 0.5, 0.5], dtype=fp32s) + result = s.ewise_union(s2, monoid.plus, 0, 0).new() # a Monoid lifts as its BinaryOp + np.testing.assert_array_equal(result.value, [1.5, 2.5, 3.5]) + with pytest.raises(DomainMismatch, match="ewise_add cannot use binary.first"): + s.ewise_add(s2, binary.first) + # Infix ``+`` takes the union with each operand's zero for a missing value + # instead, as ``-`` takes 0. + for left, right in [(v, w), (w, v)]: + result = (left + right).new() + assert result.dtype == fp32s + assert result.isequal(left.ewise_union(right, binary.plus, False, False).new()) + assert result.to_coo()[1].tolist() == [[1, 2, 3], [4.5, 5.5, 6.5], [7, 8, 9]] + result = (A + B).new() + assert result.dtype == fp32s + assert result.to_coo()[2].tolist() == [[1, 2, 3], [0.5, 0.5, 0.5]] + assert (A + B.T.T).new().isequal(result) + result = (A + Vector.from_coo([1], [[0.5, 0.5, 0.5]], dtype=fp32s, size=2)).new() + assert result.to_coo()[2].tolist() == [[1, 2, 3], [0.5, 0.5, 0.5]] + # That zero is -0.0 in a float field, so an entry of one operand alone keeps + # its value, as ewise_add would copy it; beside an int field, 0 turns -0.0 + # into 0.0. + zeros = Vector(fp32s, size=3) + zeros[2] = [-0.0, 0.0, -1.5] + for left, right in [(zeros, Vector(f64s, size=3)), (Vector(f64s, size=3), zeros)]: + result = (left + right).new() + assert result.dtype == f64s + np.testing.assert_array_equal(np.signbit(result[2].new().value), [True, False, True]) + result = (v + zeros).new() + np.testing.assert_array_equal(np.signbit(result[2].new().value), [False, False, True]) + x = v.dup() + x += w # an accumulating store: cast, then added, as for built-in types + assert x.dtype == int8s + assert x.to_coo()[1].tolist() == [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + # The same type in and out still goes through ewise_add itself... + assert (v + v).new().dtype == int8s + # ...but not int truediv, whose result is float. + with pytest.raises(DomainMismatch, match="binary.truediv"): + v.ewise_add(v, binary.truediv) + udf = BinaryOp.register_anonymous(lambda x, y: x, "_uni_first", is_udt=True) + with pytest.raises(DomainMismatch, match="ewise_add cannot use .* on _UniPlusI8 and"): + v.ewise_add(Vector(f64s, size=3), udf) + # A size or broadcast error is reported first, as for a same-type pair. + with pytest.raises(DimensionMismatch): + v.ewise_add(Vector(fp32s, size=2), binary.plus) + with pytest.raises(DimensionMismatch, match="Matrix.nrows"): + v.ewise_add(B, binary.plus) + with pytest.raises(DimensionMismatch, match="Matrix.ncols"): + B.ewise_add(v, binary.plus) + # ewise_union computes an entry alone with the op, so it needs no cast. + result = v.ewise_union(w, binary.plus, 0, 0).new() + assert result.dtype == fp32s + assert result.to_coo()[1].tolist() == [[1, 2, 3], [4.5, 5.5, 6.5], [7, 8, 9]] + # Accumulating it into the int8 object casts as it stores, as built-ins do. + x = v.dup() + x(accum=binary.plus) << x.ewise_union(w, binary.plus, 0, 0) + assert x.dtype == int8s + assert x.to_coo()[1].tolist() == [[2, 4, 6], [8, 10, 12], [7, 8, 9]] + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_literal_converts_only_when_exact(): + """A literal that has to become an element of a UDT must be of that type. + + A literal beside a user-defined op, an IndexUnaryOp thunk and an + ``ewise_union`` default are converted into the UDT, which numpy does + silently: ``(0.5, 1.5, 2.5)`` added ``(0, 1, 2)`` to an int8 array UDT, and + a default of ``0.5`` stood in as ``0``. The literal's type decides, as + numpy's same_kind rule decides ``ints += x``: a Python number or sequence + is weak, so ``1`` and ``(1, 2, 3)`` are int8 and ``0.5`` and ``2.0`` are + float64; a numpy value is strong and must cast safely. A Monoid in + ``apply`` is typed as its BinaryOp, so its literal promotes instead. + """ + import datetime + import re + import warnings + + arr = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_ExactArr") + rec = dtypes.register_anonymous( + np.dtype([("ex_i", np.int16), ("ex_f", np.float32)], align=True), "_ExactRec" + ) + floats = dtypes.register_anonymous(np.dtype((np.float32, (2,))), "_ExactF32") + nested = dtypes.register_anonymous( + np.dtype([("ex_v", np.int8, (2,)), ("ex_w", np.int8)]), "_ExactNested" + ) + text = dtypes.register_anonymous(np.dtype([("ex_s", "S4"), ("ex_n", np.int8)]), "_ExactText") + char = dtypes.register_anonymous(np.dtype([("ex_c", "S1"), ("ex_k", np.int8)]), "_ExactChar") + bare = dtypes.register_anonymous(np.dtype("S3"), "_ExactBareText") + seconds = dtypes.register_anonymous(np.dtype("m8[s]"), "_ExactBareSeconds") + dates = dtypes.register_anonymous( + np.dtype([("ex_d", "M8[s]"), ("ex_m", np.int8)]), "_ExactDate" + ) + v = Vector(arr, size=3) + v[0] = [1, 2, 3] + w = Vector(arr, size=3) + w[1] = [4, 5, 6] + r = Vector(rec, size=1) + r[0] = (1, 1.5) + f = Vector(floats, size=1) + f[0] = [1, 2] + n = Vector(nested, size=1) + n[0] = ([1, 2], 3) + t = Vector(text, size=1) + t[0] = (b"ab", 1) + c = Vector(char, size=1) + c[0] = (b"a", 1) + b = Vector(bare, size=1) + b[0] = b"0.5" + d = Vector(seconds, size=1) + d[0] = np.timedelta64(5, "s") + cut_text = np.array((b"abcdef", 1), dtype=[("ex_s", "S6"), ("ex_n", np.int64)])[()] + when = Vector(dates, size=1) + when[0] = (np.datetime64(0, "s"), 1) + udf = BinaryOp.register_anonymous(lambda x, y: x + y, "_exact_plus", is_udt=True) + second = BinaryOp.register_anonymous(lambda x, y: y, "_exact_second", is_udt=True) + same = BinaryOp.register_anonymous(lambda x, y: x == y, "_exact_same", is_udt=True) + # Literals of the UDT's type convert, a float rounding into a float32 field. + assert v.apply(udf, right=(1, 2, 3)).new()[0].new().value.tolist() == [2, 4, 6] + assert v.apply(udf, right=[True, 2, 3]).new().dtype == arr + assert v.apply(udf, right=np.ones(3, dtype=np.int8)).new().dtype == arr + assert v.apply(udf, right=np.int8(1)).new().dtype == arr + assert r.apply(second, right=(1, 0.1)).new()[0].new().value.tolist() == (1, np.float32(0.1)) + assert r.apply(second, right={"ex_i": 2, "ex_f": 0.5}).new()[0].new().value.tolist() == (2, 0.5) + assert r.apply(second, right=np.int8(2)).new()[0].new().value.tolist() == (2, 2.0) + # A numpy number in a record is strong, as alone: np.int16 fits the int16 field. + assert r.apply(second, right=(np.int16(2), 0.5)).new()[0].new().value.tolist() == (2, 0.5) + # A sequence holding numpy values is strong, as the array numpy makes of it + # (int64, or int32 on Windows with numpy 1). + made = np.asarray((np.int8(1), 2, 3)).dtype + assert binary.plus(v, (np.int8(1), 2, 3)).new().dtype.np_type == np.dtype((made, (3,))) + # The error names the literal's type; a layout's name is whatever UDT + # registered it first, so only the built-in scalar types are pinned here. + for expr, typed_as in [ + (lambda: v.apply(udf, right=(0.5, 1.5, 2.5)), "_ExactArr"), + (lambda: v.apply(udf, left=np.array([0.5, 1, 1])), "_ExactArr"), + (lambda: v.apply(udf, right=(np.float64(0.5), 1, 1)), "_ExactArr"), + (lambda: r.apply(second, left=(0.5, 1)), "_ExactRec"), + (lambda: r.apply(second, right=(np.float64(1.5), 1)), "_ExactRec"), + (lambda: r.apply(second, right={"ex_i": np.float64(1.5), "ex_f": 1}), "_ExactRec"), + (lambda: n.apply(second, right=((np.float64(1.5), 2), 3)), "_ExactNested"), + # Text or a date in a record is typed too, so the number beside it is checked, + # and text longer than its field does not fit. + (lambda: t.apply(second, right=(b"ab", 1.5)), "_ExactText"), + (lambda: t.apply(second, right=(b"ab", np.float64(1.5))), "_ExactText"), + (lambda: t.apply(second, right=(b"abcdef", 1)), "_ExactText"), + (lambda: when.apply(second, right=(np.datetime64(0, "s"), 1.5)), "_ExactDate"), + (lambda: when.apply(second, right=(datetime.date(1970, 1, 1), 1.5)), "_ExactDate"), + (lambda: when.apply(second, right=("1970-01-01", 1.5)), "_ExactDate"), + (lambda: when.apply(second, right=(np.datetime64(5, "ms"), 1)), "_ExactDate"), + (lambda: when.apply(second, right=(np.timedelta64(5, "s"), 1)), "_ExactDate"), + (lambda: v.ewise_union(w, udf, 0.5, 0), "_ExactArr"), + (lambda: v.ewise_union(w, udf, 0, 0.5), "_ExactArr"), + ( + lambda: gb.Scalar.from_value([1, 2, 3], arr).ewise_union(gb.Scalar(arr), udf, 0, 0.5), + "_ExactArr", + ), + # Converting would lose a number: an imaginary part, a finite float that + # float32 makes infinite, or a fraction in an array field of a record. + (lambda: f.apply(udf, right=1 + 2j), "_ExactF32"), + (lambda: f.apply(udf, right=np.complex128(1 + 2j)), "_ExactF32"), + (lambda: f.apply(udf, right=np.float64(1e300)), "_ExactF32: it is typed as FP64,"), + (lambda: n.apply(second, right=([1, 2.5], 3)), "_ExactNested"), + # A number does not become text exactly (0.5 would be b"0.5"), nor a string + # cut short, nor a duration in a coarser unit (5 ms is 0 s). + (lambda: t.apply(second, right=(1.5, 1)), "_ExactText"), + (lambda: b.apply(same, right=0.5), "_ExactBareText: it is typed as FP64,"), + (lambda: b.apply(same, right=np.True_), "_ExactBareText: it is typed as BOOL,"), + (lambda: b.apply(same, right=b"0.5000"), "_ExactBareText"), # was cut to b"0.5" + (lambda: t.apply(second, right=cut_text), "_ExactText"), + # NaT as b"N" + (lambda: c.apply(second, right=(np.datetime64("NaT", "s"), 1)), "_ExactChar"), + (lambda: d.apply(same, right=np.timedelta64(5, "ms")), "_ExactBareSeconds"), + ]: + with pytest.raises(ValueError, match=re.escape(f"does not fit {typed_as}")): + expr() + # One that does not fit by type but converts exactly still converts, as it + # did on main, with a DeprecationWarning. A field that is not a number, such + # as bytes, is converted as it is. + text_literal = np.array((b"cd", 2), dtype=[("ex_s", "S4"), ("ex_n", np.int64)])[()] + for expr, typed_as in [ + (lambda: v.apply(udf, right=[2.0, 2.0, 2.0]), "_ExactArr"), + (lambda: v.apply(udf, right=np.ones(3, dtype=np.int64)), "_ExactArr"), + (lambda: v.apply(udf, right=np.int64(1)), "_ExactArr: it is typed as INT64,"), + (lambda: v.apply(udf, right=(np.int8(1), 2, 3)), "_ExactArr"), + (lambda: r.apply(second, right=np.int32(2)), "_ExactRec: it is typed as INT32,"), + (lambda: r.apply(second, right=(np.int64(2), 0.5)), "_ExactRec"), + (lambda: n.apply(second, right=([1, 2.0], 3)), "_ExactNested"), + (lambda: t.apply(second, right=text_literal), "_ExactText"), + (lambda: when.apply(second, right=("1970-01-01", 2.0)), "_ExactDate"), # a date as text + ]: + with pytest.warns(DeprecationWarning, match=re.escape(f"does not fit {typed_as}")): + expr() + with pytest.warns(DeprecationWarning, match="does not fit _ExactText"): + assert t.apply(second, right=text_literal).new()[0].new().value.tolist() == (b"cd", 2) + # Text for a text leaf converts, bytes or str. + for literal in [b"cd", "cd"]: + assert t.apply(second, right=(literal, 1)).new()[0].new().value.tolist() == (b"cd", 1) + assert t.apply(second, right=(b"c", 1)).new()[0].new().value.tolist() == (b"c", 1) + # Text or a date that converts into its field and back fits it, with no warning: + # trailing NULs, a finer date that is whole seconds, NaT. + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert t.apply(second, right=(b"ab\0\0\0", 1)).new()[0].new().value.tolist() == (b"ab", 1) + fine = when.apply(second, right=(np.datetime64(5000, "ms"), 1)).new()[0].new().value + assert fine["ex_d"] == np.datetime64(5, "s") + nat = when.apply(second, right=(np.datetime64("NaT", "ms"), 1)).new()[0].new().value + assert np.isnat(nat["ex_d"]) + # A UDT of durations alone too; NaT stays NaT in any unit, and equals nothing. + assert d.apply(same, right=np.timedelta64(5000, "ms")).new()[0].new().value + assert np.isnat(d.apply(second, right=np.timedelta64("NaT", "ms")).new()[0].new().value) + assert not d.apply(same, right=np.timedelta64("NaT", "ms")).new()[0].new().value + assert np.isnat(d.apply(second, right=np.timedelta64("NaT", "m")).new()[0].new().value) + assert (t == (b"ab", 1)).new()[0].new().value + assert not (t == (b"ab", 1.5)).new()[0].new().value + assert (when == (np.datetime64(0, "s"), 1)).new()[0].new().value + assert not (when == (np.datetime64(0, "s"), 1.5)).new()[0].new().value + for literal in [b"0.5", "0.5"]: + assert b.apply(same, right=literal).new()[0].new().value + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + v.apply(udf, right=(1, 2, 300)) + # Defaults that fit still work, and fill in for the missing side. + result = v.ewise_union(w, udf, 10, 20).new() + assert result.to_coo()[1].tolist() == [[21, 22, 23], [14, 15, 16]] + # apply with a Monoid promotes as its BinaryOp does, on either side. + for expr in [v.apply(monoid.plus, right=0.5), v.apply(monoid.plus, left=0.5)]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + np.testing.assert_array_equal(result[0].new().value, [1.5, 2.5, 3.5]) + assert v.apply(monoid.max, right=2).new().dtype == arr + + +# _SeqNest's numpy repr is over 128 chars on numpy 1; see test_udt_eq_nested_record_with_nan_leaf. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +@pytest.mark.skipif("not supports_udfs") +def test_udt_sequence_and_array_literals_type_like_numbers(): + """A tuple, list or array beside a UDT under a lifted op is typed, not converted into the UDT. + + A tuple or list of Python numbers is weak leaf by leaf, as a Python number + is: ``int8_udt + (1, 2, 3)`` stays int8 and ``+ (0.5, 1.5, 2.5)`` is + float64, where converting it raised (and before that added ``(0, 1, 2)``). + A numpy array is strong, as its own array type. ``eq`` and ``ne`` compare + the values. A literal beside a user-defined op is still converted + (``test_udt_literal_converts_only_when_exact``). + """ + int8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_SeqI8") + fp32s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_SeqF32") + nested = dtypes.register_anonymous( + np.dtype( + [("sq_a", np.int16), ("sq_in", [("sq_b", np.int8), ("sq_c", np.float32)])], + align=True, + ), + "_SeqNest", + ) + f64s = np.dtype((np.float64, (3,))) + v = Vector(int8s, size=2) + v[0] = [1, 2, 3] + v[1] = [4, 5, 6] + # Weak: the leaf's dtype when its kind holds the numbers, of the widest kind among them. + for literal in [(1, 2, 3), [1, 2, 3], (True, 2, 3), (1,)]: + result = (v + literal).new() + assert result.dtype == int8s + np.testing.assert_array_equal(result[0].new().value, np.add([1, 2, 3], literal)) + halves = (0.5, 1.5, 2.5) + for expr in [v + halves, halves + v, v.apply(binary.plus, right=list(halves))]: + result = expr.new() + assert result.dtype.np_type == f64s + np.testing.assert_array_equal(result[0].new().value, [1.5, 3.5, 5.5]) + result = binary.plus(v, (1, 0.5, 3)).new() # one float anywhere makes the leaf float + assert result.dtype.np_type == f64s + np.testing.assert_array_equal(result[0].new().value, [2, 2.5, 6]) + result = binary.times(v, (1j, 1, 1)).new() + assert result.dtype.np_type == np.dtype((np.complex128, (3,))) + np.testing.assert_array_equal(result[0].new().value, [1j, 2, 3]) + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + binary.plus(v, (1, 2, 300)) + v += (1, 1, 1) + np.testing.assert_array_equal(v[0].new().value, [2, 3, 4]) + # The float64 result is cast as it is stored, as a built-in int8 vector's would be. + v += (0.5, 1, 1) + np.testing.assert_array_equal(v[0].new().value, [2, 4, 5]) + # A float list beside float32 stays float32, as 0.5 does. + x = Vector(fp32s, size=1) + x[0] = [1, 2, 3] + assert binary.plus(x, [0.5, 0.5, 0.5]).new().dtype == fp32s + # Strong: an array is its own dtype, in its shape, or broadcast to the UDT's. + for literal, expected in [ + (np.array([0.5, 1, 2]), f64s), + (np.array([1, 2, 3]), np.dtype((np.array([1]).dtype, (3,)))), # int32 on numpy 1 Windows + (np.array([1], np.int8), int8s.np_type), + (np.array([0.5]), f64s), + (np.ones((1, 3)), np.dtype((np.float64, (1, 3)))), + ]: + result = (v + literal).new() + assert result.dtype.np_type == expected + np.testing.assert_array_equal(result[0].new().value, np.add([2, 4, 5], literal)) + # Records: by position, nested for a nested record, a dict by name, or a + # 0-d structured array of another record type. + n = Vector(nested, size=1) + n[0] = (1, (2, 3.5)) + assert binary.plus(n, (1, (1, 1))).new().dtype == nested + result = binary.plus(n, (1, (0.5, 1))).new() + assert result.dtype.np_type == np.dtype( + [("sq_a", np.int16), ("sq_in", [("sq_b", np.float64), ("sq_c", np.float32)])], + align=True, + ) + assert result[0].new().value.tolist() == (2, (2.5, 4.5)) + assert (n + {"sq_a": 0.5, "sq_in": {"sq_b": 1, "sq_c": 1}}).new()[0].new().value.tolist() == ( + 1.5, + (3, 4.5), + ) + strong = np.array( + (1, (0.5, 1)), + dtype=[("sq_a", np.int8), ("sq_in", [("sq_b", np.float64), ("sq_c", np.int8)])], + ) + assert (n + strong).new()[0].new().value.tolist() == (2, (2.5, 4.5)) + # A numpy number in a record literal keeps its dtype, as a 0-d array does. + result = binary.plus(n, (np.float64(0.5), (1, 1))).new() + assert result.dtype.np_type["sq_a"] == np.float64 + assert result[0].new().value.tolist() == (1.5, (3, 4.5)) + assert (n == (np.float64(1), (2, 3.5))).new()[0].new().value + assert not (n == (np.float64(1.5), (2, 3.5))).new()[0].new().value + # Numba cannot compute on float16, so a float16 leaf is refused by name. + for expr in [ + lambda: binary.plus(n, (np.float16(0.5), (1, 1))), + lambda: n == (np.float16(1), (2, 3.5)), + ]: + with pytest.raises(KeyError, match="Numba cannot compute on float16"): + expr() + # A record's array field takes a sequence as an array UDT does: an array, or a + # sequence holding a numpy value, is the array numpy makes of it. + arr_rec = dtypes.register_anonymous( + np.dtype([("sq_av", np.int8, (3,)), ("sq_an", np.int8)]), "_SeqArrRec" + ) + ar = Vector(arr_rec, size=1) + ar[0] = ([1, 2, 3], 4) + for literal in [(np.int8(1), 300, 3), np.array([1, 300, 3])]: + made = np.dtype((np.asarray(literal).dtype, (3,))) + assert binary.plus(v, literal).new().dtype.np_type == made + result = binary.plus(ar, (literal, 1)).new() + assert result.dtype.np_type["sq_av"] == made + assert result[0].new().value["sq_av"].tolist() == [2, 302, 6] + # A comparison keeps a Python int beyond int8 exact, there too: it equals nothing. + big = (np.int8(1), 2**70, 3) + for x, literal in [(ar, (big, 4)), (v, big)]: + assert not (x == literal).new()[0].new().value + assert (x != literal).new()[0].new().value + result = binary.plus(n, (1, np.float64(1.5))).new() # one number for the nested record + assert result.dtype.np_type["sq_in"]["sq_b"] == np.float64 + assert result[0].new().value.tolist() == (2, (3.5, 5.0)) + # Text or a duration that converts into its field and back takes the field's type, + # so a lifted op pairs the record with itself; one that does not keeps its own. + text = dtypes.register_anonymous(np.dtype([("sq_tx", "S4"), ("sq_tn", np.int8)]), "_SeqText") + tv = Vector(text, size=1) + tv[0] = (b"ab", 1) + assert binary.plus(tv, (b"ab", 1)).new()[0].new().value.tolist() == (b"abab", 2) + assert tv.apply(binary.max, right=(b"ac", 0)).new()[0].new().value.tolist() == (b"ac", 1) + assert tv.apply(binary.min, right=(b"ac", 0)).new()[0].new().value.tolist() == (b"ab", 0) + assert not (tv == (b"ab", 300)).new()[0].new().value # 300 is no int8 + secs = dtypes.register_anonymous( + np.dtype([("sq_sd", "m8[s]"), ("sq_sn", np.int8)]), "_SeqSeconds" + ) + sv = Vector(secs, size=1) + sv[0] = (np.timedelta64(10, "s"), 1) + for delta, total in [(np.timedelta64(5, "m"), 310), (np.timedelta64(5000, "ms"), 15)]: + result = binary.plus(sv, (delta, 1)).new() + assert result[0].new().value["sq_sd"] == np.timedelta64(total, "s") + with pytest.raises(KeyError, match="no common type"): + binary.plus(sv, (np.timedelta64(5, "ms"), 1)) # 5 ms is no whole number of seconds + for nat in [np.timedelta64("NaT", "ms"), np.timedelta64("NaT", "m")]: + result = binary.plus(sv, (nat, 1)).new() # NaT is NaT in seconds + assert result.dtype == secs + assert np.isnat(result[0].new().value["sq_sd"]) + first = dtypes.register_anonymous( + np.dtype([("sq_fin", [("sq_fb", np.int8), ("sq_fc", np.float32)]), ("sq_fa", np.int16)]) + ) + m = Vector(first, size=1) + m[0] = ((2, 3.5), 1) + assert binary.plus(m, ((0.5, 1), 1)).new()[0].new().value.tolist() == ((2.5, 4.5), 2) + # eq and ne take the literal as it is, like numpy, instead of raising. + assert (v == (0.5, 4, 5)).new().isequal(Vector.from_coo([0, 1], [False, False])) + assert (v == (2, 4, 5)).new().isequal(Vector.from_coo([0, 1], [True, False])) + assert (v != (2, 4, 500)).new().isequal(Vector.from_coo([0, 1], [True, True])) + assert (v == np.array([2.0, 4, 5])).new()[0].new().value + # Matrix and Scalar take the same path. + A = Matrix(int8s, nrows=1, ncols=1) + A[0, 0] = [1, 2, 3] + ones = [1, 1, 1] + assert (A + halves).new().dtype.np_type == f64s + assert (A + ones).new().dtype == int8s + s = gb.Scalar.from_value([1, 2, 3], int8s) + assert (s + halves).new().dtype.np_type == f64s + assert (s * ones).new().dtype == int8s + + +def test_udt_record_array_field_refuses_literal_of_another_shape(): + """A literal for a record's array field must broadcast to the field's shape. + + An array UDT's element takes the shape a literal broadcasts it to, as numpy + does: ``INT8[3]`` plus ``np.array([[1, 300, 3]])`` is ``INT64[1, 3]``. A + record's field keeps its shape, so the record refuses such a literal and + names both shapes. Before, its typing gave up and the literal was converted + into the field as numpy assigns it, dropping the leading axis and wrapping + 300 to 44, through ``+`` and through a user's op alike; a nested list + raised numpy's ValueError. A store still converts as numpy assigns. + """ + arr = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_ShapeArrI8") + rec = dtypes.register_anonymous( + np.dtype([("shp_a", np.int8, (3,)), ("shp_b", np.int64)]), "_ShapeRec" + ) + nest = dtypes.register_anonymous( + np.dtype([("shp_n", np.int64), ("shp_in", rec.np_type)]), "_ShapeNest" + ) + v = Vector(arr, size=1) + v[0] = [1, 2, 3] + r = Vector(rec, size=1) + r[0] = ([1, 2, 3], 0) + n = Vector(nest, size=1) + n[0] = (0, ([1, 2, 3], 0)) + exprs = [ + lambda L: r == (L, 0), + lambda L: r.apply(binary.second, right=(L, 0)), + lambda L: r.apply(binary.first, left={"shp_a": L, "shp_b": 0}), + lambda L: r[0].new() == (L, 0), + lambda L: n == (0, (L, 0)), + ] + if supports_udfs: + udf = BinaryOp.register_anonymous(lambda x, y: y, "_shape_right", is_udt=True) + exprs += [ + lambda L: binary.plus(r, (L, 0)), + lambda L: r.apply(udf, right=(L, 0)), + lambda L: r.ewise_union(r, udf, (L, 0), (L, 0)), + ] + for literal in [ + np.array([[1, 300, 3]]), + [[1, 300, 3]], + np.array([[1, 2, 3]], np.int8), + [1, 2], + np.zeros((2, 3), np.int8), + ]: + reason = f"has shape {np.shape(literal)}, which does not fit field 'shp_a'" + for expr in exprs: + with pytest.raises(ValueError, match="of shape \\(3,\\): a literal for an array") as e: + expr(literal) + assert reason in str(e.value) + # A store converts as numpy assigns, on both sides + v[0] = np.array([[1, 300, 3]]) + r[0] = (np.array([[1, 300, 3]]), 0) + assert v[0].new().value.tolist() == [1, 44, 3] + assert r[0].new().value["shp_a"].tolist() == [1, 44, 3] + if supports_udfs: + # The array UDT keeps numpy's arithmetic shape + result = (v + np.array([[1, 300, 3]])).new() + assert result.dtype.np_type == np.dtype((np.array([1]).dtype, (1, 3))) + assert result[0].new().value.tolist() == [[2, 344, 6]] + # and a literal that broadcasts to the field is typed as before + value = binary.plus(r, (np.array([5]), 1)).new()[0].new().value + assert value["shp_a"].tolist() == [6, 49, 8] + assert value["shp_b"] == 1 + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_literal_broadcasts_with_an_array_udt(): + """An array literal pairs with an array UDT when their shapes broadcast, as a UDT operand does. + + The result has the shape numpy's would: ``INT8[3]`` and a 2-by-3 array + give ``[2, 3]``, as ``INT8[3]`` and an ``INT64[2, 3]`` UDT do. Before, an + array paired only when it broadcast to the UDT's shape; any other was + converted into the UDT and raised numpy's "could not broadcast" error. + """ + i8x3 = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_BcastI8x3") + i8x1 = dtypes.register_anonymous(np.dtype((np.int8, (1,))), "_BcastI8x1") + int64 = np.array([1]).dtype # int32 on Windows with numpy 1 + v = Vector(i8x3, size=1) + v[0] = [1, 2, 3] + u = Vector(i8x1, size=1) + u[0] = [1] + rows = np.array([[1, 2, 3], [4, 5, 6]]) + for expr in [lambda: v + rows, lambda: v.apply(binary.plus, left=rows)]: + result = expr().new() + assert result.dtype.np_type == np.dtype((int64, (2, 3))) + assert result[0].new().value.tolist() == [[2, 4, 6], [5, 7, 9]] + result = (u + np.array([1, 2, 3])).new() + assert result.dtype.np_type == np.dtype((int64, (3,))) + assert result[0].new().value.tolist() == [2, 3, 4] + assert (v + np.array([[1], [2]])).new()[0].new().value.tolist() == [[2, 3, 4], [3, 4, 5]] + assert (v == np.array([[1, 2, 3], [1, 2, 3]])).new()[0].new().value + assert not (v == rows).new()[0].new().value + assert v[0].new() == np.array([[1, 2, 3], [1, 2, 3]]) + # A literal that must become an element must still fit it. + with pytest.raises(ValueError, match="does not fit _BcastI8x3"): + v.ewise_union(v, binary.plus, rows, 0) + + +def test_udt_record_date_literals_compare_as_numpy(): + """A Python date or duration, or a date as text, in a record literal is typed as numpy types it. + + It takes its field's type only when it converts into the field and back + unchanged, as a numpy date does; otherwise it keeps numpy's type for it (a + ``datetime`` is in microseconds), so ``==`` compares it whole: noon is not + the day, and 5 s and 1 us is not 5 s. Before, it always took the field's + type, so it was cut to the field's unit and compared equal. + """ + import datetime + + rec = dtypes.register_anonymous(np.dtype([("dl_d", "M8[D]"), ("dl_t", "m8[s]")]), "_DateLit") + day, secs = np.datetime64("2020-01-01"), np.timedelta64(5, "s") + v = Vector(rec, size=1) + v[0] = (day, secs) + s = v[0].new() + # numpy's dates have no time zone, so neither do these. + midnight = datetime.datetime(2020, 1, 1) # noqa: DTZ001 + noon = datetime.datetime(2020, 1, 1, 12) # noqa: DTZ001 + for literal, equal in [ + ((noon, secs), False), + ((day, datetime.timedelta(seconds=5, microseconds=1)), False), + (("2020-01-01T12:00", secs), False), + ((datetime.date(2020, 1, 1), datetime.timedelta(seconds=5)), True), + ((midnight, secs), True), + (("2020-01-01", secs), True), + # Text in a duration field takes the field's type, as before. + ((day, "5"), True), + ]: + assert (s == literal) is equal + if supports_udfs: + assert (v == literal).new()[0].new().value == equal + assert (v != literal).new()[0].new().value != equal + # Text numpy reads as NaT is NaT in the field's unit (numpy 2.5 warns for NaT with no + # unit), and NaT equals nothing, as in numpy. + bare = dtypes.register_anonymous(np.dtype("M8[s]"), "_DateLitBare") + b = gb.Scalar.from_value(np.datetime64(0, "s"), dtype=bare) + for nat in ["NaT", "nat", "", b"NaT"]: + assert (s == (nat, secs)) is False + assert (b == nat) is False + if supports_udfs: + assert not (v == (nat, secs)).new()[0].new().value + # So does text numpy cannot read as a date or duration, and converting it says why. + for literal, reason in [ + (("not a date", secs), "Error parsing datetime string"), + ((day, "abc"), "Could not convert object to NumPy timedelta"), + ]: + with pytest.raises(ValueError, match=reason): + s == literal # noqa: B015 + if supports_udfs: + with pytest.raises(ValueError, match=reason): + (v == literal).new() + if supports_udfs: + # A user's op refuses what its field would cut short, as for a numpy date. + second = BinaryOp.register_anonymous(lambda x, y: y, "_date_lit_second", is_udt=True) + with pytest.raises(ValueError, match="does not fit _DateLit"): + v.apply(second, right=(noon, secs)) + result = v.apply(second, right=(midnight, secs)).new() + assert result[0].new().value.tolist() == (day, secs) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_typed_op_takes_a_literal_in_its_own_input_type(): + """A literal beside a typed op is typed by the op's input type on its side, not by the UDT's. + + ``op[A, INT64]`` takes its second input or thunk as INT64, so a literal + ``2`` there is an INT64, as a Scalar would be, and ``2.5`` is converted to + it by GraphBLAS, as for a typed op on built-in types. Before, the literal + became an element of A, whose bytes the op read as an int64. + """ + A = dtypes.register_anonymous(np.dtype([("tl_a", "f8"), ("tl_b", "f8")]), "_TypedLitA") + B = dtypes.register_anonymous(np.dtype([("tl_a", "i8"), ("tl_b", "i8")]), "_TypedLitB") + v = Vector(A, size=2) + v[0] = (1.0, 2.0) + v[1] = (3.0, 4.0) + M = Matrix(A, nrows=2, ncols=1) + M[0, 0] = (1.0, 2.0) + M[1, 0] = (3.0, 4.0) + iu = IndexUnaryOp.register_anonymous( + lambda x, i, j, th: x.tl_a > th, "_typed_lit_iu", is_udt=True + ) + iu_rec = IndexUnaryOp.register_anonymous( + lambda x, i, j, th: x.tl_a > th.tl_a, "_typed_lit_iu_rec", is_udt=True + ) + se = SelectOp.register_anonymous(lambda x, i, j, th: x.tl_a > th, "_typed_lit_se", is_udt=True) + right = BinaryOp.register_anonymous(lambda x, y: x.tl_a + y, "_typed_lit_right", is_udt=True) + left = BinaryOp.register_anonymous(lambda x, y: x + y.tl_a, "_typed_lit_left", is_udt=True) + for thunk in [2, 2.5, np.int8(2), gb.Scalar.from_value(2)]: + assert v.apply(iu[A, INT64], thunk).new().to_coo()[1].tolist() == [False, True] + assert v.select(se[A, INT64], thunk).new().to_coo()[0].tolist() == [1] + assert M.select(se[A, INT64], thunk).new().to_coo()[0].tolist() == [1] + assert v[0].new().apply(iu[A, INT64], thunk).new().value is False + assert v.apply(iu_rec[A, B], (2, 0)).new().to_coo()[1].tolist() == [False, True] + assert v.apply(right[A, INT64], right=2).new().to_coo()[1].tolist() == [3, 5] + assert v.apply(left[INT64, A], left=2).new().to_coo()[1].tolist() == [3, 5] + assert M.apply(left[INT64, A], left=2).new().to_coo()[2].tolist() == [3, 5] + assert v[1].new().ewise_mult(2, right[A, INT64]).new().value == 5 + # A UDT input type still takes only a literal that fits it. + with pytest.raises(ValueError, match="does not fit _TypedLitB"): + v.apply(iu_rec[A, B], (2.5, 0)) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ewise_union_defaults_are_typed_not_converted(): + """An ``ewise_union`` default beside a UDT is typed as a literal and must fit the operand. + + GraphBLAS casts a default to the op's input type, which for a UDT operand + is its own type, and converting the operands instead would copy them. So + a default is typed by the literal rules and accepted only when that type + fits: ``1`` and ``(1, 1, 1)`` fit an int8 UDT, ``0.5``, ``np.int64(1)`` and + a float64 Scalar do not. A Scalar (one element) of a type that casts + safely is converted. + """ + import re + + int8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_UnionI8") + fp32s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_UnionF32") + v = Vector(int8s, size=3) + v[0] = [0, 1, 2] + v[1] = [4, 5, 6] + w = Vector(int8s, size=3) + w[1] = [1, 1, 1] + w[2] = [7, 8, 9] + plus = binary.plus + expected = [[1, 2, 3], [5, 6, 7], [7, 8, 9]] + for left, right in [ + (0, 1), + ((0, 0, 0), [1, 1, 1]), + (np.int8(0), np.array([1, 1, 1], np.int8)), + (gb.Scalar.from_value(0, "INT8"), gb.Scalar.from_value([1, 1, 1], int8s)), + (False, True), + ]: + for expr in [ + v.ewise_union(w, plus, left, right), + v.ewise_union(w, monoid.plus, left, right), + ]: + result = expr.new() + assert result.dtype == int8s + assert result.to_coo()[1].tolist() == expected + + def union_forms(d): + return [ + lambda: v.ewise_union(w, plus, d, 0), + lambda: v.ewise_union(w, plus, 0, d), + lambda: plus(v | w, left_default=d, right_default=0), + lambda: v.ewise_union(w, monoid.plus, d, 0), + ] + + for default, message in [ + (0.5, "0.5 does not fit _UnionI8"), + ((0.5, 0.5, 0.5), "(0.5, 0.5, 0.5) does not fit _UnionI8"), + (np.float32(0.5), "does not fit _UnionI8: it is typed as FP32,"), + (gb.Scalar.from_value(0.5), "Scalar of type FP64 does not fit _UnionI8"), + ]: + for expr in union_forms(default): + with pytest.raises(ValueError, match=re.escape(message)): + expr() + # Exact values of another type still convert, as on main, with a + # DeprecationWarning. + for default, message in [ + (2.0, "2.0 does not fit _UnionI8"), + (np.int64(1), "does not fit _UnionI8: it is typed as INT64,"), + (np.array([1, 1, 1]), "array([1, 1, 1]) does not fit _UnionI8"), + (gb.Scalar.from_value(1), "Scalar of type INT64 does not fit _UnionI8"), + ]: + for expr in union_forms(default): + with pytest.warns(DeprecationWarning, match=re.escape(message)): + expr() + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + v.ewise_union(w, plus, 300, 0) + # Each side is typed beside its own operand, so a mixed pair takes a + # default of each type, and an int8 Scalar casts safely into float32. + x = Vector(fp32s, size=3) + x[2] = [0.5, 0.5, 0.5] + result = v.ewise_union(x, plus, 0, 0.5).new() + assert result.dtype == fp32s + assert result.to_coo()[1].tolist() == [[0.5, 1.5, 2.5], [4.5, 5.5, 6.5], [0.5, 0.5, 0.5]] + result = v.ewise_union(x, plus, 0, gb.Scalar.from_value([1, 1, 1], int8s)).new() + assert result.to_coo()[1].tolist() == [[1, 2, 3], [5, 6, 7], [0.5, 0.5, 0.5]] + with pytest.raises(ValueError, match="0.5 does not fit _UnionI8"): + v.ewise_union(x, plus, 0.5, 0) + # Matrix and Scalar check the same way. + A = Matrix(int8s, nrows=2, ncols=1) + A[0, 0] = [0, 1, 2] + B = Matrix(int8s, nrows=2, ncols=1) + B[1, 0] = [7, 8, 9] + assert A.ewise_union(B, plus, 1, 1).new().to_coo()[2].tolist() == [[1, 2, 3], [8, 9, 10]] + s = gb.Scalar.from_value([1, 2, 3], int8s) + assert s.ewise_union(gb.Scalar(int8s), plus, 0, 1).new().value.tolist() == [2, 3, 4] + for expr in [ + lambda: A.ewise_union(B, plus, 0.5, 0), + lambda: A.ewise_union(v, plus, 0, 0.5), + lambda: s.ewise_union(gb.Scalar(int8s), plus, 0, 0.5), + ]: + with pytest.raises(ValueError, match="0.5 does not fit _UnionI8"): + expr() + + +@pytest.mark.skipif("not supports_udfs") +# SS < 9 names a UDT by its numpy repr and warns when that is over 128 chars, +# as the records with a field per dtype pair are. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_store_casts_each_element_as_builtin_types_do(): + """A UDT result stored in an object of another UDT type is cast element by element. + + GraphBLAS cannot cast a UDT, so storing across UDT types raised, and + setting one element from a Scalar of another UDT type read its bytes as the + object's type. Each array element or record leaf is now cast as GraphBLAS + casts built-in types on store (compared below), without a temporary copy: + an object stores through the cast op (``<<``, a mask, an accumulator, + ``replace``, ``dup``), an element-wise expression of a lifted op casts as + it computes, and a Scalar or one element is converted. Anything that would + need a converted copy of a whole object raises, as do layouts that do not + correspond and a built-in type on either side. + """ + # One record casts every pair of leaf dtypes, so it compiles once. + pairs = [ + ("f8", "i1"), + ("f8", "u1"), + ("f8", "i8"), + ("f8", "u8"), + ("f8", "?"), + ("f8", "f4"), + ("f4", "u4"), + ("i8", "i1"), + ("u8", "i8"), + ("i1", "u8"), + ("?", "f8"), + ] + if dtypes._supports_complex: + pairs += [("c16", "f8"), ("c16", "i2"), ("c16", "?"), ("i4", "c8")] + samples = { + "f": [np.nan, np.inf, -np.inf, 300.7, -1.7, 2.0**63, -(2.0**64)], + "c": [complex(np.nan, 1), 1 + 2j, -1.7 - 3j, 300.5j, 2.0**64, 0j, -2.5], + "i": [2**63 - 1, -(2**63), 128, -129, -1, 0, 255], + "u": [2**64 - 1, 2**63, 255, 256, 1, 0, 65535], + "b": [True, False, True, True, False, True, False], + } + src_type = np.dtype([(f"sc{i}", src) for i, (src, _dst) in enumerate(pairs)]) + dst_type = np.dtype([(f"sc{i}", dst) for i, (_src, dst) in enumerate(pairs)]) + values = np.zeros(7, dtype=src_type) + for i, (src, _dst) in enumerate(pairs): + kind = np.dtype(src).kind + values[f"sc{i}"] = np.array(samples[kind], np.uint64 if kind == "u" else None).astype(src) + x = Vector.from_coo(np.arange(7), values, dtype=dtypes.register_anonymous(src_type)) + z = Vector(dtypes.register_anonymous(dst_type), size=7) + z << x + cast = z.to_coo()[1] + for i, (src, dst) in enumerate(pairs): + expected = Vector(dst, size=7) + expected << Vector.from_coo(np.arange(7), values[f"sc{i}"], dtype=src) + np.testing.assert_array_equal(cast[f"sc{i}"], expected.to_coo()[1], err_msg=f"{src} {dst}") + + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3, 2))), "_StoreCastI8") + f64s = dtypes.register_anonymous(np.dtype((np.float64, (3, 2))), "_StoreCastF64") + v = Vector(f64s, size=3) + v[0] = [[1.7, -1.7], [300.0, np.nan], [-np.inf, 2.5]] + v[2] = np.full((3, 2), 9.5) + first = [[1, -1], [127, 0], [-128, 2]] + nines = [[9, 9]] * 3 + w = Vector(i8s, size=3) + w << v + assert w.to_coo()[1].tolist() == [first, nines] + assert (v * 1).new(dtype=i8s).isequal(w) + # The mask, accumulator and replace act on the cast values, as for built-in + # types: ``1 + 127`` wraps in int8. + mask = Vector.from_coo([0, 1], [True, True], size=3) + ones = [[1, 1]] * 3 + w = Vector.from_coo([0, 1, 2], [ones] * 3, dtype=i8s) + w(mask.S, accum=binary.plus) << v + assert w.to_coo()[1].tolist() == [[[2, 0], [-128, 1], [-127, 3]], ones, ones] + w(mask.S) << v * 1 # an expression of a lifted op casts as it computes + assert w.to_coo()[1].tolist() == [first, ones] + w(mask.S, replace=True) << v + assert w.to_coo()[1].tolist() == [first] + # Scalars, as GraphBLAS scalars or not, and assignment of one element or many. + for is_cscalar in [True, False]: + s = gb.Scalar(i8s, is_cscalar=is_cscalar) + s << v[0].new() + assert s.value.tolist() == first + s << v.reduce(monoid.plus) # summed in float64, then cast + assert s.value.tolist() == [[11, 7], [127, 0], [-128, 12]] + assert v[0].new(is_cscalar=is_cscalar).dup(i8s).value.tolist() == first + assert v.dup(i8s).isequal((v * 1).new(dtype=i8s)) + w = Vector(i8s, size=3) + w[1] = v[0].new() # read as int8 bytes before + w[[0, 2]] = v[2].new() + assert w.to_coo()[1].tolist() == [nines, first, nines] + A = Matrix.from_coo([0], [1], [v[0].new().value], dtype=f64s, nrows=2, ncols=2) + C = Matrix(i8s, nrows=2, ncols=2) + C << A.T + assert C.to_coo()[2].tolist() == [first] + # A whole Vector of another type would need a converted copy to assign. + with pytest.raises(DomainMismatch, match="converted copy of the whole value"): + C[0, :] = v[[0, 1]].new() + C[0, :] = v[[0, 1]].new().dup(dtype=i8s) + assert C[0, 0].new().value.tolist() == first + # So would a computation that cannot write another type as it computes. + with pytest.raises(DomainMismatch, match="would need a temporary copy"): + C << A.mxm(A, semiring.plus_times) + with pytest.raises(DomainMismatch, match="would need a temporary copy"): + w << v.ewise_add(v, binary.plus) + # Only a built-in op that lifts to UDTs can, so not first or a user's op. + udf = BinaryOp.register_anonymous(lambda x, y: x, "_store_cast_first", is_udt=True) + for not_lifted in [binary.first, udf]: + with pytest.raises(DomainMismatch, match="would need a temporary copy"): + w << v.ewise_mult(v, not_lifted) + with pytest.raises(DomainMismatch, match="extraction cannot"): + w << v[[0, 1, 2]] + # The computed result, stored as an existing object, casts in one pass. + C << A.mxm(A, semiring.plus_times).new() + # power raises before it computes any product; n of 0 and 1 store as above. + with pytest.raises(DomainMismatch, match="mxm with plus_times cannot"): + C << A.power(3) + C << A.power(1) + assert C.to_coo()[2].tolist() == [first] + # A semiring whose result is another type (int truediv is float) gives that + # type, which stores as it is, and each later product takes the semiring + # typed for its operands, so the powers are those of a float64 copy, as for + # built-in types. + P = Matrix(i8s, nrows=2, ncols=2) + P[0, 0] = np.full((3, 2), 2) + P[0, 1] = np.full((3, 2), 1) + P[1, 1] = [[4, 2], [1, -1], [3, 5]] + Q = Matrix(f64s, nrows=2, ncols=2) + for n in [1, 2, 3, 4, 5]: + Q << P.power(n, semiring.plus_truediv) + assert Q.isequal(P.dup(dtype=f64s).power(n, semiring.plus_truediv).new()), n + # Leading axes of length 1 do not change the layout. + lead = dtypes.register_anonymous(np.dtype((np.float32, (1, 3, 2))), "_StoreCastLead") + u = Vector(lead, size=3) + u << w + assert u.to_coo()[1].tolist() == [[nines], [first], [nines]] + + other_shape = Vector(dtypes.register_anonymous(np.dtype((np.int16, (6,)))), size=3) + with pytest.raises(DomainMismatch, match="same shape, apart from leading axes"): + other_shape << v + renamed = np.dtype([(f"sr{i}", dst) for i, (_src, dst) in enumerate(pairs)]) + with pytest.raises(DomainMismatch, match="same field names"): + Vector(dtypes.register_anonymous(renamed), size=7) << x + # A record field keeps its shape, and one that is not a number casts only to + # its own type, as for array UDTs and built-in types. + narrow = dtypes.register_anonymous(np.dtype([("sf_v", np.int8, (2,)), ("sf_s", "S4")])) + wide = dtypes.register_anonymous(np.dtype([("sf_v", np.int8, (3,)), ("sf_s", "S4")])) + longer = dtypes.register_anonymous(np.dtype([("sf_v", np.int8, (2,)), ("sf_s", "S8")])) + y = Vector(narrow, size=1) + y[0] = ([1, 2], b"ab") + with pytest.raises(DomainMismatch, match="same shape for each field"): + Vector(wide, size=1) << y + with pytest.raises(DomainMismatch, match=r"elements of \|S4 do not cast to \|S8"): + Vector(longer, size=1) << y + raw = Vector(dtypes.register_anonymous(np.dtype("S48"), "_StoreCastBytes"), size=3) + with pytest.raises(DomainMismatch, match="only record and array UDTs cast"): + raw << v + for expr in [ + lambda: Vector(f64s, size=3) << Vector.from_coo([0], [1.5], size=3), + lambda: Vector(dtypes.FP64, size=3) << v, + lambda: w.__setitem__(0, gb.Scalar.from_value(1.5)), + ]: + with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): + expr() + + +def test_udt_store_casts_need_numba(monkeypatch): + """Without Numba no UDT cast compiles, so every store across types with a UDT raises.""" + i8s = dtypes.register_anonymous(np.dtype((np.int8, (2,))), "_NoNumbaCastI8") + f64s = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_NoNumbaCastF64") + v = Vector(i8s, size=1) + v[0] = [1, 2] + monkeypatch.setattr(gb.core.udt, "_has_numba", False) + for out in [Vector(f64s, size=1), Vector(dtypes.FP64, size=1)]: + with pytest.raises(DomainMismatch, match=r"_NoNumbaCastI8 in an object .*: UDTs do not"): + out << v + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_extracted_element_of_another_type_is_cast(): + """One element extracted into another UDT type is cast, not copied as raw bytes. + + GraphBLAS's extractElement for a UDT copies the bytes of the object's type + into the destination, so ``s << v[0]`` with ``s`` of another UDT read + garbage, and wrote past the end of a smaller element (also on main). + """ + f64s = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_ExtractF64") + i16s = dtypes.register_anonymous(np.dtype((np.int16, (2,))), "_ExtractI16") + rec = dtypes.register_anonymous(np.dtype([("ex_a", np.int8), ("ex_b", np.int8)]), "_ExtractRec") + v = Vector(f64s, size=3) + v[0] = [1.7, -300.2] + A = Matrix(f64s, nrows=2, ncols=3) + A[1, 2] = [2.5, -1e30] + s = gb.Scalar(i16s) + s << v[0] + assert s.value.tolist() == [1, -300] + s << A[1, 2] + assert s.value.tolist() == [2, -32768] # saturates + assert v[0].new(dtype=i16s).value.tolist() == [1, -300] + assert A.T[2, 1].new(dtype=i16s).value.tolist() == [2, -32768] # saturates + w = Vector(i16s, size=3) + C = Matrix(i16s, nrows=2, ncols=2) + with gb.config.set(autocompute=True): + w[1] = v[0] + C[1, 0] = A[1, 2] + assert w[1].new().value.tolist() == [1, -300] + assert C[1, 0].new().value.tolist() == [2, -32768] + # A larger destination, and a missing element. + x = Vector(i16s, size=2) + x[0] = [7, -8] + assert x[0].new(dtype=f64s).value.tolist() == [7.0, -8.0] + assert x[1].new(dtype=f64s).is_empty + for dtype in [rec, dtypes.FP64]: + with pytest.raises(DomainMismatch, match="cannot store"): + v[0].new(dtype=dtype) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_fused_store_matches_compute_then_cast(): + """Storing an element-wise result into another UDT type casts as the op computes. + + The op computes each element as it would for its own result type and casts + it into the object's type as it writes it, so the values are those of + computing the result and then casting it (``.new().dup(dtype=)``), with no + temporary of the result's type: special values, records with array + fields, literals on either side, unary ops, masks and accumulators. + """ + i16s = dtypes.register_anonymous(np.dtype((np.int16, (2, 2))), "_FuseI16") + f32s = dtypes.register_anonymous(np.dtype((np.float32, (2, 2))), "_FuseF32") + rec_i = dtypes.register_anonymous( + np.dtype([("fz_a", np.int8), ("fz_b", np.uint16, (3,))]), "_FuseRecI" + ) + rec_f = dtypes.register_anonymous( + np.dtype([("fz_a", np.float64), ("fz_b", np.float64, (3,))]), "_FuseRecF" + ) + i = Vector(i16s, size=4) + i[0] = [[1, -2], [32767, -32768]] + i[2] = [[7, 0], [-7, 3]] + f = Vector(f32s, size=4) + f[0] = [[np.nan, -np.inf], [1e10, -0.5]] + f[1] = [[2.5, -2.5], [65535.9, 1.5]] + f[2] = [[0.25, 3.75], [-1.0, 2.0]] + r = Vector(rec_f, size=4) + r[0] = (300.7, [-1.5, 70000.2, np.nan]) + r[1] = (-129.9, [2.5, 3.5, 1e300]) + cases = [ + (lambda: i * 2.5, i16s), + (lambda: 2.5 - i, i16s), + (lambda: i / 4, i16s), + (lambda: i.ewise_mult(f, binary.times), i16s), + (lambda: f.ewise_mult(i, binary.minus), i16s), + (lambda: i.ewise_union(f, binary.plus, 0, 0), i16s), + (lambda: f.apply(unary.ainv), i16s), + (lambda: f.apply(monoid.max, right=3), i16s), + (lambda: f * f, i16s), + (lambda: f.ewise_mult(f, monoid.plus), i16s), + (lambda: r * 1.5, rec_i), + (lambda: r.apply(binary.minus, left=gb.Scalar.from_value(0.5)), rec_i), + (lambda: r.apply(unary.abs), rec_i), + ] + mask = Vector.from_coo([0, 2], [True, True], size=4) + for make, out in cases: + expected = make().new().dup(dtype=out) + fused = Vector(out, size=expected.size) + fused << make() + assert fused.isequal(expected), make + # A mask and an accumulator act on the cast values, as they would on + # values that were already of the object's type. + start = expected.dup() + explicit = start.dup() + explicit(mask.S, accum=binary.plus) << make().new().dup(dtype=out) + start(mask.S, accum=binary.plus) << make() + assert start.isequal(explicit), make + A = Matrix(f32s, nrows=2, ncols=2) + A[0, 1] = [[1.5, -2.5], [np.nan, 1e9]] + C = Matrix(i16s, nrows=2, ncols=2) + C << A.T.apply(binary.times, right=2) + assert C.isequal(A.T.apply(binary.times, right=2).new().dup(dtype=i16s)) + + +def _assert_same_floats(got, expected, msg): + """Equal, NaN equal to NaN, and with the same sign on every zero (of each complex part).""" + got, expected = np.asarray(got), np.asarray(expected) + np.testing.assert_array_equal(got, expected, err_msg=msg) + for part in (np.real, np.imag) if got.dtype.kind == "c" else (np.asarray,): + np.testing.assert_array_equal( + np.signbit(part(got)), np.signbit(part(expected)), err_msg=msg + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_record_array_fields_compute_element_by_element(): + """An array-valued record field computes each element as a scalar field does. + + As a whole-array expression it ran through Numba's ufuncs: ``min`` and + ``max`` did not compile (and the scalar NaN rule, applied to an array, left + garbage), a zero complex divisor raised inside the cfunc and left the + element unwritten, and ``int8 / int8`` ran in float32. Each element now + takes the scalar expression, so a field agrees with numpy's ``fmin``, + ``fmax``, ``floor_divide`` and ``true_divide`` element for element, with a + record or a number on the other side. ``INT64_MIN // -1`` is the one + integer quotient Numba gets wrong by itself (0, where numpy wraps). + """ + nan, inf = float("nan"), float("inf") + rec = dtypes.register_anonymous( + np.dtype( + [ + ("aef_s", np.float64), + ("aef_f", np.float64, (2, 3)), + ("aef_i", np.int8, (4,)), + ("aef_l", np.int64, (3,)), + ("aef_m", np.int64), + ], + align=True, + ), + "_ArrElemFields", + ) + fx = np.array([[nan, 1.0, -0.0], [inf, 1.0, 7.5]]) + fy = np.array([[1.0, nan, 2.0], [2.0, 0.1, -2.5]]) + ix = np.array([-128, 7, -7, 100], np.int8) + iy = np.array([-1, 0, 2, 3], np.int8) + int64_min = np.iinfo(np.int64).min + lx = np.array([int64_min, -7, 9], np.int64) + ly = np.array([-1, 2, 0], np.int64) + v = Vector(rec, size=1) + v[0] = (nan, fx, ix, lx, int64_min) + w = Vector(rec, size=1) + w[0] = (1.0, fy, iy, ly, -1) + # The references overflow (``-128 // -1``) and divide by zero on purpose. + with np.errstate(all="ignore"): + for gb_op, reference in [ + (binary.min, np.fmin), + (binary.max, np.fmax), + (binary.floordiv, np.floor_divide), + (binary.truediv, np.true_divide), + ]: + got = gb_op(v & w).new()[0].new().value + _assert_same_floats(got["aef_s"], reference(nan, 1.0), gb_op.name) + _assert_same_floats(got["aef_f"], reference(fx, fy), gb_op.name) + np.testing.assert_array_equal(got["aef_i"], reference(ix, iy), err_msg=gb_op.name) + np.testing.assert_array_equal(got["aef_l"], reference(lx, ly), err_msg=gb_op.name) + np.testing.assert_array_equal( + got["aef_m"], reference(np.int64(int64_min), np.int64(-1)), err_msg=gb_op.name + ) + # A number on the other side reaches every element of the field. + for other in [np.float64(0.5), np.float64(nan)]: + got = v.apply(gb_op, right=other).new()[0].new().value + _assert_same_floats(got["aef_f"], reference(fx, other), f"{gb_op.name} {other}") + got = v.apply(gb_op, right=np.int8(-1)).new()[0].new().value + np.testing.assert_array_equal( + got["aef_i"], reference(ix, np.int8(-1)), err_msg=gb_op.name + ) + np.testing.assert_array_equal( + got["aef_l"], reference(lx, np.int8(-1)), err_msg=gb_op.name + ) + + cplx = dtypes.register_anonymous( + np.dtype([("aef_c", np.complex128, (3,))], align=True), "_ArrElemComplex" + ) + cx = np.array([3 + 4j, 1j, 2 - 2j]) + cy = np.array([0j, 2, 0j]) + v = Vector(cplx, size=1) + v[0] = (cx,) + w = Vector(cplx, size=1) + w[0] = (cy,) + with np.errstate(divide="ignore", invalid="ignore"): + expected = cx / cy + np.testing.assert_array_equal(binary.truediv(v & w).new()[0].new().value["aef_c"], expected) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_min_max_and_division_match_numpy(udt_op_path): + """Array UDT elements take ``fmin``, ``fmax`` and numpy's divisions on both paths. + + NaN sits on each side in turn, and ``INT64_MIN // -1`` is the integer + quotient Numba gets wrong by itself. + """ + nan, inf = float("nan"), float("inf") + f64 = dtypes.register_anonymous(np.dtype((np.float64, (2, 3))), "_ArrNumpyF64") + i64 = dtypes.register_anonymous(np.dtype((np.int64, (3,))), "_ArrNumpyI64") + fx = np.array([[nan, 1.0, -7.5], [inf, 1.0, 0.0]]) + fy = np.array([[1.0, nan, 2.0], [2.0, 0.1, 0.0]]) + int64_min = np.iinfo(np.int64).min + ix = np.array([int64_min, -7, 9], np.int64) + iy = np.array([-1, 2, 0], np.int64) + # Two entries each, with the values swapped in the second, so neither + # vector is iso-valued (SuiteSparse answers those without a kernel). + v = Vector(f64, size=2) + w = Vector(f64, size=2) + v[0], w[0], v[1], w[1] = fx, fy, fy, fx + a = Vector(i64, size=2) + b = Vector(i64, size=2) + a[0], b[0], a[1], b[1] = ix, iy, ix[::-1], iy[::-1] + # The references divide by zero and overflow on purpose. + with np.errstate(all="ignore"): + for gb_op, reference in [ + (binary.min, np.fmin), + (binary.max, np.fmax), + (binary.floordiv, np.floor_divide), + (binary.truediv, np.true_divide), + ]: + got = gb_op(v & w).new() + for k, (x, y) in enumerate([(fx, fy), (fy, fx)]): + msg = f"{udt_op_path} {gb_op.name} [{k}]" + _assert_same_floats(got[k].new().value, reference(x, y), msg) + got = gb_op(a & b).new() + for k, (x, y) in enumerate([(ix, iy), (ix[::-1], iy[::-1])]): + msg = f"{udt_op_path} {gb_op.name} [{k}]" + np.testing.assert_array_equal(got[k].new().value, reference(x, y), err_msg=msg) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_fused_store_keeps_min_max_and_division_semantics(udt_op_path): + """A fused store computes ``min``, ``max`` and the divisions as the op does, then casts. + + The fused op is cfunc only, so with the C JIT on, the plain op runs the C + kernel and the two paths are compared here directly: NaN in ``min`` and + ``max``, zero divisors, ``INT_MIN // -1``, a complex zero divisor, and + array-valued record fields. + """ + nan, inf = float("nan"), float("inf") + f64 = dtypes.register_anonymous(np.dtype((np.float64, (2, 2))), "_FuseSemF64") + f32 = dtypes.register_anonymous(np.dtype((np.float32, (2, 2))), "_FuseSemF32") + i8s = dtypes.register_anonymous(np.dtype((np.int8, (4,))), "_FuseSemI8") + i32s = dtypes.register_anonymous(np.dtype((np.int32, (4,))), "_FuseSemI32") + rec64 = dtypes.register_anonymous( + np.dtype([("fss_a", np.float64), ("fss_b", np.float64, (3,))], align=True), "_FuseSemR64" + ) + rec32 = dtypes.register_anonymous( + np.dtype([("fss_a", np.float32), ("fss_b", np.float32, (3,))], align=True), "_FuseSemR32" + ) + c128 = dtypes.register_anonymous( + np.dtype([("fss_c", np.complex128)], align=True), "_FuseSemC128" + ) + c64 = dtypes.register_anonymous(np.dtype([("fss_c", np.complex64)], align=True), "_FuseSemC64") + f, g = _udt_vectors(f64, [nan, 1.0, -7.0, inf, 0.0], [1.0, nan, 2.0, 2.0, 0.0]) + i, j = _udt_vectors(i8s, [-128, 7, -7, 100, 5], [-1, 0, 2, 3, -2]) + r = Vector(rec64, size=3) + s = Vector(rec64, size=3) + for k, (x, y) in enumerate([(nan, 1.0), (1.0, nan), (-7.5, 2.0)]): + r[k] = (x, [x, y, x * 2]) + s[k] = (y, [y, x, 0.0]) + c = Vector(c128, size=2) + d = Vector(c128, size=2) + c[0], d[0] = (3 + 4j,), (0j,) + c[1], d[1] = (1j,), (2 + 0j,) + cases = [(f, g, f32, op) for op in ("min", "max", "floordiv", "truediv")] + cases += [(i, j, i32s, "floordiv"), (r, s, rec32, "min"), (r, s, rec32, "max")] + cases += [(r, s, rec32, "floordiv"), (c, d, c64, "truediv")] + for x, y, out, name in cases: + gb_op = getattr(binary, name) + expected = gb_op(x & y).new().dup(dtype=out) + fused = Vector(out, size=x.size) + fused << gb_op(x & y) + assert fused.nvals == expected.nvals, (udt_op_path, name, out) + for k in range(x.size): + got, want = fused[k].new().value, expected[k].new().value + msg = f"{udt_op_path} {name} into {out} at {k}" + if out.np_type.names is None: + _assert_same_floats(got, want, msg) + else: + for field in out.np_type.names: + _assert_same_floats(got[field], want[field], msg) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_records_differing_only_in_layout_combine(): + """A packed and an aligned record with the same leaves combine, either way round. + + The result is the left operand's type, since it holds every leaf. Converting + the other operand to it for a Scalar's ``ewise_add`` subtracted a zero, + which came back as the other operand's own type, so the conversion recursed + forever; Vectors raised RecursionError the same way before ``ewise_add`` + refused operands of another type than the result's. + """ + fields = [("lay_f", np.float64), ("lay_i", np.int16), ("lay_n", [("lay_k", np.int32)])] + packed = dtypes.register_anonymous(np.dtype(fields), "_LayPacked") + aligned = dtypes.register_anonymous(np.dtype(fields, align=True), "_LayAligned") + x = Vector(packed, size=3) + x[0] = (1.5, 2, (3,)) + y = Vector(aligned, size=3) + y[0] = (-0.0, 1, (1,)) + y[2] = (10.0, 20, (30,)) + for a, b in [(x, y), (y, x)]: + result = a.ewise_union(b, binary.plus, 0, 0).new() + assert result.dtype == a.dtype + assert result.to_coo()[1].tolist() == [(1.5, 3, (4,)), (10.0, 20, (30,))] + assert (a * b).new().dtype == a.dtype + with pytest.raises(DomainMismatch, match="ewise_add cannot use"): + a.ewise_add(b) + assert (a + b).new().isequal(result) + s = gb.Scalar.from_value((1.5, 2, (3,)), dtype=packed) + t = gb.Scalar.from_value((1.0, 1, (1,)), dtype=aligned) + assert (s + t).new().value.tolist() == (2.5, 3, (4,)) + # Converting keeps a negative zero where there is no partner. + z = gb.Scalar.from_value((-0.0, 0, (0,)), dtype=aligned) + assert np.signbit((gb.Scalar(packed) + z).new().value["lay_f"]) + # Either layout stores into the other through the cast op. + x << y + assert x.to_coo()[1].tolist() == [(-0.0, 1, (1,)), (10.0, 20, (30,))] + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_array_field_truediv_is_float64(): + """``truediv`` on a small integer array field divides in float64, as a scalar field does. + + Numba runs an array field through a ufunc whose loop it picks itself, and + for int8 or int16 it picked float32, so ``1 / 46`` came back as + 0.021739130839705467 in an FP64 field. + """ + for np_type in [np.int8, np.uint8, np.int16, np.uint16]: + name = f"_TdArr{np.dtype(np_type).name}" + rec = dtypes.register_anonymous( + np.dtype([(f"td_a_{name}", np_type, (2,)), (f"td_s_{name}", np_type)]), name + ) + v = Vector(rec, size=1) + v[0] = ([1, 46], 46) + w = Vector(rec, size=1) + w[0] = ([46, 7], 3) + result = (v / w).new()[0].new().value + np.testing.assert_array_equal(result[0], np.array([1, 46]) / np.array([46, 7])) + assert result[1] == 46 / 3 + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_typed_scalar_operands_and_defaults(): + """Typed Scalars keep their values beside a UDT, as literals do. + + ``Scalar.ewise_union`` converted its other operand into the left UDT + (``s8 - sf`` was int8), and an ``ewise_union`` default given as a typed + Scalar was converted into the operand's UDT however much that changed it. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_TsI8") + f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_TsF8") + s8 = gb.Scalar.from_value([1, 2, 3], dtype=i8s) + sf = gb.Scalar.from_value([0.5, 0.5, 0.5], dtype=f8s) + for expr, expected in [(s8 - sf, [0.5, 1.5, 2.5]), (sf - s8, [-0.5, -1.5, -2.5])]: + result = expr.new() + assert result.dtype == f8s + np.testing.assert_array_equal(result.value, expected) + v = Vector(i8s, size=2) + v[0] = [1, 2, 3] + w = Vector(i8s, size=2) + w[1] = [4, 5, 6] + A = Matrix(i8s, nrows=1, ncols=2) + A[0, 0] = [1, 2, 3] + B = Matrix(i8s, nrows=1, ncols=2) + B[0, 1] = [4, 5, 6] + for a, b in [(v, w), (A, B)]: + # A default must have the operand's type, and FP64 does not cast safely to + # int8; 2.0 still converts, being exact, but that is deprecated. + with pytest.raises(ValueError, match="Scalar of type FP64 does not fit _TsI8"): + a.ewise_union(b, binary.plus, gb.Scalar.from_value(0.5), 0) + with pytest.warns(DeprecationWarning, match="Scalar of type FP64 does not fit _TsI8"): + a.ewise_union(b, binary.plus, gb.Scalar.from_value(2.0), 0) + with pytest.raises(ValueError, match="Scalar of type _TsF8 does not fit _TsI8"): + a.ewise_union(b, binary.plus, 0, sf) + assert a.ewise_union(b, binary.plus, gb.Scalar.from_value(2, "INT8"), 0).new().dtype == i8s + # Matrix.apply with a Monoid promotes as Vector.apply does. + result = A.apply(monoid.plus, right=0.5).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_errors_name_the_real_problem(): + """Errors that the UDT rules had made misleading. + + Infix on two UDTs of different sizes said ``ewise_add cannot use + binary.first``, the op it uses to provoke the size error, and an op that + does not take UDTs at all reported the literal check's ValueError. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_ErrI8") + f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_ErrF8") + v = Vector(i8s, size=2) + w = Vector(f8s, size=3) + for expr in [lambda: binary.plus(v | w), lambda: binary.first(v | w), lambda: v + w]: + with pytest.raises(DimensionMismatch): + expr().new() + with pytest.raises(KeyError, match="lt does not work with _ErrI8"): + v.apply(binary.lt, right=2.5) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_literal_rules_edge_cases(): + """Edges of the weak rule, each pinned on its own. + + numpy 2 raises OverflowError itself when ``300`` is converted to int8, which + hid that the weak rule's own range check matters on numpy 1, where the + conversion wraps. A ``Fraction`` or ``Decimal`` is typed as the int or + float it is worth; numpy would have truncated either into an int field. + """ + from decimal import Decimal + from fractions import Fraction + + from graphblas.core.literals import _weak_literal_udt + + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_EdgeI8") + f4s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_EdgeF4") + c8s = dtypes.register_anonymous(np.dtype((np.complex64, (3,))), "_EdgeC8") + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + _weak_literal_udt(i8s, 300) + assert _weak_literal_udt(i8s, 1) is i8s + assert _weak_literal_udt(f4s, 1j).np_type == np.dtype((np.complex64, (3,))) + v = Vector(i8s, size=1) + v[0] = [1, 2, 3] + for value in [Fraction(1, 2), Decimal("0.5")]: + result = (v + value).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + assert result[0].new().value.tolist() == [1.5, 2.5, 3.5] + for value in [Fraction(2), Decimal(2)]: + result = (v + value).new() + assert result.dtype == i8s + assert result[0].new().value.tolist() == [3, 4, 5] + # An infinite or NaN Decimal is the float it stands for. + for value in [Decimal("Infinity"), Decimal("NaN")]: + result = (v + value).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + np.testing.assert_array_equal(result[0].new().value, [float(value)] * 3) + udf = BinaryOp.register_anonymous(lambda x, y: x + y, "_edge_plus", is_udt=True) + with pytest.raises(ValueError, match=r"Fraction\(1, 2\) does not fit _EdgeI8"): + v.apply(udf, right=Fraction(1, 2)) + # A literal not laid out like an element (ragged, too long or too short for + # a field, not broadcastable), or not a number, has no type of its own, so + # it is converted into the UDT, and numpy says why it cannot be. A record + # says so itself for a part that does not broadcast to its array field. + # So too for a comparison, which first looks for an int beyond its leaf, a + # part holding numpy values, and text a bytes field cannot hold. + rec = dtypes.register_anonymous( + np.dtype([("edge_v", np.float64, (3,)), ("edge_n", np.int8)]), "_EdgeRec" + ) + r = Vector(rec, size=1) + r[0] = ([1, 2, 3], 4) + nested = dtypes.register_anonymous( + np.dtype([("edge_in", [("edge_p", np.int8), ("edge_q", np.int8)]), ("edge_m", np.int8)]), + "_EdgeNested", + ) + text = dtypes.register_anonymous(np.dtype([("edge_s", "S3"), ("edge_k", np.int8)]), "_EdgeText") + for expr, reason in [ + (lambda: binary.plus(v, [1, [2, 3]]), "with a sequence"), + (lambda: v + np.array([1, 2]), "could not broadcast"), + (lambda: v + "abc", "invalid literal"), + (lambda: v.apply(udf, right="abc"), "invalid literal"), + (lambda: binary.plus(r, (1, 2, 3)), "tuple of length 3"), + (lambda: binary.plus(r, ((1, 2), 3)), "does not fit field 'edge_v'"), + (lambda: r == (1, 2, 3), "tuple of length 3"), + (lambda: Vector(nested, size=1) == ((1, 2, 3), 300), "tuple of length 3"), + (lambda: binary.plus(r, ([np.float64(1), [2, 3]], 4)), "with a sequence"), + (lambda: binary.plus(r, ([np.zeros((3, 2)), np.zeros((3, 1))], 4)), "with a sequence"), + (lambda: binary.plus(r, (np.array(["a", "b", "c"]), 4)), "could not convert string"), + (lambda: Vector(text, size=1) == ("\xe9", 1), "can't encode"), + ]: + with pytest.raises(ValueError, match=reason): + expr() + c = Vector(c8s, size=1) + c[0] = [1, 2, 3] + assert c.apply(binary.plus, right=(1 + 2j, 0, 0)).new()[0].new().value[0] == 2 + 2j + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_store_casts_reach_every_store_path(): + """Each way of storing into a UDT object casts as ``<<`` does, or says why it cannot. + + ``agg.first`` into a Scalar of another UDT handed GraphBLAS the output as it + was, which read the int8 bytes as float64 (24 bytes from a 3-byte + element). Indexed assignment of an expression skipped the cast, and + ``dup`` failed for a dtype given as a numpy dtype or a string. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_PathI8") + f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_PathF8") + v = Vector(i8s, size=3) + v[0] = [1, 2, 3] + v[2] = [4, 5, 6] + aggregators = [(agg.sum, [5.0, 7.0, 9.0])] + if gb.backend == "suitesparse": + aggregators += [(agg.ss.first, [1.0, 2.0, 3.0]), (agg.ss.last, [4.0, 5.0, 6.0])] + for aggregator, expected in aggregators: + s = gb.Scalar(f8s) + s << v.reduce(aggregator) + np.testing.assert_array_equal(s.value, expected) + # A Scalar expression assigned to many elements is computed, then cast. + u = Vector(f8s, size=3) + u[[0, 2]] << unary.ainv(gb.Scalar.from_value([1, 2, 3], dtype=i8s)) + assert u.to_coo()[1].tolist() == [[-1, -2, -3]] * 2 + with pytest.raises(DomainMismatch, match="result of type .*casts only to another UDT. Use"): + gb.Scalar(FP64) << v.reduce(agg.sum) + # Assigned to part of an object, a computed result is still a result. + with pytest.raises(DomainMismatch, match="result of type FP64 .*another UDT. Use"): + v[:] << Vector.from_coo([0], [0.5], size=3).apply(unary.ainv) + w = Vector(f8s, size=3) + w[0] = [0.5, 0.5, 0.5] + w[1] = [1.5, 1.5, 1.5] + # Assigning part of an object from a whole object of another type would + # need a converted copy, an expression's (extract) as well; so would + # ewise_add, which copies an entry without a partner as it is. + x = v.dup() + for expr in [lambda: x[[0, 1]] << w[[0, 1]], lambda: x[[0, 2]] << (v + 0.5).new()[[0, 2]]]: + with pytest.raises(DomainMismatch, match="converted copy of the whole value"): + expr() + x[[0, 1]] << w[[0, 1]].new().dup(dtype=i8s) + assert x.to_coo()[1].tolist() == [[0, 0, 0], [1, 1, 1], [4, 5, 6]] + s = gb.Scalar.from_value([1, 2, 3], dtype=i8s) + s += 1 + np.testing.assert_array_equal(s.value, [2, 3, 4]) + result = s.dup(dtype=np.dtype((np.float64, (3,)))) + assert result.dtype == f8s + np.testing.assert_array_equal(result.value, [2, 3, 4]) + with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): + s.dup(dtype="FP64") + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_scalar_accumulates_its_value_converted_to_its_type(): + """A UDT Scalar's accumulator adds a value converted to the Scalar's type, as a built-in's does. + + The accumulator is typed on the Scalar's dtype, and GraphBLAS converts the + value to that type first: an INT8 Scalar of -1 ``+= 0.5`` stays -1, and so + does a Vector's ``v(accum=binary.plus) << 0.5``. A UDT Scalar used the + untyped op instead, so it computed -1 + 0.5 and truncated the sum to 0. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_AccumI8") + f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_AccumF8") + halves = Vector(f8s, size=2) + halves[1] = [0.5, 1.5, 2.5] + + def added(value): + s = gb.Scalar.from_value([-1, -1, -1], dtype=i8s) + s(accum=binary.plus) << value + return s.value.tolist() + + def added_to_int8(value): + s = gb.Scalar.from_value(-1, dtype=INT8) + s(accum=binary.plus) << value + return s.value + + assert added(0.5) == [added_to_int8(0.5)] * 3 == [-1, -1, -1] + assert added(2.5) == [added_to_int8(2.5)] * 3 == [1, 1, 1] + assert added((0.5, 1.5, 2.5)) == [-1, 0, 1] + assert added(halves[1]) == [-1, 0, 1] + assert added(gb.Scalar.from_value([0.5, 1.5, 2.5], dtype=f8s)) == [-1, 0, 1] + v = Vector(i8s, size=1) + v[0] = [-1, -1, -1] + v(accum=binary.plus) << (0.5, 1.5, 2.5) + assert v[0].new().value.tolist() == [-1, 0, 1] + s = gb.Scalar.from_value([-1, -1, -1], dtype=i8s) + s += 0.5 + assert s.value.tolist() == [-1, -1, -1] + s.clear() + s(accum=binary.plus) << 0.5 + assert s.value.tolist() == [0, 0, 0] + # A built-in Scalar casts only into a built-in type, as for a store. The error + # calls it a value, not a result (it said "a result of type FP64"). + with pytest.raises(DomainMismatch, match="value of type FP64 .*casts only to another UDT.$"): + added(gb.Scalar.from_value(0.5)) + # A Vector is not a value to accumulate, into a UDT Scalar or a built-in one. + for accumulate in [added, added_to_int8]: + with pytest.raises(TypeError, match="Bad type for keyword argument `other=`"): + accumulate(halves) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_literal_fits_in_select_bool_arrays_and_overflow(): + """Three edges of the literal rules found by review. + + A user IndexUnaryOp's thunk in ``select`` skipped the must-fit check (1.5 + became 1), a bool numpy array beside an array UDT failed to look up a + ``BOOL[3]`` type (main converted it), and a float literal too large for a + float32 field became infinity where it must become an element. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_EdgeSelI8") + f4s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_EdgeSelF4") + + def ge_first(x, i, j, t): # pragma: no cover (numba) + return x[0] >= t[0] + + ge = IndexUnaryOp.register_anonymous(ge_first, "_edge_sel_ge", is_udt=True) + v = Vector(i8s, size=3) + v[0] = [1, 2, 3] + v[1] = [2, 3, 4] + assert v.select(ge, 2).new().to_coo()[0].tolist() == [1] + for thunk in [1.5, (1.5, 0, 0), np.float64(2.5)]: + with pytest.raises(ValueError, match="does not fit _EdgeSelI8"): + v.select(ge, thunk) + A = Matrix(i8s, nrows=1, ncols=2) + A[0, 0] = [1, 2, 3] + with pytest.raises(ValueError, match="does not fit _EdgeSelI8"): + A.select(ge, 1.5) + # A bool array converts into any numeric element, as it did, and bools in a + # sequence compare as 0 and 1 (a BOOL[3] UDT cannot be registered). + result = v.apply(binary.plus, right=np.array([True, False, True])).new() + assert result.dtype == i8s + assert result[0].new().value.tolist() == [2, 2, 4] + ones = Vector(i8s, size=2) + ones[0] = [1, 0, 1] + assert (ones == (True, False, True)).new()[0].new().value + assert not (ones != [True, False, True]).new()[0].new().value + # A float too large for a float32 field raises where it must become an element. + f = Vector(f4s, size=2) + f[0] = [1, 2, 3] + with pytest.raises(OverflowError, match="overflows _EdgeSelF4"): + f.ewise_union(f, binary.plus, 1e300, 0) + with pytest.raises(OverflowError, match="overflows _EdgeSelF4"): + f.apply(binary.register_anonymous(lambda x, y: x, "_edge_sel_first", is_udt=True), 1e300) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_exact_literal_of_another_type_is_deprecated(): + """A must-fit value that fits only by value converts, with a DeprecationWarning. + + On main any literal, thunk or default was converted into the UDT, and what + did not fit was lost. Now it must fit by type; one whose values convert + exactly, such as 0 for a bool field, still converts for the deprecation + period, and the warning points at the caller's line. + """ + rec = dtypes.register_anonymous( + np.dtype([("dp_b", np.bool_), ("dp_f", np.float64), ("dp_i", np.int8)]), "_DeprecRec" + ) + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_DeprecI8") + r = Vector(rec, size=2) + r[0] = (True, 1.5, 2) + s = Vector(rec, size=2) + s[1] = (False, 2.5, 3) + expected = r.ewise_union(s, binary.plus, False, False).new() + for default in [0, (0, 0.0, 0), {"dp_b": 0, "dp_f": 0.0, "dp_i": 0}]: + with pytest.warns(DeprecationWarning, match="does not fit _DeprecRec") as record: + result = r.ewise_union(s, binary.plus, default, False).new() + assert result.isequal(expected) + assert {w.filename for w in record} == {__file__} + with pytest.raises(ValueError, match="3 does not fit _DeprecRec"): + r.ewise_union(s, binary.plus, 3, False) + + def ge_first(x, i, j, t): # pragma: no cover (numba) + return x[0] >= t[0] + + ge = IndexUnaryOp.register_anonymous(ge_first, "_deprec_ge", is_udt=True) + v = Vector(i8s, size=3) + v[0] = [1, 2, 3] + v[1] = [2, 3, 4] + A = Matrix(i8s, nrows=1, ncols=2) + A[0, 1] = [2, 3, 4] + for x in [v, A]: + with pytest.warns(DeprecationWarning, match="2.0 does not fit _DeprecI8"): + assert x.select(ge, 2.0).new().isequal(x.select(ge, 2).new()) + for thunk in [1.5, float("nan"), np.int64(300)]: + with pytest.raises(ValueError, match="does not fit _DeprecI8"): + x.select(ge, thunk) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_isequal_is_array_equal_and_eq_broadcasts(): + """Isequal compares array elements as np.array_equal; == broadcasts, on Scalars too. + + eq on array UDTs broadcasts, as numpy's == does, and isequal used it, so an + FP64[3] of [1, 1, 1] was isequal to an FP64[1] of [1]. A Scalar's == was + isequal itself, and a literal was converted into the UDT first, so an + INT8[3] of [0, 0, 0] equaled 0.5 (truncated to 0) by either spelling. + """ + f3 = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_IeqF3") + f1 = dtypes.register_anonymous(np.dtype((np.float64, (1,))), "_IeqF1") + f13 = dtypes.register_anonymous(np.dtype((np.float64, (1, 3))), "_IeqF13") + i3 = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_IeqI3") + rec = dtypes.register_anonymous(np.dtype([("ieq_a", np.int8), ("ieq_b", np.float32)]), "_IeqR") + + def vec(dtype, value): + v = Vector(dtype, size=2) + v[1] = value + return v + + ones = vec(f3, [1, 1, 1]) + assert not ones.isequal(vec(f1, [1])) + assert (ones == vec(f1, [1])).new().reduce(monoid.land).new().value # == broadcasts + assert ones.isequal(vec(f13, [[1, 1, 1]])) # leading axes of length 1 do not count + assert ones.isequal(vec(i3, [1, 1, 1])) + assert not ones.isequal(Vector.from_coo([1], [1.0], size=2)) + assert vec(f1, [5]).isequal(Vector.from_coo([1], [5.0], size=2)) + A = Matrix.from_coo([0], [1], [[1, 1, 1]], dtype=f3, nrows=1, ncols=2) + assert not A.isequal(Matrix.from_coo([0], [1], [[1]], dtype=f1, nrows=1, ncols=2)) + assert A.isequal(Matrix.from_coo([0], [1], [[1, 1, 1]], dtype=i3, nrows=1, ncols=2)) + + u = gb.Scalar.from_value([1, 1, 1], i3) + zeros = gb.Scalar.from_value([0, 0, 0], i3) + assert u == 1 + assert u != 0 + assert u == gb.Scalar.from_value(1) + assert gb.Scalar.from_value(1) == u + assert u == gb.Scalar.from_value([1.0], f1) + assert not u.isequal(1) + assert not u.isequal([1]) + assert not u.isequal(gb.Scalar.from_value([1.0], f1)) + assert u.isequal((1, 1, 1)) + assert u.isequal(np.ones((1, 3), dtype=np.int8)) + with pytest.raises(ValueError, match="with a sequence"): # ragged: it has no shape + u.isequal([1, [1, 1]]) + assert gb.Scalar.from_value([5], f1).isequal(5) + # A literal is compared as given, not converted into the UDT first. + assert zeros != 0.5 + assert not zeros.isequal(0.5) + assert not zeros.isequal((0, 0, 0.5)) + r = gb.Scalar.from_value((0, 0.1), rec) + assert r == (0, 0.1) # 0.1 beside a float32 field is float32, as == types it + assert r.isequal((0, 0.1)) + assert r != (0.5, 0.1) + assert not r.isequal((0.5, 0.1)) + # Empty Scalars are equal to each other, and to None. + assert gb.Scalar(i3) == gb.Scalar(f3) + assert gb.Scalar(i3) == None # noqa: E711 + assert gb.Scalar(i3) != u + assert u != gb.Scalar(i3) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_monoid_literal_is_typed_as_its_binaryop(): + """A literal beside a Monoid is typed as beside its BinaryOp, everywhere. + + ``apply`` already used the BinaryOp, but a Scalar's ``ewise_add`` and + ``ewise_mult`` with a Monoid converted the literal into the UDT (so ``0.5`` + had to fit), where the same call with ``binary.plus`` promoted. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_MonoLitI8") + s = gb.Scalar.from_value([1, 2, 3], dtype=i8s) + for expr in [s.ewise_add(0.5, monoid.plus), s.ewise_add(0.5, binary.plus)]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + np.testing.assert_array_equal(result.value, [1.5, 2.5, 3.5]) + assert s.ewise_mult(2, monoid.times).new().dtype == i8s + # ewise_union defaults still have to fit the op's input types. + v = Vector(i8s, size=2) + v[0] = [1, 2, 3] + with pytest.raises(ValueError, match="does not fit _MonoLitI8"): + v.ewise_union(v, monoid.plus, 0.5, 0) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_outer_on_mixed_types(): + """``outer`` builds its semiring from the op, typed on both input types. + + ``get_semiring`` typed the semiring by the op's first input type only, so + ``outer`` on two different UDTs, or a UDT and a built-in dtype, was a bare + GrB_DOMAIN_MISMATCH. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_OuterI8") + f4s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_OuterF4") + v = Vector(i8s, size=2) + v[0] = [1, 2, 3] + u = Vector(f4s, size=2) + u[1] = [0.5, 0.5, 0.5] + for a, b in [(v, u), (u, v)]: + result = a.outer(b).new() + assert result.dtype == f4s + assert result.to_coo()[2].tolist() == [[0.5, 1.0, 1.5]] + result = v.outer(Vector.from_coo([0], [2.0], size=2)).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + assert result.to_coo()[2].tolist() == [[2.0, 4.0, 6.0]] + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_jit_kernel_only_for_a_same_type_result(): + """A same-type pair gets the arithmetic JIT kernel only when its result is that type too. + + The kernel declares its result with the operands' type, so for ``truediv`` + on an integer UDT, whose result is float64, it would write integer + quotients into a float64 buffer. That op runs through the cfunc instead. + 64 entries keep SuiteSparse off its iso and short-vector shortcuts, so a + JIT kernel would run. + """ + from graphblas.core.operator.udt_codegen import _has_jit_set + + ints = dtypes.register_anonymous(np.dtype((np.int32, (11,))), "_JitDivI11") + N = 64 + v = Vector(ints, size=N) + w = Vector(ints, size=N) + for i in range(N): + v[i] = np.arange(11) + i + w[i] = np.full(11, 4) + result = v.ewise_mult(w, binary.truediv).new() + assert result.dtype.np_type == np.dtype((np.float64, (11,))) + for i in [0, 1, N - 1]: + np.testing.assert_array_equal(result[i].new().value, (np.arange(11) + i) / 4) + if suitesparse and _has_jit_set: + assert binary.truediv[ints].jit_c_source is None + assert binary.plus[ints].jit_c_source is not None + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_on_leaves_of_other_kinds_is_a_keyerror(): + """``eq`` refuses leaves that hold different kinds of value with a KeyError, as ``plus`` does. + + Text, a date, a duration and a number compare only with their own kind, so a + bytes field against a float field, or a number given for a bytes field, is + binary.eq's KeyError, as it is for ``plus``: bytes and float64 have no + common dtype for the result field. + """ + text = dtypes.register_anonymous(np.dtype([("ueq_a", "S4"), ("ueq_b", np.float64)]), "_UeqS") + num = dtypes.register_anonymous( + np.dtype([("ueq_a", np.float64), ("ueq_b", np.float64)]), "_UeqF" + ) + v = Vector(text, size=1) + v[0] = (b"ab", 1.0) + w = Vector(num, size=1) + w[0] = (1.0, 1.0) + with pytest.raises(KeyError, match=r"elements of \|S4 and float64 have no common type"): + binary.plus(v & w).new() + # A number beside the bytes field is an int64 there, which pairs no better. + with pytest.raises(KeyError, match=r"elements of \|S4 and int64 have no common type"): + v + 1 + nested = dtypes.register_anonymous( + np.dtype([("ueq_x", "S2"), ("ueq_in", [("ueq_y", "S2"), ("ueq_n", np.int8)])]), "_UeqN" + ) + n = Vector(nested, size=1) + n[0] = (b"x", (b"y", 1)) + for expr in [ + lambda: binary.eq(v & w), + lambda: v == 1, + lambda: v != (1.0, 1.0), + lambda: gb.Scalar.from_value((b"ab", 1.0), text) == (1.0, 1.0), + lambda: n == (b"x", np.float64(1)), + ]: + with pytest.raises(KeyError, match=r"elements of \|S\d and \w+ do not compare"): + expr() + assert (n == (b"x", (b"y", 1))).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_compares_a_udt_that_is_neither_record_nor_array(): + """``==``, ``!=`` and ``isequal`` compare UDT elements of bytes or dates as numpy does. + + Such a UDT holds one value per element, which ``binary.eq`` and + ``binary.ne`` compare as a record field of that type compares: bytes by + value, and NaT equal to nothing. Only two elements of the one type + compare, or text of one kind at two widths, or dates or durations in two + units; any other pair raises binary.eq's KeyError. + """ + text = dtypes.register_anonymous(np.dtype("S4"), "_EqOpaqueBytes") + v = Vector(text, size=3) + v[0] = b"ab" + v[1] = b"cd" + w = Vector(text, size=3) + w[0] = b"ab" + w[1] = b"ce" + w[2] = b"x" + assert binary.eq(v & w).new().isequal(Vector.from_coo([0, 1], [True, False], size=3)) + assert (v != w).new().isequal(Vector.from_coo([0, 1], [False, True], size=3)) + assert (v == b"ab").new().isequal(Vector.from_coo([0, 1], [True, False], size=3)) + assert v.isequal(v.dup()) + assert not v.isequal(w) + A = Matrix(text, nrows=2, ncols=2) + A[0, 1] = b"ab" + A[1, 0] = b"cd" + B = A.dup() + assert A.isequal(B) + B[1, 0] = b"cc" + assert not A.isequal(B) + s = gb.Scalar.from_value(b"ab", text) + assert s == gb.Scalar.from_value(b"ab", text) + assert s != gb.Scalar.from_value(b"cd", text) + assert s == b"ab" + assert s != b"cd" + # Under eq and ne, text pairs with text of its kind at another width, as numpy + # compares it, so a literal longer than the element is unequal, not cut short. + wide = dtypes.register_anonymous(np.dtype("S6"), "_EqOpaqueWide") + u = Vector(wide, size=3) + u[0] = b"ab" + u[1] = b"cdefgh" + for left, right in [(v, u), (u, v)]: + assert binary.eq(left & right).new().to_coo()[1].tolist() == [True, False] + assert binary.ne(left & right).new().to_coo()[1].tolist() == [False, True] + four = Vector(text, size=1) + four[0] = b"abcd" + assert not (four == b"abcdef").new()[0].new().value + assert (four != b"abcdef").new()[0].new().value + assert gb.Scalar.from_value(b"abcd", text) != b"abcdef" + assert gb.Scalar.from_value(b"ab", text) == gb.Scalar.from_value(b"ab", wide) + # A Scalar's == compares in numpy, so an element Numba has no type for compares too. + half = dtypes.register_anonymous(np.dtype(np.float16), "_EqOpaqueHalf") + h = gb.Scalar.from_value(1.5, half) + assert h == gb.Scalar.from_value(1.5, half) + assert h == 1.5 + assert gb.Scalar.from_value(np.nan, half) != gb.Scalar.from_value(np.nan, half) + # Numba cannot compute on float16, so an op on it is refused by name, for a + # record with a float16 field too; a Scalar's == compares in numpy. + halves = Vector(half, size=1) + halves[0] = 1.5 + half_rec = dtypes.register_anonymous(np.dtype([("eqo_h", np.float16)]), "_EqOpaqueHalfRec") + assert gb.Scalar.from_value((1.5,), half_rec) == (1.5,) + for expr in [ + lambda: halves == halves, + lambda: halves.isequal(halves), + lambda: binary.plus(Vector(half_rec, size=1), (1.0,)), + lambda: -Vector(half_rec, size=1), + ]: + with pytest.raises(KeyError, match="Numba cannot compute on float16"): + expr() + # A store cast is Numba code too, so it raises the DomainMismatch of a store that cannot cast. + single_rec = dtypes.register_anonymous(np.dtype([("eqo_h", np.float32)]), "_EqOpaqueSingleRec") + for expr in [ + lambda: Vector(single_rec, size=1) << Vector(half_rec, size=1), + lambda: Vector(half_rec, size=1).dup(dtype=single_rec), + ]: + with pytest.raises(DomainMismatch, match="Numba cannot compute on float16"): + expr() + dates = dtypes.register_anonymous(np.dtype("M8[s]"), "_EqOpaqueDates") + d = Vector(dates, size=2) + d[0] = np.datetime64("NaT", "s") + d[1] = np.datetime64(0, "s") + assert (d == d).new().to_coo()[1].tolist() == [False, True] + # Under eq and ne, a date or duration pairs with one in another unit too, as + # numpy compares them, so a literal in a finer unit is not cut to the element's. + import datetime + + days = dtypes.register_anonymous(np.dtype("M8[D]"), "_EqOpaqueDays") + day = Vector(days, size=1) + day[0] = np.datetime64("2020-01-01") + seconds = dtypes.register_anonymous(np.dtype("m8[s]"), "_EqOpaqueSeconds") + five = Vector(seconds, size=1) + five[0] = np.timedelta64(5, "s") + for x, literal, equal in [ + (day, np.datetime64("2020-01-01T12:00"), False), + (day, datetime.datetime(2020, 1, 1, 12), False), # noqa: DTZ001 (numpy's are naive) + (day, "2020-01-01T12:00", False), + (day, np.datetime64("2020-01-01T00"), True), + (day, datetime.date(2020, 1, 1), True), + (day, "2020-01-01", True), + (five, np.timedelta64(5001, "ms"), False), + (five, datetime.timedelta(seconds=5, microseconds=1), False), + (five, np.timedelta64(5000, "ms"), True), + (five, datetime.timedelta(seconds=5), True), + ]: + assert (x == literal).new()[0].new().value == equal + assert (x != literal).new()[0].new().value != equal + assert (x[0].new() == literal) is equal + in_ms = Vector(dtypes.register_anonymous(np.dtype("m8[ms]"), "_EqOpaqueMillis"), size=1) + in_ms[0] = np.timedelta64(5000, "ms") + assert binary.eq(five & in_ms).new()[0].new().value + assert five.isequal(in_ms) + with pytest.raises(KeyError, match="pairs only with its own type"): + day == np.timedelta64(5, "s") # noqa: B015 + # Nor with a record or an array UDT, for a Scalar as for a Vector. + pair = dtypes.register_anonymous(np.dtype((np.int8, (5,))), "_EqOpaqueArr5") + rec = dtypes.register_anonymous(np.dtype([("eqo_a", np.int8)]), "_EqOpaqueRec") + for expr in [ + lambda: v == 1, + lambda: binary.eq(v & Vector.from_coo([0], [1], size=3)), + lambda: binary.eq(v & Vector(rec, size=3)), + lambda: binary.eq(v & Vector(dtypes.register_anonymous(np.dtype("U4")), size=3)), + lambda: gb.Scalar.from_value([1, 2, 3, 4, 5], pair) == h, + lambda: gb.Scalar.from_value((1,), rec) == s, + ]: + with pytest.raises(KeyError, match="pairs only with its own type"): + expr() + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_numba_error_summary_is_never_empty(): + """A Numba error's one-line summary falls back to its first line, then to a sentence. + + The summary is Numba's diagnostic line when one is there, with the + signature from the line below when the diagnostic ends in a colon. A + message without one gives its first non-empty line, and an empty message + a fixed sentence, so a UdfParseError always says something. + """ + from graphblas.core.operator.udf import _summarize_numba_typing_error as summarize + + line = "No implementation of function f found for signature:" + colored = f"\x1b[1m{line}\x1b[0m" + assert summarize(ValueError(f"{colored}\n\n >>> f(int64)")) == f"{line} >>> f(int64)" + assert summarize(ValueError(f"Failed in nopython mode\n{colored}\n\n")) == line + assert summarize(ValueError("\n \nSomething else went wrong\nat line 2")) == ( + "Something else went wrong" + ) + assert summarize(ValueError("")) == "Numba could not compile the function for these input types" + # From a lifted op: Numba cannot add two records, so not an array of records either + arr_rec = dtypes.register_anonymous( + np.dtype(([("sum_p", np.int32), ("sum_q", np.float64)], (2,))), "_SumArrRec" + ) + with pytest.raises(UdfParseError, match=r"plus does not work with \(_SumArrRec, _SumArrRec\)"): + binary.plus[arr_rec] + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_pair_is_never_unified(): + """Two different UDTs reach the op as they are, never unified to one type. + + numpy promotes two records of the same layout to one of them or to a third, + which ``unify`` registered as a new type. An op that does not lift to UDTs + then named that third type in its KeyError, and ``ewise_add`` with a monoid + passed GraphBLAS two operands it rejected with GrB_DOMAIN_MISMATCH. A + Monoid takes one type, so on a mixed pair it runs as its BinaryOp, and + indexed with the pair it takes the type the pair promotes to, as + ``monoid.plus["INT8", "INT16"]`` is ``monoid.plus["INT16"]``. + """ + X = dtypes.register_anonymous(np.dtype([("uni_a", np.float32), ("uni_b", np.int64)]), "_UniX") + Y = dtypes.register_anonymous(np.dtype([("uni_a", np.int64), ("uni_b", np.float32)]), "_UniY") + with pytest.raises(KeyError, match=r"lt does not work with \(_UniX, _UniY\)"): + binary.lt[X, Y] + v = Vector(X, size=1) + v[0] = (1.0, 2) + w = Vector(Y, size=1) + w[0] = (1, 2.0) + for expr in [v.ewise_mult(w, monoid.plus), v.ewise_union(w, monoid.plus, 0, 0)]: + result = expr.new() + assert result.dtype.np_type == np.dtype([("uni_a", np.float64), ("uni_b", np.float64)]) + assert result[0].new().value.tolist() == (2.0, 4.0) + assert v.ewise_mult(w, monoid.max).new()[0].new().value.tolist() == (1.0, 2.0) + # ewise_add, the Monoid's default, would need converted copies of the operands. + for expr in [lambda: v.ewise_add(w), lambda: monoid.plus(v | w)]: + with pytest.raises(DomainMismatch, match="ewise_add cannot use binary.plus"): + expr() + # Indexed with the pair: TypeError (a Monoid takes one type) before; main and + # built-in types promote. + XY = binary.plus[X, Y].return_type + assert monoid.plus[X, Y] is monoid.plus[Y, X] is monoid.plus[XY] + assert monoid.max[X, Y] is monoid.max[XY] + R8 = dtypes.register_anonymous(np.dtype([("uni_c", np.int8)]), "_UniR8") + R16 = dtypes.register_anonymous(np.dtype([("uni_c", np.int16)]), "_UniR16") + assert monoid.plus[R8, R16] is monoid.plus[R16] + assert monoid.plus[R8, "INT16"] is monoid.plus[R16] + # binary.any gives the second operand's type, but monoid.any promotes the pair in + # either order, as main did + for pair in [(R8, R16), (R16, R8), (R8, "INT16"), ("INT16", R8)]: + assert monoid.any[pair] is monoid.any[R16], pair + assert monoid.any[X, Y] is monoid.any[Y, X] is monoid.any[XY] + with pytest.raises(KeyError, match="plus does not work with"): + monoid.plus[X, R8] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_aggregators(): + """Monoid-based aggregators must auto-extend to UDTs the underlying monoid supports. + + Before: ``v.reduce(agg.sum)`` on a UDT vector failed with + ``KeyError: 'sum does not work with '``. Now ``Aggregator.__getitem__`` + triggers UDT compilation of the underlying monoid for monoid-based aggs. + """ + record = np.dtype([("a", np.int64), ("b", np.float64)], align=True) + udt = dtypes.register_anonymous(record, "_AggUdt") + v = Vector(udt, size=3) + v[0] = (1, 2.0) + v[1] = (3, 4.0) + v[2] = (5, 6.0) + + assert v.reduce(agg.sum).new() == (9, 12.0) + assert v.reduce(agg.prod).new() == (15, 48.0) + assert v.reduce(agg.min).new() == (1, 2.0) + assert v.reduce(agg.max).new() == (5, 6.0) + # any_value uses any_dtype=True and works on any input + assert v.reduce(agg.any_value).new() in [(1, 2.0), (3, 4.0), (5, 6.0)] + # count is dtype-agnostic + assert v.reduce(agg.count).new() == 3 + + # __contains__ should agree + assert udt in agg.sum + assert udt in agg.prod + assert udt in agg.min + assert udt in agg.max + # Composite/semiring-based aggregators aren't auto-lifted + assert udt not in agg.hypot + + if suitesparse: + # agg.ss.first / agg.ss.last are positional aggregators (any_dtype=True): + # they pick an existing entry rather than combining values, so they + # work on any dtype, UDT included, without per-UDT compilation. + assert tuple(v.reduce(agg.ss.first).new().value) == (1, 2.0) + assert tuple(v.reduce(agg.ss.last).new().value) == (5, 6.0) + + # Array UDT path + adt = np.dtype((np.float64, (3,))) + audt = dtypes.register_anonymous(adt, "_AggArrUdt") + a = Vector(audt, size=2) + a[0] = [1.0, 2.0, 3.0] + a[1] = [4.0, 5.0, 6.0] + np.testing.assert_array_equal(a.reduce(agg.sum).new().value, [5.0, 7.0, 9.0]) + np.testing.assert_array_equal(a.reduce(agg.min).new().value, [1.0, 2.0, 3.0]) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_lazy_registration(): + """``lazy=True`` must preserve ``is_udt`` so registration succeeds when fired. + + Regression: the ``module._delayed[funcname]`` kwargs dict didn't include + ``is_udt``, so the delayed callback compiled the function for standard + types only and failed with ``UdfParseError``. + """ + from graphblas.core.operator import BinaryOp, IndexUnaryOp, UnaryOp + + record = np.dtype([("a", np.int64), ("b", np.float64)], align=True) + udt = dtypes.register_anonymous(record, "_LazyUdt") + + BinaryOp.register_new("_lazy_udt_add", _pkl_udt_add, is_udt=True, lazy=True) + UnaryOp.register_new("_lazy_udt_neg", _pkl_udt_neg, is_udt=True, lazy=True) + IndexUnaryOp.register_new("_lazy_udt_iu", _pkl_udt_get_a, is_udt=True, lazy=True) + try: + # Trigger the delayed registration by attribute access; the failure + # mode was UdfParseError raised inside the delayed callback. + add_op = binary._lazy_udt_add + neg_op = unary._lazy_udt_neg + iu_op = indexunary._lazy_udt_iu + assert add_op._is_udt + assert neg_op._is_udt + assert iu_op._is_udt + + v = Vector(udt, 2) + v[0] = (1, 2.0) + v[1] = (3, 4.0) + assert add_op(v & v).new()[0].new() == (2, 4.0) + assert v.apply(neg_op).new()[0].new() == (-1, -2.0) + finally: + # Remove these UDT-only ops from every namespace they leaked into, + # including the combined ``op`` (binary and unary register there too). + # ``test_dir`` enumerates module names, and test_op_namespace iterates + # ``op._delayed`` and would fail to resolve an op whose per-type entry + # is gone. Pop ``__dict__`` and ``_delayed`` directly to avoid + # re-triggering ``__getattr__``. + for module, name in [ + (binary, "_lazy_udt_add"), + (unary, "_lazy_udt_neg"), + (indexunary, "_lazy_udt_iu"), + (op, "_lazy_udt_add"), + (op, "_lazy_udt_neg"), + ]: + vars(module).pop(name, None) + module._delayed.pop(name, None) + + +def _pkl_udt_add(x, y): # pragma: no cover (numba) + return (x["a"] + y["a"], x["b"] + y["b"]) + + +def _pkl_udt_neg(x): # pragma: no cover (numba) + return (-x["a"], -x["b"]) + + +def _pkl_udt_get_a(x, ix, jx, t): # pragma: no cover (numba) + return x["a"] + + +def _pkl_udt_big_a(x, ix, jx, t): # pragma: no cover (numba) + return x["a"] > t["a"] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_op_pickle(): + """Pickle round-trip for typed and untyped UDT operators. + + Catches regressions like: + - TypedUserMonoid / TypedUserSemiring being built with the cffi ``GrB_*`` + pointer as ``parent`` instead of the Monoid / Semiring Python object. + - Anonymous reduce paths losing ``is_udt`` when re-registering. + """ + import pickle + + from graphblas.core.operator import ( + BinaryOp, + IndexUnaryOp, + Monoid, + SelectOp, + Semiring, + UnaryOp, + ) + + record = np.dtype([("a", np.int64), ("b", np.float64)], align=True) + udt = dtypes.register_anonymous(record, "_PickleUdt") + + bin_op = BinaryOp.register_anonymous(_pkl_udt_add, "_pkl_b", is_udt=True) + un_op = UnaryOp.register_anonymous(_pkl_udt_neg, "_pkl_u", is_udt=True) + iu_op = IndexUnaryOp.register_anonymous(_pkl_udt_get_a, "_pkl_iu", is_udt=True) + sel_op = SelectOp.register_anonymous(_pkl_udt_big_a, "_pkl_s", is_udt=True) + mon_op = Monoid.register_anonymous(bin_op, (0, 0.0), "_pkl_m") + sr_op = Semiring.register_anonymous(mon_op, bin_op, "_pkl_sr") + + # Anonymous-op round-trip; verifies `is_udt` flows through `__reduce__`. + for anon_op in [bin_op, un_op, iu_op, sel_op, mon_op, sr_op]: + op2 = pickle.loads(pickle.dumps(anon_op)) + assert op2._is_udt is True, f"is_udt lost on {anon_op.name}" + + # Typed UDT instances on user-defined parents. + for anon_op in [bin_op, un_op, iu_op, sel_op, mon_op, sr_op]: + typed = anon_op[udt] + typed2 = pickle.loads(pickle.dumps(typed)) + assert typed2.name == typed.name + # parent must be the Python op object, not a cffi pointer + assert isinstance( + typed2.parent, type(anon_op) + ), f"{anon_op.name} typed parent had wrong type: {type(typed2.parent).__name__}" + + # Typed UDT instances on built-in monoid or semiring used to fail with + # ``cannot pickle '_cffi_backend.__CDataOwn'`` because the parent slot + # held the raw GrB pointer instead of the Python Monoid object. + pickle.loads(pickle.dumps(monoid.plus[udt])) + pickle.loads(pickle.dumps(monoid.times[udt])) + pickle.loads(pickle.dumps(semiring.plus_times[udt])) + pickle.loads(pickle.dumps(semiring.min_plus[udt])) + + # Vector round-trip with UDT (was already working; this is a sanity check). + v = Vector(udt, 2) + v[0] = (1, 2.0) + v[1] = (3, 4.0) + v2 = pickle.loads(pickle.dumps(v)) + assert v.isequal(v2) + + +def test_dir(): + for mod in [unary, binary, monoid, semiring, op]: + assert not set(mod._delayed) - set(dir(mod)) + + +def test_semiring_commute_exists(): + from .conftest import orig_semirings + + vals = { + semiring._deprecated[key] if key in semiring._deprecated else getattr(semiring, key) + for key in orig_semirings + } + missing = set() + for key in orig_semirings: + val = semiring._deprecated[key] if key in semiring._deprecated else getattr(semiring, key) + commutes_to = val.commutes_to + if commutes_to is not None and commutes_to not in vals: # pragma: no cover (debug) + missing.add(commutes_to.name) + if missing: + raise AssertionError("Missing semirings: " + ", ".join(sorted(missing))) + + +def test_binaryop_commute_exists(): + from .conftest import orig_binaryops + + vals = { + binary._deprecated[key] if key in binary._deprecated else getattr(binary, key) + for key in orig_binaryops + } + missing = set() + for key in orig_binaryops: + val = binary._deprecated[key] if key in binary._deprecated else getattr(binary, key) + commutes_to = val.commutes_to + if commutes_to is not None and commutes_to not in vals: # pragma: no cover (debug) + missing.add(commutes_to.name) + if missing: + raise AssertionError("Missing binaryops: " + ", ".join(sorted(missing))) + + +@pytest.mark.skipif("not supports_udfs") +def test_binom(): + v = Vector.from_coo([0, 1, 2], [3, 4, 5]) + result = v.apply(binary.binom, 2).new() + expected = Vector.from_coo([0, 1, 2], [3, 6, 10]) + assert result.isequal(expected) + assert op.binom is binary.binom + + +def test_builtins(): + v1 = Vector.from_coo([0, 1, 2], [1, 2, 3]) + v2 = Vector.from_coo([0, 1, 2], [3, 2, 1]) + result = v1.ewise_mult(v2, min).new() + expected = Vector.from_coo([0, 1, 2], [1, 2, 1]) + assert result.isequal(expected) + v1(max) << v2 + expected = Vector.from_coo([0, 1, 2], [3, 2, 3]) + assert v1.isequal(expected) + + +def test_op_ss(): + if suitesparse: + gb.unary.ss.positioni + gb.binary.ss.firsti + gb.semiring.ss.max_secondj + gb.op.ss.positionj + gb.agg.ss.argmin + else: + with pytest.raises(AttributeError, match="suitesparse"): + gb.unary.ss + with pytest.raises(AttributeError, match="suitesparse"): + gb.binary.ss + with pytest.raises(AttributeError, match="suitesparse"): + gb.semiring.ss + with pytest.raises(AttributeError, match="suitesparse"): + gb.op.ss + with pytest.raises(AttributeError, match="suitesparse"): + gb.agg.ss + + +def test_deprecated(): + with pytest.warns(DeprecationWarning, match="please use"): + gb.unary.erf + with pytest.warns(DeprecationWarning, match="please use `gb.indexunary.rowindex`"): + gb.unary.positioni + with pytest.warns(DeprecationWarning, match="please use"): + gb.binary.firsti + with pytest.warns(DeprecationWarning, match="please use"): + gb.semiring.min_firsti + with pytest.warns(DeprecationWarning, match="please use"): + gb.op.secondj + with pytest.warns(DeprecationWarning, match="please use"): + gb.agg.argmin + + +@pytest.mark.slow +def test_is_idempotent(): + assert monoid.min.is_idempotent + assert monoid.max[int].is_idempotent + assert monoid.lor.is_idempotent + assert monoid.band.is_idempotent + if shouldhave(monoid.numpy, "gcd"): + assert monoid.numpy.gcd.is_idempotent + assert not monoid.plus.is_idempotent + assert not monoid.times[float].is_idempotent + if config["mapnumpy"] or shouldhave(monoid.numpy, "equal"): + assert not monoid.numpy.equal.is_idempotent + with pytest.raises(AttributeError): + binary.min.is_idempotent + + +def _isidem_factory(scale): # pragma: no cover (called by Numba) + # Plain arithmetic so the binop compiles for every builtin type that + # ``BinaryOp._build`` samples, including complex (no ``>`` lowering). + def inner(x, y): + return (x + y) * scale + + return inner + + +def _isidem_pick_first(x, y): # pragma: no cover (called by Numba) + # Returning ``x`` is trivially idempotent (``op(x, x) == x``) and + # compiles for every builtin type, including complex where ``max`` + # has no Numba lowering. + return x + + +@pytest.mark.skipif("not supports_udfs") +def test_monoid_pickle_preserves_is_idempotent(): + """Regression: anonymous ``Monoid`` and ``ParameterizedMonoid`` both + dropped ``is_idempotent`` from the pickle round-trip, silently turning + a known-idempotent op into a non-idempotent one. Both + ``Monoid.__reduce__`` and ``ParameterizedMonoid.__reduce__`` now carry + the flag explicitly so ``_deserialize`` can pass it to + ``register_anonymous`` / ``register_new``. + """ + import pickle + + # Plain Monoid (non-parameterized) over an anonymous BinaryOp. + bin_op = BinaryOp.register_anonymous(_isidem_pick_first, "isidem_pick_first") + mon = Monoid.register_anonymous(bin_op, 0, "isidem_monoid", is_idempotent=True) + assert mon.is_idempotent is True + assert pickle.loads(pickle.dumps(mon)).is_idempotent is True + + # ParameterizedMonoid wrapping a ParameterizedBinaryOp. + pbin = BinaryOp.register_anonymous(_isidem_factory, parameterized=True) + pmon = Monoid.register_anonymous(pbin, 0, "isidem_param_monoid", is_idempotent=True) + assert pmon.is_idempotent is True + assert pickle.loads(pickle.dumps(pmon)).is_idempotent is True + + +def _parameterized_is_udt_factory(scale): # pragma: no cover (called by Numba inside the op) + def inner(x, y): + return x * scale + y * scale + + return inner + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize( + "module_name", + ["unary", "binary", "indexunary", "select", "indexbinary"], +) +def test_parameterized_is_udt_pickle_roundtrip(module_name): + """Parameterized + ``is_udt=True`` propagates through ``__reduce__``. + + Regression: ``Parameterized{Unary,Binary,IndexUnary,Select,IndexBinary}Op.__reduce__`` + used to emit ``(name, func, anonymous)``, so ``_deserialize`` invoked + ``register_*(..., parameterized=True)`` without ``is_udt``. Cross-process + re-register lost the flag, then dispatch took the non-UDT compile path + and failed at first use. + """ + import pickle + + module = getattr(gb, module_name) + op = module.register_anonymous(_parameterized_is_udt_factory, parameterized=True, is_udt=True) + assert op._is_udt is True + op2 = pickle.loads(pickle.dumps(op)) + assert op2._is_udt is True + + +def test_ops_have_ss(): + modules = [unary, binary, monoid, semiring, indexunary, select, op] + if suitesparse: + for mod in modules: + assert mod.ss is not None + else: + for mod in modules: + with pytest.raises(AttributeError): + mod.ss + + +def test_compile_codegen_helper(): + """The ``_compile_codegen`` helper validates source and surfaces typos clearly. + + Codegen bugs used to surface as a cryptic ``SyntaxError`` from ``exec`` + or, worse, as a Numba ``TypingError`` at first use of the generated + function. The helper catches them at the call site with the offending + source attached, and registers each generated function with + ``linecache`` so any later traceback shows real lines instead of + ````. + """ + import linecache + + from graphblas.core.operator.udt_codegen import _compile_codegen + + fn = _compile_codegen( + "def _op(x, y):\n return x + y\n", + func_name="_op", + source_label="", + ) + assert fn(2, 3) == 5 + # The synthetic filename is registered with linecache so a traceback + # raised from inside the generated function points at real source. + co_filename = fn.__code__.co_filename + assert co_filename.startswith(" #") + assert "x + y" in "".join(linecache.cache[co_filename][2]) + + # A bad source surfaces as RuntimeError with the offending source attached + # and the underlying SyntaxError as ``__cause__``. + bad_src = "def _op(x, y):\n return (x + y\n" # missing close paren + with pytest.raises(RuntimeError) as exc_info: + _compile_codegen( + bad_src, + func_name="_op", + source_label="", + ) + msg = str(exc_info.value) + assert "" in msg + assert "not valid Python" in msg + assert "Source:" in msg + assert bad_src in msg + assert isinstance(exc_info.value.__cause__, SyntaxError) + + +def test_operator_namespace_typo_suggestions(): + # A typo in an operator namespace should suggest close matches (via difflib), + # drawn from __dir__() so lazily-registered operators are offered without + # forcing them to build. + with pytest.raises(AttributeError, match="has no attribute 'pluss'.*Did you mean 'plus'"): + binary.pluss + with pytest.raises(AttributeError, match="Did you mean 'plus'"): + monoid.pluss + with pytest.raises(AttributeError, match="plus_times"): + semiring.plus_time + with pytest.raises(AttributeError, match="Did you mean 'sum'"): + agg.summ + with pytest.raises(AttributeError, match="Did you mean"): + unary.expp + with pytest.raises(AttributeError, match="rowindex"): + indexunary.rowindexx + with pytest.raises(AttributeError, match="triu"): + select.triu_typo + with pytest.raises(AttributeError, match="Did you mean 'plus'"): + op.pluss + + # No close match -> plain message, no suggestion appended + with pytest.raises(AttributeError) as exc_info: + binary.zzzzzz + assert "has no attribute 'zzzzzz'" in str(exc_info.value) + assert "Did you mean" not in str(exc_info.value) + + # Building suggestions must not force lazy operators to compile + before = set(binary._delayed) + with pytest.raises(AttributeError): + binary.pluss + assert set(binary._delayed) == before + + +# Touching an operator namespace must not compile any lazily-registered UDF. +# Run in a subprocess: BinaryOp._initialize runs once per process, so by the +# time any test executes, the import-time behavior under test is long past. +_LAZY_UDF_PROBE = """ +import graphblas as gb + +gb.binary.plus # forces BinaryOp._initialize + +lazy_udfs = ("floordiv", "rfloordiv", "absfirst", "abssecond", "rpow") +print("package: " + gb.__file__) +print("still lazy: " + " ".join(n for n in lazy_udfs if n in gb.binary._delayed)) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_anonymous_op_named_like_a_builtin_is_not_lifted_to_udts(): + # Only the built-in ops lift field by field to UDTs. An anonymous op that + # shares a built-in's name is a user's op, as before ops were seeded in + # BinaryOp.__init__. + def plus_like(x, y): # pragma: no cover (numba) + return x + y + + op = BinaryOp.register_anonymous(plus_like, "plus") + assert op._custom_dtype is None + assert binary.plus._custom_dtype is not None + rec = dtypes.register_anonymous(np.dtype([("x", np.int32), ("y", np.float64)])) + v = Vector.from_coo([0], [(1, 2.0)], dtype=rec) + with pytest.raises(KeyError, match="does not work with"): + v.ewise_mult(v, op).new() + + +@pytest.mark.skipif("not supports_udfs") +def test_initialize_does_not_build_lazy_udfs(): + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + result = subprocess.run( + [sys.executable, "-c", _LAZY_UDF_PROBE], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + # Assert the child probed this tree before trusting what it reports about it. + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + + still_lazy = lines.get("still lazy", "").split() + assert still_lazy == ["floordiv", "rfloordiv", "absfirst", "abssecond", "rpow"], report + + +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udfs_are_disk_cached(): + # The built-in UDF binops are module-level functions, so numba can persist + # their compilation across processes. Anything a user registers cannot be: + # numba keys a cache entry on a stable on-disk source, which a lambda or an + # interactively defined function does not have. + from numba.core.caching import NullCache + + for name in ["floordiv", "rfloordiv", "absfirst", "abssecond", "rpow"]: + numba_func = getattr(binary, name)._numba_func + assert not isinstance(numba_func._cache, NullCache), name + + def _uncached_probe(x, y): # pragma: no cover (numba) + return x + y + + user_op = BinaryOp.register_anonymous(_uncached_probe) + assert isinstance(user_op._numba_func._cache, NullCache) + + +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udf_compiles_uncached_where_numba_cannot_cache(monkeypatch): + # A read-only install whose user has no writable cache directory either + # leaves Numba no cache locator, and njit(cache=True) raises RuntimeError. + # The op then compiles uncached instead of failing. + from numba.core import caching + + from graphblas.core.operator.binary import _floordiv + + monkeypatch.setattr(caching.CacheImpl, "_locator_classes", []) + op = BinaryOp._build("floordiv", _floordiv, anonymous=True, cache=True) + assert isinstance(op._numba_func._cache, caching.NullCache) + v = Vector.from_coo([0, 1], [-3, 4]) + assert v.ewise_mult(v, op).new().isequal(Vector.from_coo([0, 1], [1, 1])) diff --git a/graphblas/tests/test_operator_types.py b/graphblas/tests/test_operator_types.py index 027f02fcc..9192b37ae 100644 --- a/graphblas/tests/test_operator_types.py +++ b/graphblas/tests/test_operator_types.py @@ -185,6 +185,7 @@ "plus_copysign", "lazy", "lazy2", + "lazy_idem", "lazy_lazy", # UDFs created during tests "is_positive", @@ -204,6 +205,7 @@ "binary_pickle", "monoid_pickle", "semiring_pickle", + "_test_weak_lt", # numpy-graphblas commutation (can we clean this up?) "band_land", "band_lor", diff --git a/graphblas/tests/test_pickle.py b/graphblas/tests/test_pickle.py index 724f43d76..3659b666c 100644 --- a/graphblas/tests/test_pickle.py +++ b/graphblas/tests/test_pickle.py @@ -1,4 +1,7 @@ +import multiprocessing import pickle +import platform +import sys from pathlib import Path import numpy as np @@ -312,6 +315,20 @@ def test_udt(extra): assert udt2._is_anonymous assert pickle.loads(pickle.dumps(udt2)).np_type == udt2.np_type + if ( + np.__version__.startswith("1.") + and sys.platform == "linux" + and platform.machine() == "aarch64" + ): + # pickle3*.pkl was written on x86, where the align=True struct dtype's flags byte 0x90 + # (NPY_ALIGNED_STRUCT | 0x10) pickled from a signed char as -112. numpy 1.x + # dtype.__setstate__ stores flags back into a char and rejects values that do not + # round-trip; on Linux aarch64 char is unsigned, so (char)-112 == 144 and loading raises + # ValueError("incorrect value for flags variable (overflow)"). numpy 2.0 widened flags to + # uint64 and normalizes signed pickled values, so it loads fine everywhere. Regenerating + # the pickle cannot help under numpy 1.x: positive flags would fail the same check on + # signed-char platforms. + pytest.skip("pickle3 dtype flags cannot round-trip on unsigned-char platforms in numpy 1.x") path = Path(__file__).parent / f"pickle3{extra}.pkl" with path.open("rb") as f: d = pickle.load(f) @@ -338,3 +355,180 @@ def test_udt(extra): any_udt = d["any[udt]"] assert any_udt is gb.binary.any[udt3] assert pickle.loads(pickle.dumps(gb.binary.first[udt, int])) is gb.binary.first[udt, int] + + +# Module-level so the spawn worker can pickle it; local functions can't cross processes. +def _crossproc_ibo_add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + +# Operations dispatched by ``_crossproc_worker``. Each takes the unpickled dict +# and returns a JSON-friendly value to send through the Pipe. +def _op_udt_reduce(d): + return d["vector"].reduce(d["op"]).new().value + + +def _op_bound_ibo_ewise(d): + rows, cols, vals = d["A"].ewise_mult(d["B"], d["bound"]).new().to_coo() + return list(rows), list(cols), [int(v) for v in vals] + + +_CROSSPROC_OPS = { + "udt_reduce": _op_udt_reduce, + "bound_ibo_ewise": _op_bound_ibo_ewise, +} + + +def _crossproc_worker(backend, op_name, payload, conn): + """Run ``_CROSSPROC_OPS[op_name]`` on the unpickled payload in a fresh interpreter.""" + try: + import graphblas as gb_ + + gb_.init(backend) + unpickled = pickle.loads(payload) + conn.send(("ok", _CROSSPROC_OPS[op_name](unpickled))) + except Exception as exc: # pragma: no cover (only on failure) + conn.send(("err", f"{type(exc).__name__}: {exc}")) + finally: + conn.close() + + +def _run_crossproc(op_name, payload): + """Spawn a child running ``_crossproc_worker`` and return ``rest`` from ``("ok", *rest)``.""" + ctx = multiprocessing.get_context("spawn") + parent_conn, child_conn = ctx.Pipe(duplex=False) + proc = ctx.Process( + target=_crossproc_worker, + args=(gb.backend, op_name, payload, child_conn), + ) + proc.start() + child_conn.close() + try: + assert parent_conn.poll(timeout=120), "child process timed out" + result = parent_conn.recv() + finally: + proc.join(timeout=5) + if proc.is_alive(): # pragma: no cover (defensive) + proc.terminate() + proc.join(timeout=5) + status, *rest = result + assert status == "ok", f"child failed: {rest}" + return rest + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_pickle_crossprocess(): + """Verify a typed UDT monoid round-trips through a spawned child process. + + Same-process pickle already exercises the ``__reduce__`` and + ``_deserialize`` paths. This test additionally proves the child can + re-register the UDT and auto-lift the parent monoid for it without + inheriting any state from the parent process; that's the case + ``dask`` and ``multiprocessing`` actually hit. + """ + record_dtype = np.dtype([("x", np.int64), ("y", np.int64)], align=True) + udt_name = "CrossProcPickleUDT" + # Use register_new so the child can re-register under the same name. + if hasattr(gb.dtypes, udt_name): + delattr(gb.dtypes, udt_name) + gb.core.dtypes._registry.pop(udt_name, None) + udt = gb.dtypes.register_new(udt_name, record_dtype) + try: + monoid = gb.monoid.plus[udt] + v = gb.Vector(udt, size=3) + v[0] = (1, 10) + v[1] = (2, 20) + v[2] = (3, 30) + payload = pickle.dumps({"op": monoid, "vector": v}) + (value,) = _run_crossproc("udt_reduce", payload) + # Reduce of (1,10), (2,20), (3,30) with plus = (6, 60). + assert tuple(value) == (6, 60) + finally: + delattr(gb.dtypes, udt_name) + gb.core.dtypes._registry.pop(udt_name, None) + + +def _crossproc_parameterized_is_udt_factory(scale): # pragma: no cover (called by Numba) + def inner(x, y): + return x * scale + y * scale + + return inner + + +def _op_compile_parameterized_is_udt(d): + op = d["op"] + # ``op._is_udt`` must survive the spawn so the right compile path runs in the child. + return bool(op._is_udt) + + +_CROSSPROC_OPS["compile_parameterized_is_udt"] = _op_compile_parameterized_is_udt + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +@pytest.mark.parametrize( + "module_name", + ["unary", "binary", "indexunary", "select", "indexbinary"], +) +def test_parameterized_is_udt_pickle_crossprocess(module_name): + """``is_udt`` survives a spawn -> unpickle -> re-register chain. + + Same-process pickle is covered by ``test_parameterized_is_udt_pickle_roundtrip`` + in ``test_op.py``; this exercises the cross-process path that wave-5 + fixed (parameterized ``__reduce__`` used to drop ``is_udt`` and the + child silently took the non-UDT compile path). + """ + module = getattr(gb, module_name) + op_name = f"_cpu_param_is_udt_{module_name}" + if hasattr(module, op_name): + delattr(module, op_name) + module.register_new( + op_name, _crossproc_parameterized_is_udt_factory, parameterized=True, is_udt=True + ) + try: + op = getattr(module, op_name) + assert op._is_udt is True + payload = pickle.dumps({"op": op}) + (got,) = _run_crossproc("compile_parameterized_is_udt", payload) + assert got is True + finally: + delattr(module, op_name) + # binary and unary ops also register into the combined ``op`` namespace, + # which ``delattr(module, ...)`` doesn't touch; clean it so + # test_op_namespace doesn't see an op with no per-type home. + vars(gb.op).pop(op_name, None) + gb.op._delayed.pop(op_name, None) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_bound_ibo_pickle_crossprocess(): + """A bound IBO must round-trip through a spawned child process. + + Bound IBOs pickle as ``(_rebind_indexbinaryop, (parent_ibo, type, theta))``. + The parent IBO must re-register in the child (via the registered name) so + the rebind can resolve it. Same-process pickle is already covered by + ``test_pickle_bound`` in ``test_indexbinary.py``. + """ + from graphblas.core.operator.indexbinary import _has_idxbinop + + if not _has_idxbinop: + pytest.skip("requires SuiteSparse:GraphBLAS 9.4+") + + op_name = "CrossProcBoundIBO" + if hasattr(gb.indexbinary, op_name): + delattr(gb.indexbinary, op_name) + gb.indexbinary.register_new(op_name, _crossproc_ibo_add_theta) + try: + bound = getattr(gb.indexbinary, op_name)[int](100) + A = gb.Matrix.from_coo([0, 1], [0, 1], [3, 7]) + B = gb.Matrix.from_coo([0, 1], [0, 1], [5, 2]) + payload = pickle.dumps({"A": A, "B": B, "bound": bound}) + ((rows, cols, vals),) = _run_crossproc("bound_ibo_ewise", payload) + # (0,0): 3 + 5 + 100 = 108 ; (1,1): 7 + 2 + 100 = 109 + assert rows == [0, 1] + assert cols == [0, 1] + assert vals == [108, 109] + finally: + delattr(gb.indexbinary, op_name) diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index e93511914..35439f34b 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -1,4 +1,6 @@ import inspect +import itertools +import math import pickle import random import sys @@ -10,6 +12,7 @@ import graphblas as gb from graphblas import backend, binary, dtypes, monoid, replace, select, unary +from graphblas.core.ss import version_major as ss_version_major # noqa: F401 (used in skipif) from graphblas.exceptions import EmptyObject from .conftest import autocompute, compute, pypy @@ -198,6 +201,180 @@ def test_isclose(): assert Scalar.from_value(None, dtype="FP64").isclose(Scalar.from_value(None, dtype="FP32")) +def test_isclose_infinity(): + # As in math.isclose, an infinity is close only to itself. It was close to every + # number (abs(inf - y) is inf, and so is rel_tol * inf), in both of isclose's paths. + values = [math.inf, -math.inf, 1.0, 0.0, math.nan] + for x, y in itertools.product(values, repeat=2): + for abs_tol in [0.0, math.inf]: + expected = math.isclose(x, y, rel_tol=1e-7, abs_tol=abs_tol) + assert Scalar.from_value(x).isclose(y, abs_tol=abs_tol) is expected, (x, y) + other = Scalar.from_value(y) + assert Scalar.from_value(x).isclose(other, abs_tol=abs_tol) is expected, (x, y) + + +def test_isclose_integers_do_not_wrap(): + # x - y or abs(x) wrapped in an integer type, and a UINT64 value above INT64's + # maximum raised OverflowError with Numba; see test_vector.ISCLOSE_WRAPPING_CASES + from .test_vector import ISCLOSE_WRAPPING_CASES + + for xtype, ytype, x, y, rel_tol in ISCLOSE_WRAPPING_CASES: + expected = math.isclose(x, y, rel_tol=rel_tol) + for (lt, lv), (rt, rv) in [((xtype, x), (ytype, y)), ((ytype, y), (xtype, x))]: + left = Scalar.from_value(lv, dtype=lt) + right = Scalar.from_value(rv, dtype=rt) + assert left.isclose(right, rel_tol=rel_tol) is expected, (lt, lv, rt, rv) + + +def test_isclose_in_the_type_the_dtypes_unify_to(): + # INT64 and UINT64 unify to FP64, where 2**62 + 1 rounds to 2**62, so they are close + # even with no tolerance, as in Vector.isclose; with Numba and without + x = Scalar.from_value(2**62 + 1, dtype="INT64") + y = Scalar.from_value(2**62, dtype="UINT64") + assert x.isclose(y, rel_tol=0) + assert y.isclose(x, rel_tol=0) + v = Vector.from_coo([0], [2**62 + 1]) + assert v.isclose(Vector.from_coo([0], [2**62], "UINT64"), rel_tol=0) + # Two INT64 values are compared exactly + assert not x.isclose(Scalar.from_value(2**62, dtype="INT64"), rel_tol=0) + + +def test_isclose_udt_raises(): + """Reject ``isclose`` on UDTs with a clear message that points to isequal.""" + arr_udt = dtypes.register_anonymous(np.dtype((np.int64, (3,))), "_IscloseUdt3") + s = gb.Scalar(arr_udt) + s.value = np.array([1, 2, 3], dtype=np.int64) + other = gb.Scalar(arr_udt) + other.value = np.array([1, 2, 3], dtype=np.int64) + # A raw value and a typed UDT Scalar both get the same clear message. + for arg in (np.array([1, 2, 3], dtype=np.int64), other): + with pytest.raises(TypeError, match="isclose is not defined for user-defined types"): + s.isclose(arg) + # isequal (the suggested alternative) does work for UDTs. + assert s.isequal(other) + + +def test_udt_scalar_padding_is_zeroed(): + """A record UDT's alignment padding is zeroed, not left as process memory. + + The whole record is copied into the GrB_Scalar and read back out of + ``Scalar.value``, so anything that compares or hashes those raw bytes would + see an equal-valued scalar as unequal. numpy promises nothing about which + bytes an assignment writes: it may fill the fields and leave the padding + alone, or copy whole items out of a temporary it never initialized (numpy + 2.5 switched ``arr[:] = record`` to the latter). + """ + # int8 then float64 leaves 7 bytes of padding under C alignment rules. + udt = dtypes.register_anonymous( + np.dtype([("pad_a", np.int8), ("pad_b", np.float64)], align=True), "_PaddingUdt" + ) + assert udt.np_type.itemsize == 16 # 1 + 7 padding + 8 + + # numpy serves a small buffer from a cache of recently freed blocks, so + # dirtying one of the right size first is what makes an uninitialized + # allocation read back as junk rather than as incidentally-zero memory. + junk = np.full(udt.np_type.itemsize, 0xAA, dtype=np.uint8) + del junk + raw = Scalar.from_value((3, 4.0), dtype=udt).value.tobytes() + assert raw[1:8] == bytes(7), f"alignment padding leaked into the scalar: {raw!r}" + + # Same guarantee without relying on what the allocator happened to hand + # back: a source array whose padding is deliberately 0xAA must not carry + # those bytes through, no matter how numpy chooses to copy it. + dirty = np.frombuffer(bytearray(b"\xaa" * udt.np_type.itemsize), dtype=udt.np_type) + dirty["pad_a"] = 3 + dirty["pad_b"] = 4.0 + raw = Scalar.from_value(dirty, dtype=udt).value.tobytes() + assert raw[1:8] == bytes(7), f"a value's own padding leaked into the scalar: {raw!r}" + assert raw == Scalar.from_value((3, 4.0), dtype=udt).value.tobytes() + + +# A SuiteSparse compiled without variable-length arrays (MSVC, so the Windows +# builds) rejects a user-defined type larger than GB_VLA_MAXSIZE with +# GrB_INVALID_VALUE. That ceiling was 128 bytes through SS 8.x and is 1024 as +# of 9.0, so skipping on SS < 9 covers the affected builds and costs one test +# on the rest. +@pytest.mark.skipif( + "ss_version_major < 9", + reason="SuiteSparse < 9 rejects a 240-byte UDT on builds without VLA support", +) +def test_udt_scalar_nested_array_roundtrip(): + """A UDT declared as an array of arrays reads back at full extent. + + numpy keeps subarray dtypes layered, so a 6-long array of 5-long FP64 + arrays reports ``subdtype`` as a 40-byte inner dtype with shape ``(6,)``. + Viewing all 240 bytes as that inner dtype raised ValueError: "Changing the + dtype to a subarray type is only supported if the total itemsize is + unchanged". Registration now flattens it to ``FP64[6, 5]``. + """ + inner = np.dtype((np.float64, (5,))) + udt = dtypes.register_anonymous(np.dtype((inner, (6,))), "_NestedArrayUdt") + assert udt.np_type == np.dtype((np.float64, (6, 5))) + + val = np.arange(30, dtype=np.float64).reshape(6, 5) + got = Scalar.from_value(val, dtype=udt).value + # Assert the values, not just that it did not raise: a wrong collapse + # produces a correctly-shaped but transposed array. + assert got.shape == (6, 5) + np.testing.assert_array_equal(got, val) + + +def test_udt_nat_duration_in_another_unit_stays_nat(): + """A NaT duration stored into a duration of another unit is NaT, not a number. + + numpy before 2.3 converts a NaT ``timedelta64`` scalar into another unit as + the integer it is stored as: NaT ms became 9223372036854774 s, and NaT + minutes 0 s. A date or an array of durations was not affected. + """ + seconds = dtypes.register_anonymous(np.dtype("m8[s]"), "_NatSeconds") + rec = dtypes.register_anonymous(np.dtype([("nat_d", "m8[s]"), ("nat_n", np.int8)]), "_NatRec") + days = dtypes.register_anonymous(np.dtype([("nat_t", "M8[D]"), ("nat_n", np.int8)])) + nested = dtypes.register_anonymous( + np.dtype([("nat_o", [("nat_d", "m8[s]")]), ("nat_n", np.int8)]), "_NatNested" + ) + # Numba cannot type an array of durations, so registering one needs it absent. + pair = pairs = None + if not gb.core._has_numba: + pair = dtypes.register_anonymous(np.dtype(("m8[s]", 2))) + pairs = dtypes.register_anonymous(np.dtype([("nat_p", "m8[s]", 2), ("nat_n", np.int8)])) + for unit in ["ns", "ms", "m", "D"]: + nat = np.timedelta64("NaT", unit) + assert np.isnat(Scalar.from_value(nat, seconds).value) + s = Scalar(seconds) + s.value = nat + assert np.isnat(s.value) + v = Vector(seconds, size=2) + v[0] = nat + assert np.isnat(v[0].new().value) + v[:] = nat + assert np.isnat(v.to_coo()[1]).all() + r = Vector(rec, size=2) + r[0] = (nat, 1) + r[1] = {"nat_d": nat, "nat_n": 1} + assert np.isnat(r.to_coo()[1]["nat_d"]).all() + # Off SuiteSparse, ``build`` uses ``binary.any`` as its ``dup_op``, which a + # UDT gets only by compiling it with Numba. + if suitesparse or gb.core._supports_udfs: + v = Vector.from_coo([0, 1], [nat, nat], dtype=seconds) + assert np.isnat(v.to_coo()[1]).all() + r = Vector.from_coo([0, 1], [(nat, 1), (nat, 2)], dtype=rec) + assert np.isnat(r.to_coo()[1]["nat_d"]).all() + assert np.isnat(Scalar.from_value(((nat,), 1), nested).value["nat_o"]["nat_d"]) + if pair is not None: + got = Scalar.from_value([nat, np.timedelta64(5000, "ms")], pair).value + assert np.isnat(got[0]) + assert got[1] == np.timedelta64(5, "s") + got = Scalar.from_value(([nat, np.timedelta64(5000, "ms")], 1), pairs).value + assert np.isnat(got["nat_p"][0]) + assert got["nat_p"][1] == np.timedelta64(5, "s") + r = Vector(days, size=1) + r[0] = (np.datetime64("NaT", unit), 1) + assert np.isnat(r[0].new().value["nat_t"]) + # Only NaT changes how it converts; a value converts as numpy converts it. + assert Scalar.from_value(np.timedelta64(5000, "ms"), seconds).value == np.timedelta64(5, "s") + assert Scalar.from_value(np.timedelta64(2, "m"), seconds).value == np.timedelta64(120, "s") + + def test_nvals(s): assert s.nvals == 1 s.clear() @@ -286,13 +463,71 @@ def test_scalar_to_numpy(s): assert a.shape == b.shape +def test_uint64_top_half_from_value(): + """An int from 2**63 to 2**64 - 1 goes into a UINT64 GrB_Scalar, Vector or Matrix. + + A Python int was typed INT64 by itself, so it raised OverflowError. + """ + for value in [2**63, 2**64 - 1]: + for is_cscalar in [False, True]: + s = Scalar.from_value(value, dtype=dtypes.UINT64, is_cscalar=is_cscalar) + assert s.value == value + assert s == value + v = Vector(dtypes.UINT64, size=3) + v[0] = value + v[1:] = value + assert v.to_coo()[1].tolist() == [value] * 3 + v << value + assert v.to_coo()[1].tolist() == [value] * 3 + A = Matrix(dtypes.UINT64, nrows=2, ncols=2) + A[0, 1] = value + A[1, :] = value + assert A.to_coo()[2].tolist() == [value] * 3 + f = Vector(dtypes.FP64, size=1) + f[0] = value + assert f[0].new().value == float(value) + with pytest.raises(OverflowError): + Scalar.from_value(2**63, dtype=dtypes.INT8, is_cscalar=False) + i64 = Vector(dtypes.INT64, size=1) + with pytest.raises(OverflowError): + i64[0] = 2**63 # not wrapped + u64 = Vector(dtypes.UINT64, size=1) + with pytest.raises(OverflowError): + u64[0] = 2**64 + + +def test_eq_types_a_literal_as_binary_eq_does(): + """``==`` on a Scalar compares as ``binary.eq`` does, as on a Vector, and returns a bool.""" + assert Scalar.from_value(0.1, dtype=dtypes.FP32) == 0.1 # in FP32, as numpy 2 compares it + assert Scalar.from_value(1, dtype=dtypes.INT8) != 300 # exact, as a comparison + assert Scalar.from_value(1, dtype=dtypes.INT8) == Scalar.from_value(1.0) + not_a_number = float("nan") + assert Scalar.from_value(not_a_number) != not_a_number + assert not (Scalar.from_value(not_a_number) == not_a_number) # noqa: SIM201 + empty = Scalar(dtypes.INT64) + assert empty == Scalar(dtypes.FP64) + assert empty == None # noqa: E711 + assert empty != 0 + # A value that is not a literal names the operator used, not Scalar.ewise_mult. + s = Scalar.from_value(1, dtype=dtypes.INT8) + with pytest.raises(TypeError, match=r"argument `other` in Scalar\.__eq__\(\.\.\.\)"): + s == [1] # noqa: B015 + with pytest.raises(TypeError, match=r"argument `other` in Scalar\.__ne__\(\.\.\.\)"): + s != object() # noqa: B015 + # A dict is not an element of an array UDT (AttributeError from reading its fields) + a = Scalar.from_value([1, 2, 3], dtypes.register_anonymous(np.dtype((np.int8, 3)))) + match = r"Scalar\.__eq__" if gb.core._supports_udfs else None + with pytest.raises(TypeError, match=match): + a == {} # noqa: B015 + + @autocompute def test_neg(): for dtype in sorted( ( dtype for attr, dtype in vars(dtypes).items() - if isinstance(dtype, dtypes.DataType) and attr not in {"_INDEX"} + if isinstance(dtype, dtypes.DataType) and attr != "_INDEX" ), key=lambda x: x.name, reverse=random.choice([False, True]), # used to segfault when False @@ -300,16 +535,14 @@ def test_neg(): s = Scalar.from_value(1, dtype=dtype) empty = Scalar(dtype) if dtype._is_udt: - with pytest.raises(KeyError, match="ainv does not work with"): - -s - with pytest.raises(KeyError, match="ainv does not work with"): - -empty - else: - minus_s = Scalar.from_value(-1, dtype=dtype, is_cscalar=False) # pragma: is_grbscalar - assert s == -minus_s - assert (-s).value == minus_s.value - assert empty == -empty - assert compute((-empty).value) is None + # ainv works on UDTs with numeric record/array fields; skip UDTs + # where the negation result is hard to validate generically + continue + minus_s = Scalar.from_value(-1, dtype=dtype, is_cscalar=False) # pragma: is_grbscalar + assert s == -minus_s + assert (-s).value == minus_s.value + assert empty == -empty + assert compute((-empty).value) is None @autocompute @@ -349,6 +582,7 @@ def test_expr_is_like_scalar(s): expected = { "__call__", "__del__", + "__getattr__", "__imatmul__", "__lshift__", "_carg", @@ -393,6 +627,7 @@ def test_index_expr_is_like_scalar(s): # Should we make any of these raise informative errors? expected = { "__del__", + "__getattr__", "__imatmul__", "_carg", "_deserialize", @@ -541,6 +776,19 @@ def test_ewise_union(s): s.ewise_union(t, binary.plus, 10, v) +def test_select_row_compares_its_index_with_the_bound_given(s): + # A Scalar's only index is 0. The bound is the number given, made a whole number + # (s > -0.5 dropped it, and s < 1 raised AttributeError), not the comparison's thunk + # typed beside the value (300 beside INT8 is infinity). + assert select.row(s > -0.5).new() == s + assert select.row(s < 0.5).new() == s + assert select.row(s < 1).new() == s + assert select.row(s >= 0.5).new().is_empty + s8 = Scalar.from_value(1, dtype=dtypes.INT8) + assert select.row(s8 < 300).new() == s8 + assert select.row(s8 > 299).new().is_empty + + def test_ewise_mult_add(s): assert s.ewise_add(s).new() == 10 assert s.ewise_mult(s).new() == 25 @@ -549,6 +797,9 @@ def test_ewise_mult_add(s): s.ewise_add(v) with pytest.raises(TypeError, match="Literal scalars also"): s.ewise_mult(v) + # A literal is typed beside a string op that names no BinaryOp, so the op reports it. + with pytest.raises(ValueError, match="Unknown binary string: 'plus_times'"): + s.ewise_mult(1, "plus_times") def test_select(s): @@ -581,6 +832,69 @@ def test_record_from_dict(): assert s == (1, (2, 3)) +@pytest.mark.skipif("not gb.core._has_numba") +def test_udt_eq_without_numba_matches_binary_eq(monkeypatch): + """Without Numba, ``==`` on a UDT Scalar compares in numpy, as ``binary.eq`` would. + + ``binary.eq`` on a UDT needs Numba; ``==`` on a Scalar used to compare in + Python, so it works without Numba still. Here both ways must agree. + """ + rec = dtypes.register_anonymous(np.dtype([("x", np.int64), ("y", np.float64)], align=True)) + arr = dtypes.register_anonymous(np.dtype((np.int8, (3,)))) + f32 = dtypes.register_anonymous(np.dtype((np.float32, (2,)))) + r = Scalar.from_value((1, 2.5), rec) + a = Scalar.from_value([1, 1, 1], arr) + f = Scalar.from_value([0.1, 0.1], f32) + text = dtypes.register_anonymous(np.dtype("S6")) + t = Scalar.from_value(b"ab", text) + cases = [ + (r, (1, 2.5)), + (r, (1, 2)), + (r, {"x": 1, "y": 2.5}), + (a, 1), + (a, [1, 1, 1]), + (a, [1]), + (a, [1, 2, 1]), + (a, 0.5), + (a, 300), + (a, 2**63), + (a, Scalar.from_value([1, 1, 1], arr)), + (f, 0.1), + (f, np.float64(0.1)), + (t, b"ab"), + (t, b"cd"), + (t, Scalar.from_value(b"ab", text)), + (t, Scalar.from_value(b"ab", dtypes.register_anonymous(np.dtype("S7")))), + (t, b"abcdefgh"), + ] + # UDTs that binary.eq does not pair raise its KeyError both ways. + other_names = dtypes.register_anonymous( + np.dtype([("p", np.int64), ("q", np.float64)], align=True) + ) + arr2 = dtypes.register_anonymous(np.dtype((np.int8, (2,)))) + unpaired = [ + (r, Scalar.from_value((1, 2.5), other_names)), + (a, Scalar.from_value([1, 1], arr2)), + (r, a), + (t, Scalar.from_value("ab", dtypes.register_anonymous(np.dtype("U7")))), + (t, 1), + # A number for a bytes field does not compare. + (Scalar.from_value((b"ab", 1), np.dtype([("tx", "S2"), ("tn", np.int8)])), (1, 1)), + ] + + def compare_all(): + for s, other in unpaired: + with pytest.raises(KeyError, match="does not work with"): + s == other # noqa: B015 + return [(s == other, s != other) for s, other in cases] + + with_numba = compare_all() + monkeypatch.setattr(gb.core.scalar, "_has_numba", False) + without_numba = compare_all() + assert without_numba == with_numba + assert with_numba[:3] == [(True, False), (False, True), (True, False)] + + def test_get(s): assert s.get() == 5 assert s.get("mittens") == 5 @@ -633,3 +947,44 @@ def test_scalar_operators(s): unary.sin(object()) with pytest.raises(TypeError): binary.plus(object(), object()) + + +@autocompute +def test_scalar_operators_literal_on_the_left(): + # `3 ** s` is pow(3, s), as `binary.pow(3, s)` is, with no reflected op: numpy.mod has + # none, rpow is a Numba function, and none takes a UDT + s = Scalar.from_value(5, dtype=dtypes.INT8) + assert 2 - s == -3 + assert 2**s == 32 + result = (3**s).new() + assert result.dtype == dtypes.INT8 + assert result.value == 127 # GraphBLAS's pow saturates + t = Scalar.from_value(True) + result = (False**t).new() + assert result.dtype == dtypes.BOOL + assert result.value is False + if gb.core._supports_udfs: + assert 7 % s == 2 + + +@pytest.mark.skipif("not gb.core._supports_udfs") +def test_scalar_operators_literal_on_the_left_udt(): + udt = dtypes.register_anonymous(np.dtype((np.int8, (2,))), "_LeftLiteralI8") + u = Scalar.from_value([5, 7], udt) + for result, expected in [ + (1 - u, [-4, -6]), + ((1, 2) - u, [-4, -5]), + (1 // u, [0, 0]), + (10 / u, [2.0, 10 / 7]), + ]: + result = result.new() + np.testing.assert_array_equal(result.value, expected) + assert (1 - u).new().dtype == udt + assert (10 / u).new().dtype.np_type == np.dtype((np.float64, (2,))) + # An op typed for its two sides takes a literal on the left in its left type + rec = dtypes.register_anonymous(np.dtype([("ll_a", "f8"), ("ll_b", "f8")]), "_LeftLiteralRec") + add_a = gb.core.operator.BinaryOp.register_anonymous( + lambda x, y: x + y.ll_a, "_left_literal_add_a", is_udt=True + ) + s = Scalar.from_value((3.0, 4.0), rec) + assert add_a[dtypes.INT64, rec](2 & s).new().value == 5.0 diff --git a/graphblas/tests/test_ss_refcount.py b/graphblas/tests/test_ss_refcount.py new file mode 100644 index 000000000..8b64d4bac --- /dev/null +++ b/graphblas/tests/test_ss_refcount.py @@ -0,0 +1,144 @@ +"""Regression tests for gh-559: Matrix/Vector must not be born in a reference cycle. + +Historically each Matrix and Vector stored ``self.ss = ss(self)``, and the ss +object (and its config) held ``_parent`` back to the object, so instances died +only via the cyclic garbage collector, not by reference counting. For large +matrices in a tight loop this deferred the release of the C-side GrB buffer +until gc happened to run. ``.ss`` is now built lazily per access and stored +nowhere, so the cycle never forms. These tests pin that behavior; they are +suitesparse-only because ``.ss`` exists only on that backend. +""" + +import gc +import weakref + +import numpy as np +import pytest + +from graphblas import Matrix, Vector, backend, semiring +from graphblas.core.ss.matrix import ss as matrix_ss_class +from graphblas.core.ss.vector import ss as vector_ss_class + +if backend != "suitesparse": + pytest.skip("A.ss only available with suitesparse backend", allow_module_level=True) + + +def _dies_by_reference_count(factory): + """True if the object dies the instant its last strong ref drops, with gc off.""" + gc.collect() + was_enabled = gc.isenabled() + gc.disable() + try: + obj = factory() + wref = weakref.ref(obj) + del obj + return wref() is None + finally: + if was_enabled: + gc.enable() + + +def test_matrix_dies_by_reference_count(): + assert _dies_by_reference_count(lambda: Matrix(float, 5, 5)) + + +def test_vector_dies_by_reference_count(): + assert _dies_by_reference_count(lambda: Vector(float, 5)) + + +def test_mxm_result_dies_by_reference_count(): + A = Matrix.from_dense(np.arange(9.0).reshape(3, 3) + 1) + B = Matrix.from_dense(np.arange(9.0).reshape(3, 3) + 1) + assert _dies_by_reference_count(lambda: A.mxm(B, semiring.min_plus).new()) + + +def test_ss_namespace_is_functional(): + A = Matrix.from_coo([0, 1, 2], [0, 1, 2], [1.0, 2.0, 3.0], nrows=3, ncols=3) + # Built fresh each access (nothing is stored on the instance). + assert A.ss is not A.ss + # Introspection still works through a fresh access. + assert A.ss.nbytes > 0 + matrix_formats = { + "csr", + "csc", + "hypercsr", + "hypercsc", + "bitmapr", + "bitmapc", + "fullr", + "fullc", + "coor", + "cooc", + } + assert A.ss.export()["format"] in matrix_formats + # Config get then set then get, each through an independent `.ss`. + assert A.ss.config["format"] in {"by_row", "by_col"} + A.ss.config["format"] = "by_col" + assert A.ss.config["format"] == "by_col" + + v = Vector.from_coo([0, 2], [1.0, 3.0], size=4) + assert v.ss is not v.ss + assert v.ss.nbytes > 0 + assert v.ss.export()["format"] in {"sparse", "bitmap", "full"} + + +def test_ss_class_access_returns_namespace_class(): + # Class-level access must still yield the ss class so its import_* classmethods work. + assert Matrix.ss is matrix_ss_class + assert Vector.ss is vector_ss_class + assert hasattr(Matrix.ss, "import_any") + assert hasattr(Vector.ss, "import_any") + + +def test_ss_attribute_is_read_only(): + A = Matrix(float, 3, 3) + with pytest.raises(AttributeError): + A.ss = 5 + v = Vector(float, 3) + with pytest.raises(AttributeError): + v.ss = 5 + + +def test_views_have_working_ss(): + # A single-column Matrix cast to a Vector (_as_vector) is a view with _parent set. + A = Matrix.from_coo([0, 1], [0, 0], [1.0, 2.0], nrows=3, ncols=1) + v = A._as_vector() + assert v._parent is A + assert v.ss.nbytes > 0 + # A Vector cast to a Matrix (_as_matrix) is likewise a view. + w = Vector.from_coo([0, 2], [1.0, 3.0], size=4) + M = w._as_matrix() + assert M._parent is w + assert M.ss.nbytes > 0 + + +def test_batched_mxm_loop_does_not_accumulate_matrices(): + # gh-559: with the cyclic collector switched off, a batched + # mxm -> to_dense -> discard loop must not pile up Matrix objects. + rng = np.random.default_rng(0) + A = Matrix.from_dense(rng.random((16, 8)) + 0.1) + B = Matrix.from_dense(rng.random((8, 24)) + 0.1) + + def live_matrices(): + return sum(1 for o in gc.get_objects() if type(o) is Matrix) + + gc.collect() + was_enabled = gc.isenabled() + gc.disable() + try: + # Warm up one iteration so any one-time caches are populated first. + C = Matrix(float, 16, 24) + C << A.mxm(B, semiring.min_plus) + C.to_dense(0.0) + del C + baseline = live_matrices() + for _ in range(50): + C = Matrix(float, 16, 24) + C << A.mxm(B, semiring.min_plus) + C.to_dense(0.0) + del C + # Without the fix this would be baseline + 50. + assert live_matrices() <= baseline + finally: + if was_enabled: # pragma: no branch (safety) + gc.enable() diff --git a/graphblas/tests/test_ss_utils.py b/graphblas/tests/test_ss_utils.py index 2df7ab939..e76f6e45d 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -1,14 +1,46 @@ +import ctypes +import os +import subprocess +import sys +import tracemalloc +from pathlib import Path + import numpy as np import pytest from numpy.testing import assert_array_equal import graphblas as gb from graphblas import Matrix, Vector, backend +from graphblas.core.utils import _distinct_buffers, values_to_numpy_buffer from graphblas.exceptions import InvalidValue if backend != "suitesparse": pytest.skip("gb.ss and A.ss only available with suitesparse backend", allow_module_level=True) +try: + from suitesparse_graphblas.utils import can_unclaim_buffer, empty +except ImportError: # pragma: no cover (import) + # suitesparse-graphblas < 10.5.1.1 + can_unclaim_buffer = empty = None + + +# ``from_coo`` with all-equal values used to come back iso, because SuiteSparse +# checked after a build whether every value was the same. 10.5.0 dropped that +# check ("GrB_Matrix_build and GrB_Vector_build: no longer do a post-iso +# check"), so the iso arm of these tests has to ask for iso outright. +# ``ss.build_scalar`` does, and produces the same storage formats on every +# version, which is what the tests below are really selecting for. +def _iso_vector(indices, value, size): + v = Vector(int, size) + v.ss.build_scalar(indices, value) + return v + + +def _iso_matrix(rows, cols, value, nrows, ncols): + A = Matrix(int, nrows, ncols) + A.ss.build_scalar(rows, cols, value) + return A + @pytest.mark.parametrize("do_iso", [False, True]) def test_vector_head(do_iso): @@ -19,9 +51,14 @@ def test_vector_head(do_iso): values1 = [10, 20, 30] values2 = [2, 4, 6] values3 = [1, 2, 3] - v1 = Vector.from_coo([0, 1, 2], values1) # full - v2 = Vector.from_coo([1, 3, 5], values2) # bitmap - v3 = Vector.from_coo([100, 200, 300], values3) # sparse + if do_iso: + v1 = _iso_vector([0, 1, 2], values1[0], 3) # full + v2 = _iso_vector([1, 3, 5], values2[0], 6) # bitmap + v3 = _iso_vector([100, 200, 300], values3[0], 301) # sparse + else: + v1 = Vector.from_coo([0, 1, 2], values1) # full + v2 = Vector.from_coo([1, 3, 5], values2) # bitmap + v3 = Vector.from_coo([100, 200, 300], values3) # sparse assert v1.ss.export()["format"] == "full" assert v2.ss.export()["format"] == "bitmap" assert v3.ss.export()["format"] == "sparse" @@ -72,10 +109,16 @@ def test_matrix_head(do_iso): values2 = [1, 2, 4] values3 = values4 = [1, 2, 3] - A1 = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], values1) # fullr - A2 = Matrix.from_coo([0, 0, 1], [0, 1, 1], values2) # Bitmap - A3 = Matrix.from_coo([5, 5, 10], [4, 5, 10], values3) # CSR - A4 = Matrix.from_coo([500, 500, 1000], [400, 500, 1000], values4) # HyperCSR + if do_iso: + A1 = _iso_matrix([0, 0, 1, 1], [0, 1, 0, 1], values1[0], 2, 2) # fullr + A2 = _iso_matrix([0, 0, 1], [0, 1, 1], values2[0], 2, 2) # Bitmap + A3 = _iso_matrix([5, 5, 10], [4, 5, 10], values3[0], 11, 11) # CSR + A4 = _iso_matrix([500, 500, 1000], [400, 500, 1000], values4[0], 1001, 1001) # HyperCSR + else: + A1 = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], values1) # fullr + A2 = Matrix.from_coo([0, 0, 1], [0, 1, 1], values2) # Bitmap + A3 = Matrix.from_coo([5, 5, 10], [4, 5, 10], values3) # CSR + A4 = Matrix.from_coo([500, 500, 1000], [400, 500, 1000], values4) # HyperCSR d = A1.ss.export(raw=True) assert d["format"] == "fullr" d["format"] = "fullc" @@ -194,19 +237,37 @@ def test_about(): with pytest.raises(KeyError): about["badkey"] assert "SuiteSparse" in about["library_name"] + # SS:GraphBLAS 10.5.2 made the GPU modes aliases of these; neither may read as a GPU mode + assert about["mode"] == ("blocking" if gb._init_params["blocking"] else "nonblocking") with pytest.raises(TypeError): del about["library_name"] # pylint: disable=unsupported-delete-operation assert "library_name" in repr(about) def test_openmp_enabled(): - # SuiteSparse:GraphBLAS without OpenMP enabled is very undesirable - assert gb.ss.about["openmp"] + # A SuiteSparse:GraphBLAS built without OpenMP is single-threaded, which is + # a large and entirely silent performance loss: nothing else about the + # library reports it, and the conda package still depends on an OpenMP + # runtime whether or not the build actually found one. Keep this assertion + # hard; it is how the conda-forge osx-arm64 regression in graphblas 10.5.0 + # was noticed. + assert gb.ss.about["openmp"], ( + "libgraphblas was built without OpenMP and will run single-threaded. " + "This is a defect in the installed build, not in python-graphblas. " + "Reinstall or rebuild graphblas with OpenMP enabled " + "(SuiteSparse only warns when its CMake cannot find OpenMP, so a " + "serial library can ship looking normal)." + ) -def test_global_config(): +def test_global_config(monkeypatch): d = {} config = gb.ss.config + if "jit_c_control" in config: # SuiteSparse:GraphBLAS 8+ + from graphblas.core.ss import jit_config + + # Writing jit_c_control also sets whether UDT work keeps a "run"; undo that after + monkeypatch.setattr(jit_config, "_user_set_run", jit_config._user_set_run) for k in config: d[k] = config[k] assert d == config @@ -229,6 +290,13 @@ def test_global_config(): for k in config: if k in config._defaults: config[k] = None + elif k == "jit_c_control": + # None is the default "run" whatever graphblas.config has, and not a "run" + # the user chose, so UDT work may still raise it to "on" + config[k] = None + assert config[k] == "run" + assert jit_config._user_set_run is False + config[k] = d[k] else: with pytest.raises(ValueError, match="Unable to set default value for"): config[k] = None @@ -237,6 +305,80 @@ def test_global_config(): assert "format" in repr(config) +def test_global_config_key_completions(): + # IPython offers these for `config[]`, so they must be the keys that + # actually resolve, not the attributes of the mapping object. + config = gb.ss.config + completions = config._ipython_key_completions_() + assert set(completions) == set(config._options) + for key in completions: + config[key] + # About aliases the same hook onto its own __iter__ (gb.ss.about[]) + about = gb.ss.about + completions = about._ipython_key_completions_() + assert set(completions) == set(about) + for key in completions: + about[key] + + +def test_global_config_attribute_assignment_raises(): + # Attribute assignment used to create a dead instance attribute and leave + # the real config unchanged; now it raises and points at item assignment. + config = gb.ss.config + before = config["nthreads"] + with pytest.raises(AttributeError, match="item assignment"): + config.nthreads = before + 1 + assert config["nthreads"] == before + assert "nthreads" not in config.__dict__ + with pytest.raises(AttributeError, match="Unknown config option 'nthread'"): + config.nthread = before + 1 + # Item assignment is the supported write and still works + config["nthreads"] = before + 1 + assert config["nthreads"] == before + 1 + config["nthreads"] = before + assert config["nthreads"] == before + # Reads, iteration, and key completions are unaffected by the guard + assert dict(config) == {k: config[k] for k in config} + assert set(config._ipython_key_completions_()) == set(config._options) + + +def test_object_config_attribute_assignment_raises(): + v = Vector(int, 3) + with pytest.raises(AttributeError, match="item assignment"): + v.ss.config.sparsity_control = "bitmap" + assert v.ss.config["sparsity_control"] == {"auto"} + with pytest.raises(AttributeError, match="read-only"): + v.ss.config.sparsity_status = "bitmap" + A = Matrix(int, 2, 2) + with pytest.raises(AttributeError, match="item assignment"): + A.ss.config.format = "by_col" + assert A.ss.config["format"] == "by_row" + # Item assignment is the supported write and still works + v.ss.config["sparsity_control"] = "bitmap" + assert v.ss.config["sparsity_control"] == {"bitmap"} + v.ss.config["sparsity_control"] = "auto" + assert v.ss.config["sparsity_control"] == {"auto"} + + +def test_about_attribute_assignment_raises(): + about = gb.ss.about + mode = about["mode"] + with pytest.raises(AttributeError, match="read-only"): + about.mode = "junk" + assert about["mode"] == mode + assert "mode" not in about.__dict__ + with pytest.raises(AttributeError, match="Unknown About option"): + about.junkattr = 1 + assert dict(about) == {k: about[k] for k in about} + + +@pytest.mark.skipif("gb.core.ss._IS_SSGB7") +def test_context_attribute_assignment_raises(): + context = gb.ss.Context(engage=False) + with pytest.raises(AttributeError, match="item assignment"): + context.nthreads = 4 + + @pytest.mark.skipif("gb.core.ss._IS_SSGB7") def test_context(): context = gb.ss.Context() @@ -256,8 +398,8 @@ def test_context(): assert context4["nthreads"] == context["nthreads"] + 1 assert context == context.dup() assert context4 == context.dup(chunk=context["chunk"] + 1, nthreads=context["nthreads"] + 1) - assert context.dup(gpu_id=-1)["gpu_id"] == -1 - + if "gpu_id" in gb.ss.Context._options: + assert context.dup(gpu_id=-1)["gpu_id"] == -1 context.engage() assert gb.core.ss.context.threadlocal.context is context with gb.ss.Context(nthreads=1) as ctx: @@ -294,3 +436,398 @@ def test_context(): context._context = context # This is allowed to work with config with pytest.raises(AttributeError, match="_context"): context._context = ctx # This is not + + +def _libc_arena(): + """Give the last GraphBLAS arena libc's allocator, and return its number. + + Arrays from ``empty(..., arena=k)`` then carry a NumPy memory handler of + their own, as arrays from any custom handler (NEP 49) do. + """ + from graphblas.core import ffi, lib + + k = lib.GxB_NARENAS - 1 + initialized = ffi.new("int*") + assert lib.GxB_arena_initialized(initialized, k) == lib.GrB_SUCCESS + # GxB_arena_init refuses an arena that is already initialized, which it is + # only if the test runs again in the same process + if not initialized[0]: # pragma: no branch (safety) + libc = ctypes.CDLL(None) + funcs = [ + ffi.cast(sig, ctypes.cast(getattr(libc, name), ctypes.c_void_p).value) + for name, sig in [ + ("malloc", "void *(*)(size_t)"), + ("calloc", "void *(*)(size_t, size_t)"), + ("realloc", "void *(*)(void *, size_t)"), + ("free", "void (*)(void *)"), + ] + ] + assert lib.GxB_arena_init(k, *funcs) == lib.GrB_SUCCESS + return k + + +@pytest.mark.skipif("empty is None") +@pytest.mark.skipif("sys.platform == 'win32'") # no ``ctypes.CDLL(None)`` +def test_take_ownership_of_foreign_memory(): + # These arrays own their data, but GraphBLAS would free it with NumPy's + # allocator, so it must copy them. It used to take them, and then the + # process aborted when GraphBLAS freed one (NumPy >=2.5, PYTHONMALLOC=debug). + k = _libc_arena() + + def foreign(values, dtype): + array = empty(len(values), dtype, arena=k) + array[:] = values + assert array.flags.owndata + assert not can_unclaim_buffer(array) + return array + + indptr = foreign([0, 2, 3], np.uint64) + col_indices = foreign([0, 2, 1], np.uint64) + values = foreign([1.0, 2.0, 3.0], np.float64) + A = Matrix.ss.import_csr( + nrows=2, + ncols=3, + indptr=indptr, + col_indices=col_indices, + values=values, + take_ownership=True, + ) + expected = Matrix.from_coo([0, 0, 1], [0, 2, 1], [1.0, 2.0, 3.0], nrows=2, ncols=3) + assert A.isequal(expected) + for array in [indptr, col_indices, values]: + assert array.flags.owndata + assert array.flags.writeable + values[:] = 0 + assert A.isequal(expected) + + values = foreign([1, 2, 3], np.int64) + v = Vector.ss.import_full(values=values, take_ownership=True) + assert v.isequal(Vector.from_coo([0, 1, 2], [1, 2, 3])) + assert values.flags.owndata + del A, v + + +def test_take_ownership_asks_suitesparse_graphblas(monkeypatch): + # Whether GraphBLAS may own an array is up to ``can_unclaim_buffer`` + # (``owndata`` alone can't tell), even on a suitesparse-graphblas without it + monkeypatch.setattr(gb.core.utils, "can_unclaim_buffer", lambda array: False) + v = Vector.from_coo([1, 3], [10, 20], size=5) + d = v.ss.export("sparse") + w = Vector.ss.import_sparse(take_ownership=True, **d) + assert w.isequal(v) + assert d["indices"].flags.owndata + assert d["values"].flags.owndata + + +def test_take_ownership_without_can_unclaim_buffer(monkeypatch): + # suitesparse-graphblas < 10.5.1.1 has neither function. GraphBLAS then takes an + # array that owns its data and is writeable, and gets a NumPy copy of any other. + monkeypatch.setattr(gb.core.utils, "can_unclaim_buffer", None) + monkeypatch.setattr(gb.core.utils, "empty", None) + v = Vector.from_coo([1, 3], [10, 20], size=5) + d = v.ss.export("sparse") + d["indices"].flags.writeable = False + values = np.array([10, 20, 0, 0]) + w = Vector.ss.import_sparse(take_ownership=True, **dict(d, values=values[:2])) + assert w.isequal(v) + assert d["indices"].flags.owndata + values[:] = 0 # w has a copy of the view + assert w.isequal(v) + + d = v.ss.export("sparse") + w = Vector.ss.import_sparse(take_ownership=True, **d) + assert w.isequal(v) + assert not d["indices"].flags.owndata + assert not d["values"].flags.owndata + + +def test_distinct_buffers(): + # (Importing an array given twice is tested in a subprocess, below.) + a = np.arange(3, dtype=np.uint64) + b = np.asfortranarray(np.ones((2, 3), np.bool_)) + rv = _distinct_buffers(a, b, a, b) + assert rv[0] is a + assert rv[1] is b + for orig, copy in [(a, rv[2]), (b, rv[3])]: + assert copy is not orig + assert copy.flags.owndata + assert copy.flags.f_contiguous == orig.flags.f_contiguous + assert_array_equal(copy, orig) + + +def test_ownable_copy_of_nested_subarray_values_is_made_once(): + # Values whose record dtype nests subarrays are viewed with the flattened + # dtype. A view of the copy would not own its data, so it was copied again. + point = np.dtype((np.float64, (3,))) + layered = np.dtype([("oc_k", np.int64), ("oc_pts", point, (2,))], align=True) + vals = np.zeros(2**16, layered) + vals["oc_pts"] = np.arange(vals.size * 6.0).reshape(-1, 2, 3) + started = not tracemalloc.is_tracing() + if started: # pragma: no branch (safety) + tracemalloc.start() + try: + tracemalloc.reset_peak() + before = tracemalloc.get_traced_memory()[0] + arr, dtype = values_to_numpy_buffer(vals, copy=True, ownable=True) + numpy_bytes = tracemalloc.get_traced_memory()[1] - before + finally: + if started: # pragma: no branch (safety) + tracemalloc.stop() + assert numpy_bytes < 1.5 * vals.nbytes + assert dtype.np_type != layered + assert arr.dtype == dtype.np_type + assert arr.flags.owndata + assert_array_equal(arr["oc_pts"], vals["oc_pts"]) + + +# Run in a subprocess: GraphBLAS is initialized once per process, and this needs +# it initialized with libc's allocator before python-graphblas starts. +_LIBC_ALLOCATOR_PROBE = """ +import tracemalloc + +import numpy as np +import suitesparse_graphblas as ssgb + +ssgb.initialize(memory_manager="c") +import graphblas as gb +from suitesparse_graphblas.utils import empty + +print("package: " + gb.__file__) +for take_ownership in [False, True]: + values = np.arange(1000, dtype=np.float64) + v = gb.Vector.ss.import_full(values=values, take_ownership=take_ownership) + assert v.reduce().new().value == values.sum() + del v + +# GraphBLAS can't own memory NumPy allocates here, so each array is copied, or +# converted, once, straight into libc's memory, which tracemalloc doesn't see. +# A copy NumPy made first would show up as an array's worth of NumPy memory. +n = 2**20 +indptr = np.arange(n + 1) # int64, converted to uint64 +indices = np.arange(n) +values = np.arange(n, dtype=np.float64) +tracemalloc.start() +for dtype in [None, np.float32]: + for kwargs in [{}, {"take_ownership": True}, {"take_ownership": True, "secure_import": True}]: + tracemalloc.reset_peak() + before = tracemalloc.get_traced_memory()[0] + A = gb.Matrix.ss.import_csr( + nrows=n, ncols=n, indptr=indptr, col_indices=indices, values=values, dtype=dtype, + **kwargs, + ) + numpy_bytes = tracemalloc.get_traced_memory()[1] - before + assert numpy_bytes < values.nbytes // 4, (dtype, kwargs, numpy_bytes) + assert A.nvals == n, (dtype, kwargs) + assert A[n - 1, n - 1].new().value == n - 1, (dtype, kwargs) + assert values.flags.owndata, (dtype, kwargs) + del A +tracemalloc.stop() + +# Memory from GraphBLAS's allocator is still taken, not copied +values = empty(n) +values[:] = 1 +v = gb.Vector.ss.import_full(values=values, take_ownership=True) +assert not values.flags.owndata +assert v.reduce().new().value == n +del v +print("done") +""" + + +def _run_probe(source, **env): + """Run ``source`` in a new process that imports this tree's graphblas, and check it ends.""" + repo_root = Path(__file__).resolve().parents[2] + pythonpath = os.pathsep.join(filter(None, [str(repo_root), os.environ.get("PYTHONPATH")])) + env = dict(os.environ, PYTHONPATH=pythonpath, PYTHONUNBUFFERED="1", **env) + result = subprocess.run( + [sys.executable, "-c", source], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + assert result.stdout.splitlines()[-1] == "done", report + + +@pytest.mark.slow +@pytest.mark.skipif("empty is None") +def test_import_with_libc_allocator(): + # With ``memory_manager="c"`` GraphBLAS frees with libc's ``free``, so not + # even the copy that ``take_ownership=False`` makes may come from NumPy + _run_probe(_LIBC_ALLOCATOR_PROBE, PYTHONMALLOC="debug") + + +# Run in a subprocess: a regression frees an array twice, which aborts the +# process instead of failing the test. +_ONE_ARRAY_TWICE_PROBE = """ +import numpy as np + +import graphblas as gb +from graphblas import Matrix + +print("package: " + gb.__file__) +a = np.array([0, 2], np.uint64) +A = Matrix.ss.import_csr( + nrows=1, ncols=3, indptr=a, col_indices=a, values=[1.0, 2.0], take_ownership=True +) +assert A.isequal(Matrix.from_coo([0, 0], [0, 2], [1.0, 2.0], nrows=1, ncols=3)) +assert not a.flags.owndata # taken once, copied once + +a = np.array([0, 2], np.uint64) +A = Matrix.ss.import_hypercsr( + nrows=3, + ncols=3, + indptr=[0, 1, 2], + rows=a, + col_indices=a, + values=[1.0, 2.0], + take_ownership=True, +) +assert A.isequal(Matrix.from_coo([0, 2], [0, 2], [1.0, 2.0], nrows=3, ncols=3)) + +a = np.array([0, 2], np.uint64) +A = Matrix.ss.import_hypercsc( + nrows=3, + ncols=3, + indptr=[0, 1, 2], + cols=a, + row_indices=a, + values=[1.0, 2.0], + take_ownership=True, +) +assert A.isequal(Matrix.from_coo([0, 2], [0, 2], [1.0, 2.0], nrows=3, ncols=3)) +del A +print("done") +""" + + +def test_take_ownership_of_one_array_twice(): + # GraphBLAS frees every buffer it takes, so an array given as two arguments + # has to be copied for one of them. Only an index array given as the values + # was, and an index array given twice was freed twice, crashing the process. + _run_probe(_ONE_ARRAY_TWICE_PROBE) + + +# Run in a subprocess: a regression frees memory twice, which aborts the process +# at the second free or later. Arrays over 1 KiB: NumPy caches smaller blocks +# instead of freeing them, so a double free of those is silent. +_FAILED_SECURE_IMPORT_PROBE = """ +import numpy as np + +import graphblas as gb +from graphblas import Matrix, Vector +from graphblas.exceptions import GraphblasException + +print("package: " + gb.__file__) +n = 200 + + +def indices(): + # An index out of range: only the secure check reads the indices. + rv = np.zeros(n, np.uint64) + rv[-1] = n + 5 + return rv + + +def bitmap(order="C"): + # nvals does not match the bitmap: only the secure check counts it. + return {"bitmap": np.ones((n, n), np.bool_, order=order), + "values": np.ones((n, n), order=order), "nvals": 5} + + +indptr = lambda: np.arange(n + 1, dtype=np.uint64) # noqa: E731 +formats = { + "csr": lambda: {"indptr": indptr(), "col_indices": indices(), "values": np.ones(n)}, + "csc": lambda: {"indptr": indptr(), "row_indices": indices(), "values": np.ones(n)}, + "hypercsr": lambda: {"indptr": indptr(), "rows": np.arange(n, dtype=np.uint64), + "col_indices": indices(), "values": np.ones(n)}, + "hypercsc": lambda: {"indptr": indptr(), "cols": np.arange(n, dtype=np.uint64), + "row_indices": indices(), "values": np.ones(n)}, + "bitmapr": bitmap, + "bitmapc": lambda: bitmap("F"), + "sparse": lambda: {"indices": indices(), "values": np.ones(n)}, + "bitmap": lambda: {"bitmap": np.ones(n * n, np.bool_), "values": np.ones(n * n), + "nvals": 5}, +} +A = Matrix.from_coo([0, 1], [1, 0], [1.0, 2.0], nrows=n, ncols=n) +v = Vector.from_coo([1], [1.0], size=n * n) +for take_ownership in [False, True]: + for format, make in formats.items(): + for method in ["import", "pack"]: + data = make() + if format in {"sparse", "bitmap"}: + obj = v.dup() + func = getattr(obj.ss, f"{method}_{format}") + if method == "import": + data["size"] = n if format == "sparse" else n * n + else: + obj = A.dup() + func = getattr(obj.ss, f"{method}_{format}") + if method == "import": + data.update(nrows=n, ncols=n) + arrays = [val for val in data.values() if isinstance(val, np.ndarray)] + try: + func(**data, secure_import=True, take_ownership=take_ownership) + except GraphblasException: + pass + else: + raise AssertionError(f"the secure check passed: {method}_{format}") + # A secure import copies the arrays, even with take_ownership=True, so + # GraphBLAS freed only the copies: the caller's arrays still own their + # data and still hold it. + owndata = [a.flags.owndata for a in arrays] + assert all(owndata), (method, format, owndata) + fresh = [val for val in make().values() if isinstance(val, np.ndarray)] + assert all(np.array_equal(a, b) for a, b in zip(arrays, fresh)), (method, format) + # A pack leaves the object as it was. GraphBLAS freed it. + assert obj.isequal(v if format in {"sparse", "bitmap"} else A), (method, format) + del data, arrays, obj, func +for _ in range(100): # use the heap, so a corrupted one aborts here + gb.Vector.from_dense(np.ones(3000)) +print("done") +""" + + +def test_failed_secure_import_frees_once(): + # GraphBLAS moves the arrays into the object before a secure check, and + # frees them with the object when the check fails. The arrays, or the copies + # made of them, were unclaimed only after a successful call, so NumPy freed + # them a second time, and the process aborted. A secure pack also freed the + # packed object itself (a SuiteSparse bug), whose handle Python then used and + # freed again; it now checks the data with a secure import into a new object. + _run_probe(_FAILED_SECURE_IMPORT_PROBE) + + +@pytest.mark.parametrize("take_ownership", [False, True]) +def test_secure_import(take_ownership): + # Valid data passes a secure import's or pack's check, in every format. + sparse = Matrix.from_coo([0, 1, 1], [1, 0, 2], [1.0, 2.0, 3.0], nrows=2, ncols=3) + dense = Matrix.from_dense(np.arange(6.0).reshape(2, 3)) + for format in ["csr", "csc", "hypercsr", "hypercsc", "bitmapr", "bitmapc", "fullr", "fullc"]: + A = dense if format.startswith("full") else sparse + info = A.ss.export(format) + arrays = [val for val in info.values() if isinstance(val, np.ndarray)] + owndata = [array.flags.owndata for array in arrays] + B = Matrix.ss.import_any(**info, secure_import=True, take_ownership=take_ownership) + assert B.isequal(A, check_dtype=True), format + # A secure import copies, so the arrays keep their data either way. + assert [array.flags.owndata for array in arrays] == owndata, format + B.ss.pack_any(**B.ss.export(format), secure_import=True, take_ownership=take_ownership) + assert B.isequal(A, check_dtype=True), format + sparse = Vector.from_coo([1, 3], [1.0, 2.0], size=5) + dense = Vector.from_dense(np.arange(5.0)) + for format in ["sparse", "bitmap", "full"]: + v = dense if format == "full" else sparse + info = v.ss.export(format) + arrays = [val for val in info.values() if isinstance(val, np.ndarray)] + owndata = [array.flags.owndata for array in arrays] + w = Vector.ss.import_any(**info, secure_import=True, take_ownership=take_ownership) + assert w.isequal(v, check_dtype=True), format + # A secure import copies, so the arrays keep their data either way. + assert [array.flags.owndata for array in arrays] == owndata, format + w.ss.pack_any(**w.ss.export(format), secure_import=True, take_ownership=take_ownership) + assert w.isequal(v, check_dtype=True), format diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 4cea0b563..573d0f236 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1,17 +1,34 @@ +import contextlib import os -import pathlib -import platform +import subprocess import sys import sysconfig +import uuid +from pathlib import Path import numpy as np import pytest from numpy.testing import assert_array_equal import graphblas as gb -from graphblas import backend, binary, dtypes, indexunary, select, unary +from graphblas import ( + agg, + backend, + binary, + dtypes, + indexbinary, + indexunary, + monoid, + select, + semiring, + unary, +) from graphblas.core import _supports_udfs as supports_udfs +from graphblas.core.operator.indexbinary import _has_idxbinop +from graphblas.core.operator.udt_codegen import _has_jit_set from graphblas.core.ss import _IS_SSGB7 +from graphblas.core.ss.jit_config import _enable_jit_for_udt +from graphblas.exceptions import EmptyObject from .conftest import autocompute, burble @@ -26,114 +43,784 @@ pytest.skip("not suitesparse backend", allow_module_level=True) +# The fix-conda-build-paths logic lives in ``gb.ss.fix_jit_config`` so users +# can call it themselves. Tests use it through the public surface. +_fix_jit_config = gb.ss.fix_jit_config if not _IS_SSGB7 else (lambda: None) + + +def _jit_works_here(): + """Whether this environment can JIT-compile at all, whatever control the session has. + + The backend's init set ``jit_c_control`` from ``gb.config`` (``'run'`` by + default) and, under ``'on'``, probed the compiler. ``_enable_jit_for_udt`` + raises the default ``'run'`` and probes, but answers no under ``'off'``, + ``'pause'``, ``'load'`` or a ``'run'`` the user set without probing, so ask + it under ``'on'``, and put the control back. It caches the answer for the + rest of the process. Asked once here, before the autouse ``_setup_jit`` + fixture mutates the live config. + """ + if _IS_SSGB7 or backend != "suitesparse": + return False + with _jit_mode("on"): + return _enable_jit_for_udt() + + @pytest.fixture(scope="module", autouse=True) def _setup_jit(): - """Set up the SuiteSparse:GraphBLAS JIT.""" + """Set up the SuiteSparse:GraphBLAS JIT. + + Strategy: + 1. _fix_jit_config(): fix conda-baked compiler paths and probe. + - Returns True: JIT works, proceed. + - Returns False: probe failed. Whatever ``jit_c_control`` + SuiteSparse was left at is left untouched. + - Returns None: no conda env, try sysconfig instead. + 2. Sysconfig fallback: for non-conda installs (pure pip). + """ if _IS_SSGB7: # SuiteSparse JIT was added in SSGB 8 yield return - if not os.environ.get("GITHUB_ACTIONS"): - # Try to run the tests with defaults from sysconfig if not running in CI - prev = gb.ss.config["jit_c_control"] + prev = gb.ss.config["jit_c_control"] + + result = _fix_jit_config() + # ``True``: conda JIT configured and verified. ``False``: probe failed; + # don't try sysconfig, since if the conda compiler can't build GraphBLAS + # JIT kernels, Python's sysconfig compiler won't either. Both done. + if result is None: + # No conda env. Try sysconfig for non-conda installs. cc = sysconfig.get_config_var("CC") cflags = sysconfig.get_config_var("CFLAGS") include = sysconfig.get_path("include") libs = sysconfig.get_config_var("LIBS") - if not (cc is None or cflags is None or include is None or libs is None): + if cc and cflags and include: gb.ss.config["jit_c_control"] = "on" gb.ss.config["jit_c_compiler_name"] = cc gb.ss.config["jit_c_compiler_flags"] = f"{cflags} -I{include}" - gb.ss.config["jit_c_libraries"] = libs - else: - # Should we skip or try to run if sysconfig vars aren't set? - gb.ss.config["jit_c_control"] = "on" # "off" - try: - yield - finally: - gb.ss.config["jit_c_control"] = prev - return + if libs: + gb.ss.config["jit_c_libraries"] = libs - if ( - sys.platform == "darwin" - or sys.platform == "linux" - and "conda" not in gb.ss.config["jit_c_compiler_name"] - ): - # XXX TODO: tests for SuiteSparse JIT are not passing on linux when using wheels or on osx - # This should be understood and fixed! - gb.ss.config["jit_c_control"] = "off" + try: yield - return + finally: + gb.ss.config["jit_c_control"] = prev - # Configuration values below were obtained from the output of the JIT config - # in CI, but with paths changed to use `{conda_prefix}` where appropriate. - conda_prefix = os.environ["CONDA_PREFIX"] - prev = gb.ss.config["jit_c_control"] - gb.ss.config["jit_c_control"] = "on" - if sys.platform == "linux": - gb.ss.config["jit_c_compiler_name"] = f"{conda_prefix}/bin/x86_64-conda-linux-gnu-cc" - gb.ss.config["jit_c_compiler_flags"] = ( - "-march=nocona -mtune=haswell -ftree-vectorize -fPIC -fstack-protector-strong " - f"-fno-plt -O2 -ffunction-sections -pipe -isystem {conda_prefix}/include -Wundef " - "-std=c11 -lm -Wno-pragmas -fexcess-precision=fast -fcx-limited-range " - "-fno-math-errno -fwrapv -O3 -DNDEBUG -fopenmp -fPIC" - ) - gb.ss.config["jit_c_linker_flags"] = ( - "-Wl,-O2 -Wl,--sort-common -Wl,--as-needed -Wl,-z,relro -Wl,-z,now " - "-Wl,--disable-new-dtags -Wl,--gc-sections -Wl,--allow-shlib-undefined " - f"-Wl,-rpath,{conda_prefix}/lib -Wl,-rpath-link,{conda_prefix}/lib " - f"-L{conda_prefix}/lib -shared" - ) - gb.ss.config["jit_c_libraries"] = ( - f"-lm -ldl {conda_prefix}/lib/libgomp.so " - f"{conda_prefix}/x86_64-conda-linux-gnu/sysroot/usr/lib/libpthread.so" - ) - gb.ss.config["jit_c_cmake_libs"] = ( - f"m;dl;{conda_prefix}/lib/libgomp.so;" - f"{conda_prefix}/x86_64-conda-linux-gnu/sysroot/usr/lib/libpthread.so" - ) - elif sys.platform == "darwin": - gb.ss.config["jit_c_compiler_name"] = f"{conda_prefix}/bin/clang" - gb.ss.config["jit_c_compiler_flags"] = ( - "-march=core2 -mtune=haswell -mssse3 -ftree-vectorize -fPIC -fPIE " - f"-fstack-protector-strong -O2 -pipe -isystem {conda_prefix}/include -DGBNCPUFEAT " - f"-Wno-pointer-sign -O3 -DNDEBUG -fopenmp=libomp -fPIC -arch {platform.machine()}" - ) - gb.ss.config["jit_c_linker_flags"] = ( - "-Wl,-pie -Wl,-headerpad_max_install_names -Wl,-dead_strip_dylibs " - f"-Wl,-rpath,{conda_prefix}/lib -L{conda_prefix}/lib -dynamiclib" - ) - gb.ss.config["jit_c_libraries"] = f"-lm -ldl {conda_prefix}/lib/libomp.dylib" - gb.ss.config["jit_c_cmake_libs"] = f"m;dl;{conda_prefix}/lib/libomp.dylib" - elif sys.platform == "win32": # pragma: no branch (sanity) - if "mingw" in gb.ss.config["jit_c_libraries"]: - # This probably means we're testing a `python-suitesparse-graphblas` wheel - # in a conda environment. This is not yet working. - gb.ss.config["jit_c_control"] = "off" - yield - return - - gb.ss.config["jit_c_compiler_name"] = f"{conda_prefix}/bin/cc" - gb.ss.config["jit_c_compiler_flags"] = ( - '/DWIN32 /D_WINDOWS -DGBNCPUFEAT /O2 -wd"4244" -wd"4146" -wd"4018" ' - '-wd"4996" -wd"4047" -wd"4554" /O2 /Ob2 /DNDEBUG -openmp' - ) - gb.ss.config["jit_c_linker_flags"] = "/machine:x64" - gb.ss.config["jit_c_libraries"] = "" - gb.ss.config["jit_c_cmake_libs"] = "" - if not pathlib.Path(gb.ss.config["jit_c_compiler_name"]).exists(): - # Can't use the JIT if we don't have a compiler! - gb.ss.config["jit_c_control"] = "off" - yield - return +def _require_jit_on(): + """Skip the test if the SuiteSparse JIT can't compile in this environment.""" + if gb.ss.config["jit_c_control"] != "on": + pytest.skip("JIT compilation not available (probe failed or compiler missing)") + + +@contextlib.contextmanager +def _jit_mode(mode): + """Temporarily set ``jit_c_control`` to ``mode``; restore on exit.""" + prev = gb.ss.config["jit_c_control"] + gb.ss.config["jit_c_control"] = mode try: yield finally: gb.ss.config["jit_c_control"] = prev +_JIT_WORKS_AT_IMPORT = _jit_works_here() + + +def _run_in_fresh_process(script, jit_c_control=None, **extra_env): + """Run ``script`` in a new interpreter on this tree; return its ``key: value`` lines. + + The script must print ``package: `` so the child can be shown + to have imported this tree before anything it reports is trusted. The + child gets ``GRAPHBLAS_JIT_C_CONTROL=jit_c_control``, or none at all, and + ``extra_env``. + """ + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root), **extra_env) + env.pop("GRAPHBLAS_JIT_C_CONTROL", None) + if jit_c_control is not None: + env["GRAPHBLAS_JIT_C_CONTROL"] = jit_c_control + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + return lines, report + + +# Run in a subprocess: the import under test happens once per process, and by +# the time any test executes it is long past. ``jit_c_control`` has to be read +# through the C API rather than ``gb.ss.config``, because reaching for +# ``gb.ss`` is itself the thing under test. +_SS_IMPORT_PROBE = """ +import graphblas as gb + +gb.Vector # initialize the backend, but not gb.ss +from graphblas.core import ffi, lib + +def control(): + val_ptr = ffi.new("int32_t*") + assert lib.GxB_Global_Option_get_INT32(lib.GxB_JIT_C_CONTROL, val_ptr) == lib.GrB_SUCCESS + return val_ptr[0] + +before = control() +gb.ss.about["library_version"] # the line a user writes for a bug report +print("package: " + gb.__file__) +print("before: " + str(before)) +print("after: " + str(control())) +""" + + +@pytest.mark.skipif("_IS_SSGB7") +@pytest.mark.parametrize("jit_c_control", [None, "run", "on", "pause"]) +def test_ss_import_leaves_jit_c_control_alone(jit_c_control): + """Reading from ``gb.ss`` must not change the numerical machinery. + + ``gb.ss`` is reached by attribute access, so importing it has to be inert. + It used to raise ``jit_c_control`` to ``'on'``, which decides whether + SuiteSparse loads kernels from its on-disk cache, so a line as harmless as + ``print(gb.ss.about["library_version"])`` changed how every later + operation in the process was computed. The control is set by the + backend's init, from ``gb.config`` (here ``GRAPHBLAS_JIT_C_CONTROL``), + and ``'on'`` becomes ``'run'`` where the JIT cannot compile. Not set, it + is ``'run'``. + """ + lines, report = _run_in_fresh_process(_SS_IMPORT_PROBE, jit_c_control) + assert lines["after"] == lines["before"], report + if jit_c_control != "on": + expected = getattr(gb.core.lib, f"GxB_JIT_{(jit_c_control or 'run').upper()}") + assert lines["before"] == str(expected), report + elif _JIT_WORKS_AT_IMPORT: + assert lines["before"] == str(gb.core.lib.GxB_JIT_ON), report + else: + # psg's wheels (built without the JIT), or no working C compiler + assert lines["before"] == str(gb.core.lib.GxB_JIT_RUN), report + + +_BAD_JIT_C_CONTROL_PROBE = """ +import graphblas as gb +import suitesparse_graphblas as ssgb + +print("package: " + gb.__file__) +try: + gb.Vector +except ValueError as exc: + print("error: " + str(exc)) +print("initialized: " + str(ssgb.is_initialized())) +""" + + +def test_bad_jit_c_control_config_fails_init_before_it_starts(): + """A bad ``GRAPHBLAS_JIT_C_CONTROL`` says so, and leaves SuiteSparse uninitialized.""" + lines, report = _run_in_fresh_process(_BAD_JIT_C_CONTROL_PROBE, "of") + assert "GRAPHBLAS_JIT_C_CONTROL" in lines.get("error", ""), report + assert "'of'" in lines["error"], report + assert lines["initialized"] == "False", report + + +@pytest.mark.skipif("_IS_SSGB7") +def test_set_jit_c_control_at_init(monkeypatch): + """Init sets each configured value as given, and ``'on'`` only where the JIT compiles. + + Where it cannot, it sets ``'run'`` with no error: a library built without + the JIT, no compiler on disk, a cache it cannot write to, or a compiler + that fails the probe. Nothing configured (None) sets ``'run'`` too, and + only that ``'run'`` is left for UDT work to raise. Called directly here, + as init calls it. + """ + from graphblas.core.ss import jit_config + from graphblas.core.ss.jit_config import _set_jit_c_control_at_init + + cfg = gb.ss.config + prev_control = cfg["jit_c_control"] + prev_cc = cfg["jit_c_compiler_name"] + for name in ["_jit_enabled_for_udt", "_configured_jit_c_control", "_user_set_run"]: + monkeypatch.setattr(jit_config, name, getattr(jit_config, name)) + try: + for control in ["load", "run", "pause", "off"]: + _set_jit_c_control_at_init(control) + # A library built without the JIT clamps 'load' down to 'run'. + expected = {control, "run"} if control == "load" else {control} + assert cfg["jit_c_control"] in expected + assert jit_config._configured_jit_c_control == control + assert jit_config._user_set_run is (control == "run") + _set_jit_c_control_at_init(None) + assert cfg["jit_c_control"] == "run" + assert jit_config._configured_jit_c_control is None + assert jit_config._user_set_run is False + _set_jit_c_control_at_init("on") + assert jit_config._configured_jit_c_control == "on" + assert cfg["jit_c_control"] == ("on" if jit_config._jit_enabled_for_udt else "run") + if _JIT_WORKS_AT_IMPORT: + assert cfg["jit_c_control"] == "on" + # A cache it cannot write to: the cached probe would load, and the first + # new kernel would fail to compile. + with monkeypatch.context() as m: + m.setattr(jit_config, "_cache_is_writable", lambda cfg: False) + _set_jit_c_control_at_init("on") + assert cfg["jit_c_control"] == "run" + # A "compiler" that exists but cannot compile anything. The probe's type + # is named after the toolchain, so the probe kernel this process has + # already loaded does not answer for this one. SuiteSparse built with + # MSVC always compiles through cmake, whose Visual Studio generator + # chooses the compiler itself, so this setting cannot fail the probe. + if not cfg["jit_use_cmake"]: + cfg["jit_c_compiler_name"] = str(Path(__file__).parent) + _set_jit_c_control_at_init("on") + assert cfg["jit_c_control"] == "run" + assert jit_config._jit_enabled_for_udt is False + # No compiler on disk and nothing to repair it from: no probe at all. + cfg["jit_c_compiler_name"] = "/nonexistent/path/to/cc" + probes = [] + monkeypatch.setattr(jit_config, "_repair_jit_compiler", lambda cfg: None) + monkeypatch.setattr(jit_config, "_probe_jit", probes.append) + _set_jit_c_control_at_init("on") + assert cfg["jit_c_control"] == "run" + assert jit_config._jit_enabled_for_udt is False + assert probes == [] + finally: + cfg["jit_c_compiler_name"] = prev_cc + cfg["jit_c_control"] = prev_control + + +# Run in a subprocess: this is about the very first call in a process, and by +# the time any test executes, plenty of JIT setup has already happened. +_FIX_JIT_TWICE_PROBE = """ +import graphblas as gb + +print("package: " + gb.__file__) +print("first: " + str(gb.ss.fix_jit_config())) +print("second: " + str(gb.ss.fix_jit_config())) +""" + + +@pytest.mark.skipif("_IS_SSGB7") +def test_fix_jit_config_first_call_agrees_with_the_second(): + """``fix_jit_config()`` must report on the environment, not on itself. + + Its probe once registered a UDT, and registering a UDT asks to enable the + JIT, which started a second probe inside the first. The inner probe + registered the probe type and the outer one then died redeclaring it, so a + working environment answered ``False`` the first time and ``True`` every + time after. The probe now goes through the C API and registers nothing. + """ + lines, report = _run_in_fresh_process(_FIX_JIT_TWICE_PROBE) + assert lines["first"] == lines["second"], report + + +# Run in a subprocess: this process has already probed the JIT and cached the +# answer, and a fresh one has to probe at its first registration. Each +# registration starts from the default ``'run'`` (None), so one registrar +# raising the control cannot cover for another that does not. +_SS_REGISTER_NEW_PROBE = """ +import graphblas as gb +from graphblas import binary, indexbinary, indexunary, select, unary +from graphblas.core.operator.indexbinary import _has_idxbinop + +print("package: " + gb.__file__) +registrars = { + "unary": lambda: unary.ss.register_new( + "rn_neg", "void rn_neg (double *z, double *x) { (*z) = -(*x) ; }", "FP64", "FP64" + ), + "binary": lambda: binary.ss.register_new( + "rn_absdiff", + "void rn_absdiff (double *z, double *x, double *y) { (*z) = fabs ((*x) - (*y)) ; }", + "FP64", "FP64", "FP64", + ), + "indexunary": lambda: indexunary.ss.register_new( + "rn_diffy", + "void rn_diffy (double *z, double *x, GrB_Index i, GrB_Index j, double *y) " + "{ (*z) = (i + j) * fabs ((*x) - (*y)) ; }", + "FP64", "FP64", "FP64", + ), + "select": lambda: select.ss.register_new( + "rn_woot", + "void rn_woot (bool *z, const int32_t *x, GrB_Index i, GrB_Index j, int32_t *y) " + "{ (*z) = ((*x) + i + j == (*y)) ; }", + "INT32", "INT32", + ), +} +if _has_idxbinop: + registrars["indexbinary"] = lambda: indexbinary.ss.register_new( + "rn_add_theta", + "void rn_add_theta (double *z, double *x, GrB_Index ix, GrB_Index jx, " + "double *y, GrB_Index iy, GrB_Index jy, double *theta) " + "{ (*z) = (*x) + (*y) + (*theta) ; }", + "FP64", "FP64", "FP64", "FP64", + ) +for kind, register in registrars.items(): + gb.ss.config["jit_c_control"] = None + try: + register() + except Exception as exc: + print(kind + ": " + type(exc).__name__) + else: + print(kind + ": ok") +""" + + +@pytest.mark.skipif("_IS_SSGB7") +def test_ss_register_new_enables_the_jit(): + """Registering an op from C source raises the default ``'run'``. + + SuiteSparse compiles the C definition, or loads it from its cache, when the + op is created. ``'run'`` does neither, so without the raise the + registration fails. Each registrar has to ask for the JIT itself: since + ``gb.ss`` stopped raising the control at import, nothing else does it + first. + """ + if not _JIT_WORKS_AT_IMPORT: + pytest.skip("JIT not compiled in or not working here; nothing can compile") + lines, report = _run_in_fresh_process(_SS_REGISTER_NEW_PROBE) + kinds = ["unary", "binary", "indexunary", "select"] + kinds += ["indexbinary"] if _has_idxbinop else [] + assert {kind: lines.get(kind) for kind in kinds} == dict.fromkeys(kinds, "ok"), report + + +@pytest.mark.skipif("_IS_SSGB7") +def test_ss_register_new_says_why_the_jit_cannot_build_it(monkeypatch): + """A C-source registration the JIT cannot build names the reason. + + For each of these SuiteSparse reports a bare ``NullPointer`` (op) or + ``InvalidValue`` (type) with no message. None of them compiles anything, so + this runs with or without a working compiler; a library built without the + JIT skips ``'load'`` and the failed probe. + """ + from graphblas.core.ss import jit_config + from graphblas.exceptions import InvalidValue, NullPointer + + # New names every run, so no cached kernel can satisfy ``'load'``. + suffix = uuid.uuid4().hex[:8] + op_name = f"why_neg_{suffix}" + cdef = f"void {op_name} (double *z, double *x) {{ (*z) = -(*x) ; }}" + type_name = f"why_t_{suffix}" + tdef = f"typedef struct {{ int64_t why_a ; }} {type_name} ;" + prev = gb.ss.config["jit_c_control"] + try: + for control, match in [ + ("off", "jit_c_control is 'off', so SuiteSparse will not compile it"), + ("load", "jit_c_control is 'load', so SuiteSparse only loads kernels"), + ]: + gb.ss.config["jit_c_control"] = control + if gb.ss.config["jit_c_control"] != control: + # A library built without the JIT clamps 'load' down to 'run'. + continue + with pytest.raises(NullPointer, match=match): + unary.ss.register_new(op_name, cdef, "FP64", "FP64") + with pytest.raises(InvalidValue, match=match): + dtypes.ss.register_new(type_name, tdef) + # A process that cannot compile leaves the default 'run' in place. + monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", False) + gb.ss.config["jit_c_control"] = None + with pytest.raises(NullPointer, match="jit_c_control is still 'run'"): + unary.ss.register_new(op_name, cdef, "FP64", "FP64") + # A probe that fails in this very call leaves 'load' before the op is created. + gb.ss.config["jit_c_control"] = "on" + if gb.ss.config["jit_c_control"] == "on": + + def failed_probe(cfg): + cfg["jit_c_control"] = "load" + return False + + monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", None) + monkeypatch.setattr(jit_config, "jit_compiler_is_usable", lambda: True) + monkeypatch.setattr(jit_config, "_probe_jit", failed_probe) + gb.ss.config["jit_c_control"] = None + with pytest.raises(NullPointer, match="the C compiler does not work here"): + unary.ss.register_new(op_name, cdef, "FP64", "FP64") + finally: + gb.ss.config["jit_c_control"] = prev + + +@pytest.mark.skipif("_IS_SSGB7") +def test_ss_register_new_under_a_chosen_run_says_how_to_opt_in(monkeypatch): + """C source cannot be built under a ``'run'`` the user set, and the error says so. + + UDT work raises only the default ``'run'``. One that was set, often to keep + processes from compiling into one cache at the same time, is kept, so the + registration fails, and the message has to say why and what to set instead. + Here the user configured ``'on'`` before init and writes ``'run'`` later. + """ + from graphblas.core.ss import jit_config + from graphblas.exceptions import InvalidValue, NullPointer + + monkeypatch.setattr(jit_config, "_configured_jit_c_control", "on") + monkeypatch.setattr(jit_config, "_user_set_run", jit_config._user_set_run) + suffix = uuid.uuid4().hex[:8] + op_name = f"chosen_neg_{suffix}" + type_name = f"chosen_t_{suffix}" + prev = gb.ss.config["jit_c_control"] + try: + gb.ss.config["jit_c_control"] = "run" + assert jit_config._user_set_run is True + match = "jit_c_control was set to 'run' .*GRAPHBLAS_JIT_C_CONTROL.* = 'on', or None" + with pytest.raises(NullPointer, match=match): + unary.ss.register_new( + op_name, + f"void {op_name} (double *z, double *x) {{ (*z) = -(*x) ; }}", + "FP64", + "FP64", + ) + with pytest.raises(InvalidValue, match=match): + dtypes.ss.register_new(type_name, f"typedef struct {{ int64_t c_a ; }} {type_name} ;") + assert gb.ss.config["jit_c_control"] == "run" + finally: + gb.ss.config["jit_c_control"] = prev + + +# Run in a subprocess: the forms that set the control are read once, at init. +# ``form`` is defined by a line put in front of this. +_UDT_WORK_PROBE = """ +import uuid + +import numpy as np + +import graphblas as gb + +if form == "gb.config.set": + gb.config.set(jit_c_control="run") +elif form == "gb.ss.config": + gb.ss.config["jit_c_control"] = "run" +print("package: " + gb.__file__) +print("before: " + gb.ss.config["jit_c_control"]) +name = "uw_" + uuid.uuid4().hex[:8] +try: + gb.dtypes.ss.register_new(name, "typedef struct { int64_t uw_a ; } " + name + " ;") +except Exception as exc: + print("register_new: " + type(exc).__name__ + ": " + str(exc)) +else: + print("register_new: ok") +print("after register_new: " + gb.ss.config["jit_c_control"]) +try: + import numba +except ImportError: + numba = None +if numba is not None: + udt = gb.dtypes.register_anonymous( + np.dtype([("uw_x", np.int64), ("uw_y", np.int64)], align=True), "UdtWorkProbe" + ) + v = gb.Vector.from_coo([0, 1], np.array([(1, 2), (3, 4)], dtype=udt.np_type), dtype=udt) + w = v.ewise_add(v, gb.binary.plus).new() + print("udt op: " + str(w[1].new().value.tolist())) + print("after udt op: " + gb.ss.config["jit_c_control"]) +""" + + +@pytest.mark.skipif("_IS_SSGB7") +@pytest.mark.parametrize( + ("form", "env_control", "kept"), + [ + pytest.param("default", None, False, id="default"), + pytest.param("GRAPHBLAS_JIT_C_CONTROL", "run", True, id="env-run"), + pytest.param("gb.config.set", None, True, id="config-set-run"), + pytest.param("YAML file", None, True, id="yaml-run"), + # With nothing configured, a 'run' written later is the default again. + pytest.param("gb.ss.config", None, False, id="ss-config-run"), + # Written over a configured 'on', it changes the control, so it counts. + pytest.param("gb.ss.config", "on", True, id="env-on-then-ss-config-run"), + pytest.param("GRAPHBLAS_JIT_C_CONTROL", "off", True, id="env-off"), + pytest.param("GRAPHBLAS_JIT_C_CONTROL", "pause", True, id="env-pause"), + pytest.param("GRAPHBLAS_JIT_C_CONTROL", "load", True, id="env-load"), + ], +) +def test_udt_work_raises_only_the_default_run(form, env_control, kept, tmp_path): + """UDT work raises the default ``'run'`` to ``'on'``, and keeps a control the user set. + + A ``'run'`` set before init (``GRAPHBLAS_JIT_C_CONTROL``, ``gb.config.set`` + or a YAML config file) keeps processes from compiling, at the cost of UDT + ops running their Numba functions and C source failing to register. One + written to ``gb.ss.config`` later counts only when it changes the configured + control. ``'off'``, ``'pause'`` and ``'load'`` were kept before the default + was told apart. + """ + if not kept and not _JIT_WORKS_AT_IMPORT: + pytest.skip("JIT not compiled in or not working here; nothing raises to 'on'") + extra_env = {} + if form == "YAML file": + (tmp_path / "jit.yaml").write_text("jit_c_control: run\n") + extra_env["GRAPHBLAS_CONFIG"] = str(tmp_path) + script = f"form = {form!r}\n" + _UDT_WORK_PROBE + lines, report = _run_in_fresh_process(script, env_control, **extra_env) + control = env_control if form == "GRAPHBLAS_JIT_C_CONTROL" else "run" + if not kept: + assert lines["before"] == "run", report + assert lines["register_new"] == "ok", report + assert lines["after register_new"] == "on", report + expected = "on" + else: + # A library built without the JIT clamps 'load' down to 'run'. + expected = lines["before"] + assert expected in {control, "run"}, report + assert lines["register_new"].startswith("InvalidValue: Cannot create"), report + if control == "run": + assert "was set to 'run'" in lines["register_new"], report + assert lines["after register_new"] == expected, report + if numba is not None: + assert lines["udt op"] == "(6, 8)", report + assert lines["after udt op"] == expected, report + + +@pytest.mark.skipif("_IS_SSGB7") +def test_ss_register_new_says_the_compile_failed(): + """C source that does not compile says so, and that ``jit_c_control`` is now ``'load'``.""" + if not _JIT_WORKS_AT_IMPORT: + pytest.skip("JIT not compiled in or not working here; nothing compiles to fail") + from graphblas.exceptions import JitError, NullPointer + + cdef = "void why_bad (double *z, double *x) { (*z) = why_undeclared ; }" + prev = gb.ss.config["jit_c_control"] + try: + gb.ss.config["jit_c_control"] = "on" + # SuiteSparse reports a failed compile as JitError from 9.4 on. + with pytest.raises((NullPointer, JitError), match="SuiteSparse failed to compile it"): + unary.ss.register_new("why_bad", cdef, "FP64", "FP64") + assert gb.ss.config["jit_c_control"] == "load" + finally: + gb.ss.config["jit_c_control"] = prev + + +@pytest.mark.skipif("_IS_SSGB7") +def test_auto_fix_jit_at_import_left_compiler_usable(): + """The import-time repair leaves a compiler the UDT path can actually use.""" + if not gb.ss.jit_compiler_is_usable(): + pytest.skip("sandboxed env without a usable compiler; auto-fix had nothing to repair") + if not _JIT_WORKS_AT_IMPORT: + # Compiler file exists but the probe failed (e.g., the baked-in flags + # target a different arch than the host). After a compile failure + # SuiteSparse drops ``jit_c_control`` to a non-compiling mode so + # downstream ops punt to generic cleanly; the probe absorbs that first + # failure. Which mode it lands in varies by SuiteSparse version. + assert gb.ss.config["jit_c_control"] in {"load", "run"} + pytest.skip("compiler present but JIT probe failed in this env") + assert _JIT_WORKS_AT_IMPORT + assert gb.ss.config["jit_c_control"] == "on" + + +@pytest.mark.skipif("_IS_SSGB7") +def test_enable_jit_for_udt_re_raises_a_lowered_control(monkeypatch): + """A cached "yes" must still raise a lowered ``'run'`` before arming a kernel. + + The cache answers "can this process compile?", which is settled once. The + control is separate state that anything may have moved since, so returning + the cached answer without re-raising it would leave an op armed with C + source that SuiteSparse never builds, and nothing would warn: codegen + succeeded. + + Driven through a stand-in config, so the assertion holds on every host: a + library built without the JIT clamps writes down to ``'run'``, and would + never read back the ``'on'`` this is about. The live round trip is pinned + separately below, where the clamp is not in the way. + """ + from graphblas.core.ss import jit_config + + fake = {} + monkeypatch.setattr(jit_config, "_ss_config", lambda: fake) + monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", True) + monkeypatch.setattr(jit_config, "_user_set_run", False) + # An explicit opt-out is a choice, not the default to raise, and ``'load'`` is + # also where SuiteSparse parks the control after a compile fails. + for start, expected in (("run", "on"), ("load", "load"), ("off", "off"), ("pause", "pause")): + fake["jit_c_control"] = start + assert _enable_jit_for_udt() is (expected == "on") + assert fake["jit_c_control"] == expected + # A 'run' the user set is an opt-out too. + jit_config._user_set_run = True + fake["jit_c_control"] = "run" + assert _enable_jit_for_udt() is False + assert fake["jit_c_control"] == "run" + + +@pytest.mark.skipif("_IS_SSGB7") +def test_enable_jit_for_udt_does_not_cache_an_opt_out(monkeypatch): + """An ``'off'`` at the first UDT op must not decide the rest of the process. + + It says nothing about whether this process can compile. Caching a "no" + for it would leave the default ``'run'``, restored afterwards, unraised for + good, so every later UDT op would run without the JIT and nothing would + say so. The same goes for ``'pause'``, ``'load'``, and a ``'run'`` the user + chose. + """ + from graphblas.core.ss import jit_config + + fake = {} + monkeypatch.setattr(jit_config, "_ss_config", lambda: fake) + monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", None) + monkeypatch.setattr(jit_config, "jit_compiler_is_usable", lambda: True) + monkeypatch.setattr(jit_config, "_probe_jit", lambda cfg: True) + monkeypatch.setattr(jit_config, "_cache_is_writable", lambda cfg: True) + for start in ("off", "pause", "load", "run"): + # The last is a 'run' the user set. + monkeypatch.setattr(jit_config, "_user_set_run", start == "run") + fake["jit_c_control"] = start + assert _enable_jit_for_udt() is False + assert fake["jit_c_control"] == start + assert jit_config._jit_enabled_for_udt is None + monkeypatch.setattr(jit_config, "_user_set_run", False) + # A cache it cannot write to: the default 'run' stays, and the answer is cached. + fake["jit_c_control"] = "run" + with monkeypatch.context() as m: + m.setattr(jit_config, "_cache_is_writable", lambda cfg: False) + assert _enable_jit_for_udt() is False + assert fake["jit_c_control"] == "run" + assert jit_config._jit_enabled_for_udt is False + monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", None) + fake["jit_c_control"] = "run" + assert _enable_jit_for_udt() is True + assert fake["jit_c_control"] == "on" + + +@pytest.mark.skipif("_IS_SSGB7") +def test_enable_jit_for_udt_re_raises_the_live_control(monkeypatch): + """The same raise, through SuiteSparse's own config rather than a stand-in. + + A ``'run'`` written to ``gb.ss.config`` is the default, and raised, while + nothing is configured; written over a configured ``'on'``, it is the user's, + and kept. None is the default either way. + """ + if not _JIT_WORKS_AT_IMPORT: + pytest.skip("JIT not compiled in or not working here; the control need not hold 'on'") + + from graphblas.core.ss import jit_config + + prev_control = gb.ss.config["jit_c_control"] + prev_enabled = jit_config._jit_enabled_for_udt + try: + jit_config._jit_enabled_for_udt = True + monkeypatch.setattr(jit_config, "_configured_jit_c_control", None) + monkeypatch.setattr(jit_config, "_user_set_run", jit_config._user_set_run) + gb.ss.config["jit_c_control"] = "run" + assert _enable_jit_for_udt() is True + assert gb.ss.config["jit_c_control"] == "on" + monkeypatch.setattr(jit_config, "_configured_jit_c_control", "on") + gb.ss.config["jit_c_control"] = "RUN" + assert jit_config._user_set_run is True + assert _enable_jit_for_udt() is False + assert gb.ss.config["jit_c_control"] == "run" + gb.ss.config["jit_c_control"] = None + assert gb.ss.config["jit_c_control"] == "run" + assert jit_config._user_set_run is False + assert _enable_jit_for_udt() is True + assert gb.ss.config["jit_c_control"] == "on" + gb.ss.config["jit_c_control"] = gb.core.lib.GxB_JIT_RUN + assert jit_config._user_set_run is True + finally: + gb.ss.config["jit_c_control"] = prev_control + jit_config._jit_enabled_for_udt = prev_enabled + + +@pytest.mark.skipif("not supports_udfs") +# The op this test arms, binary.plus on a UDT, carries C source only on SS 9+ +# (see test_udt_jit_c_source_introspection). On 8.x the stand-in config tests +# above still check that _enable_jit_for_udt leaves 'load' alone. +@pytest.mark.skipif("not _has_jit_set") +def test_failed_compile_is_not_retried_at_the_next_arming(): + """A compile that fails is tried once, not again at every later UDT op. + + SuiteSparse lowers ``jit_c_control`` to ``'load'`` after a failed compile + so the failure is not repeated. Arming another op must leave it there: + raising it back to ``'on'`` would run the failing compiler again at the + next kernel, and on SuiteSparse 9.4 and later raise ``JitError`` again, + with ``'off'`` the only way out. + """ + if not _JIT_WORKS_AT_IMPORT: + pytest.skip("JIT not compiled in or not working here; nothing compiles to fail") + from graphblas.core.ss import jit_config + from graphblas.exceptions import JitError + + # Field names unique to this test: ``register_anonymous`` caches one + # DataType per ``np.dtype``. + udt = dtypes.register_anonymous( + np.dtype([("nrt_a", np.int64), ("nrt_b", np.int64)], align=True), "_NoRetryUdt" + ) + v = Vector.from_coo([0, 1], np.array([(1, 2), (3, 4)], dtype=udt.np_type), dtype=udt) + # Before 9.4 SuiteSparse does not report a failed compile: it returns + # ``GrB_NO_VALUE`` ("use the generic kernel"), though it still lowers the + # control. + reported = gb.ss.about["library_version"] >= (9, 4) + prev_cc = gb.ss.config["jit_c_compiler_name"] + prev_control = gb.ss.config["jit_c_control"] + prev_enabled = jit_config._jit_enabled_for_udt + try: + # Start from the default, so the first arming shows the raise. + gb.ss.config["jit_c_control"] = None + # A compiler that does not exist fails under every toolchain, where a + # flag it does not know is only a warning to MSVC. The cache is pinned + # to yes (true here) so the arming does not repair the path first. + gb.ss.config["jit_c_compiler_name"] = "/nonexistent/path/to/cc" + jit_config._jit_enabled_for_udt = True + plus = binary.plus[udt] + assert plus.jit_c_source is not None + assert gb.ss.config["jit_c_control"] == "on" + with pytest.raises(JitError) if reported else contextlib.nullcontext(): + v.ewise_add(v, plus).new() + assert gb.ss.config["jit_c_control"] == "load" + assert binary.minus[udt].jit_c_source is not None + assert gb.ss.config["jit_c_control"] == "load" + result = v.ewise_add(v, plus).new() # the generic kernel, without raising + assert tuple(result[1].new().value) == (6, 8) + finally: + gb.ss.config["jit_c_compiler_name"] = prev_cc + gb.ss.config["jit_c_control"] = prev_control + jit_config._jit_enabled_for_udt = prev_enabled + + +@pytest.mark.skipif("_IS_SSGB7") +def test_public_fix_jit_config_repairs_a_broken_compiler(): + """``gb.ss.fix_jit_config()`` repairs a clobbered compiler path and returns ``True``.""" + if not _JIT_WORKS_AT_IMPORT: + pytest.skip("env JIT doesn't actually compile; nothing to repair to") + prev = { + "control": gb.ss.config["jit_c_control"], + "cc": gb.ss.config["jit_c_compiler_name"], + "flags": gb.ss.config["jit_c_compiler_flags"], + } + if not os.environ.get("CONDA_PREFIX"): + pytest.skip("test requires a CONDA_PREFIX to repair from") + try: + # Force a broken path; the auto-fix must put back something usable. + gb.ss.config["jit_c_compiler_name"] = "/nonexistent/path/to/cc" + assert not gb.ss.jit_compiler_is_usable() + result = gb.ss.fix_jit_config() + assert result is True + assert gb.ss.jit_compiler_is_usable() + # Calling fix_jit_config a second time must not disable JIT. The probe + # used to register a one-shot ``_jit_probe`` UDT, and took the + # duplicate-name error as a probe failure. + gb.ss.config["jit_c_compiler_name"] = "/nonexistent/path/to/cc" + result2 = gb.ss.fix_jit_config() + assert result2 is True + assert gb.ss.config["jit_c_control"] == "on" + finally: + gb.ss.config["jit_c_control"] = prev["control"] + gb.ss.config["jit_c_compiler_name"] = prev["cc"] + gb.ss.config["jit_c_compiler_flags"] = prev["flags"] + + +@pytest.mark.skipif("_IS_SSGB7") +def test_public_fix_jit_config_returns_none_without_env(): + """With no ``CONDA_PREFIX`` and ``use_sysconfig=False``, the helper signals no-environment.""" + if os.environ.get("CONDA_PREFIX"): + # Drop the env for this call only. + prev_env = os.environ.pop("CONDA_PREFIX") + else: + prev_env = None + prev_cc = gb.ss.config["jit_c_compiler_name"] + try: + gb.ss.config["jit_c_compiler_name"] = "/nonexistent/path/to/cc" + result = gb.ss.fix_jit_config(use_sysconfig=False) + assert result is None + finally: + if prev_env is not None: + os.environ["CONDA_PREFIX"] = prev_env + gb.ss.config["jit_c_compiler_name"] = prev_cc + + @pytest.fixture def v(): return Vector.from_coo([1, 3, 4, 6], [1, 1, 2, 0]) @@ -147,8 +834,7 @@ def test_jit_udt(): "myquaternion", "typedef struct { float x [4][4] ; int color ; } myquaternion ;" ) return - if gb.ss.config["jit_c_control"] == "off": - return + _require_jit_on() with burble(): dtype = dtypes.ss.register_new( "myquaternion", "typedef struct { float x [4][4] ; int color ; } myquaternion ;" @@ -189,21 +875,55 @@ def test_jit_udt(): assert dtype.np_type == np_type +@pytest.mark.skipif("_IS_SSGB7 or numba is None") +def test_jit_udt_np_type_numba_cannot_type(): + # The np_type gets the error register_anonymous gives for an array Numba cannot type + _require_jit_on() + np_type = np.dtype([("x", np.bool_, (2,))]) + with pytest.raises(ValueError, match="cannot type an array of bool") as exc_info: + dtypes.ss.register_new( + "_jit_bool_pair", "typedef struct { uint8_t x [2] ; } _jit_bool_pair ;", np_type=np_type + ) + assert isinstance(exc_info.value.__cause__, AttributeError) + assert not hasattr(dtypes.ss, "_jit_bool_pair") + # The failure leaves the name free, so it registers with an np_type Numba can type. + dtype = dtypes.ss.register_new( + "_jit_bool_pair", + "typedef struct { uint8_t x [2] ; } _jit_bool_pair ;", + np_type=np.dtype([("x", np.uint8, (2,))]), + ) + assert dtypes.ss._jit_bool_pair is dtype + + +@pytest.mark.skipif("_IS_SSGB7") +def test_jit_udt_name_leaves_room_for_the_nul(): + """A type name as long as ``GxB_MAX_NAME_LEN`` is rejected before SuiteSparse sees it. + + SuiteSparse keeps one character fewer, so that name would no longer match + its typedef, and the registration would fail as a compile error. + """ + from graphblas.core import lib + + limit = lib.GxB_MAX_NAME_LEN + name = "n" * limit + with pytest.raises(ValueError, match=f"Max size is {limit - 1}; got {limit}"): + dtypes.ss.register_new(name, f"typedef struct {{ int64_t a ; }} {name} ;") + + def test_jit_unary(v): cdef = "void square (float *z, float *x) { (*z) = (*x) * (*x) ; } ;" if _IS_SSGB7: with pytest.raises(RuntimeError, match="JIT was added"): unary.ss.register_new("square", cdef, "FP32", "FP32") return - if gb.ss.config["jit_c_control"] == "off": - return + _require_jit_on() with burble(): square = unary.ss.register_new("square", cdef, "FP32", "FP32") assert not hasattr(unary, "square") assert unary.ss.square is square assert square.name == "ss.square" assert square.types == {dtypes.FP32: dtypes.FP32} - # The JIT is unforgiving and does not coerce--use the correct types! + # JIT ops don't coerce: wrong dtype raises KeyError, not a silent cast. with pytest.raises(KeyError, match="square does not work with INT64"): v << square(v) v = v.dup("FP32") @@ -227,14 +947,14 @@ def test_jit_unary(v): unary.ss.register_new("nested", cdef, "FP32", "FP32") +@pytest.mark.slow def test_jit_binary(v): cdef = "void absdiff (double *z, double *x, double *y) { (*z) = fabs ((*x) - (*y)) ; }" if _IS_SSGB7: with pytest.raises(RuntimeError, match="JIT was added"): binary.ss.register_new("absdiff", cdef, "FP64", "FP64", "FP64") return - if gb.ss.config["jit_c_control"] == "off": - return + _require_jit_on() with burble(): absdiff = binary.ss.register_new( "absdiff", @@ -249,7 +969,7 @@ def test_jit_binary(v): assert absdiff.types == {(dtypes.FP64, dtypes.FP64): dtypes.FP64} # different than normal assert "FP64" in absdiff assert absdiff["FP64"].return_type == dtypes.FP64 - # The JIT is unforgiving and does not coerce--use the correct types! + # JIT ops don't coerce: wrong dtype raises KeyError, not a silent cast. with pytest.raises(KeyError, match="absdiff does not work with .INT64, INT64. types"): v << absdiff(v & v) w = (v - 1).new("FP64") @@ -292,6 +1012,7 @@ def test_jit_binary(v): assert ("FP32", "FP64") not in absdiff +@pytest.mark.slow def test_jit_indexunary(v): cdef = ( "void diffy (double *z, double *x, GrB_Index i, GrB_Index j, double *y) " @@ -301,8 +1022,7 @@ def test_jit_indexunary(v): with pytest.raises(RuntimeError, match="JIT was added"): indexunary.ss.register_new("diffy", cdef, "FP64", "FP64", "FP64") return - if gb.ss.config["jit_c_control"] == "off": - return + _require_jit_on() with burble(): diffy = indexunary.ss.register_new("diffy", cdef, "FP64", "FP64", "FP64") assert not hasattr(indexunary, "diffy") @@ -313,7 +1033,7 @@ def test_jit_indexunary(v): assert diffy.types == {(dtypes.FP64, dtypes.FP64): dtypes.FP64} assert "FP64" in diffy assert diffy["FP64"].return_type == dtypes.FP64 - # The JIT is unforgiving and does not coerce--use the correct types! + # JIT ops don't coerce: wrong dtype raises KeyError, not a silent cast. with pytest.raises(KeyError, match="diffy does not work with .INT64, INT64. types"): v << diffy(v, 1) v = v.dup("FP64") @@ -356,9 +1076,170 @@ def test_jit_indexunary(v): assert ("FP32", "FP64") not in diffy +@pytest.mark.slow +@pytest.mark.skipif(not _has_idxbinop, reason="requires SuiteSparse:GraphBLAS 9.4+") +def test_jit_indexbinary(v): + cdef = ( + "void add_theta (double *z, double *x, GrB_Index ix, GrB_Index jx, " + "double *y, GrB_Index iy, GrB_Index jy, double *theta) " + "{ (*z) = (*x) + (*y) + (*theta) ; }" + ) + _require_jit_on() + with burble(): + add_theta = indexbinary.ss.register_new("add_theta", cdef, "FP64", "FP64", "FP64", "FP64") + assert not hasattr(indexbinary, "add_theta") + assert indexbinary.ss.add_theta is add_theta + assert add_theta.name == "ss.add_theta" + assert add_theta.types == {(dtypes.FP64, dtypes.FP64, dtypes.FP64): dtypes.FP64} + assert "FP64" in add_theta + assert add_theta["FP64"].return_type == dtypes.FP64 + assert add_theta["FP64"].jit_c_definition == cdef + # SuiteSparse compiles it under a C name of its types, name and definition + c_name = add_theta["FP64"].jit_c_name + assert c_name.startswith("add_theta_4_FP64_4_FP64_4_FP64_") + assert len(c_name) == len("add_theta_4_FP64_4_FP64_4_FP64_") + 12 + assert add_theta["FP64"].jit_c_source == cdef.replace("add_theta", c_name) + # Bind theta and use as BinaryOp + v64 = v.dup("FP64") + binop = add_theta["FP64"](10.0) + assert binop.opclass == "BinaryOp" + res = v64.ewise_mult(v64, binop).new() + # v has values at [1, 3, 4, 6] with vals [1, 1, 2, 0] + # ewise_mult: x+y+theta = 2*val + 10 + expected = Vector.from_coo([1, 3, 4, 6], [12.0, 12.0, 14.0, 10.0], dtype="FP64") + assert expected.isequal(res) + # Bound with no type, it runs at the inputs' types and theta's, which a JIT op must + # have: like a JIT IndexUnaryOp's thunk, theta is not cast to the op's theta type. + assert expected.isequal(v64.ewise_mult(v64, add_theta(10.0)).new(), check_dtype=True) + with pytest.raises(KeyError, match="add_theta does not work with .FP64, FP64, INT64. types"): + v64.ewise_mult(v64, add_theta(10)) + with pytest.raises(KeyError, match="add_theta does not work with .INT64, INT64, FP64. types"): + v.ewise_mult(v, add_theta(10.0)) + # An empty theta is refused, as for any IndexBinaryOp (SuiteSparse would read its value) + for bind in [add_theta, add_theta["FP64"]]: + with pytest.raises(EmptyObject, match="theta is an empty Scalar"): + bind(gb.Scalar("FP64")) + # Test duplicate registration fails + assert "FP32" not in add_theta + with burble(): + add_theta_fp32 = indexbinary.ss.register_new( + "add_theta", + cdef.replace("double", "float"), + "FP32", + "FP32", + "FP32", + "FP32", + ) + assert add_theta_fp32 is add_theta + assert "FP32" in add_theta + with pytest.raises( + TypeError, + match="gb.indexbinary.ss.add_theta already defined for .FP64, FP64, FP64. input", + ): + indexbinary.ss.register_new("add_theta", cdef, "FP64", "FP64", "FP64", "FP64") + # Test nested names + indexbinary.ss.register_new("nested.add_theta", cdef, "FP64", "FP64", "FP64", "FP64") + with pytest.raises(AttributeError, match="nested is already defined"): + indexbinary.ss.register_new("nested", cdef, "FP64", "FP64", "FP64", "FP64") + # Test mixed types (x=FP64, y=FP64, theta=FP32) + mixed_cdef = ( + "void add_theta (double *z, double *x, GrB_Index ix, GrB_Index jx, " + "double *y, GrB_Index iy, GrB_Index jy, float *theta) " + "{ (*z) = (*x) + (*y) + 2 * (double)(*theta) ; }" + ) + add_theta_mixed = indexbinary.ss.register_new( + "add_theta", mixed_cdef, "FP64", "FP64", "FP32", "FP64" + ) + assert add_theta_mixed is add_theta + assert ("FP64", "FP64", "FP32") in add_theta + assert ("FP64", "FP32") not in add_theta # theta of y's type, FP32 + assert repr(add_theta["FP64", "FP64", "FP32"]) == "indexbinary.add_theta[FP64, FP64, FP32]" + # An FP32 theta picks it, and its kernel runs. SuiteSparse finds a compiled kernel by its + # op's name, and its encoding has no theta type, so with one C name for both theta types + # this ran the FP64 theta's kernel, which read the float as a double: [2, 2, 4, 0]. + assert add_theta["FP64", "FP64", "FP32"].jit_c_name.startswith( + "add_theta_4_FP64_4_FP64_4_FP32_" + ) + expected_mixed = Vector.from_coo([1, 3, 4, 6], [22.0, 22.0, 24.0, 20.0]) + res = v64.ewise_mult(v64, add_theta(np.float32(10))).new() + assert expected_mixed.isequal(res, check_dtype=True) + assert expected.isequal(v64.ewise_mult(v64, add_theta(10.0)).new(), check_dtype=True) + # A name too long for SuiteSparse with its types is cut, with a digest of the whole + long_name = "add_theta_" + "long" * 25 + long_op = indexbinary.ss.register_new( + long_name, cdef.replace("add_theta", long_name), "FP64", "FP64", "FP64", "FP64" + ) + c_name = long_op["FP64"].jit_c_name + assert len(c_name) == 117 + assert c_name.startswith(long_name[:100]) + assert expected.isequal(v64.ewise_mult(v64, long_op(10.0)).new(), check_dtype=True) + + +@pytest.mark.slow +def test_jit_c_name_per_name_and_definition(v): + """An op registered from C is compiled under a C name of its full name and definition. + + SuiteSparse finds a compiled kernel in memory by its op's name, and checks a definition + only when it loads a kernel from disk. When the C name was the op's last name, an op + defined again after ``delattr``, or one of the same last name in another namespace, ran + the kernel compiled first. + """ + if _IS_SSGB7: + return + _require_jit_on() + v64 = v.dup("FP64") + # One last name in two namespaces + plus = "void jcn_f (double *z, double *x, double *y) { (*z) = (*x) + (*y) ; }" + f = binary.ss.register_new("jcn_f", plus, "FP64", "FP64", "FP64") + g = binary.ss.register_new("jcn_ns.jcn_f", plus.replace("+", "-"), "FP64", "FP64", "FP64") + try: + expected = Vector.from_coo([1, 3, 4, 6], [2.0, 2.0, 4.0, 0.0]) + assert v64.ewise_mult(v64, f).new().isequal(expected, check_dtype=True) + expected = Vector.from_coo([1, 3, 4, 6], [0.0, 0.0, 0.0, 0.0]) + assert v64.ewise_mult(v64, g).new().isequal(expected, check_dtype=True) + assert f["FP64"].jit_c_name != g["FP64"].jit_c_name + finally: + delattr(binary.ss, "jcn_f") + delattr(binary.ss, "jcn_ns") + # One name defined again, after delattr, with another definition and then the first + square = "void jcn_u (double *z, double *x) { (*z) = (*x) * (*x) ; }" + add10 = square.replace("* (*x)", "+ 10") + names = [] + for cdef, values in [(square, [1, 1, 4, 0]), (add10, [11, 11, 12, 10]), (square, [1, 1, 4, 0])]: + op = unary.ss.register_new("jcn_u", cdef, "FP64", "FP64") + try: + expected = Vector.from_coo([1, 3, 4, 6], values, dtype="FP64") + assert v64.apply(op).new().isequal(expected, check_dtype=True) + names.append(op["FP64"].jit_c_name) + finally: + delattr(unary.ss, "jcn_u") + assert names[0] != names[1] + assert names[0] == names[2] # one definition, one kernel + # A struct member of the op's name is not renamed + rec = dtypes.ss.register_new( + "JcnRec", + "typedef struct { int64_t jcn_m ; double b ; } JcnRec ;", + np_type=np.dtype([("jcn_m", "introsp_a = (x->introsp_a) + (y->introsp_a)" in src + assert "z->introsp_b = (x->introsp_b) + (y->introsp_b)" in src + + # Auto-lifted record UDT unary op + ainv_udt = unary.ainv[udt] + assert ainv_udt.jit_c_name == "ainv_gbudt__IntrospectUDT" + assert "z->introsp_a = -(x->introsp_a)" in ainv_udt.jit_c_source + + # Array UDT auto-lifted op. Use a shape no other test uses: ``np.dtype`` + # identity is shared across the session, so once SS sets + # ``GxB_JIT_C_NAME`` we can't rename. A distinctive shape keeps the + # C-side name stable for this test. + arr_udt = dtypes.register_anonymous(np.dtype("(11,)float64"), "_IntrospectArr") + times_arr = binary.times[arr_udt] + src_arr = times_arr.jit_c_source + assert src_arr is not None + assert "for (int64_t i = 0 ; i < 11 ; i++) { z->v[i] = (x->v[i]) * (y->v[i]) ; }" in src_arr + + # Builtin scalar op: no JIT source + assert binary.plus[int].jit_c_source is None + assert binary.plus[int].jit_c_name is None + assert unary.abs[float].jit_c_source is None + + # Monoid / Semiring / Aggregator walk-down: introspection should follow the + # underlying binary op so users can chase the JIT'd source from any layer. + plus_binop_src = binary.plus[udt].jit_c_source + assert gb.monoid.plus[udt].jit_c_source == plus_binop_src + # Semiring delegates to the multiplier; .monoid still works on its own. + semi = gb.semiring.plus_times[udt] + assert semi.jit_c_source == binary.times[udt].jit_c_source + assert semi.monoid.jit_c_source == plus_binop_src + # Monoid-based aggregator walks to its monoid; composite agg has no kernel. + assert gb.agg.sum[udt].jit_c_source == plus_binop_src + assert gb.agg.count[udt].jit_c_source is None + + +@pytest.mark.skipif("not supports_udfs") +# jit_c_name is only ever set on SS 9+; see test_udt_jit_c_source_introspection. +@pytest.mark.skipif("not _has_jit_set") +def test_op_jit_signature_uses_pinned_type_name(): + """After a UDT is renamed, auto-lifted ops must still reference the + pinned (first-registration) C name in their signature. SS's + ``GxB_JIT_C_NAME`` on a ``GrB_Type`` is one-shot, so a mismatched op + signature would reference an undefined struct and SS would silently + fall back to the Numba cfunc. + """ + record_dtype = np.dtype([("pinned_a", np.int64), ("pinned_b", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_PinnedOrig") + udt2 = dtypes.register_anonymous(record_dtype, "_PinnedRenamed") + assert udt2 is udt + assert udt.name == "_PinnedRenamed" + assert udt.jit_c_name == "gbudt__PinnedOrig" + + op = binary.plus[udt] + # Op signature must reference the pinned struct name, not the renamed one. + assert op.jit_c_name == "plus_gbudt__PinnedOrig" + assert "gbudt__PinnedOrig *" in op.jit_c_source + assert "_PinnedRenamed" not in op.jit_c_source + + +@pytest.mark.skipif("not supports_udfs") +def test_jit_compiles_auto_udt_ops(): + """JIT must actually compile a kernel for an auto-generated UDT op. + + Regression for the bug where ``GxB_JIT_C_NAME`` wasn't being set on + UDT types, so SuiteSparse's eWise JIT template (which depends on the + type having a JIT name) failed to expand with errors like + ``use of undeclared identifier 'Bx'``, disabling JIT for the rest of + the session. Verifies both that (a) the op runs without JIT erroring + out, and (b) ``jit_c_control`` stays ``"on"`` after the op. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + _require_jit_on() + + record_dtype = np.dtype([("a", np.int64), ("b", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_JitAutoUdt") + + v = gb.Vector(udt, 3) + v[0] = (1, 2.0) + v[1] = (3, 4.0) + v[2] = (5, 6.0) + w = v.dup() + + # Force JIT on and probe. If our typedef + JIT_C_NAME wiring is correct, + # SuiteSparse should compile and load the kernel without errors. + with _jit_mode("on"): + with burble(): + result = binary.plus(v & w).new() + assert result[0].new() == (2, 4.0) + assert result[2].new() == (10, 12.0) + # JIT must remain in compile-on mode; a failed compile flips it to 'load'. + assert ( + gb.ss.config["jit_c_control"] == "on" + ), "JIT compilation got disabled (flipped from 'on' to 'load') after a built-in UDT op" + + +@pytest.mark.skipif("not supports_udfs") +def test_floordiv_udt_jit_matches_python_semantics(): + """``binary.floordiv`` on a UDT must round toward minus infinity on the JIT path. + + Regression: the JIT codegen used to lower ``//`` to plain C ``/``, which + is trunc-toward-zero for ints and true division for floats. For positive + integer operands the two agreed; for negative operands or any float they + silently disagreed with the Numba cfunc path (Python ``//`` is floor). + + Use N=50 non-iso vectors with mixed-sign operands to bypass any + iso/short-vector shortcuts and exercise the actual JIT kernel. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + _require_jit_on() + + N = 50 + + # float64 record: was returning 3.5 / -3.5 (true division) before the fix. + # Use field names unique to this test so the anonymous-UDT cache (keyed + # on np.dtype) doesn't share a DataType with another test, which would + # also share the SS-side ``jit_c_name``. + udt = dtypes.register_anonymous( + np.dtype([("fd_a", np.float64), ("fd_b", np.float64)]), "_FdJitF64" + ) + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + for i in range(N): + v[i] = (-7.0 - i, 7.0 + i) + u[i] = (2.0, 2.0) + w = v.ewise_mult(u, binary.floordiv).new() + # -7.0 // 2.0 == -4.0, 7.0 // 2.0 == 3.0 + assert w[0].new() == (-4.0, 3.0) + # -8.0 // 2.0 == -4.0, 8.0 // 2.0 == 4.0 + assert w[1].new() == (-4.0, 4.0) + + # int64 record: was returning -3 (trunc-to-zero) for -7//2 before the fix. + udt_i = dtypes.register_anonymous(np.dtype([("a", np.int64), ("b", np.int64)]), "_FdJitI64") + v = gb.Vector(udt_i, N) + u = gb.Vector(udt_i, N) + for i in range(N): + v[i] = (-7 - 2 * i, 7 + 2 * i) + u[i] = (2, 3) + w = v.ewise_mult(u, binary.floordiv).new() + # -7 // 2 == -4, 7 // 3 == 2 + assert w[0].new() == (-4, 2) + # -9 // 2 == -5, 9 // 3 == 3 + assert w[1].new() == (-5, 3) + # -11 // 2 == -6, 11 // 3 == 3 + assert w[2].new() == (-6, 3) + + # Array UDT: same hazard on the flattened ``v[i]`` path. + udt_arr = dtypes.register_anonymous(np.dtype((np.float32, (4,))), "_FdJitArrF32") + v = gb.Vector(udt_arr, N) + u = gb.Vector(udt_arr, N) + for i in range(N): + v[i] = np.array([-7.0 - i, 7.0 + i, -8.0 - i, 8.0 + i], dtype=np.float32) + u[i] = np.array([2.0, 2.0, 3.0, 3.0], dtype=np.float32) + w = v.ewise_mult(u, binary.floordiv).new() + # -7//2=-4, 7//2=3, -8//3=-3, 8//3=2 + assert_array_equal(w[0].new().value, np.array([-4.0, 3.0, -3.0, 2.0], dtype=np.float32)) + # -8//2=-4, 8//2=4, -9//3=-3, 9//3=3 + assert_array_equal(w[1].new().value, np.array([-4.0, 4.0, -3.0, 3.0], dtype=np.float32)) + + +@pytest.mark.skipif("not supports_udfs") +def test_min_max_udt_jit_ignores_nan_and_breaks_zero_ties(): + """The C JIT kernel for ``binary.min`` on a float UDT answers as the cfunc does. + + ``GrB_MIN_FP64`` is C99 ``fmin``, which ignores a NaN operand from either + side, and so does ``binary.min`` on a UDT field. C99 leaves the tie + between ``-0.0`` and ``0.0`` to the C library (glibc's x86-64 ``fmin`` + returns an operand by position), so the kernel spells the comparison out, + as the cfunc does, rather than calling ``fmin``: a tie gives ``-0.0`` for + ``min`` and ``0.0`` for ``max`` on every platform. Earlier spellings + decided NaN by position, in opposite directions: ``(a < b ? a : b)`` + dropped a NaN on the right, and ``(b < a ? b : a)`` one on the left. + + Integer fields keep the plain comparison: they have no NaN or signed zero. + """ + if not _has_jit_set: + pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + + # Field names unique to this test; see floordiv test for the cache rationale. + udt = dtypes.register_anonymous( + np.dtype([("nan_a", np.float64), ("nan_b", np.float32), ("nan_c", np.int32)]), "_NanJitMM" + ) + csrc = binary.min[udt].jit_c_source + assert "fmin" not in csrc, csrc + assert "signbit(x->nan_a)" in csrc, csrc + assert "signbit(x->nan_b)" in csrc, csrc + assert "((x->nan_c) < (y->nan_c) ? (x->nan_c) : (y->nan_c))" in csrc, csrc + assert "!signbit(x->nan_a)" in binary.max[udt].jit_c_source + + N = 100 + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + nan = float("nan") + for i in range(N): + # field nan_a: NaN on the left; field nan_b: NaN on the right at odd indices. + v[i] = (nan, 2.0 + i, i) + u[i] = (1.0 + i, nan if i % 2 else 3.0 + i, 2 * i) + + w = v.ewise_mult(u, binary.min).new() + assert w[0].new().value[0] == 1.0 # min(NaN, 1.0) -> 1.0 + assert w[1].new().value[1] == 3.0 # min(3.0, NaN) -> 3.0 + assert w[2].new().value[1] == 4.0 # min(4.0, 5.0) -> 4.0 (normal case) + assert w[3].new().value[2] == 3 # integer field is unaffected + + w = v.ewise_mult(u, binary.max).new() + assert w[0].new().value[0] == 1.0 # max(NaN, 1.0) -> 1.0 + assert w[1].new().value[1] == 3.0 # max(3.0, NaN) -> 3.0 + assert w[2].new().value[1] == 5.0 # max(4.0, 5.0) -> 5.0 (normal case) + assert w[3].new().value[2] == 6 # integer field is unaffected + + # Both operands NaN is the one case where a NaN survives, for min and max + # alike, and it is the only case ``fmin`` has no non-NaN answer for. + nan_only = gb.Vector(udt, N) + for i in range(N): + nan_only[i] = (nan, np.float32(i), i) + w = nan_only.ewise_mult(nan_only, binary.min).new() + assert np.isnan(w[0].new().value[0]) + + # Signed-zero ties, in both orders and both fields, on the C kernel and the cfunc. + left = gb.Vector(udt, N) + right = gb.Vector(udt, N) + for i in range(N): + a, b = (-0.0, 0.0) if i % 2 else (0.0, -0.0) + left[i] = (a, a, i) + right[i] = (b, b, i) + for mode in ["on", "off"]: + with _jit_mode(mode): + lo = left.ewise_mult(right, binary.min).new() + hi = left.ewise_mult(right, binary.max).new() + for i in range(2): + assert np.signbit(lo[i].new().value[0]), mode + assert np.signbit(lo[i].new().value[1]), mode + assert not np.signbit(hi[i].new().value[0]), mode + assert not np.signbit(hi[i].new().value[1]), mode + + +@pytest.mark.skipif("not supports_udfs") +def test_complex_truediv_udt_jit_matches_cfunc_at_extremes(): + """Complex division in the C JIT kernel is the cfunc's, bit for bit, at any magnitude. + + SuiteSparse compiles C JIT kernels with ``-fcx-limited-range`` under GCC, + which turns C's ``_Complex`` division into the naive formula: complex64 + ``1e20 / 1e20`` came out ``nan`` and ``(1+1j) / (1e160+1e160j)`` came out + ``0`` on Linux. The kernel divides with real arithmetic instead, by + CPython's algorithm before 3.14, so no flag changes it; the cfunc spells + out the same steps, since from Python 3.14 Numba's own ``/`` recovers + infinities as CPython's does (``(inf+nanj) / (1+1j)`` is ``inf-infj``). + """ + if not _has_jit_set: + pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + + udt = dtypes.register_anonymous( + np.dtype([("cxd_s", np.complex64), ("cxd_d", np.complex128)], align=True), "_CxDivJit" + ) + csrc = binary.truediv[udt].jit_c_source + assert "crealf" in csrc, csrc + assert "(x->cxd_d) / (y->cxd_d)" not in csrc, csrc + # (complex64 pair, complex128 pair); the first four are finite and the C + # library's division got them wrong under the limited-range flag. + pairs = [ + ((1e20, 1e20), (1e200, 1e200)), + ((1e-25 + 1e-25j, 1e-25), (1e200 + 1e200j, 1e200 + 1e200j)), + ((1 + 1j, 1e30 + 1e30j), (1 + 1j, 1e160 + 1e160j)), + ((3e30 + 4e30j, 1e30 - 2e30j), (1e-170, 1e-170 + 1e-170j)), + ((0j, -2 + 0.5j), (0j, -2 + 0.5j)), + ((-2 + 0.5j, -2 + 0.5j), (-2 + 0.5j, -2 + 0.5j)), + ((3 + 4j, complex(0.0, -0.0)), (3 + 4j, complex(0.0, -0.0))), + ((complex(np.inf, np.nan), 1 + 1j), (complex(np.inf, np.nan), 1 + 1j)), + ] + N = 64 + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + for i in range(N): + (a_s, b_s), (a_d, b_d) = pairs[i % len(pairs)] + v[i] = (a_s, a_d) + u[i] = (b_s, b_d) + with _jit_mode("on"): + jit = v.ewise_mult(u, binary.truediv).new() + with _jit_mode("off"): + cfunc = v.ewise_mult(u, binary.truediv).new() + assert jit.to_dense().tobytes() == cfunc.to_dense().tobytes() + result = jit.to_dense() + for i, ((a_s, b_s), (a_d, b_d)) in enumerate(pairs[:4]): + assert np.isclose(result[i]["cxd_s"], np.complex64(a_s) / np.complex64(b_s), rtol=1e-6) + assert np.isclose(result[i]["cxd_d"], np.complex128(a_d) / np.complex128(b_d), rtol=1e-15) + + +@pytest.mark.skipif("not supports_udfs") +def test_complex_times_udt_jit_matches_cfunc_with_infinities(): + """Complex multiplication in the C JIT kernel is the cfunc's, bit for bit. + + numpy multiplies by the textbook formula; clang's ``_Complex`` multiply + adds C99 Annex G's recovery, so ``1j * (inf+infj)`` came out + ``-inf+infj`` from a C kernel on macOS and ``nan+nanj`` from the cfunc + (GCC's ``-fcx-limited-range`` drops the recovery, so Linux agreed). The + kernel multiplies with real arithmetic instead, for record fields and + array elements alike. CPython adds the recovery from 3.14 and Numba 0.68 + follows it, so the cfunc spells the formula out too. + """ + if not _has_jit_set: + pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + + rec = dtypes.register_anonymous( + np.dtype([("cxt_s", np.complex64), ("cxt_d", np.complex128)], align=True), "_CxTimesJit" + ) + arr = dtypes.register_anonymous(np.dtype((np.complex128, (2,))), "_CxTimesJitArr") + assert "(x->cxt_d) * (y->cxt_d)" not in binary.times[rec].jit_c_source + inf, nan = np.inf, np.nan + pairs = [ + (1j, complex(inf, inf)), + (-2.5j, complex(inf, -inf)), + (complex(inf, 0.0), complex(0.0, 1.0)), + (complex(nan, inf), 2 + 1j), + (complex(1e30, 1e30), complex(1e30, -1e30)), + (3 + 4j, -1.5 + 0.25j), + (complex(-0.0, 0.0), complex(0.0, -0.0)), + ] + N = 3 * len(pairs) + for udt, element in [(rec, lambda z: (z, z)), (arr, lambda z: [z, -z])]: + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + for i in range(N): + a, b = pairs[i % len(pairs)] + v[i] = element(a) + u[i] = element(b) + with _jit_mode("on"): + jit = v.ewise_mult(u, binary.times).new() + with _jit_mode("off"): + cfunc = v.ewise_mult(u, binary.times).new() + assert jit.to_dense().tobytes() == cfunc.to_dense().tobytes() + result = jit[0].new().value # 1j * (inf+infj), numpy's nan+nanj + assert np.isnan(result[0].real) + assert np.isnan(result[0].imag) + + +@pytest.mark.skipif("not supports_udfs") +def test_abs_udt_jit_matches_python_negative_zero(): + """``unary.abs`` on a float UDT must clear the sign bit of ``-0.0``. + + Regression: the JIT codegen used to emit ``(x < 0 ? -x : x)``. For + ``x == -0.0`` the comparison ``-0.0 < 0`` is false (negative zero is + not less than zero in IEEE-754), so the ternary returned ``-0.0`` + with the sign bit intact. Python ``abs(-0.0) == 0.0`` (sign bit + cleared), and the cfunc path uses Python's ``abs``, so the two paths + silently disagreed on the sign of zero. The fix routes float fields + through ``fabs`` / ``fabsf``. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + _require_jit_on() + + import struct + + def sign_bit(f): + return struct.pack(" (a, b)). +# +# Per-variant unique field names: anonymous UDTs are keyed on ``np.dtype``, so +# same-shape dtypes across variants would share a single ``DataType`` (and its +# pinned JIT C name from the first registration), causing test-order coupling. + + +def _parity_record_i64(i): + # cross-sign division operands + return (-7 - 2 * i, 7 + 2 * i), (2, -3) + + +def _parity_record_f64(i): + nan = float("nan") + inf = float("inf") + if i == 0: + return (nan, 1.0), (2.0, 3.0) + if i == 1: + return (1.0, nan), (inf, -inf) + if i == 2: + return (-7.5, 7.5), (2.0, 2.0) + return (-7.0 - i, 7.0 + i), (2.0 + 0.1 * i, 3.0) + + +def _parity_record_u32(i): + return (10 + i, 20 + i), (2, 3) + + +def _parity_record_mixed(i): + return (5 + i, 1.5 + i), (2, 0.5) + + +def _parity_array_f32(i): + return ( + np.array([-7.0 - i, 7.0 + i, -8.0 - i, 8.0 + i], np.float32), + np.array([2.0, 2.0, 3.0, 3.0], np.float32), + ) + + +def _parity_array_i16(i): + return ( + np.array([1 + i, 2 + i, 3 + i, 4 + i], np.int16), + np.array([2, 2, 3, 3], np.int16), + ) + + +_PARITY_VARIANTS = { + "record_i64": ( + np.dtype([("p_i_x", np.int64), ("p_i_y", np.int64)]), + "_ParityI64", + _parity_record_i64, + ), + "record_f64": ( + np.dtype([("p_f_x", np.float64), ("p_f_y", np.float64)]), + "_ParityF64", + _parity_record_f64, + ), + "record_u32": ( + np.dtype([("p_u_x", np.uint32), ("p_u_y", np.uint32)]), + "_ParityU32", + _parity_record_u32, + ), + "record_mixed": ( + np.dtype([("p_m_i", np.int32), ("p_m_f", np.float64)]), + "_ParityMixed", + _parity_record_mixed, + ), + "array_f32": ( + np.dtype((np.float32, (4,))), + "_ParityArrF32", + _parity_array_f32, + ), + "array_i16": ( + np.dtype((np.int16, (4,))), + "_ParityArrI16", + _parity_array_i16, + ), +} + + +@pytest.mark.slow +@pytest.mark.parametrize("udt_kind", list(_PARITY_VARIANTS)) +def test_udt_op_jit_cfunc_parity(udt_kind): + """JIT and cfunc paths produce the same result on every auto-lifted op. + + The two paths are independent code generators (string-templated C vs Numba + njit). N=64 sidesteps SS's iso/short-vector shortcuts so the kernel actually + runs. The variants cover signed-int floor formula, float floor/NaN/inf, + unsigned trunc-div, mixed-width with C alignment, and both array shapes. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + _require_jit_on() + if numba is None: + pytest.skip("numba required for the cfunc baseline") + + np_dtype, type_name, fill = _PARITY_VARIANTS[udt_kind] + udt = dtypes.register_anonymous(np_dtype, type_name) + N = 64 + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + for i in range(N): + a, b = fill(i) + v[i] = a + u[i] = b + + binary_ops = ["plus", "minus", "times", "truediv", "floordiv", "min", "max"] + unary_ops = ["ainv", "abs"] + + with _jit_mode("on"): + jit_binary = {op: v.ewise_mult(u, getattr(binary, op)).new() for op in binary_ops} + jit_unary = {op: getattr(unary, op)(v).new() for op in unary_ops} + # cfunc path: JIT off so SS uses the registered function pointer instead of + # compiling a new kernel. + with _jit_mode("off"): + cf_binary = {op: v.ewise_mult(u, getattr(binary, op)).new() for op in binary_ops} + cf_unary = {op: getattr(unary, op)(v).new() for op in unary_ops} + + def values_equal(j, c): + # Compare raw bytes via ``to_dense`` rather than ``isequal``. With the + # IEEE-aware ``binary.eq[udt]`` fix, two NaN-bearing records compare + # unequal under ``isequal``, so ``isequal`` can't distinguish "JIT and + # cfunc agree on a NaN bit-pattern" from "they disagree". Byte-equality + # of the dense numpy view captures the parity question correctly. + if j.dtype != c.dtype or j.nvals != c.nvals: + return False + return j.to_dense().tobytes() == c.to_dense().tobytes() + + for op in binary_ops: + assert values_equal( + jit_binary[op], cf_binary[op] + ), f"binary.{op} on {udt_kind}: JIT and cfunc disagree" + for op in unary_ops: + assert values_equal( + jit_unary[op], cf_unary[op] + ), f"unary.{op} on {udt_kind}: JIT and cfunc disagree" + + +@pytest.mark.skipif("not supports_udfs") +def test_anonymous_udt_with_no_name_still_jits(): + """A UDT registered without ``name=`` should still take the JIT path. + + ``register_anonymous(np.dtype)`` gives the DataType a Python-side name + like ``"{'a': FP64, 'b': INT64}"`` that isn't a valid C identifier, so + SS can't use it as the JIT type name. ``_pick_c_type_name`` synthesizes + a ``_gbudt_NNN`` name for SuiteSparse instead, leaving ``udt.name`` + alone for Python-side display. If that fallback regresses, the op runs + through the slower Numba cfunc path and ``jit_c_source`` is ``None``. + """ + if not _has_jit_set: + pytest.skip("jit_c_* introspection requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + + spec = np.dtype([("anon_a", np.float64), ("anon_b", np.int64)], align=True) + udt = dtypes.register_anonymous(spec) # no name= + + # Python-side default name is the np.dtype repr; not a valid C identifier. + assert udt.name != udt.jit_c_name + assert udt.jit_c_name is not None + assert udt.jit_c_name.startswith("_gbudt_") + assert udt.jit_c_definition is not None + + # The synthetic name flows through to op codegen so the JIT path actually + # fires and matches the cfunc path. + plus_op = binary.plus[udt] + assert plus_op.jit_c_source is not None + assert plus_op.jit_c_name.startswith("plus__gbudt_") + + # Run the kernel to confirm the JIT path produces correct results. + v = gb.Vector(udt, 2) + v[0] = (1.0, 10) + v[1] = (2.0, 20) + w = gb.Vector(udt, 2) + w[0] = (3.0, 30) + w[1] = (4.0, 40) + result = v.ewise_mult(w, binary.plus).new() + assert result[0].new().value.tolist() == (4.0, 40) + assert result[1].new().value.tolist() == (6.0, 60) + + +@pytest.mark.slow +@pytest.mark.skipif("not supports_udfs") +def test_complex_field_udt_jits_arithmetic(): + """Complex fields take the JIT path for ``plus``/``minus``/``times``/ + ``truediv``/``abs``/``ainv``. ``min``/``max``/``floordiv`` are rejected + early with a clear error rather than a numba LLVM crash. + + The C99 ``_Complex`` types arrive via SS's GraphBLAS.h (typedef'd to + ``GxB_FC32_t`` / ``GxB_FC64_t``), so the generated kernel compiles + without extra includes. ``abs`` uses ``cabs`` / ``cabsf`` which match + Numba's ``abs(complex)`` behavior (real magnitude in the real part). + """ + if not _has_jit_set: + pytest.skip("jit_c_definition introspection requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + + spec = np.dtype([("cval", np.complex128), ("scale", np.float64)]) + udt = dtypes.register_anonymous(spec, "_JitCplx1") + assert "GxB_FC64_t" in udt.jit_c_definition + + N = 64 + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + for i in range(N): + v[i] = (complex(1.0 + i, 2.0), 3.0 + i) + u[i] = (complex(0.5, 0.25 + i), 0.1) + + # JIT must agree with cfunc bit-exactly on ops that lower to single C + # operators. ``truediv`` is intentionally excluded: C99 ``_Complex`` + # division and Numba's complex division differ in their last-bit + # rounding / overflow handling, so they're equivalent within ulps but + # not bit-identical. + for op_name in ("plus", "minus", "times"): + op = getattr(binary, op_name) + with _jit_mode("on"): + wj = v.ewise_mult(u, op).new() + with _jit_mode("off"): + wc = v.ewise_mult(u, op).new() + assert wj.isequal(wc, check_dtype=True), f"binary.{op_name} on complex UDT diverged" + for op_name in ("abs", "ainv"): + op = getattr(unary, op_name) + with _jit_mode("on"): + wj = op(v).new() + with _jit_mode("off"): + wc = op(v).new() + assert wj.isequal(wc, check_dtype=True), f"unary.{op_name} on complex UDT diverged" + + # abs source must use cabs (not the ternary, which doesn't compile on _Complex). + assert "cabs(" in unary.abs[udt].jit_c_source + + # min/max/floordiv reject early with a clear KeyError citing complex. + for op_name in ("min", "max", "floordiv"): + with pytest.raises(KeyError, match="complex fields"): + getattr(binary, op_name)[udt] + + +@pytest.mark.skipif("not supports_udfs") +def test_packed_record_layout_skips_jit(): + """A packed numpy dtype (no inter-field padding) must skip the JIT path. + + Otherwise the JIT-compiled kernel reads fields at the C-aligned offsets + while the numpy buffer holds them at packed offsets, producing silent + garbage. The cfunc path still works because Numba mirrors numpy's + layout. Users wanting JIT can re-register with ``align=True`` or use + the dict / dataclass form (which auto-aligns). + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + + # int32 + float64 packed: numpy offsets 0, 4 (itemsize 12); + # C struct would put float64 at offset 8 (itemsize 16). + packed = np.dtype([("p_flag", np.int32), ("p_val", np.float64)]) # no align=True + udt = dtypes.register_anonymous(packed, "_PackedUdt1") + assert udt.jit_c_name is None # JIT skipped + assert udt.jit_c_definition is None + + # cfunc still works (the existing fallback path). + v = gb.Vector(udt, 4) + v[0] = (1, 1.5) + v[1] = (2, 2.5) + w = v.ewise_mult(v, binary.plus).new() + assert w[0].new() == (2, 3.0) + + +@pytest.mark.skipif("not supports_udfs") +def test_packed_nested_record_layout_skips_jit(): + """Packed-inside-packed must skip the JIT path too. + + Regression: ``_is_c_compatible_layout`` originally rebuilt the + comparison dtype with ``np.dtype([(name, inner_dtype), ...], align=True)``, + but numpy's ``align=True`` respects the *inner* dtype's own alignment. + When the inner is itself packed (the default), it stays packed with + alignment 1, and the outer aligned reconstruction happens to match + the original packed layout. The JIT then produced garbage because a + C compiler aligns the inner to its natural 8-byte boundary. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + + spec = np.dtype( + [ + ("flag", np.int32), + ("coord", [("x", np.float64), ("y", np.float64)]), + ], + ) # no align=True at either level + # numpy itemsize 20 (flag@0, coord@4); C would pick 24 (flag@0, pad, coord@8). + assert spec.itemsize == 20 + udt = dtypes.register_anonymous(spec, "_PackedNestedUdt") + assert udt.jit_c_name is None # JIT must be skipped + assert udt.jit_c_definition is None + + +@pytest.mark.skipif("not supports_udfs") +def test_record_layouts_c_cannot_declare_skip_jit(): + """A record whose fields a C struct cannot lay out the same way skips the JIT. + + Explicit offsets can give a record a C struct's size with a field + somewhere else, as can such a record nested in one whose own offsets + match. A nested record with an array field has no C struct here, so the + record holding it has none either. Each computes with the cfunc. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + + # Same size as the C struct (8 bytes), but the int32 at offset 1, not 4. + offset = np.dtype( + { + "names": ["ol_a", "ol_b"], + "formats": [np.int8, np.int32], + "offsets": [0, 1], + "itemsize": 8, + } + ) + outer = np.dtype([("ol_in", offset), ("ol_z", np.float64)]) + with_array = np.dtype( + [("ol_x", np.float64), ("ol_rec", [("ol_v", np.float64, (2,)), ("ol_w", np.float64)])], + ) + for np_type, value, doubled in [ + (offset, (1, 400), (2, 800)), + (outer, ((1, 400), 2.5), ((2, 800), 5.0)), + (with_array, (1.5, ([2, 3], 4)), (3.0, ([4, 6], 8))), + ]: + udt = dtypes.register_anonymous(np_type) + assert udt.jit_c_name is None + assert udt.jit_c_definition is None + v = gb.Vector(udt, 1) + v[0] = value + assert ((v + v).new().to_coo()[1] == np.array([doubled], dtype=np_type)).all() + + +@pytest.mark.slow +@pytest.mark.skipif("not supports_udfs") +def test_nested_record_udt_jits(): + """Record-of-record UDTs take the JIT path. + + The codegen emits inner typedefs first in ``GxB_JIT_C_DEFINITION``, with + synthesized ``_gbnest_NNN`` names to avoid collisions with user types. + Op bodies use C nested-struct access (``z->outer.inner_a``), and the + Numba cfunc wrapper writes the flat result tuple leaf-by-leaf back into + the nested record fields (Numba can't ``setitem`` a tuple to a record + field, but leaf scalar writes work). + """ + if not _has_jit_set: + pytest.skip("jit_c_definition introspection requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + + # ``align=True`` is required for mixed-width fields so numpy's offsets + # match what a C compiler would produce. Without it, the JIT kernel + # would read fields at the wrong offsets and produce garbage; the + # codegen now refuses JIT in that case but the test wants the JIT path + # exercised. + spec = np.dtype( + [ + ("nest_flag", np.int32), + ("nest_coord", [("nest_x", np.float64), ("nest_y", np.float64)]), + ], + align=True, + ) + udt = dtypes.register_anonymous(spec, "_JitNested1") + + # Outer + inner typedefs both appear in the definition; the outer + # references the inner by its synthesized C name. + defn = udt.jit_c_definition + assert defn.count("typedef struct") == 2 + assert "_gbnest_" in defn + assert "_JitNested1" in defn + + plus_src = binary.plus[udt].jit_c_source + # Nested struct access uses ``.`` after the leading ``->``. + assert "z->nest_coord.nest_x = (x->nest_coord.nest_x) + (y->nest_coord.nest_x)" in plus_src + + N = 64 + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + for i in range(N): + v[i] = (i, (1.0 + i, 2.0 + i)) + u[i] = (1, (0.5, -0.5)) + + for op_name in ("plus", "minus", "times"): + op = getattr(binary, op_name) + with _jit_mode("on"): + wj = v.ewise_mult(u, op).new() + with _jit_mode("off"): + wc = v.ewise_mult(u, op).new() + assert wj.isequal(wc, check_dtype=True), f"binary.{op_name} on nested UDT diverged" + for op_name in ("abs", "ainv"): + op = getattr(unary, op_name) + with _jit_mode("on"): + wj = op(v).new() + with _jit_mode("off"): + wc = op(v).new() + assert wj.isequal(wc, check_dtype=True), f"unary.{op_name} on nested UDT diverged" + + +@pytest.mark.slow +@pytest.mark.skipif("not supports_udfs") +# SS < 9 has no GrB_NAME setter, so registration falls back to storing the +# numpy repr in the type name and warns when it does not fit in 128 chars. +# This dtype's repr is 141; how it serializes is not what the test is about. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_nested_record_monoid_semiring_agg(): + r"""Built-in monoid/semiring/agg auto-lift on nested record UDTs. + + Regression: ``_udt_identity`` flattened nested-record fields by calling + ``_scalar_identity`` directly on each field, but for a record-valued + field that returned ``numpy.void(b'\x00')`` (uninitialized bytes), and + ``Scalar.value =`` then choked. Recurse into nested records so the + identity tuple matches the dtype's nested shape. + """ + spec = np.dtype( + [ + ("agg_id", np.int32), + ("agg_pt", [("agg_x", np.float64), ("agg_y", np.float64)]), + ], + align=True, + ) + udt = dtypes.register_anonymous(spec, "_NestedMonoid1") + + plus_id = monoid.plus[udt].identity.value + assert plus_id["agg_id"] == 0 + assert plus_id["agg_pt"]["agg_x"] == 0.0 + assert plus_id["agg_pt"]["agg_y"] == 0.0 + min_id = monoid.min[udt].identity.value + assert min_id["agg_id"] == np.iinfo(np.int32).max + assert min_id["agg_pt"]["agg_x"] == np.inf + assert min_id["agg_pt"]["agg_y"] == np.inf + max_id = monoid.max[udt].identity.value + assert max_id["agg_id"] == np.iinfo(np.int32).min + assert max_id["agg_pt"]["agg_x"] == -np.inf + assert max_id["agg_pt"]["agg_y"] == -np.inf + + v = gb.Vector(udt, size=3) + v[0] = (1, (1.0, 2.0)) + v[1] = (2, (3.0, 4.0)) + v[2] = (3, (5.0, 6.0)) + s = v.reduce(agg.sum[udt]).new().value + assert s["agg_id"] == 6 + assert s["agg_pt"]["agg_x"] == 9.0 + assert s["agg_pt"]["agg_y"] == 12.0 + + A = gb.Matrix(udt, nrows=2, ncols=2) + B = gb.Matrix(udt, nrows=2, ncols=2) + A[0, 0] = (1, (1.0, 0.0)) + A[0, 1] = (2, (0.0, 1.0)) + A[1, 0] = (3, (1.0, 1.0)) + A[1, 1] = (4, (1.0, 1.0)) + B[0, 0] = (10, (1.0, 0.0)) + B[0, 1] = (20, (0.0, 1.0)) + B[1, 0] = (30, (1.0, 1.0)) + B[1, 1] = (40, (1.0, 1.0)) + C = A.mxm(B, semiring.plus_times[udt]).new() + c00 = C[0, 0].new().value + c11 = C[1, 1].new().value + assert (c00["agg_id"], c00["agg_pt"]["agg_x"], c00["agg_pt"]["agg_y"]) == (70, 1.0, 1.0) + assert (c11["agg_id"], c11["agg_pt"]["agg_x"], c11["agg_pt"]["agg_y"]) == (220, 1.0, 2.0) + + +# Per-shape fill functions for the eq/ne JIT-vs-cfunc parity test below. +# Each returns ``(v1_record, v2_record)`` for index ``i``. Variants that put +# a NaN at index 0 also exercise IEEE comparison semantics (eq(NaN,NaN) is +# False), pinned by the ``nan_at_zero`` flag in ``_EQ_NE_VARIANTS``. + + +def _eq_ne_fill_record_float(i): + nan = float("nan") + if i == 0: + return (nan, 1.0), (nan, 1.0) + return (1.0 + i, 2.0), (1.0 + i, 2.0 if i % 2 == 0 else 3.0) + + +def _eq_ne_fill_record_int(i): + return (i, 2 * i), (i, 2 * i if i % 3 == 0 else 2 * i + 1) + + +def _eq_ne_fill_record_nested(i): + return (i, (1.0, 2.0)), (i, (1.0, 2.0) if i % 3 != 0 else (1.5, 2.0)) + + +def _eq_ne_fill_record_complex(i): + return ( + (complex(1.0 + i, 2.0), 3.0), + (complex(1.0 + i, 2.0 if i % 2 == 0 else 3.0), 3.0), + ) + + +def _eq_ne_fill_array_f64(i): + return ( + np.array([1.0 + i, 2.0, 3.0]), + np.array([1.0 + i, 2.0 if i % 2 == 0 else -2.0, 3.0]), + ) + + +_EQ_NE_VARIANTS = { + "record_float": ( + np.dtype([("eq_f_a", np.float64), ("eq_f_b", np.float64)], align=True), + "_EqJitF", + _eq_ne_fill_record_float, + True, + ), + "record_int": ( + np.dtype([("eq_i_a", np.int64), ("eq_i_b", np.int32)], align=True), + "_EqJitI", + _eq_ne_fill_record_int, + False, + ), + "record_nested": ( + np.dtype( + [("eq_n_id", np.int32), ("eq_n_pt", [("nx", np.float64), ("ny", np.float64)])], + align=True, + ), + "_EqJitN", + _eq_ne_fill_record_nested, + False, + ), + "record_complex": ( + np.dtype([("eq_c_z", np.complex128), ("eq_c_s", np.float64)]), + "_EqJitC", + _eq_ne_fill_record_complex, + False, + ), + "array_f64": ( + np.dtype((np.float64, (3,))), + "_EqJitArr", + _eq_ne_fill_array_f64, + False, + ), +} + + +@pytest.mark.slow +@pytest.mark.parametrize("shape", list(_EQ_NE_VARIANTS)) +def test_eq_ne_udt_jit_matches_cfunc(shape): + """``binary.eq[udt]`` / ``binary.ne[udt]`` JIT-compile a leaf-wise + comparison kernel and produce the same result as the cfunc fallback. + + Verifies the JIT kernel for several shapes (flat record, nested + record, complex, array UDT), and that NaN propagation matches IEEE + (``eq(NaN, NaN) == False``) on variants flagged ``nan_at_zero``. + """ + if not _has_jit_set: + pytest.skip("attaching the UDT JIT kernel requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + if numba is None: + pytest.skip("numba required for the cfunc baseline") + + np_dtype, type_name, fill, nan_at_zero = _EQ_NE_VARIANTS[shape] + udt = dtypes.register_anonymous(np_dtype, type_name) + N = 64 + v1 = gb.Vector(udt, N) + v2 = gb.Vector(udt, N) + for i in range(N): + a, b = fill(i) + v1[i] = a + v2[i] = b + + # The C kernel must actually be set, not None (i.e. JIT path is engaged). + assert binary.eq[udt].jit_c_source is not None + assert binary.ne[udt].jit_c_source is not None + + with _jit_mode("on"): + eq_jit = v1.ewise_mult(v2, binary.eq).new() + ne_jit = v1.ewise_mult(v2, binary.ne).new() + with _jit_mode("off"): + eq_cf = v1.ewise_mult(v2, binary.eq).new() + ne_cf = v1.ewise_mult(v2, binary.ne).new() + + # Byte-equal across paths (the result is BOOL, so this is exact). + assert ( + eq_jit.to_dense().tobytes() == eq_cf.to_dense().tobytes() + ), f"eq on {shape}: JIT and cfunc disagree" + assert ( + ne_jit.to_dense().tobytes() == ne_cf.to_dense().tobytes() + ), f"ne on {shape}: JIT and cfunc disagree" + + if nan_at_zero: + # Both records carry NaN at index 0; IEEE ``a == a`` is False. + assert eq_jit[0].new().value is False + assert ne_jit[0].new().value is True + + +@pytest.mark.skipif("not supports_udfs") +def test_eq_ne_udt_cfunc_works_for_unrepresentable_dtypes(): + """``binary.eq`` cfunc path produces correct results on UDTs that can't take JIT. + + The udt-level skip is pinned by ``test_packed_record_layout_skips_jit`` + and ``test_udt_with_c_reserved_field_name_falls_back_to_cfunc``; this + test pins that the eq codegen (separate from plus codegen) still + handles those UDTs end-to-end. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + # Packed mixed-width: numpy offset 4 for float, C would use 8. + packed = np.dtype([("eq_pk_a", np.int32), ("eq_pk_b", np.float64)]) + udt = dtypes.register_anonymous(packed, "_EqJitPacked") + assert binary.eq[udt].jit_c_source is None + v1 = gb.Vector(udt, 2) + v2 = gb.Vector(udt, 2) + v1[0] = (1, 2.5) + v2[0] = (1, 2.5) + v1[1] = (1, 2.5) + v2[1] = (1, 3.5) + eq = v1.ewise_mult(v2, binary.eq).new() + assert eq[0].new().value is True + assert eq[1].new().value is False + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_scalar_broadcast_jit_routing(): + """Mixed UDT and scalar arithmetic gets a C definition with a type for each of z, x + and y; eq/ne on a mixed pair, and ops that have no C form, keep the cfunc. + + This test pins that routing so a future change doesn't accidentally regress it (either + direction is fine, but it should be a conscious decision). + """ + if not _has_jit_set: + # Without the SS 9 string setters, jit_c_source is None for every op, + # so there is no routing to observe. + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + record = dtypes.register_anonymous( + np.dtype([("bcast_u", np.float64), ("bcast_v", np.float64)], align=True), + name="_BcastJitUV", + ) + # Same-type pairs do JIT. + assert binary.plus[record, record].jit_c_source is not None + assert binary.eq[record, record].jit_c_source is not None + # Mixed arithmetic does, with the scalar read once for every leaf + assert binary.plus[record, dtypes.INT64].jit_c_source == ( + "void plus_17_gbudt__BcastJitUV_int64_t (gbudt__BcastJitUV *z, " + "const gbudt__BcastJitUV *x, const int64_t *y) { " + "z->bcast_u = (x->bcast_u) + (((double) ((*y)))) ; " + "z->bcast_v = (x->bcast_v) + (((double) ((*y)))) ; }" + ) + assert binary.plus[dtypes.INT64, record].jit_c_source is not None + # Comparisons on a mixed pair keep the cfunc + assert binary.eq[record, dtypes.INT64].jit_c_source is None + assert binary.ne[record, dtypes.FP64].jit_c_source is None + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_jit_c_source(): + """The C definitions of lifted ops on a scalar and an array UDT, and the cases with no C + form (an array operand that broadcasts, integer truediv, an array leaf in a record). + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + arr = dtypes.register_anonymous(np.dtype((np.float64, (13,))), "_MixedArr13") + # (the C name of whichever test registered this numpy type first: ``test_op`` has it too) + c = arr.jit_c_name + assert binary.times[dtypes.FP64, arr].jit_c_source == ( + f"void times_6_double_{c} ({c} *z, const double *x, " + f"const {c} *y) {{ for (int64_t i = 0 ; i < 13 ; i++) " + "{ z->v[i] = ((*x)) * (y->v[i]) ; } }" + ) + assert binary.times[dtypes.FP64, arr].jit_c_name == f"times_6_double_{c}" + assert binary.times[arr, dtypes.FP64].jit_c_name == f"times_{len(c)}_{c}_double" + # plus_times on these types: its monoid is the same-type plus + semi = gb.semiring.plus_times[dtypes.FP64, arr] + assert semi.jit_c_source == binary.times[dtypes.FP64, arr].jit_c_source + assert semi.monoid.jit_c_source is not None + # Operands converted to the type the op computes in, as the cfunc converts them + assert "(((double) ((*x)))) - (y->v[i])" in binary.minus[dtypes.INT32, arr].jit_c_source + arr_i8 = dtypes.register_anonymous(np.dtype((np.int8, (13,))), "_MixedArrI8") + assert "(((double) (x->v[i]))) + (y->v[i])" in binary.plus[arr_i8, arr].jit_c_source + # No C form: these keep the cfunc + arr1 = dtypes.register_anonymous(np.dtype((np.float64, (1,))), "_MixedArr1") + assert binary.plus[arr1, arr].jit_c_source is None # (broadcasts) + assert binary.truediv[dtypes.INT64, arr_i8].jit_c_source is None + nested = dtypes.register_anonymous( + np.dtype([("mx_w", np.float64, (2,)), ("mx_c", np.int64)]), "_MixedArrLeaf" + ) + assert binary.plus[nested, dtypes.INT64].jit_c_source is None + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_jit_c_name_is_its_own(): + """A mixed op's C name is not a one-type op's on a UDT named like the pair, nor another + pair's. SuiteSparse finds a compiled kernel by its op's name, so the second op ran the + first's kernel: wrong values, or writes past the end of a smaller element. + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + a = dtypes.register_anonymous(np.dtype((np.float64, (19,))), "_CnA") + b = dtypes.register_anonymous(np.dtype((np.int64, (19,))), "_CnB") + ab = dtypes.register_anonymous(np.dtype((np.int8, (19,))), f"{a.jit_c_name}_{b.jit_c_name}") + assert binary.plus[ab, ab].jit_c_name != binary.plus[a, b].jit_c_name + n = 50 + # Values that vary, or SuiteSparse before 10.5 builds iso vectors and computes an iso + # result without the kernel + k = np.broadcast_to(np.arange(n)[:, None] % 7, (n, 19)) + vab = Vector.from_coo(np.arange(n), k.astype(np.int8), dtype=ab) + va = Vector.from_coo(np.arange(n), k + 1.5, dtype=a) + vb = Vector.from_coo(np.arange(n), 10 * k, dtype=b) + with _jit_mode("on"): + same = vab.ewise_mult(vab, binary.plus).new() + mixed = va.ewise_mult(vb, binary.plus).new() + assert (same.to_dense() == 2 * k).all() + assert (mixed.to_dense() == 11 * k + 1.5).all() + + +def _mixed_parity_cases(): + rng = np.random.default_rng(17) + arr_f = dtypes.register_anonymous(np.dtype((np.float64, (5,))), "_MixedParF") + arr_i = dtypes.register_anonymous(np.dtype((np.int64, (5,))), "_MixedParI") + arr_c = dtypes.register_anonymous(np.dtype((np.complex128, (5,))), "_MixedParC") + rec = dtypes.register_anonymous( + np.dtype([("mp_a", np.float64), ("mp_b", np.int32)], align=True), "_MixedParRec" + ) + n = 64 + vf = rng.standard_normal((n, 5)) + vf[0, :2] = [np.nan, np.inf] + vi = rng.integers(-9, 10, (n, 5)) + vc = rng.standard_normal((n, 5)) + 1j * rng.standard_normal((n, 5)) + vr = np.zeros(n, rec.np_type) + vr["mp_a"] = rng.standard_normal(n) + vr["mp_b"] = rng.integers(-9, 10, n) + sf = rng.standard_normal(n) + sf[1] = 0.0 + si = rng.integers(-3, 4, n) + su = rng.integers(0, 2**63, n, dtype=np.uint64) + idx = np.arange(n) + u = { + "FP64[5]": Vector.from_coo(idx, vf, dtype=arr_f), + "INT64[5]": Vector.from_coo(idx, vi, dtype=arr_i), + "FC64[5]": Vector.from_coo(idx, vc, dtype=arr_c), + "record": Vector.from_coo(idx, vr, dtype=rec), + } + s = { + "FP64": Vector.from_coo(idx, sf), + "INT64": Vector.from_coo(idx, si), + "UINT64": Vector.from_coo(idx, su), + } + real_ops = ["plus", "minus", "times", "truediv", "floordiv", "min", "max"] + return ( + [ + ("FP64[5]", "FP64", real_ops), + ("FP64[5]", "INT64", real_ops), + ("INT64[5]", "UINT64", real_ops), # (computed in FP64, as on built-in vectors) + ("INT64[5]", "INT64", ["plus", "minus", "times", "floordiv", "min", "max"]), + ("FC64[5]", "FP64", ["plus", "minus", "times", "truediv"]), + ("record", "FP64", real_ops), + ("record", "INT64", ["plus", "minus", "times", "min", "max"]), + ], + u, + s, + ) + + +@pytest.mark.slow +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_jit_cfunc_parity(): + """A mixed op's C kernel and its Numba cfunc give the same bytes, the UDT on either side.""" + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + cases, u, s = _mixed_parity_cases() + checked = 0 + for ukey, skey, ops in cases: + for op in ops: + for left, right in [(u[ukey], s[skey]), (s[skey], u[ukey])]: + typed = getattr(binary, op)[left.dtype, right.dtype] + if typed.jit_c_source is None: + continue + with _jit_mode("on"): + jit = left.ewise_mult(right, getattr(binary, op)).new() + with _jit_mode("off"): + cf = left.ewise_mult(right, getattr(binary, op)).new() + assert jit.dtype == cf.dtype + assert jit.nvals == cf.nvals + assert ( + jit.to_dense().tobytes() == cf.to_dense().tobytes() + ), f"binary.{op}[{left.dtype}, {right.dtype}]: JIT and cfunc disagree" + checked += 1 + assert checked > 60 + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_semiring_runs_jit(capfd): + """The scalar x array semiring compiles as one JIT kernel (burble), with the result of the + generic kernel. + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + arr = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_MixedSemi6") + rng = np.random.default_rng(23) + A = gb.Matrix.from_coo( + rng.integers(0, 30, 200), + rng.integers(0, 20, 200), + rng.random(200), + nrows=30, + ncols=20, + dup_op=binary.plus, + ) + U = rng.random((20, 6)) + u = Vector.from_coo(np.arange(20), U, dtype=arr) + semi = gb.semiring.plus_times[dtypes.FP64, arr] + capfd.readouterr() + with burble(): + w = A.mxv(u, semi).new() + w.wait() + out = capfd.readouterr().out + assert "GrB_mxv" in out + assert "generic" not in out + # (and burble names the generic kernel when it runs) + with _jit_mode("off"), burble(): + w_generic = A.mxv(u, semi).new() + w_generic.wait() + assert "generic" in capfd.readouterr().out + idx, vals = w.to_coo() + gidx, gvals = w_generic.to_coo() + assert_array_equal(idx, gidx) + np.testing.assert_allclose(vals, gvals, rtol=1e-13) + np.testing.assert_allclose(vals, (A.to_dense(0) @ U)[idx], rtol=1e-13) + + def test_context_importable(): if _IS_SSGB7: with pytest.raises(ImportError, match="Context was added"): @@ -436,3 +2546,226 @@ def test_context_importable(): from graphblas.core.ss.context import global_context assert gb.ss.global_context is global_context + + +def test_fix_from_conda_finds_windows_compilers(tmp_path): + """The conda repair looks in ``bin`` and, as Windows conda installs them, ``Library/bin`` + with ``.exe`` names; a baked-in Windows path's basename is found too. Nothing there: None. + """ + from graphblas.core.ss.jit_config import _fix_from_conda + + def cfg_for(name): + return {"jit_c_compiler_name": name, "jit_c_compiler_flags": "-O3"} + + win = tmp_path / "win" + (win / "Library" / "bin").mkdir(parents=True) + (win / "Library" / "bin" / "clang.exe").write_text("") + cfg = cfg_for(r"C:\bld\_h_env\Library\bin\cl.exe") + assert _fix_from_conda(cfg, str(win)) is True + assert cfg["jit_c_compiler_name"] == str(win / "Library" / "bin" / "clang.exe") + (win / "Library" / "bin" / "cl.exe").write_text("") + cfg = cfg_for(r"C:\bld\_h_env\Library\bin\cl.exe") + assert _fix_from_conda(cfg, str(win)) is True + assert cfg["jit_c_compiler_name"] == str(win / "Library" / "bin" / "cl.exe") + posix = tmp_path / "posix" + (posix / "bin").mkdir(parents=True) + (posix / "bin" / "gcc").write_text("") + cfg = cfg_for("/home/conda/feedstock_root/build_artifacts/bin/x86_64-conda-linux-gnu-cc") + assert _fix_from_conda(cfg, str(posix)) is True + assert cfg["jit_c_compiler_name"] == str(posix / "bin" / "gcc") + assert _fix_from_conda(cfg_for("/nowhere/cc"), str(tmp_path / "empty")) is None + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_jit_c_name_fits_suitesparse(): + """A mixed op's C name from two long UDT names still fits SuiteSparse's 127 characters. + + SuiteSparse refuses a longer name and takes the definition under an empty one; the kernel + it compiles then calls no function and leaves its output uninitialized. + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + rec = [("ln_x", np.float64), ("ln_y", np.float64)] + rec32 = [("ln_x", np.float32), ("ln_y", np.float32)] + a = dtypes.register_new("LongNameA" + "a" * 60, np.dtype(rec, align=True)) + b = dtypes.register_new("LongNameB" + "b" * 60, np.dtype(rec32, align=True)) + op = binary.plus[a, b] + assert op.jit_c_source is not None + assert len(op.jit_c_name) < gb.core.lib.GxB_MAX_NAME_LEN + n = 40 + va = np.zeros(n, a.np_type) + va["ln_x"] = np.arange(n) + va["ln_y"] = 2 * np.arange(n) + vb = np.zeros(n, b.np_type) + vb["ln_x"] = 0.5 + vb["ln_y"] = np.arange(n) % 3 + u = Vector.from_coo(np.arange(n), va, dtype=a) + w = Vector.from_coo(np.arange(n), vb, dtype=b) + with _jit_mode("on"): + r = u.ewise_mult(w, binary.plus).new() + _, vals = r.to_coo() + np.testing.assert_array_equal(vals["ln_x"], np.arange(n) + 0.5) + np.testing.assert_array_equal(vals["ln_y"], 2 * np.arange(n) + np.arange(n) % 3) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_jit_c_name_too_long_keeps_the_cfunc(): + """An op whose C name SuiteSparse refuses (128 characters or more) runs its cfunc. + + Setting the definition after a refused name would compile a kernel that calls no + function, so the op gets no C form at all. + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + rec = np.dtype([("tl_x", np.float64), ("tl_y", np.int64)], align=True) + # A numpy UDT's C name leaves room for its ops' names, but one registered from C + # source keeps the name given: 125 characters fits, plus_ or eq_ and that does not. + name = "TooLongName" + "t" * 114 + t = dtypes.ss.register_new( + name, f"typedef struct {{ double tl_x ; int64_t tl_y ; }} {name} ;", np_type=rec + ) + assert len(t.jit_c_name) == 125 + n = 40 + vals = np.zeros(n, t.np_type) + vals["tl_x"] = np.arange(n) / 4 + vals["tl_y"] = np.arange(n) + v = Vector.from_coo(np.arange(n), vals, dtype=t) + plus_op = binary.plus[t] + eq_op = binary.eq[t] + assert plus_op.jit_c_source is None + assert eq_op.jit_c_source is None + with _jit_mode("on"): + r = v.ewise_mult(v, binary.plus).new() + same = v.ewise_mult(v, binary.eq).new() + _, got = r.to_coo() + np.testing.assert_array_equal(got["tl_x"], np.arange(n) / 2) + np.testing.assert_array_equal(got["tl_y"], 2 * np.arange(n)) + assert same.reduce(monoid.land).new() + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not _has_jit_set") +def test_udt_field_named_like_a_standard_macro_skips_the_jit(): + """A field named like a macro of the C standard headers the JIT includes skips the JIT. + + The preprocessor expands it inside the struct, so the kernel did not compile: + SuiteSparse raised ``JitError`` (with no message) for that op and set + ``jit_c_control`` to 'load' for the rest of the process. Asked of the C name + only, so nothing compiles here. + """ + from graphblas.exceptions import NoJITWarning + + fields = ["INT_MAX", "UINT8_MAX", "SIZE_MAX", "BUFSIZ", "EXIT_FAILURE", "PRId64", "EDOM"] + # and GraphBLAS's own macros, as GrB_INDEX_MAX, failed the same way + fields += ["SIGINT", "SIG_DFL", "GrB_INDEX_MAX", "GxB_IMPLEMENTATION_MAJOR", "GB_ANY"] + for i, field in enumerate(fields): + record = np.dtype([(field, np.float64), (f"macro_other_{i}", np.float64)], align=True) + udt = dtypes.register_anonymous(record) + assert udt.jit_c_name is None, field + v = Vector.from_coo([0], np.array([(1.0, 2.0)], dtype=record), dtype=udt, size=2) + with pytest.warns(NoJITWarning, match="without JIT"): + result = v.ewise_mult(v, binary.plus).new() + assert result[0].new().value.tolist() == (2.0, 4.0), field + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not _has_jit_set") +def test_udt_named_like_an_identifier_in_the_kernel_takes_the_jit(): + """A UDT named like an identifier in a JIT kernel compiles, under its ``gbudt_`` C name. + + SuiteSparse's kernels name the type beside their own identifiers. A UDT named like a + parameter of our op definitions (``z``), a local of its templates (``A``, ``p``, ``i``, + ``result``) or something its headers declare (``div``) failed to compile: an empty + ``JitError``, and 'load' for the rest of the process. + """ + _require_jit_on() + for name in ["z", "A", "p", "i", "div", "result"]: + record = np.dtype([(f"kn_{name}_a", np.float64), (f"kn_{name}_b", np.int64)], align=True) + udt = dtypes.register_anonymous(record, name) + assert udt.jit_c_name == f"gbudt_{name}" + vals = np.array([(1.0, 2), (3.0, 4)], dtype=record) + v = Vector.from_coo([0, 1], vals, dtype=udt) + A = gb.Matrix.from_coo([0, 1], [1, 0], vals, dtype=udt) + with _jit_mode("on"): + plus = v.ewise_mult(v, binary.plus).new() + biggest = v.reduce(monoid.max).new() + product = A.mxv(v, semiring.plus_times).new() + assert gb.ss.config["jit_c_control"] == "on", name + assert plus.to_coo()[1].tolist() == [(2.0, 4), (6.0, 8)], name + assert biggest.value.tolist() == (3.0, 4), name + assert product.to_coo()[1].tolist() == [(3.0, 8), (3.0, 8)], name + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not _has_jit_set") +def test_udt_long_name_gets_a_c_name_that_fits(): + """A UDT whose ``gbudt_`` C name would not leave room for its ops' names gets a synthetic one. + + SuiteSparse refuses a name of 128 characters or more. A type with such a name was kept + with an empty one, so its first kernel, a transpose say, raised ``JitError``; and with a + C name of 119 to 127 characters, ``floordiv_`` or ``plus_`` and it did not fit, so the + op lost its JIT with a NoJITWarning that blamed the fields. + """ + from graphblas.core.operator.udt_codegen import _max_type_c_name_len + + _require_jit_on() + assert _max_type_c_name_len() == 118 # 127 less floordiv_ + for length in [112, 113, 119, 200]: + record = np.dtype([(f"fit{length}_a", np.float64), (f"fit{length}_b", np.int64)]) + udt = dtypes.register_anonymous(record, "n" * length) + if length == 112: + assert udt.jit_c_name == "gbudt_" + "n" * 112 # 118 characters + else: + assert udt.jit_c_name.startswith("_gbudt_"), length + for op in [binary.floordiv[udt], binary.plus[udt], binary.eq[udt], unary.ainv[udt]]: + assert op.jit_c_source is not None, (length, op) + vals = np.array([(1.0, 2), (3.0, 4)], dtype=record) + A = gb.Matrix.from_coo([0, 1], [1, 0], vals, dtype=udt) + with _jit_mode("on"): + B = A.T.new() + C = A.ewise_mult(A, binary.plus).new() + assert B.to_coo()[2].tolist() == [(3.0, 4), (1.0, 2)], length + assert C.to_coo()[2].tolist() == [(2.0, 4), (6.0, 8)], length + + +def test_udt_c_name_keeps_python_graphblas_own_names(): + # A type python-graphblas makes itself has a name starting with _gb, which no kernel + # identifier starts with either, so it keeps that name as its C name; a user's name gets + # the gbudt_ prefix, even one that starts with it. + from graphblas.core.operator.udt_codegen import _pick_c_type_name + + assert _pick_c_type_name("Edge") == "gbudt_Edge" + assert _pick_c_type_name("gbudt_Edge") == "gbudt_gbudt_Edge" + for name in ["_gbudt_7", "_gbnest_3", "_gbinternal_fp64x16"]: + assert _pick_c_type_name(name) == name + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not _has_jit_set") +def test_udt_from_c_source_lifts_ops_under_its_own_c_name(): + """An op lifted to a type registered from C source names the type as its C source does. + + The prefix that keeps a numpy UDT's C name clear of kernel identifiers was also put on + a C-source type's name, so ``binary.plus[t]`` defined ``plus_gbudt_cpt_t`` on a type + ``gbudt_cpt_t`` that SuiteSparse does not know: ``JitError``, then 'load'. + """ + _require_jit_on() + np_type = np.dtype([("cs_x", np.float64), ("cs_y", np.float64)], align=True) + t = dtypes.ss.register_new( + "cs_lift_t", "typedef struct { double cs_x ; double cs_y ; } cs_lift_t ;", np_type=np_type + ) + assert t.jit_c_name == "cs_lift_t" + assert binary.plus[t].jit_c_name == "plus_cs_lift_t" + assert binary.eq[t].jit_c_name == "eq_cs_lift_t" + v = Vector.from_coo([0, 1], np.array([(1.0, 2.0), (3.0, 4.0)], dtype=np_type), dtype=t) + with _jit_mode("on"): + r = v.ewise_mult(v, binary.plus).new() + same = (v == v).new() + assert gb.ss.config["jit_c_control"] == "on" + assert r.to_coo()[1].tolist() == [(2.0, 4.0), (6.0, 8.0)] + assert same.to_coo()[1].tolist() == [True, True] + # A numpy UDT's own name always gets the prefix, even one that starts with it + udt = dtypes.register_anonymous(np.dtype([("cs_pre", np.float64)]), "gbudt_CsPre") + assert udt.jit_c_name == "gbudt_gbudt_CsPre" diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index db80cdf71..509b74173 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1,8 +1,11 @@ +import cmath import inspect import itertools +import math import pickle import sys import types +import warnings import weakref import numpy as np @@ -12,6 +15,7 @@ import graphblas as gb from graphblas import agg, backend, binary, dtypes, indexunary, monoid, select, semiring, unary from graphblas.core import _supports_udfs as supports_udfs +from graphblas.core import lib from graphblas.exceptions import ( DimensionMismatch, DomainMismatch, @@ -27,7 +31,6 @@ from graphblas import Matrix, Scalar, Vector # isort:skip (for dask-graphblas) - suitesparse = backend == "suitesparse" if suitesparse: ss_version_major = gb.core.ss.version_major @@ -167,6 +170,36 @@ def test_from_coo_scalar(): u.ss.iso_value +@pytest.mark.skipif("not suitesparse") +def test_ss_iso_value_of_an_empty_iso_vector(): + # As for Matrix: an empty Vector SuiteSparse marks iso has an empty iso_value. + u = Vector.from_coo([], 7, size=3) + w = Vector.from_coo([1], [1.0], size=3).select("==", 2).new() + x = Vector(int, 3) + assert u.ss.is_iso + assert not x.ss.is_iso + for v in [u, w, x]: + if v.ss.is_iso: + value = v.ss.iso_value + assert isinstance(value, Scalar) + assert value.dtype == v.dtype + assert value.is_empty + else: + with pytest.raises(ValueError, match="not iso"): + v.ss.iso_value + + +@pytest.mark.skipif("not suitesparse") +@pytest.mark.parametrize("method", ["export", "unpack"]) +def test_ss_export_full_of_a_vector_that_is_not_full(method): + # SuiteSparse refuses with GrB_INVALID_VALUE and sets no message, so the + # error used to be blank. + v = Vector.from_coo([1], [2.0], size=3) + with pytest.raises(InvalidValue, match="1 of 3 entries is not full"): + getattr(v.ss, method)("full") + assert v.isequal(Vector.from_coo([1], [2.0], size=3)) + + def test_from_coo_default_values(): u = Vector.from_coo([1, 3]) v = Vector.from_coo([1, 3], [1.0, 1.0]) @@ -750,6 +783,11 @@ def test_select(v): w8 = select.index(v >= 4).new() expected = Vector.from_coo([4, 6], [2, 0], size=7) assert w8.isequal(expected) + # The index is compared with the bound given, not with 300 typed beside INT8 (inf) + w = Vector.from_coo([0, 200, 399], 1, dtype=dtypes.INT8, size=400) + assert select.index(w < 300).new().to_coo()[0].tolist() == [0, 200] + assert select.index(w >= 300).new().to_coo()[0].tolist() == [399] + assert select.index(w > 2.5).new().to_coo()[0].tolist() == [200, 399] @autocompute @@ -1101,6 +1139,77 @@ def test_isclose(v): assert not u4.isclose(u4b) +def test_isclose_infinity(): + # As in math.isclose, an infinity is close only to itself. It was close to every + # number (abs(inf - y) is inf, and so is rel_tol * inf), with Numba and without. + from graphblas.core.vector import _isclose_recipe + + values = [math.inf, -math.inf, 1.0, 0.0, math.nan] + for x, y in itertools.product(values, repeat=2): + left = Vector.from_coo([0, 2], [x, 5.0], size=3) + right = Vector.from_coo([0, 2], [y, 5.0], size=3) + for abs_tol in [0.0, math.inf]: + expected = math.isclose(x, y, rel_tol=1e-7, abs_tol=abs_tol) + if supports_udfs: + assert left.isclose(right, abs_tol=abs_tol) is expected, (x, y) + assert _isclose_recipe(left, right, 1e-7, abs_tol) is expected, (x, y) + + +# (x's dtype, y's dtype, x, y, rel_tol): each was wrong in at least one of isclose's +# paths, because x - y or abs(x) wrapped in an integer type. The values are far enough +# apart (or exactly on the tolerance) that FP64 gives math.isclose's answer too. +ISCLOSE_WRAPPING_CASES = [ + ("UINT64", "UINT64", 0, 2**64 - 1, 1e-7), + ("UINT64", "UINT64", 2**63, 2**64 - 1, 1e-7), + ("UINT32", "UINT32", 5, 6, 0.5), + ("UINT8", "UINT8", 5, 6, 0.5), + ("UINT8", "UINT64", 0, 2**64 - 1, 1e-7), + ("INT64", "INT64", -(2**63), 2**63 - 1, 1e-7), + ("INT64", "INT64", -(2**63), 0, 1.0), + ("INT8", "INT8", -128, 127, 0.5), + ("INT8", "INT8", -128, 0, 1.0), + ("INT8", "INT64", 0, -(2**63), 1e-7), + ("INT8", "UINT64", -1, 2**64 - 1, 1e-7), +] + + +def test_isclose_integers_do_not_wrap(): + from graphblas.core.vector import _isclose_recipe + + for xtype, ytype, x, y, rel_tol in ISCLOSE_WRAPPING_CASES: + expected = math.isclose(x, y, rel_tol=rel_tol) + for (lt, lv), (rt, rv) in [((xtype, x), (ytype, y)), ((ytype, y), (xtype, x))]: + left = Vector.from_coo([0, 2], [lv, 5], dtype=lt, size=3) + right = Vector.from_coo([0, 2], [rv, 5], dtype=rt, size=3) + if supports_udfs: + assert left.isclose(right, rel_tol=rel_tol) is expected, (lt, lv, rt, rv) + assert _isclose_recipe(left, right, rel_tol, 0.0) is expected, (lt, lv, rt, rv) + + +@pytest.mark.skipif("not dtypes._supports_complex") +def test_isclose_complex(): + # Without Numba, isclose kept only the real part of x - y, so 0 was close to 1j + from graphblas.core.vector import _isclose_recipe + + values = [0j, 1j, -1j, 1 + 1j, 3 + 4j, 3 + 4.000001j] + for x, y in itertools.product(values, repeat=2): + expected = cmath.isclose(x, y, rel_tol=1e-6) + left = Vector.from_coo([0, 2], [x, 5], dtype="FC64", size=3) + right = Vector.from_coo([0, 2], [y, 5], dtype="FC64", size=3) + if supports_udfs: + assert left.isclose(right, rel_tol=1e-6) is expected, (x, y) + assert _isclose_recipe(left, right, 1e-6, 0.0) is expected, (x, y) + # abs(x) and abs(y) are in the type x - y is in, here FC64; abs(1 + 1j) in FC32 is + # smaller than in FC64, so the case exactly on the tolerance was not close + left = Vector.from_coo([0], [0.0]) + right = Vector.from_coo([0], [1 + 1j], dtype="FC32") + expected = cmath.isclose(0.0, complex(np.complex64(1 + 1j)), rel_tol=1.0) + assert expected + if supports_udfs: + assert left.isclose(right, rel_tol=1.0) is expected + assert _isclose_recipe(left, right, 1.0, 0.0) is expected + + def test_binary_op(v): v2 = v.dup() v2[1] = 0 @@ -1493,6 +1602,345 @@ def test_ss_nbytes(v): assert v.ss.nbytes > 0 +def test_python_number_is_weak_beside_builtin_dtypes(): + """A Python number takes a built-in collection's dtype when that kind holds it (NEP 50). + + How a literal was typed used to depend on the installed numpy: on numpy 2 a + Python int was INT64 and a float FP64 beside anything, and on numpy 1 the + collection's dtype won by value, so ``int8_vec + 300`` wrapped and + ``int8_vec < 300`` compared with 44. python-graphblas now applies numpy 2's + rule itself, the table UDT elements use, on numpy 1 and 2 alike. numpy + scalars and Scalars keep their dtype. + """ + i8 = Vector.from_coo([0, 1], [100, -100], dtype=dtypes.INT8) + u8 = Vector.from_coo([0, 1], [0, 200], dtype=dtypes.UINT8) + f32 = Vector.from_coo([0, 1], [0.1, 1.5], dtype=dtypes.FP32) + u64 = Vector.from_coo([0, 1], [0, 2**63], dtype=dtypes.UINT64) + bools = Vector.from_coo([0, 1], [True, False]) + for expr, dtype, values in [ + (i8 + 1, dtypes.INT8, [101, -99]), + (1 + i8, dtypes.INT8, [101, -99]), + (i8 + True, dtypes.INT8, [101, -99]), + (binary.plus(i8, 1), dtypes.INT8, [101, -99]), + (i8.apply(binary.minus, left=1), dtypes.INT8, [-99, 101]), + (i8.apply(monoid.max, right=1), dtypes.INT8, [100, 1]), + (i8.apply("+", right=1), dtypes.INT8, [101, -99]), + (1 - u8, dtypes.UINT8, [1, 57]), # wraps, as numpy 2 does + (i8 * 0.5, dtypes.FP64, [50, -50]), + (u64 + 1, dtypes.UINT64, [1, 2**63 + 1]), + (bools + 1, dtypes.INT64, [2, 1]), + (f32 * 2, dtypes.FP32, [np.float32(0.2), 3]), + # Strong: numpy scalars, Scalars, and the literal of a typed op + (i8 + np.int64(1), dtypes.INT64, [101, -99]), + (i8 + Scalar.from_value(1), dtypes.INT64, [101, -99]), + (binary.plus["INT64"](i8, 1), dtypes.INT64, [101, -99]), + (binary.plus["FP64"](i8, 300), dtypes.FP64, [400, 200]), + ]: + assert expr.dtype == dtype + assert expr.new().to_coo()[1].tolist() == values + assert (f32 * 0.5).dtype == dtypes.FP32 + if dtypes._supports_complex: + assert (f32 + 1j).dtype == dtypes.FC32 + else: + # Without complex types, a complex number keeps its own type, which plus does not + # take, and which compares only with itself. + for expr in [lambda: f32 + 1j, lambda: f32 == 1j]: + with pytest.raises(KeyError, match="does not work with"): + expr() + s8 = Scalar.from_value(100, dtype=dtypes.INT8) + assert (s8 + 1).new().dtype == (1 + s8).new().dtype == dtypes.INT8 + assert (Scalar.from_value(0.5, dtype=dtypes.FP32) * 0.5).new().dtype == dtypes.FP32 + A = Matrix.from_coo([0], [0], [100], dtype=dtypes.INT8) + assert (A * 2).dtype == dtypes.INT8 + # Beside a bool, an int is INT64, and must fit it, as in numpy's ``bools & 2**70``. + exprs = [ + lambda: i8 + 300, + lambda: u8 + -1, + lambda: s8 * 300, + lambda: bools + 2**70, + lambda: bools.apply(binary.band, right=2**70), + ] + if supports_udfs: # integer floordiv is a Numba op + exprs.append(lambda: i8 // 300) + for expr in exprs: + with pytest.raises(OverflowError, match="out of bounds"): + expr() + # truediv gives FP64 whatever the int, so an int out of range does not raise: + # numpy divides integers in float64 too. + assert (i8 / 300).dtype == (300 / i8).dtype == dtypes.FP64 + assert (i8 / 300).new().to_coo()[1].tolist() == [100 / 300, -100 / 300] + # A float too large for FP32 becomes infinity with a warning, as in numpy 2. + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + result = (f32 + 1e300).new() + assert result.dtype == dtypes.FP32 + assert np.isinf(result.to_coo()[1]).all() + assert {w.filename for w in record} == {__file__} + with pytest.warns(RuntimeWarning, match="overflow encountered in cast"): + assert np.isinf((f32 + 2**200).new().to_coo()[1]).all() # an int too + # A comparison keeps no result type, so an int out of range compares exactly. + assert (i8 < 300).new().to_coo()[1].tolist() == [True, True] + assert (i8 == 300).new().to_coo()[1].tolist() == [False, False] + assert (u8 > -1).new().to_coo()[1].tolist() == [True, True] + assert i8.apply(indexunary.valuelt, 300).new().to_coo()[1].tolist() == [True, True] + assert (f32 == 0.1).new().to_coo()[1].tolist() == [True, False] # compared as float32 + # A float too large for FP32 is infinity in a comparison too, with a warning, as in + # numpy 2 (an int out of range is exact; a float has no type that is both FP32 and + # holds it). The Scalar's == warns once, at the caller's line. + finf = Vector.from_coo([0, 1], [np.inf, 1], dtype=dtypes.FP32) + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + assert (finf == 1e300).new().to_coo()[1].tolist() == [True, False] + assert {w.filename for w in record} == {__file__} + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + assert Scalar.from_value(np.inf, dtype=dtypes.FP32) == 1e300 + assert [w.filename for w in record] == [__file__] + # ewise_union defaults and fill values follow the rule too. + w8 = Vector.from_coo([1], [1], dtype=dtypes.INT8, size=2) + assert i8.ewise_union(w8, binary.plus, 0, 0).dtype == dtypes.INT8 + assert i8.ewise_union(w8, binary.plus, 0.5, 0).dtype == dtypes.FP64 + assert i8.ewise_union(w8, "<", 0, 300).new().to_coo()[1].tolist() == [True, True] + sparse = Vector.from_coo([0], [5], dtype=dtypes.INT8, size=2) + assert sparse.to_dense(fill_value=0).dtype == np.int8 + assert sparse.to_dense(fill_value=300).tolist() == [5, 300] # widened, not lost + sparse32 = Vector.from_coo([0], [5], dtype=dtypes.FP32, size=2) + assert sparse32.to_dense(fill_value=0.5).dtype == np.float32 + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert sparse32.to_dense(fill_value=1e300).tolist() == [5, 1e300] # widened + assert sparse32.to_dense(fill_value=2**200).tolist() == [5, 2.0**200] + # A thunk that is not compared with the values keeps its own type. + assert i8.select(select.rowle, 2**40).new().nvals == 2 + if supports_udfs: + above = indexunary.register_anonymous(lambda x, i, j, thunk: x > thunk) + assert i8.apply(above, 300).new().to_coo()[1].tolist() == [False, False] + # A user's BinaryOp computes in the collection's type, as plus does. + plus = binary.register_anonymous(lambda x, y: x + y) + with pytest.raises(OverflowError, match="out of bounds"): + i8.apply(plus, right=300) + # Any op whose result is another type, as truediv's and a comparison's are, takes + # an int out of range exactly: a parameterized op, a user's, and the instances a + # parameterized op makes (these raised OverflowError; main and numpy give False). + gt = binary.register_anonymous(lambda x, y: x > y) + for op, values in [ + (binary.isclose, [False, False]), + (binary.isclose(0.1), [False, False]), + (gt, [False, False]), + ]: + assert i8.apply(op, right=300).new().to_coo()[1].tolist() == values + assert i8.apply(op, left=-300).new().to_coo()[1].tolist() == values + result = i8.ewise_union(w8, binary.isclose, 0, 300).new() + assert result.to_coo()[1].tolist() == [False, False] + # A user's op named like a comparison is not one: the int is not made infinity + # (it gave [-inf, -inf] and [True, True]; main gives these, as numpy would) + + def ge(x, y): + return x * 1.0 - y + + def eq(x, y): + return y - x > 250 + + result = i8.apply(binary.register_anonymous(ge), right=300).new() + assert result.to_coo()[1].tolist() == [-200.0, -400.0] + result = i8.apply(binary.register_anonymous(eq), right=300).new() + assert result.to_coo()[1].tolist() == [False, True] + # Where the op and the int would compute in a float type, which rounds an int + # beyond 64 bits, a user's op raises (it gave [False, True]: 2**63 - 1 < 2**63 is + # False in FP64). A built-in op that computes in float anyway does not. + i64 = Vector.from_coo([0, 1], [2**63 - 1, 0], dtype=dtypes.INT64) + u64 = Vector.from_coo([0, 1], [2**64 - 1, 0], dtype=dtypes.UINT64) + lt = binary.register_anonymous(lambda x, y: x < y) + if not hasattr(binary, "_test_weak_lt"): + binary.register_new("_test_weak_lt", lambda x, y: x < y) + for vec, value in [(i64, 2**63), (i64, -(2**63) - 1), (u64, 2**64)]: + for op in [lt, binary._test_weak_lt]: + with pytest.raises(OverflowError, match="out of bounds"): + vec.apply(op, right=value) + assert vec.apply(binary.lt, right=value).new().to_coo()[1].tolist() == [value > 0] * 2 + result = i64.apply(binary.isclose, right=2**63).new() + assert result.to_coo()[1].tolist() == [True, False] # as numpy's isclose + # Built-in ops made outside register_new (truediv) or on first use (binom) count + # as built-in too: truediv computes beside a 64-bit int, as numpy divides + # (these raised OverflowError). + assert (i64 / 2**63).new().to_coo()[1].tolist() == [1.0, 0.0] + assert (u64 / -1).new().to_coo()[1].tolist() == [-(2.0**64), 0.0] + assert (i8 / 2**70).new().to_coo()[1].tolist() == [100 / 2**70, -100 / 2**70] + from graphblas.core.literals import _is_builtin_op + + assert all( + _is_builtin_op(op) + for op in [binary.truediv, binary.rtruediv, binary.cdiv, binary.binom] + ) + # Typed scalars promote as numpy 2 promotes them, on numpy 1 too. + assert dtypes.unify(dtypes.INT8, dtypes.INT64, is_right_scalar=True) == dtypes.INT64 + + +def test_python_int_beyond_64_bit_element_compares_exactly(): + """A Python int beyond a 64-bit integer element's range compares exactly, as in numpy 2. + + It compares as infinity of its sign. A type that held both would round + them: in FP64, the INT64 ``2**63 - 1`` and the int ``2**63`` are equal, so + ``int64_vec < 2**63`` was False there (main raised OverflowError). + """ + compare = { + "eq": lambda x, y: x == y, + "ne": lambda x, y: x != y, + "lt": lambda x, y: x < y, + "le": lambda x, y: x <= y, + "gt": lambda x, y: x > y, + "ge": lambda x, y: x >= y, + } + i64_values = [-(2**63), -(2**63) + 1, 0, 2**63 - 2, 2**63 - 1] + u64_values = [0, 2**63, 2**64 - 2, 2**64 - 1] + i64 = Vector.from_coo(range(5), i64_values, dtype=dtypes.INT64) + u64 = Vector.from_coo(range(4), u64_values, dtype=dtypes.UINT64) + for vec, values, numbers in [ + (i64, i64_values, [2**63, -(2**63) - 1, 2**64, -(2**64), 10**400]), + (u64, u64_values, [2**64, 2**64 + 1, -1, -(10**400)]), + ]: + empty = Vector(vec.dtype, size=vec.size) + for name, py_op in compare.items(): + for number in numbers: + expected = [py_op(x, number) for x in values] + for expr in [ + py_op(vec, number), + vec.apply(getattr(binary, name), right=number), + vec.apply(getattr(indexunary, f"value{name}"), number), + vec.ewise_union(empty, getattr(binary, name), 0, number), + ]: + assert expr.new().to_coo()[1].tolist() == expected + kept = vec.select(getattr(select, f"value{name}"), number).new() + assert kept.to_coo()[0].tolist() == [i for i, x in enumerate(expected) if x] + flipped = py_op(number, vec).new() + assert flipped.to_coo()[1].tolist() == [py_op(number, x) for x in values] + A = Matrix.from_coo([0, 1], [0, 1], [2**63 - 1, -(2**63)]) + assert (A < 2**63).new().to_coo()[2].tolist() == [True, True] + assert (A == -(2**63) - 1).new().to_coo()[2].tolist() == [False, False] + s = Scalar.from_value(2**63 - 1) + assert (s == 2**63) is False + assert s != 2**63 + assert (s < 2**63).new().value + assert (Scalar.from_value(2**64 - 1, dtype=dtypes.UINT64) == 2**64) is False + # A float element compares with the int as a float, as numpy 2 does. + assert (Vector.from_coo([0], [2.0**63]) == 2**63 + 1).new().to_coo()[1].tolist() == [True] + + +def test_select_value_thunk_of_another_type(): + """Select with a value op compares exactly when the thunk's type differs from the values'. + + SuiteSparse:GraphBLAS's select before 10.5.2 read such a thunk's bytes as the + values' type: on an INT8 vector, ``valueeq`` with 1.0 compared with 0 and kept + nothing, ``valuelt`` with 300 compared with 44, and ``valueeq`` with 356 + kept 100. python-graphblas now compares in the values' own type with a + thunk that keeps the same elements. 10.5.2 casts both to the op's type, + which rounds the 2**53 + 1 thunks below to 2**53 and keeps that element. + """ + i8 = Vector.from_coo([0, 1, 2, 3], [-5, 1, 2, 100], dtype=dtypes.INT8) + u8 = Vector.from_coo([0, 1], [0, 255], dtype=dtypes.UINT8) + f32 = Vector.from_coo([0, 1, 2], [np.float32(0.1), 1.0, np.nan], dtype=dtypes.FP32) + bools = Vector.from_coo([0, 1], [False, True]) + for x, op, thunk, kept in [ + (i8, select.valueeq, 1.0, [1]), + (i8, select.valueeq, 356, []), + (i8, select.valuene, 356, [0, 1, 2, 3]), + (i8, select.valuelt, 300, [0, 1, 2, 3]), + (i8, select.valuege, 2.5, [3]), + (i8, select.valuele, 2.5, [0, 1, 2]), + (i8, select.valuegt, -300, [0, 1, 2, 3]), + (i8, select.valuegt, np.nan, []), + (i8, select.valuene, np.nan, [0, 1, 2, 3]), + (i8, select.valuelt, np.inf, [0, 1, 2, 3]), + (i8, select.valuege, np.inf, []), + (i8, select.valuele, -np.inf, []), + (i8, indexunary.valueeq, np.float64(2.0), [2]), + (i8, select.valueeq, Scalar.from_value(1.0), [1]), + (i8, select.valuelt, Scalar.from_value(300, dtype=dtypes.INT64), [0, 1, 2, 3]), + (u8, select.valuelt, np.int64(-1), []), + (u8, select.valuegt, np.int64(-1), [0, 1]), + (bools, select.valueeq, 256, []), + (bools, select.valueeq, 1, [1]), + # A strong FP64 thunk compares in FP64: float32(0.1) is not 0.1. + (f32, select.valueeq, np.float64(0.1), []), + (f32, select.valuegt, np.float64(0.1), [0, 1]), + (f32, select.valuene, np.float64(0.1), [0, 1, 2]), + (f32, select.valuelt, np.float64(np.nan), []), + (f32, select.valuene, np.float64(np.nan), [0, 1, 2]), + # A Python float beside FP32 is FP32 (NEP 50), as numpy compares it. + (f32, select.valueeq, 0.1, [0]), + ]: + assert x.select(op, thunk).new().to_coo()[0].tolist() == kept, (x.dtype, op, thunk) + if dtypes._supports_complex: + # A real value equals a complex one only when its imaginary part is 0, and + # a complex value equals a thunk of another type only when that converts + # exactly. Complex values have no order, so only valueeq and valuene apply. + fc32 = Vector.from_coo( + [0, 1, 2, 3, 4, 5, 6], + [1, 2 + 1j, 3, np.nan, np.complex64(0.1), np.inf, np.complex64(1 + 0.1j)], + dtype=dtypes.FC32, + ) + fc64 = Vector.from_coo([0, 1, 2], [1, 2 + 1j, 0.1], dtype=dtypes.FC64) + # Compared exactly, as floats are: 2**53 + 1 is not the float it rounds to. + big32 = Vector.from_coo([0, 1], [2**53, 2**24], dtype=dtypes.FC32) + bigf32 = Vector.from_coo([0, 1], [2**53, 2**24], dtype=dtypes.FP32) + for x, op, thunk, kept in [ + (i8, select.valueeq, 1 + 0j, [1]), + (i8, select.valueeq, 1 + 2j, []), + (i8, select.valuene, 1 + 2j, [0, 1, 2, 3]), + (f32, select.valueeq, np.complex64(1), [1]), + (f32, select.valueeq, np.complex128(0.1), []), # float32(0.1) is not 0.1 + (f32, select.valuene, np.complex128(1 + 2j), [0, 1, 2]), + (fc32, select.valueeq, np.complex128(1), [0]), + (fc32, select.valueeq, np.complex128(2 + 1j), [1]), + (fc32, select.valueeq, np.complex64(2 + 1j), [1]), + (fc32, select.valueeq, 3, [2]), + (fc32, select.valuene, np.complex128(2 + 1j), [0, 2, 3, 4, 5, 6]), + (fc32, select.valueeq, np.float64(3), [2]), + (fc32, select.valueeq, np.int64(3), [2]), + (fc32, select.valuene, np.int64(3), [0, 1, 3, 4, 5, 6]), + (fc32, select.valueeq, np.complex64(0.1), [4]), + (fc32, select.valueeq, np.complex128(0.1), []), # complex64(0.1) is not 0.1 + (fc32, select.valuene, np.complex128(0.1), [0, 1, 2, 3, 4, 5, 6]), + (fc32, select.valueeq, np.complex128(complex(np.nan, 0)), []), + (fc32, select.valueeq, np.float64(np.inf), [5]), + (fc32, select.valueeq, np.complex128(1e300), []), # inf as complex64, not 1e300 + (fc32, select.valuene, np.complex128(1e300), [0, 1, 2, 3, 4, 5, 6]), + (fc32, select.valueeq, np.complex64(1 + 0.1j), [6]), + (fc32, select.valueeq, np.complex128(1 + 0.1j), []), # in the imaginary part + (fc32, indexunary.valueeq, np.complex128(2 + 1j), [1]), + (fc32, select.valueeq, Scalar.from_value(np.complex128(2 + 1j)), [1]), + (big32, select.valueeq, np.int64(2**53), [0]), + (big32, select.valueeq, np.int64(2**53 + 1), []), + (bigf32, select.valueeq, np.int64(2**53 + 1), []), + (fc64, select.valueeq, np.complex128(2 + 1j), [1]), + (fc64, select.valueeq, np.complex64(0.1), []), # complex64(0.1) is not 0.1 + (fc64, select.valueeq, np.float64(0.1), [2]), + (fc64, select.valuene, np.int64(1), [1, 2]), + (fc64, select.valueeq, 1, [0]), + ]: + assert x.select(op, thunk).new().to_coo()[0].tolist() == kept, (x.dtype, op, thunk) + with pytest.raises(KeyError, match="valuelt does not work with FC64"): + fc32.select(select.valuelt, np.complex128(1)) + C = Matrix.from_coo([0, 0, 1], [0, 1, 1], [1, 2 + 1j, 3], dtype=dtypes.FC32) + equal = C.select(select.valueeq, np.complex128(2 + 1j)).new() + assert equal.to_coo()[2].tolist() == [2 + 1j] + assert C.select(select.valuene, np.int64(3)).new().nvals == 2 + # A float64 too large for FP32 is compared as it is, without numpy's overflow warning. + with warnings.catch_warnings(): + warnings.simplefilter("error") + for op, thunk, kept in [ + (select.valueeq, 1e300, []), + (select.valuene, 1e300, [0, 1, 2]), + (select.valuelt, 1e300, [0, 1]), + (select.valuegt, 1e300, []), + (select.valuegt, -1e300, [0, 1]), + ]: + selected = f32.select(op, np.float64(thunk)).new().to_coo()[0].tolist() + assert selected == kept, (op, thunk) + with pytest.raises(EmptyObject): + i8.select(select.valueeq, Scalar(dtypes.INT64)).new() + A = Matrix.from_coo([0, 0, 1], [0, 1, 1], [1, 100, -5], dtype=dtypes.INT8) + assert A.select(select.valueeq, 1.0).new().to_coo()[2].tolist() == [1] + assert A.select(select.valuelt, 300).new().nvals == 3 + + def test_inner(v): R = Matrix(v.dtype, nrows=1, ncols=v.size) # row vector C = Matrix(v.dtype, nrows=v.size, ncols=1) # column vector @@ -1601,7 +2049,7 @@ def test_auto(v): # print(type(expr).__name__, method) val1 = getattr(expected, method)(expected).new() if method in {"__or__", "__ror__"} and type(expr) is VectorEwiseMultExpr: - # Doing e.g. `plus(x & y | z)` isn't allowed--make user be explicit + # Doing e.g. `plus(x & y | z)` isn't allowed; make user be explicit with pytest.raises(TypeError): val2 = getattr(expected, method)(expr) with pytest.raises(TypeError): @@ -1658,6 +2106,7 @@ def test_expr_is_like_vector(v): "__call__", "__del__", "__delitem__", + "__getattr__", "__lshift__", "__setitem__", "_assign_element", @@ -1708,6 +2157,7 @@ def test_index_expr_is_like_vector(v): expected = { "__del__", "__delitem__", + "__getattr__", "__setitem__", "_assign_element", "_delete_element", @@ -2097,11 +2547,12 @@ def test_udt(): assert s == t t = v.reduce(monoid.any, allow_empty=False).new() assert s == t - v << binary.first(1, v) + # A literal fills every field by type; an int does not fit the bool field, True does. + v << binary.first(True, v) assert v.isequal(ones) v << binary.first(s, v) assert v.isequal(zeros) - v << binary.second(v, (1, 1)) + v << binary.second(v, (True, 1)) assert v.isequal(ones) v << binary.second(v, s) assert v.isequal(zeros) @@ -2221,11 +2672,13 @@ def test_udt(): if suitesparse: vv = Vector.ss.deserialize(v.ss.serialize(), dtype=long_udt) assert v.isequal(vv, check_dtype=True) - if ss_version_major < 9: - with pytest.raises(SyntaxError): - # The size of the UDT name is limited - Vector.ss.deserialize(v.ss.serialize()) - else: + # The dtype name is too long to embed in the blob; dtype= must be provided. + # Whichever error the caller sees comes from parsing the name that is in + # the blob. Builds with GxB_Serialized_get_String let that SyntaxError + # through; older ones have no second name to try, so they wrap it. + with pytest.raises( + SyntaxError if hasattr(lib, "GxB_Serialized_get_String") else ValueError + ): Vector.ss.deserialize(v.ss.serialize()) # May be able to look up non-anonymous dtypes by name if their names are too long named_long_dtype = np.dtype([("x", np.bool_), ("y" * 1000, np.float64)], align=False) @@ -2639,6 +3092,17 @@ def test_ss_sort(v): w, p = v.ss.sort(monoid.lxor) # Weird, but user-defined monoids may not commute, so okay +def test_subarray_dtypes_iso(): + # An iso array UDT exports its one value with a leading axis of length 1 + v = Vector("INT64[4]", size=3) + v[:] = [0, 1, 2, 3] + if suitesparse: + assert v.ss.is_iso + expected = np.broadcast_to([0, 1, 2, 3], (3, 4)) + np.testing.assert_array_equal(v.to_dense(), expected) + np.testing.assert_array_equal(v.to_coo()[1], expected) + + @pytest.mark.skipif("not supports_udfs") def test_subarray_dtypes(): a = np.arange(3 * 4, dtype=np.int64).reshape(3, 4) @@ -2686,3 +3150,99 @@ def test_subarray_dtypes(): assert full1.isequal(full2, check_dtype=True) full2 = Vector.ss.import_bitmap(values=a, bitmap=[True, True, True]) assert full1.isequal(full2, check_dtype=True) + + +def test_constructor_rejects_arraylike_first_arg(): + # The first positional arg is the dtype; passing data instead used to raise a + # confusing "Unknown dtype" ValueError. It should now raise TypeError pointing + # at the from_* constructors. + with pytest.raises(TypeError, match="Vector.*dtype.*from_coo.*from_dense"): + Vector([1, 2, 3]) + with pytest.raises(TypeError, match="Vector.*expects a dtype"): + Vector((1, 2, 3)) + with pytest.raises(TypeError, match="Vector.*expects a dtype"): + Vector(np.array([1, 2, 3])) + # Valid dtype-first signatures still work, including list/tuple dtype specs + assert Vector(int, size=3).dtype == dtypes.INT64 + assert Vector("INT64", size=3).dtype == dtypes.INT64 + assert Vector(np.dtype("int64"), size=3).dtype == dtypes.INT64 + assert Vector([("x", "i8"), ("y", "f8")], size=3).dtype._is_udt + # A non-array-like bad dtype keeps the original ValueError + with pytest.raises(ValueError, match="Unknown dtype"): + Vector("not_a_dtype", size=3) + # A malformed structured dtype spec gets the hint, and the lookup's reason stays + # attached as the cause + with pytest.raises(TypeError, match="If you meant to pass values") as excinfo: + Vector([("a", "zzz")], 3) + assert excinfo.value.__cause__ is not None + + +def test_property_attribute_error_keeps_its_message(): + # A property that raises AttributeError sends Python to BaseType.__getattr__, + # which must not replace the property's own message with its generic one. + with pytest.raises(AttributeError, match="only has `__index__` for integral dtypes"): + Scalar.from_value(1.5).__index__ + if not suitesparse: + with pytest.raises(AttributeError, match='only available with "suitesparse" backend'): + Vector(int, 2).ss + + +# Without SuiteSparse, building a UDT uses dup_op=binary.any, which needs Numba. +@pytest.mark.skipif("not suitesparse and not supports_udfs") +def test_subarray_dtype_values_shape(): + # Each of these has one entry per index but not two numbers per entry. Building + # from too few numbers read past the end of the array. (The dtype's name is + # whichever another test gave this anonymous dtype.) + match = r"values must have shape \(2, 2\) for dtype \S+, not \(2," + for values in [np.ones(2), [1.0, 2.0], np.ones((2, 1)), np.arange(6.0).reshape(2, 3)]: + with pytest.raises(ValueError, match=match): + Vector.from_coo([0, 1], values, "FP64[2]") + v = Vector("FP64[2]", 3) + with pytest.raises(ValueError, match=match): + v.build([0, 1], values) + with pytest.raises(ValueError, match=r"shape \(3, 2\) for dtype \S+, not \(3, 1\)"): + Vector.from_dense(np.ones((3, 1)), dtype="FP64[2]") + v = Vector("FP64[2]", 3) + with pytest.raises(ValueError, match="shape mismatch"): + v[[0, 1]] = np.ones((2, 1)) + with pytest.raises(ValueError, match=r"shape \(2,\) for dtype FP64, not \(2, 0\)"): + Vector.from_coo([0, 1], np.ones((2, 0)), "FP64") + # A record is one element, so a float fills each field, and a 2d array is too many + rec = dtypes.register_anonymous(np.dtype([("vs_x", np.float64), ("vs_y", np.int32)])) + v = Vector.from_coo([0, 1], np.ones(2), rec) + assert v.to_coo()[1].tolist() == [(1.0, 1), (1.0, 1)] + with pytest.raises(ValueError, match=r"shape \(2,\) for dtype"): + Vector.from_coo([0, 1], np.arange(4.0).reshape(2, 2), rec) + # A scalar fills whole elements + v = Vector.from_coo([0, 1, 2], 1.5, "FP64[2]") + assert v.to_coo()[1].tolist() == [[1.5, 1.5]] * 3 + # The lengths compared are the number of elements, not of numbers (it said 2 != 6) + with pytest.raises(ValueError, match="lengths must match: 2 != 3$"): + Vector.from_coo([0, 1], np.ones((3, 2)), "FP64[2]") + + +@pytest.mark.skipif("not suitesparse") +def test_ss_import_values_size(): + # GraphBLAS before 10.5.2 compares the byte size of imported values with the number + # of entries, not the bytes they need, so it read past the end of each of these + with pytest.raises(ValueError, match="have 4 or more elements of dtype FP64, not 2"): + Vector.ss.import_sparse(size=5, indices=[0, 1, 2, 3], values=np.ones(2), nvals=4) + with pytest.raises(ValueError, match="have 5 or more elements of dtype FP64, not 3"): + Vector.ss.import_full(values=np.ones(3), size=5) + v = Vector(float, 5) + with pytest.raises(ValueError, match="have 5 or more elements of dtype FP64, not 3"): + v.ss.pack_full(values=np.ones(3)) + # The dtype's name is whichever another test gave this anonymous dtype + with pytest.raises(ValueError, match=r"have 2 or more elements of dtype \S+, not 1"): + Vector.ss.import_sparse(size=3, indices=[0, 1], values=np.ones(2), dtype="FP64[2]") + with pytest.raises(ValueError, match=r"have 3 or more elements of dtype \S+, not 1"): + Vector.ss.import_bitmap(bitmap=[True, True, True], values=np.ones(3), dtype="FP64[2]") + # One element is enough for iso + v = Vector.ss.import_full(values=[4.0], size=5, is_iso=True) + assert v.ss.is_iso + assert v.to_coo()[1].tolist() == [4.0] * 5 + v = Vector.ss.import_sparse( + size=3, indices=[0, 1], values=[2.0, 3.0], dtype="FP64[2]", is_iso=True + ) + assert v.ss.is_iso + assert v.to_coo()[1].tolist() == [[2.0, 3.0]] * 2 diff --git a/graphblas/unary/__init__.py b/graphblas/unary/__init__.py index b83ea3b8b..e48a6a387 100644 --- a/graphblas/unary/__init__.py +++ b/graphblas/unary/__init__.py @@ -42,7 +42,9 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/viz.py b/graphblas/viz.py index b6d5f6ba7..8e2a53228 100644 --- a/graphblas/viz.py +++ b/graphblas/viz.py @@ -81,7 +81,7 @@ def spy(M, *, centered=False, show=True, figure=None, axes=None, figsize=None, * datashade """ - mpl, plt, ss = _get_imports(["mpl", "plt", "ss"], "spy") + mpl, plt, _ss = _get_imports(["mpl", "plt", "ss"], "spy") A = to_scipy_sparse(M, "coo") if show: plt.ion() @@ -132,7 +132,7 @@ def datashade(M, agg="count", *, width=None, height=None, opts_kwargs=None, **kw spy """ - np, pd, bk, hv, hp, ds = _get_imports(["np", "pd", "bk", "hv", "hp", "ds"], "datashade") + np, pd, bk, hv, _hp, _ds = _get_imports(["np", "pd", "bk", "hv", "hp", "ds"], "datashade") if "df" not in kwargs: rows, cols, vals = M.to_coo() max_int = np.iinfo(np.int64).max diff --git a/pyproject.toml b/pyproject.toml index 1bad95118..cf0251968 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,14 +1,15 @@ [build-system] build-backend = "setuptools.build_meta" -requires = ["setuptools >=64", "setuptools-git-versioning"] +requires = ["setuptools >=77", "setuptools-git-versioning"] [project] name = "python-graphblas" dynamic = ["version"] description = "Python library for GraphBLAS: high-performance sparse linear algebra for scalable graph analytics" readme = "README.md" -requires-python = ">=3.10" -license = { file = "LICENSE" } +requires-python = ">=3.11" +license = "Apache-2.0" +license-files = ["LICENSE"] authors = [ { name = "Erik Welch", email = "erik.n.welch@gmail.com" }, { name = "Jim Kitchen" }, @@ -35,16 +36,16 @@ keywords = [ ] classifiers = [ "Development Status :: 5 - Production/Stable", - "License :: OSI Approved :: Apache Software License", "Operating System :: MacOS :: MacOS X", "Operating System :: POSIX :: Linux", "Operating System :: Microsoft :: Windows", "Programming Language :: Python", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: 3.15", "Programming Language :: Python :: 3 :: Only", "Intended Audience :: Developers", "Intended Audience :: Other Audience", @@ -55,14 +56,14 @@ classifiers = [ "Topic :: Software Development :: Libraries :: Python Modules", ] dependencies = [ - "numpy >=1.23", + "numpy >=1.24", "donfig >=0.6", "pyyaml >=5.4", - # These won't be installed by default after 2024.3.0 - # once pep-771 is supported: https://peps.python.org/pep-0771/ + # These should be optional once PEP 771 is accepted and supported: + # https://peps.python.org/pep-0771/ (still draft as of 2026-03) # Use e.g. "python-graphblas[suitesparse]" or "python-graphblas[default]" instead - "suitesparse-graphblas >=7.4.0.0, <10", - "numba >=0.55; python_version<'3.14'", # make optional where numba is not supported + "suitesparse-graphblas >=7.4.0.0", + "numba >=0.57; python_version<'3.16'", # make optional where numba is not supported ] [project.urls] @@ -72,9 +73,9 @@ repository = "https://github.com/python-graphblas/python-graphblas" changelog = "https://github.com/python-graphblas/python-graphblas/releases" [project.optional-dependencies] -suitesparse = ["suitesparse-graphblas >=7.4.0.0, <10"] +suitesparse = ["suitesparse-graphblas >=7.4.0.0"] networkx = ["networkx >=2.8"] -numba = ["numba >=0.55"] +numba = ["numba >=0.57"] pandas = ["pandas >=1.5"] scipy = ["scipy >=1.9"] suitesparse-udf = [ # udf requires numba @@ -83,10 +84,12 @@ suitesparse-udf = [ # udf requires numba repr = ["python-graphblas[pandas]"] io = [ "python-graphblas[networkx,scipy]", - "python-graphblas[numba]; python_version<'3.14'", + "python-graphblas[numba]; python_version<'3.16'", # make optional where numba is not supported "awkward >=2.0", - "sparse >=0.14; python_version<'3.13'", # make optional, b/c sparse needs numba - "fast-matrix-market >=1.4.5; python_version<'3.13'", # py3.13 not supported yet + # sparse needs numba, so it stays optional and takes numba's upper guard from above: + # it has to drop out wherever numba does. + "sparse >=0.14; python_version<'3.16'", + # fast-matrix-market is no longer maintained; last supported Python is 3.12 ] viz = ["python-graphblas[networkx,scipy]", "matplotlib >=3.6"] datashade = [ # datashade requires numba @@ -94,19 +97,14 @@ datashade = [ # datashade requires numba "datashader >=0.14", "hvplot >=0.8", ] -test = [ - "python-graphblas[suitesparse,pandas,scipy]", - "packaging >=21", - "pytest >=6.2", - "tomli >=1", -] +test = ["python-graphblas[suitesparse,pandas,scipy]", "packaging >=21", "pytest >=6.2", "tomli >=1"] default = [ "python-graphblas[suitesparse,pandas,scipy]", - "python-graphblas[numba]; python_version<'3.14'", # make optional where numba is not supported + "python-graphblas[numba]; python_version<'3.16'", # make optional where numba is not supported ] all = [ "python-graphblas[default,io,viz,test]", - "python-graphblas[datashade]; python_version<'3.14'", # make optional, b/c datashade needs numba + "python-graphblas[datashade]; python_version<'3.16'", # make optional, b/c datashade needs numba ] [tool.setuptools] @@ -122,6 +120,7 @@ packages = [ "graphblas.core.operator", "graphblas.core.ss", "graphblas.dtypes", + "graphblas.indexbinary", "graphblas.indexunary", "graphblas.io", "graphblas.monoid", @@ -140,7 +139,11 @@ dirty_template = "{tag}+{ccount}.g{sha}.dirty" [tool.black] line-length = 100 -target-version = ["py310", "py311", "py312", "py313"] +# MAINT: add "py315" once a black that knows it (26.5 or later) installs on Python 3.15 +# from conda-forge (2026-10-19: its pytokens >=0.4 has no 3.15 build, so 3.15 gets black +# 26.1, which rejects "py315" and stops CI's auto-generated code check). Listing it also +# makes any older black refuse to run on this repository. +target-version = ["py311", "py312", "py313", "py314"] [tool.isort] sections = ["FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER"] @@ -154,47 +157,63 @@ line_length = 100 [tool.pytest.ini_options] minversion = "6.0" testpaths = "graphblas/tests" -xfail_strict = false # 2023-07-23: awkward and numpy 1.25 sometimes conflict +xfail_strict = false addopts = [ "--strict-config", # Force error if config is mispelled "--strict-markers", # Force error if marker is mispelled (must be defined in config) "-ra", # Print summary of all fails/errors ] -markers = ["slow: Skipped unless --runslow passed"] +markers = [ + "slow: Skipped unless --runslow passed", + # Not applied by our tests. ``scipy.io.tests.test_mmio`` (which + # ``test_mmread_mmwrite`` imports for its example matrices) sets it at module + # scope for the pytest-run-parallel plugin. Importing the module evaluates the + # marker, and with ``--strict-markers`` (always set below) an unregistered + # marker fails collection outright, on pytest 8 and 9 alike. Without the + # strict flag the same import would only warn (PytestUnknownMarkWarning), + # which ``filterwarnings = error`` below would turn into a failure; with it, + # the strict route is the only one this configuration ever takes. + "thread_unsafe: Declared by scipy's test modules; unused here", + "fail_slow: Declared by scipy's test modules; unused here", +] log_cli_level = "info" filterwarnings = [ # See: https://docs.python.org/3/library/warnings.html#describing-warning-filters # and: https://docs.pytest.org/en/7.2.x/how-to/capture-warnings.html#controlling-warnings "error", - # sparse 0.14.0 (2022-02-24) began raising this warning; it has been reported and fixed upstream. + # sparse 0.14.0 began raising this warning; fixed in newer versions but CI tests with 0.14/0.15 "ignore:coords should be an ndarray. This will raise a ValueError:DeprecationWarning:sparse._coo.core", - # setuptools v67.3.0 deprecated `pkg_resources.declare_namespace` on 13 Feb 2023. See: - # https://setuptools.pypa.io/en/latest/history.html#v67-3-0 - # MAINT: check if this is still necessary in 2025 - "ignore:Deprecated call to `pkg_resources.declare_namespace:DeprecationWarning:pkg_resources", - - # This deprecation warning was added in setuptools v67.5.0 (8 Mar 2023). See: - # https://setuptools.pypa.io/en/latest/history.html#v67-5-0 - "ignore:pkg_resources is deprecated as an API:DeprecationWarning:", - - # sre_parse deprecated in 3.11; this is triggered by awkward 0.10 - "ignore:module 'sre_parse' is deprecated:DeprecationWarning:", - "ignore:module 'sre_constants' is deprecated:DeprecationWarning:", - - # numpy 1.25.0 (2023-06-17) deprecated `np.find_common_type`; many other dependencies use it. - # See if we can remove this filter in 2025. - "ignore:np.find_common_type is deprecated:DeprecationWarning:", - # pypy gives this warning "ignore:can't resolve package from __spec__ or __package__:ImportWarning:", # Python 3.12 introduced this deprecation, which is triggered by pandas 2.1.1 + # MAINT: check if still needed once we drop pandas <2.2 support "ignore:datetime.datetime.utcfromtimestamp:DeprecationWarning:dateutil", - # Pandas 2.2 warns that pyarrow will become a required dependency in pandas 3.0 - "ignore:\\nPyarrow will become a required dependency of pandas:DeprecationWarning:", + # numpy 1.25+ deprecated np.find_common_type; triggered by older scipy/networkx/pandas + # MAINT: remove once we drop numpy <2 support (numpy 2.0 removed it entirely) + "ignore:np.find_common_type is deprecated:DeprecationWarning:", + + # UDT auto-lift emits a one-time ``gb.exceptions.NoJITWarning`` (subclass + # of ``UserWarning``) when JIT setup can't be done: compiler missing, + # control mode off, or UDT not C-expressible (C-reserved field name, + # stdlib macro, unsupported field type, packed layout, etc.). The warning + # uses ``stacklevel=3`` so its traceback points at the user's code rather + # than ``jit_config``; that makes a ``module=`` qualifier here fragile + # (Python attributes the warning to whatever module is 3 frames up at + # issue time). Match on the unique message instead. Tests that exercise + # this path wrap with ``pytest.warns(NoJITWarning, match="without JIT")`` + # (or ``UserWarning``, which still matches via inheritance). + "ignore:UDT operator running without JIT:UserWarning", + + # ``pytest -n`` (xdist) makes pytest-benchmark issue this warning during + # ``pytest_configure`` to say it disabled itself. Our ``error`` filter + # turns warnings into errors, which crashes config before tests run. + # Match by message only: naming the category would make pytest import + # pytest_benchmark to resolve it, which fails where it isn't installed. + "ignore:Benchmarks are automatically disabled:", ] [tool.coverage.run] @@ -210,11 +229,7 @@ precision = 1 fail_under = 0 skip_covered = true skip_empty = true -exclude_lines = [ - "pragma: no cover", - "raise AssertionError", - "raise NotImplementedError", -] +exclude_lines = ["pragma: no cover", "raise AssertionError", "raise NotImplementedError"] [tool.codespell] ignore-words-list = "coo,ba" @@ -222,7 +237,7 @@ ignore-words-list = "coo,ba" [tool.ruff] # https://github.com/charliermarsh/ruff/ line-length = 100 -target-version = "py310" +target-version = "py311" [tool.ruff.format] exclude = ["*.ipynb"] # Consider enabling auto-formatting of notebooks @@ -268,7 +283,7 @@ select = [ "Q", # flake8-quotes "RSE", # flake8-raise "RET", # flake8-return - # "SLF", # flake8-self (We can use our own private variables--sheesh!) + # "SLF", # flake8-self (we can use our own private variables, sheesh!) "SIM", # flake8-simplify # "TID", # flake8-tidy-imports (Rely on isort and our own judgement) # "TCH", # flake8-type-checking (Note: figure out type checking later) @@ -313,6 +328,7 @@ ignore = [ "RUF012", # Mutable class attributes should be annotated with `typing.ClassVar` (Note: no annotations yet) "RUF021", # parenthesize-chained-operators (Note: results don't look good yet) "RUF023", # unsorted-dunder-slots (Note: maybe fine, but noisy changes) + "RUF043", # Unescaped metacharacters in a `match=` pattern (Note: tests pass plain strings) "PERF401", # Use a list comprehension to create a transformed list (Note: poorly implemented atm) # Intentionally ignored @@ -322,16 +338,18 @@ ignore = [ "D400", # First line should end with a period (Note: prefer D415, which also allows "?" and "!") "N801", # Class name ... should use CapWords convention (Note:we have a few exceptions to this) "N802", # Function name ... should be lowercase - "N803", # Argument name ... should be lowercase (Maybe okay--except in tests) + "N803", # Argument name ... should be lowercase (Maybe okay; except in tests) "N806", # Variable ... in function should be lowercase "N807", # Function name should not start and end with `__` "N818", # Exception name ... should be named with an Error suffix (Note: good advice) "PERF203", # `try`-`except` within a loop incurs performance overhead (Note: too strict) "PLC0205", # Class `__slots__` should be a non-string iterable (Note: string is fine) + "PLC0415", # `import` should be at the top-level of a file (Note: we use deliberate local imports) "PLR0124", # Name compared with itself, consider replacing `x == x` (Note: too strict) "PLR0911", # Too many return statements "PLR0912", # Too many branches "PLR0913", # Too many arguments to function call + "PLR0917", # Too many positional arguments (Stable since ruff 0.16; same call for PLR0913) "PLR0915", # Too many statements "PLR2004", # Magic number used in comparison, consider replacing magic with a constant variable "PLW0603", # Using the global statement to update ... is discouraged (Note: yeah, discouraged, but too strict) @@ -340,6 +358,7 @@ ignore = [ "RET502", # Do not implicitly `return None` in function able to return non-`None` value "RET503", # Missing explicit `return` at the end of function able to return non-`None` value "RET504", # Unnecessary variable assignment before `return` statement + "RUF028", # Invalid formatter-suppression comment (Note: our `# fmt:` comments are for black, not ruff) "S110", # `try`-`except`-`pass` detected, consider logging the exception (Note: good advice, but we don't log) "S112", # `try`-`except`-`continue` detected, consider logging the exception (Note: good advice, but we don't log) "S603", # `subprocess` call: check for execution of untrusted input (Note: not important for us) @@ -348,10 +367,10 @@ ignore = [ "SIM105", # Use contextlib.suppress(...) instead of try-except-pass (Note: try-except-pass is much faster) "SIM108", # Use ternary operator ... instead of if-else-block (Note: if-else better for coverage and sometimes clearer) "TRY003", # Avoid specifying long messages outside the exception class (Note: why?) - "UP038", # Use `X | Y` in `isinstance` call instead of `(X, Y)` (Note: using `|` is slower atm) # Ignored categories "C90", # mccabe (Too strict, but maybe we should make things less complex) + "CPY", # flake8-copyright (Stable since ruff 0.16; we don't put a header in every file) "I", # isort (Should we replace `isort` with this?) "ANN", # flake8-annotations (We don't use annotations yet) "BLE", # flake8-blind-except (Maybe consider) @@ -360,7 +379,7 @@ ignore = [ "EM", # flake8-errmsg (Perhaps nicer, but too much work) "ICN", # flake8-import-conventions (Doesn't allow "_" prefix such as `_np`) "PYI", # flake8-pyi (We don't have stub files yet) - "SLF", # flake8-self (We can use our own private variables--sheesh!) + "SLF", # flake8-self (we can use our own private variables, sheesh!) "TID", # flake8-tidy-imports (Rely on isort and our own judgement) "TCH", # flake8-type-checking (Note: figure out type checking later) "ARG", # flake8-unused-arguments (Sometimes helpful, but too strict) @@ -374,6 +393,8 @@ ignore = [ "graphblas/core/operator/__init__.py" = ["A005"] "graphblas/io/__init__.py" = ["A005"] # shadows a standard-library module "graphblas/core/operator/base.py" = ["S102"] # exec is used for UDF +"graphblas/core/operator/udt_codegen.py" = ["S102"] # exec is used for UDT op codegen +"graphblas/monoid/numpy.py" = ["PLW0108"] # lambda is needed for numba.njit "graphblas/core/ss/matrix.py" = [ "NPY002", # numba doesn't support rng generator yet "PLR1730", @@ -382,23 +403,16 @@ ignore = [ "NPY002", # numba doesn't support rng generator yet ] "graphblas/core/utils.py" = ["PLE0302"] # `__set__` is used as a property +"graphblas/core/formatting.py" = ["ISC004"] # Adjacent literals form one HTML fragment +"scripts/ci_pick_versions.py" = ["S311"] # random.choice is intentional for CI version selection "graphblas/ss/_core.py" = ["N999"] # We want _core.py to be underscopre # Allow useless expressions, assert, pickle, RNG, print, no docstring, and yoda in tests -"graphblas/tests/*py" = [ - "B018", - "S101", - "S301", - "S311", - "T201", - "D103", - "D100", - "SIM300", -] +"graphblas/tests/*py" = ["B018", "S101", "S301", "S311", "T201", "D103", "D100", "SIM300"] "graphblas/tests/test_formatting.py" = ["E501"] # Allow long lines "graphblas/**/__init__.py" = [ "F401", # Allow unused imports (w/o defining `__all__`) ] -"scripts/*.py" = ["INP001"] # Not a package +"scripts/*.py" = ["INP001", "T201"] # Not a package; allow print "scripts/create_pickle.py" = ["F403", "F405"] # Allow `from foo import *` "docs/*.py" = ["INP001"] # Not a package @@ -411,7 +425,7 @@ builtins-allowed-modules = ["select"] fixture-parentheses = false mark-parentheses = false -[tool.lint.ruff.pydocstyle] +[tool.ruff.lint.pydocstyle] convention = "numpy" [tool.bandit] @@ -423,7 +437,7 @@ skips = [ [tool.pylint.messages_control] # To run a single check, do: pylint graphblas --disable E,W,R,C,I --enable assignment-from-no-return max-line-length = 100 -py-version = "3.10" +py-version = "3.11" enable = ["I"] disable = [ # Error diff --git a/scripts/audit_conda_versions.py b/scripts/audit_conda_versions.py new file mode 100755 index 000000000..5b837622d --- /dev/null +++ b/scripts/audit_conda_versions.py @@ -0,0 +1,397 @@ +#!/usr/bin/env python +"""Audit the CI dependency pools in `ci_pick_versions.py` against conda-forge. + +CI installs its dependencies from conda-forge, so this is the authoritative companion +to `audit_pypi_versions.py`. It can see two things PyPI cannot: + +* which Pythons a version was actually built for. conda-forge builds more than PyPI + ships (it has a py311 build of pyyaml 5.4.1 where PyPI's newest wheel is cp39), and + sometimes fewer, so "no wheel" and "no build" are different questions. +* `constrains` metadata (conda's run_constrained). Several rules in the picker exist + only because a package constrains something it does not depend on, which is what + makes a pinned pair fail to solve. + +Usage: + python scripts/audit_conda_versions.py [--samples N] + +Requires `conda` on PATH and network access. Slower than the PyPI audit: one +`conda search` per package, a few seconds each. Exits non-zero if any check fails. +See also scripts/check_versions.sh for the by-hand version of the same questions. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import random +import re +import shutil +import subprocess +import sys +from pathlib import Path + +from packaging.specifiers import InvalidSpecifier, SpecifierSet +from packaging.version import InvalidVersion, Version + +_PICKER = Path(__file__).with_name("ci_pick_versions.py") +_spec = importlib.util.spec_from_file_location("ci_pick_versions", _PICKER) +picker = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(picker) + +# Our internal names to conda-forge package names +CONDA_NAMES = {"psg": "python-suitesparse-graphblas", "matplotlib": "matplotlib-base"} + +PER_PYTHON_POOLS = { + "numpy": picker.NUMPY_VERSIONS, + "networkx": picker.NETWORKX_VERSIONS, + "pyyaml": picker.PYYAML_VERSIONS, + "sparse": picker.SPARSE_VERSIONS, +} + +SAMPLED_DEPS = ("scipy", "pandas", "awkward", "numba", "networkx", "sparse", "pyyaml", "psg") + +# Constraints the picker hard-codes because they come from conda metadata alone. If one +# of these drifts, the matching rule in `apply_constraints` needs another look. +# Each entry is (dep, version pin, field, expected constraint). +DOCUMENTED_METADATA = [ + ("networkx", "3.7", "constrains", "numpy >=1.25"), + ("networkx", "3.7", "constrains", "scipy >=1.11.2"), + ("networkx", "3.7", "constrains", "pandas >=2.0"), + ("networkx", "3.6", "constrains", "numpy >=1.25"), + ("networkx", "3.6", "constrains", "scipy >=1.11.2"), + ("networkx", "3.6", "constrains", "pandas >=2.0"), + ("networkx", "3.5", "constrains", "numpy >=1.25"), + ("pandas", "2.3", "constrains", "scipy >=1.10.0"), + # matplotlib 3.11.0's metadata understates its numpy floor; 3.11.1 fixed it. The + # picker pins matplotlib <3.11 for old numpy because of the 3.11.0 builds. + ("matplotlib", "3.11.0", "depends", "numpy >=1.23"), + ("matplotlib", "3.11.1", "depends", "numpy >=1.25"), +] + + +# --------------------------------------------------------------------------- +# conda access +# --------------------------------------------------------------------------- + +_cache: dict[tuple[str, str], list] = {} + + +def conda_search(dep, floor): + """All conda-forge builds of `dep` at or above `floor`, cached per process.""" + if (dep, floor) not in _cache: + name = CONDA_NAMES.get(dep, dep) + spec = f"{name}[channel=conda-forge]>={floor}" + result = subprocess.run( + [_CONDA, "search", "--json", spec], + capture_output=True, + text=True, + check=False, + timeout=600, + ) + try: + payload = json.loads(result.stdout) + except json.JSONDecodeError: + raise SystemExit(f"conda search failed for {spec}:\n{result.stderr.strip()}") from None + if isinstance(payload, dict) and payload.get("error"): + raise SystemExit(f"conda search failed for {spec}: {payload['error']}") + _cache[dep, floor] = [build for builds in payload.values() for build in builds] + return _cache[dep, floor] + + +def builds_for(dep, pin, floor): + """Builds of `dep` matching a pin ("" means every build at or above `floor`).""" + return [ + b + for b in conda_search(dep, floor) + if not pin or b["version"] == pin or b["version"].startswith(f"{pin}.") + ] + + +def build_supports(build, pyver): + """Whether a conda build can be installed on `pyver`.""" + minor = int(pyver.split(".")[1]) + if match := re.search(r"py3(\d\d)", build["build"]): + return int(match.group(1)) == minor + constraints = [] + for dep in build.get("depends", []): + tokens = dep.split() + if tokens[0] != "python" or len(tokens) < 2: + continue + constraints.append(f"=={tokens[1]}" if tokens[1].endswith(".*") else tokens[1]) + if not constraints: + return True # noarch with a bare `python` dep + try: + return Version(pyver) in SpecifierSet(",".join(constraints)) + except (InvalidSpecifier, InvalidVersion): + return True + + +def supports_python(dep, pin, pyver, floor): + """Whether any build matching the pin supports `pyver`.""" + return any(build_supports(b, pyver) for b in builds_for(dep, pin, floor)) + + +def numpy_cap(dep, pin, floor): + """The tightest `numpy declared: + problems.append( + f"{dep} {pin} caps numpy at {declared} on conda-forge, " + f"but the table allows up to {ceiling}" + ) + return problems + + +def note_conservative_ceilings(): + """Ceiling rows stricter than conda-forge's cap. Advisory: coverage left on the table.""" + notes = [] + for dep, table in (("scipy", picker.SCIPY_MAX_NUMPY), ("numba", picker.NUMBA_MAX_NUMPY)): + floor = _oldest_pin(dep) + for pin, ceiling in table.items(): + declared = numpy_cap(dep, pin, floor) + if declared is not None and ceiling < declared: + notes.append( + f"{dep} {pin} is capped at {ceiling} in the table, " + f"but conda-forge allows numpy up to {declared}" + ) + return notes + + +def check_psg_python_floors(): + """PSG_MIN_PYTHON's conda-forge floor must be the oldest psg built for that Python.""" + problems = [] + floor = _oldest_pin("psg") + for pyver, sources in picker.PSG_MIN_PYTHON.items(): + expected = sources.get("conda-forge") + built = [p for p in picker.PSG_VERSIONS_NP2 if supports_python("psg", p, pyver, floor)] + if built and min(built, key=picker._ver) != expected: + problems.append( + f"PSG_MIN_PYTHON says conda-forge py{pyver} starts at psg {expected}, " + f"but {min(built, key=picker._ver)} has a py{pyver} build" + ) + return problems + + +def check_documented_metadata(): + """Constraints the picker's rules quote must still be what conda-forge publishes.""" + problems = [] + for dep, pin, field, expected in DOCUMENTED_METADATA: + matching = builds_for(dep, pin, _oldest_pin(dep)) + if not matching: + problems.append(f"{dep} {pin} has no conda-forge build; cannot confirm {expected!r}") + continue + if not any(expected in build.get(field, []) for build in matching): + problems.append(f"no {dep} {pin} build declares {field} {expected!r} any more") + return problems + + +def report_newer_releases(): + """conda-forge versions newer than our newest pin. Informational, never a failure.""" + news = [] + for dep in ( + "numpy", + "scipy", + "pandas", + "awkward", + "numba", + "networkx", + "pyyaml", + "sparse", + "psg", + ): + floor = _oldest_pin(dep) + newest_pin = max( + (b["version"] for b in builds_for(dep, "", floor)), key=_version_key, default=None + ) + tested = _newest_tested(dep) + if not newest_pin or not tested: + continue + # A pin names a series, so "2.5" already covers conda-forge's 2.5.3 + in_tested_series = newest_pin == tested or newest_pin.startswith(f"{tested}.") + if not in_tested_series and _version_key(newest_pin) > _version_key(tested): + news.append(f"{dep}: newest pin is {tested}, conda-forge has {newest_pin}") + return news + + +def _newest_tested(dep): + pools = { + "numpy": picker.NUMPY_VERSIONS, + "scipy": {"": picker.SCIPY_VERSIONS["2.x"]}, + "pandas": {"": picker.PANDAS_VERSIONS["2.x"]}, + "awkward": {"": picker.AWKWARD_VERSIONS["2.x"]}, + "numba": {"": picker.NUMBA_VERSIONS["2.x"]}, + "networkx": picker.NETWORKX_VERSIONS, + "pyyaml": picker.PYYAML_VERSIONS, + "sparse": picker.SPARSE_VERSIONS, + "psg": {"": picker.PSG_VERSIONS_NP2}, + }[dep] + pins = {p for group in pools.values() for p in group if p and p != "NA"} + return max(pins, key=picker._ver) if pins else None + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +_CONDA = shutil.which("conda") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--samples", + type=int, + default=100, + help="random picks per Python/source combination (default: 100)", + ) + parser.add_argument("--seed", type=int, default=None, help="seed for reproducible sampling") + args = parser.parse_args() + if _CONDA is None: + print("conda is not on PATH; this audit needs it", file=sys.stderr) + return 2 + if args.seed is not None: + random.seed(args.seed) + + checks = ( + ("per-Python pools", check_per_python_pools), + ("sampled picks", lambda: check_sampled_picks(args.samples)), + ("numpy ceiling tables", check_numpy_ceilings), + ("psg Python floors", check_psg_python_floors), + ("documented conda metadata", check_documented_metadata), + ) + failures = 0 + for name, check in checks: + problems = check() + if problems: + failures += len(problems) + print(f"FAIL {name}:") + for problem in problems: + print(f" - {problem}") + else: + print(f"ok {name}") + + if notes := note_conservative_ceilings(): + print("\nNotes (stricter than conda-forge needs; often right, since conda's depends") + print("carry ABI caps while a package's real runtime cap lives upstream):") + for note in notes: + print(f" - {note}") + + if news := report_newer_releases(): + print("\nNewer versions exist (not a failure; update the pools when convenient):") + for line in news: + print(f" - {line}") + + print(f"\n{failures} problem(s) found", file=sys.stderr) + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/audit_pypi_versions.py b/scripts/audit_pypi_versions.py new file mode 100755 index 000000000..75cc202e9 --- /dev/null +++ b/scripts/audit_pypi_versions.py @@ -0,0 +1,424 @@ +#!/usr/bin/env python +"""Audit the CI dependency pools in `ci_pick_versions.py` against PyPI. + +`ci_pick_versions.py --validate` proves the picker agrees with the constraints it +documents. It cannot notice when a constraint has drifted from reality: a pin that +never had a build for a Python we test, or a ceiling table that a new release moved. +This script checks the pools and tables against what the packages actually declare. + +The picker's rules are written for conda-forge, which is what CI installs from, so +treat PyPI as a fast proxy and `audit_conda_versions.py` as the authority. Anything +this script reports is worth looking at in both places. + +Usage: + python scripts/audit_pypi_versions.py [--samples N] + +Requires network access and `packaging`. Exits non-zero if any check fails. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import random +import re +import sys +import urllib.error +import urllib.request +from pathlib import Path + +from packaging.requirements import Requirement +from packaging.specifiers import SpecifierSet +from packaging.version import InvalidVersion, Version + +# Import the picker as a module so the pools stay the single source of truth +_PICKER = Path(__file__).with_name("ci_pick_versions.py") +_spec = importlib.util.spec_from_file_location("ci_pick_versions", _PICKER) +picker = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(picker) + +# Our internal names to PyPI project names (conda-forge calls psg python-suitesparse-graphblas) +PYPI_NAMES = {"pyyaml": "PyYAML", "psg": "suitesparse-graphblas"} + +# Pools keyed by Python version, where every pin is offered on that Python and nowhere else +PER_PYTHON_POOLS = { + "numpy": picker.NUMPY_VERSIONS, + "networkx": picker.NETWORKX_VERSIONS, + "pyyaml": picker.PYYAML_VERSIONS, + "sparse": picker.SPARSE_VERSIONS, +} + +# Deps a pick can carry, in the order the picker fills them in +SAMPLED_DEPS = ("scipy", "pandas", "awkward", "numba", "networkx", "sparse", "pyyaml", "psg") + + +# --------------------------------------------------------------------------- +# PyPI access +# --------------------------------------------------------------------------- + +_cache: dict[str, dict] = {} + + +def _pypi(path): + """GET a PyPI JSON endpoint, cached for the life of the process.""" + if path not in _cache: + url = f"https://pypi.org/pypi/{path}/json" + request = urllib.request.Request(url, headers={"User-Agent": "python-graphblas-audit"}) + with urllib.request.urlopen(request, timeout=90) as response: # noqa: S310 + _cache[path] = json.load(response) + return _cache[path] + + +def releases(dep): + """Every non-prerelease, non-yanked release of `dep`: {Version: {...}}.""" + key = f"releases:{dep}" + if key not in _cache: + data = _pypi(PYPI_NAMES.get(dep, dep)) + out = {} + for raw, files in data["releases"].items(): + files = [f for f in files if not f["yanked"]] + if not files: + continue + try: + version = Version(raw) + except InvalidVersion: + continue + if version.is_prerelease: + continue + out[version] = { + "requires_python": files[0].get("requires_python"), + "cpython": {int(m) for f in files for m in re.findall(r"cp3(\d+)", f["filename"])}, + } + _cache[key] = out + return _cache[key] + + +def in_series(dep, pin): + """Every release matching a pin. An empty pin means "latest", so every release.""" + return [v for v in releases(dep) if not pin or str(v) == pin or str(v).startswith(f"{pin}.")] + + +def newest_in_series(dep, pin, pyver=None): + """The release a pin resolves to, preferring one that `pyver` can install. + + Returns (version, installable). When no release in the series supports `pyver` the + newest one is returned with installable=False, so callers can still read its + metadata while reporting the gap. + """ + found = in_series(dep, pin) + if not found: + return None, False + if pyver is not None: + usable = [v for v in found if supports_python(dep, v, pyver)] + if usable: + return max(usable), True + return max(found), False + return max(found), True + + +def supports_python(dep, version, pyver): + """Whether `version` of `dep` can be installed on `pyver`. + + Projects that ship compiled wheels are judged by their CPython tags, which is the + only signal that catches a release whose requires-python is wide but that simply + never built for this interpreter. Pure-Python projects are judged by requires-python. + """ + info = releases(dep)[version] + if info["cpython"]: + return int(pyver.split(".")[1]) in info["cpython"] + requires = info["requires_python"] + return not requires or Version(pyver) in SpecifierSet(requires) + + +def numpy_bound(dep, version, pyver): + """The numpy requirement `version` of `dep` declares for `pyver`, if any. + + A project may list numpy more than once (numba 0.65 declares both `numpy>=1.22` and + `numpy<2.5,>=1.22`), so combine every applicable clause rather than taking the first. + """ + data = _pypi(f"{PYPI_NAMES.get(dep, dep)}/{version}") + combined = None + for raw in data["info"].get("requires_dist") or []: + requirement = Requirement(raw) + if requirement.name.lower() != "numpy": + continue + marker = requirement.marker + if marker is None or marker.evaluate({"python_version": pyver, "extra": ""}): + combined = ( + requirement.specifier if combined is None else combined & requirement.specifier + ) + return combined + + +# --------------------------------------------------------------------------- +# Checks +# --------------------------------------------------------------------------- + + +def check_per_python_pools(): + """Every pin in a per-Python pool must exist on PyPI at all.""" + problems = [] + for dep, pools in PER_PYTHON_POOLS.items(): + for pyver, pins in pools.items(): + for pin in pins: + if pin == "NA": + continue + version, _ = newest_in_series(dep, pin) + if version is None: + problems.append(f"{dep} {pin} in the py{pyver} pool has no release on PyPI") + return problems + + +def note_pypi_python_gaps(): + """Pins with no PyPI build for a Python we offer them on. Advisory, not a failure. + + CI installs these from conda-forge, which builds more than PyPI ships: conda-forge + has a py311 build of pyyaml 5.4.1 where PyPI's newest wheel is cp39. Treat a note + here as a prompt to check `audit_conda_versions.py`, not as a broken pool. + """ + notes = [] + for dep, pools in PER_PYTHON_POOLS.items(): + for pyver, pins in pools.items(): + for pin in pins: + if pin == "NA": + continue + version, installable = newest_in_series(dep, pin, pyver) + if version is not None and not installable: + notes.append( + f"{dep} {pin or 'latest'} is offered on py{pyver}, " + f"but PyPI's {version} has no py{pyver} build" + ) + return notes + + +def check_ceiling_tables(): + """A ceiling row must never sit above the cap the release actually declares. + + A row above the real cap lets through a pairing that breaks; a row below it is + conservative, which is often deliberate, so that is reported as a note instead. + """ + problems = [] + tables = (("scipy", picker.SCIPY_MAX_NUMPY), ("numba", picker.NUMBA_MAX_NUMPY)) + for dep, table in tables: + for pin, ceiling in table.items(): + version, _ = newest_in_series(dep, pin) + if version is None: + problems.append(f"{dep} {pin} has a ceiling row but no release on PyPI") + continue + # Read the bound on the oldest Python we test, where markers are least likely + # to narrow it, then compare against the "< X.Y" the table records. + specifier = numpy_bound(dep, version, min(picker.NUMPY_VERSIONS, key=picker._ver)) + declared = _upper_bound(specifier) + if declared is not None and ceiling > declared: + problems.append( + f"{dep} {pin} ({version}) caps numpy at {declared}, " + f"but the table allows up to {ceiling}" + ) + return problems + + +def note_conservative_ceilings(): + """Ceiling rows stricter than PyPI's cap. Advisory: usually a conda-forge cap.""" + notes = [] + for dep, table in (("scipy", picker.SCIPY_MAX_NUMPY), ("numba", picker.NUMBA_MAX_NUMPY)): + for pin, ceiling in table.items(): + version, _ = newest_in_series(dep, pin) + if version is None: + continue + declared = _upper_bound( + numpy_bound(dep, version, min(picker.NUMPY_VERSIONS, key=picker._ver)) + ) + if declared is not None and ceiling < declared: + notes.append( + f"{dep} {pin} is capped at {ceiling} in the table, " + f"but PyPI's {version} allows numpy up to {declared}" + ) + return notes + + +def _upper_bound(specifier): + """The ` resolved and not str(latest).startswith(newest_pin): + news.append(f"{dep}: newest pin is {newest_pin} ({resolved}), PyPI has {latest}") + return news + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--samples", + type=int, + default=100, + help="random picks per Python/source combination (default: 100)", + ) + parser.add_argument("--seed", type=int, default=None, help="seed for reproducible sampling") + args = parser.parse_args() + if args.seed is not None: + random.seed(args.seed) + + checks = ( + ("per-Python pools", check_per_python_pools), + ("numpy ceiling tables", check_ceiling_tables), + ("Python floor tables", check_python_floor_tables), + ("sampled picks", lambda: check_sampled_picks(args.samples)), + ) + failures = 0 + for name, check in checks: + try: + problems = check() + except (urllib.error.URLError, TimeoutError) as exc: + print(f"SKIP {name}: PyPI unreachable ({exc})", file=sys.stderr) + continue + if problems: + failures += len(problems) + print(f"FAIL {name}:") + for problem in problems: + print(f" - {problem}") + else: + print(f"ok {name}") + + try: + notes = note_conservative_ceilings() + note_pypi_python_gaps() + news = report_newer_releases() + except (urllib.error.URLError, TimeoutError) as exc: + print(f"SKIP notes: PyPI unreachable ({exc})", file=sys.stderr) + notes = news = [] + + if notes: + print("\nNotes (PyPI only; conda-forge builds more, so check the conda audit):") + for note in notes: + print(f" - {note}") + + if news: + print("\nNewer releases exist (not a failure; update the pools when convenient):") + for line in news: + print(f" - {line}") + + print(f"\n{failures} problem(s) found", file=sys.stderr) + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/check_versions.sh b/scripts/check_versions.sh index 5aa88e045..e99cb7ca6 100755 --- a/scripts/check_versions.sh +++ b/scripts/check_versions.sh @@ -3,15 +3,27 @@ # Use, adjust, copy/paste, etc. as necessary to answer your questions. # This may be helpful when updating dependency versions in CI. # Tip: add `--json` for more information. -conda search 'flake8-bugbear[channel=conda-forge]>=24.12.12' -conda search 'flake8-simplify[channel=conda-forge]>=0.21.0' -conda search 'numpy[channel=conda-forge]>=2.2.3' -conda search 'pandas[channel=conda-forge]>=2.2.3' -conda search 'scipy[channel=conda-forge]>=1.15.2' -conda search 'networkx[channel=conda-forge]>=3.4.2' -conda search 'awkward[channel=conda-forge]>=2.7.4' -conda search 'sparse[channel=conda-forge]>=0.15.5' -conda search 'fast_matrix_market[channel=conda-forge]>=1.7.6' -conda search 'numba[channel=conda-forge]>=0.61.0' -conda search 'pyyaml[channel=conda-forge]>=6.0.2' -# conda search 'python[channel=conda-forge]>=3.10 *pypy*' +# +# When updating versions throughout the repo (CI, pyproject.toml, pre-commit, etc.), +# also update these version numbers to match the latest versions we currently test. +# +# scripts/audit_pypi_versions.py and scripts/audit_conda_versions.py ask the same +# questions automatically, and also check the pools and tables in ci_pick_versions.py +# against what is published. The "Audit dependency versions" workflow runs both weekly. +conda search 'flake8-bugbear[channel=conda-forge]>=26.10.27' +conda search 'flake8-simplify[channel=conda-forge]>=0.31.1' +conda search 'numpy[channel=conda-forge]>=2.5' +# ci_pick_versions.py caps an unpinned pandas below 3.1 for now (see its MAINT note). +conda search 'pandas[channel=conda-forge]>=3.0' +conda search 'scipy[channel=conda-forge]>=1.18' +conda search 'networkx[channel=conda-forge]>=3.7' +conda search 'awkward[channel=conda-forge]>=2.14' +conda search 'sparse[channel=conda-forge]>=0.19' +# fast_matrix_market is deprecated (no longer maintained; last supported Python is 3.12) +conda search 'numba[channel=conda-forge]>=0.68' +conda search 'pyyaml[channel=conda-forge]>=6.0' +# matplotlib is installed unpinned for the notebooks job; ci_pick_versions.py caps it +# below 3.11 when the numpy pick predates matplotlib's numpy >=1.25 runtime floor. +conda search 'matplotlib-base[channel=conda-forge]>=3.11' +conda search 'python-suitesparse-graphblas[channel=conda-forge]>=10.5.2' +# conda search 'python[channel=conda-forge]>=3.11 *pypy*' diff --git a/scripts/ci_pick_versions.py b/scripts/ci_pick_versions.py new file mode 100755 index 000000000..34f76e9ee --- /dev/null +++ b/scripts/ci_pick_versions.py @@ -0,0 +1,881 @@ +#!/usr/bin/env python +"""Pick random, compatible dependency versions for python-graphblas CI. + +Replaces the bash-based version selection that used to live in +test_and_build.yml. + +Usage (in GitHub Actions workflow): + eval "$(python scripts/ci_pick_versions.py --python 3.12 --source conda-forge)" + +Output: bash-eval-safe key=value lines using the same variable names as the workflow: + npver='=2.2' # conda pin + spver='=1.15' + npver='' # empty = latest (no pin) + sparsever='NA' # NA = skip this package +""" + +import argparse +import random +import sys + +# --------------------------------------------------------------------------- +# Version pools: which versions we want to test per package. +# "" means "latest" (no pin). "NA" means "don't install". +# +# When updating versions here, also update scripts/check_versions.sh +# +# To add a Python version (3.16, say): add it to `choices` in the "RNG for Python +# version" step of test_and_build.yml, give it a row in every table keyed by Python +# below (`check_python_pools` fails the run if you miss one), and revisit the floors +# in `apply_constraints`. Those floors are written as `_ver(pyver) >= (3, N)`, so a +# new Python inherits the previous one's floors instead of escaping them. That is +# enough to keep CI honest, but usually still too loose to install, so check each dep. +# --------------------------------------------------------------------------- + +# numpy 2.5 requires Python >=3.12, so py3.11 stops at 2.4. conda-forge's first py3.15 +# builds are numpy 2.5.2, which the "2.5" pin resolves to. +NUMPY_VERSIONS = { + "3.11": ["1.24", "1.25", "1.26", "2.0", "2.1", "2.2", "2.3", "2.4", ""], + "3.12": ["1.26", "2.0", "2.1", "2.2", "2.3", "2.4", "2.5", ""], + "3.13": ["2.1", "2.2", "2.3", "2.4", "2.5", ""], + "3.14": ["2.3", "2.4", "2.5", ""], + "3.15": ["2.5", ""], +} + +# Deps that depend on numpy version (1.x vs 2.x path). +# Per-Python sublists narrow to versions with available conda builds. +SCIPY_VERSIONS = { + "1.x": { + "3.11": ["1.9", "1.10", "1.11", "1.12", "1.13", "1.14", ""], + "3.12": ["1.11", "1.12", "1.13", "1.14", ""], + }, + "2.x": ["1.13", "1.14", "1.15", "1.16", "1.17", "1.18", ""], +} + +# Oldest Python each scipy pin supports (scipy's requires-python). The "2.x" pool is +# shared by every Python, so re-picks run through `_pick_scipy` to drop pins this +# Python has no build for. +# MAINT: 2026-09-16 scipy 1.18 requires Python >=3.12 +SCIPY_MIN_PYTHON = {"1.18": (3, 12)} + +# The numpy ceiling each scipy pin declares, the same idea as NUMBA_MAX_NUMPY below: +# scipy caps `numpy =2.10 pins survive numpy >=2.5; see the awkward +# constraints in `apply_constraints`. +AWKWARD_VERSIONS = { + "1.x": { + "3.11": ["2.0", "2.1", "2.2", "2.3", "2.4", "2.5", "2.6", "2.7", "2.8", "2.9", ""], + "3.12": ["2.4", "2.5", "2.6", "2.7", "2.8", "2.9", ""], + }, + "2.x": ["2.6", "2.7", "2.8", "2.9", "2.10", "2.11", "2.12", "2.13", "2.14", ""], +} + +NUMBA_VERSIONS = { + "1.x": ["0.57", "0.58", "0.59", "0.60", "0.61", ""], + "2.x": ["0.62", "0.63", "0.64", "0.65", "0.66", "0.67", "0.68", ""], +} + +# Oldest numba that supports each Python. +NUMBA_MIN_PYTHON = { + "3.11": (0, 57), + "3.12": (0, 59), + "3.13": (0, 61), + "3.14": (0, 63), + "3.15": (0, 68), +} + +# The numpy ceiling each numba pin declares. numba caps `numpy =3.12, so py3.11 stops at 3.6. +NETWORKX_VERSIONS = { + "3.11": ["2.8", "3.0", "3.1", "3.2", "3.3", "3.4", "3.5", "3.6", ""], + "3.12": ["3.2", "3.3", "3.4", "3.5", "3.6", "3.7", ""], + "3.13": ["3.4", "3.5", "3.6", "3.7", ""], + "3.14": ["3.6", "3.7", ""], + "3.15": ["3.6", "3.7", ""], +} + +PYYAML_VERSIONS = { + "3.11": ["5.4", "6.0", ""], + "3.12": ["6.0", ""], + "3.13": ["6.0", ""], + "3.14": ["6.0", ""], + "3.15": ["6.0", ""], +} + +# sparse is noarch and caps neither Python nor numpy: every release here declares numba +# >=0.49, numpy >=1.17 and a Python floor at or below our oldest (>=3.8 for 0.14/0.15, +# >=3.10 for 0.16/0.17, >=3.11 for 0.18/0.19), so every pin installs on every Python and +# the rows are all the same. It depends on numba, so the workflow drops sparse wherever it +# skips numba. To skip sparse on a Python, give it a ["NA"] pool: a bare "NA" would be +# indexed as a string by `random.choice`. +# MAINT: 2026-09-26 sparse 0.19.2 is latest. Each pin, as conda resolves it, passes the +# pydata-sparse tests on py3.13 and py3.14 with the oldest and newest numpy/numba pair +# each one draws. +SPARSE_VERSIONS = { + "3.11": ["0.14", "0.15", "0.16", "0.17", "0.18", "0.19", ""], + "3.12": ["0.14", "0.15", "0.16", "0.17", "0.18", "0.19", ""], + "3.13": ["0.14", "0.15", "0.16", "0.17", "0.18", "0.19", ""], + "3.14": ["0.14", "0.15", "0.16", "0.17", "0.18", "0.19", ""], + "3.15": ["0.14", "0.15", "0.16", "0.17", "0.18", "0.19", ""], +} + +# PSG versions to pair with numpy 1.x (only reachable on py3.11/py3.12, since +# py3.13+ have no numpy 1.x in their pools). Only "conda-forge" and "wheel" +# builds get here: "source" builds blank numpy before picking psg (so they use +# PSG_VERSIONS_NP2), and "upstream" builds always use psg from git. +PSG_VERSIONS_NP1 = { + "conda-forge": { + "3.11": [ + "7.4.0", + "7.4.1", + "7.4.2", + "7.4.3.0", + "7.4.3.1", + "7.4.3.2", + "8.0.2.1", + "8.2.0.1", + "8.2.1.0", + ], + "3.12": ["8.2.0.1", "8.2.1.0"], + }, + # The picker does not know the OS, and 7.4.3.2's wheels are manylinux x86_64 only, + # so a "wheel" draw of it would fail to install on every other runner. + "wheel": { + "3.11": ["8.0.2.1", "8.2.0.1", "8.2.1.0"], + "3.12": ["8.2.0.1", "8.2.1.0"], + }, +} + +# PSG versions to pair with numpy 2.x, oldest first. Every source type except +# "upstream" draws from this one pool; Pythons too new for the older releases draw +# from a tail of it (see PSG_MIN_PYTHON). +PSG_VERSIONS_NP2 = [ + "9.3.1.0", + "9.4.5.0", + "10.0.1.1", + "10.1.1.0", + "10.3.1.0", + "10.4.0.0", + "10.4.1.0", + "10.5.0.0", + "10.5.1.1", + "10.5.2.0", +] + +# Oldest psg release that installs on a given Python, by where psg comes from. +# Pythons older than every entry here (3.11-3.13) can use the whole pool; a Python +# newer than every entry inherits the newest one until it gets a row of its own. +# +# py3.14 needs a different floor per source: conda-forge has py314 builds from +# 10.0.1.1 on, but PyPI's first cp314 wheels are in 10.3.1.0, so the "wheel" path +# can't reach back as far. "source" builds the sdist, and the sdists before 10.3.1.0 +# predate psg's py3.14 support (they still declare requires-python >=3.9), so hold +# them to the wheel floor rather than test a build psg never claimed works. +# +# py3.15 starts at 10.5.1.1 everywhere: it is psg's first release with py315 builds on +# conda-forge and cp315 wheels on PyPI, and its sdist is the first from after psg added +# Python 3.15 support. +# MAINT: 2026-10-28 checked conda-forge build strings and PyPI wheel tags for psg 9-10 +PSG_MIN_PYTHON = { + "3.14": {"conda-forge": "10.0.1.1", "wheel": "10.3.1.0", "source": "10.3.1.0"}, + "3.15": {"conda-forge": "10.5.1.1", "wheel": "10.5.1.1", "source": "10.5.1.1"}, +} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _ver(s): + """Parse version string to tuple for comparison. "" means latest (very large).""" + if s in ("", "NA"): + return (9999,) + return tuple(int(x) for x in s.split(".")) + + +def _numba_ceiling(pin): + """The numpy ceiling a numba pin declares, or None when numba is skipped.""" + if pin == "NA": + return None + return max(NUMBA_MAX_NUMPY.values()) if pin == "" else NUMBA_MAX_NUMPY[pin] + + +def _pick_scipy(pyver, npver, pool, floor=(0,)): + """Pick a scipy pin at or above `floor` that this Python and this numpy can use. + + Filters the pool by the pin's requires-python (SCIPY_MIN_PYTHON), its numpy ceiling + (SCIPY_MAX_NUMPY) and the one numpy floor scipy imposes (>=1.15 needs numpy + >=1.26.4), so every re-pick site gets the same answer no matter where it sits in the + constraint order. "" (latest) is always a candidate: an unpinned scipy is one conda + resolves itself, so it satisfies every rule here by construction. + """ + return random.choice( + [ + s + for s in pool + if s + and _ver(s) >= floor + and SCIPY_MAX_NUMPY[s] > npver + and (_ver(s) < (1, 15) or npver >= (1, 26)) + and _ver(pyver) >= SCIPY_MIN_PYTHON.get(s, (0,)) + ] + + [""] + ) + + +def _numpy_ver(v, pyver): + """Version of numpy to compare against, resolving "" (latest) to what it installs. + + An unpinned numpy is the newest release this Python can actually get, which is the + newest pin in its pool, not something infinitely new: numba 0.67 caps numpy at <2.6 + and still qualifies today. Deriving it from the pool keeps it honest per Python + (py3.11 tops out at numpy 2.4) and leaves nothing extra to bump. + """ + if v["numpy"]: + return _ver(v["numpy"]) + return max(_ver(n) for n in NUMPY_VERSIONS[pyver] if n) + + +def _min_for_python(table, pyver): + """Look up `pyver` in a {python version: minimum} table, or None if it predates it. + + A Python newer than every row inherits the newest row: it needs at least as new a + build as the newest Python we have characterized. That guess is conservative rather + than correct, so give a new Python its own row once its real floors are known. + """ + if pyver in table: + return table[pyver] + newest = max(table, key=_ver) + return table[newest] if _ver(pyver) > _ver(newest) else None + + +# --------------------------------------------------------------------------- +# Constraints +# --------------------------------------------------------------------------- + + +def apply_constraints(v, pyver, scipy_pool, numba_pool): + """Mutate version dict to satisfy all known compatibility constraints. + + Each constraint comment documents the real-world requirement it encodes. + Order matters: numpy/scipy constraints first, then pandas (which may bump scipy/numba), + then Python-version constraints, then numba/numpy constraints. Within the awkward + section, the numpy 2.x floor runs before the Python-availability picks (which may + narrow the choices further), and the awkward >=2.10 rule runs last so it wins. + The networkx section reads numpy, scipy, and pandas, so it runs after all three + are final. + """ + # Nothing below changes the numpy pick, so derive these once. + npver = _numpy_ver(v, pyver) + np_is_1x = v["numpy"].startswith("1.") if v["numpy"] else False + + # --- scipy / numpy constraints --- + + # scipy >=1.15 requires numpy >=1.26.4 + if v["numpy"] in ("1.24", "1.25") and _ver(v["scipy"]) >= (1, 15): + candidates = [s for s in scipy_pool if s and _ver(s) < (1, 15)] + v["scipy"] = random.choice(candidates) if candidates else "1.14" + + # scipy's numpy ceiling (see SCIPY_MAX_NUMPY). This covers every scipy/numpy + # pairing at once: the pre-1.13 pins against numpy 2.x, 1.9 against numpy 1.26, + # and each later pin against the numpy that outgrew it. + if (ceiling := SCIPY_MAX_NUMPY.get(v["scipy"])) is not None and npver >= ceiling: + v["scipy"] = _pick_scipy(pyver, npver, scipy_pool) + + # --- scipy / Python version availability --- + + # scipy <1.14 has no py3.13 builds, <1.16 none for py3.14, and <1.18 none for py3.15 + if _ver(pyver) >= (3, 15) and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 18): + v["scipy"] = _pick_scipy(pyver, npver, scipy_pool, floor=(1, 18)) + elif _ver(pyver) >= (3, 14) and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 16): + v["scipy"] = _pick_scipy(pyver, npver, scipy_pool, floor=(1, 16)) + elif _ver(pyver) >= (3, 13) and v["scipy"] == "1.13": + v["scipy"] = _pick_scipy(pyver, npver, scipy_pool, floor=(1, 14)) + + # A scipy pin whose requires-python is newer than this Python has no build for it + if v["scipy"] and _ver(pyver) < SCIPY_MIN_PYTHON.get(v["scipy"], (0,)): + v["scipy"] = _pick_scipy(pyver, npver, scipy_pool) + + # --- pandas constraints --- + + # pandas <2.3 has no py3.14 builds, and <3.0 none for py3.15 + if _ver(pyver) >= (3, 15) and v["pandas"] not in ("", "NA") and _ver(v["pandas"]) < (3, 0): + v["pandas"] = random.choice(["3.0", ""]) + elif _ver(pyver) >= (3, 14) and v["pandas"] == "2.2": + v["pandas"] = random.choice(["2.3", "3.0", ""]) + + # pandas 3.0 requires scipy >=1.14.1. It also requires numba >=0.60, which needs no + # rule: pandas 3.0 only appears in the numpy 2.x pool, whose numba pins start at 0.62. + if v["pandas"] == "3.0" and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 15): + v["scipy"] = _pick_scipy(pyver, npver, scipy_pool, floor=(1, 15)) + + # conda-forge pandas >=2.2 carries run_constrained "scipy >=1.10.0" (latest 2.2.3 + # build and all 2.3 builds; 2.0/2.1 declare no scipy constraint), so a pandas + # 2.2/2.3 pin alongside a scipy 1.9 pin fails the conda solve outright. Unpinned + # pandas backs off on its own, so only pins trigger. A scipy 1.9 pick can only + # reach here with numpy 1.24/1.25 (the 1.26 conflict above already re-picked), + # so keep the replacement pinned below the scipy 1.15 ceiling those numpys impose. + # MAINT: 2026-08-06 verified with conda search --info against conda-forge + if ( + v["pandas"] not in ("", "NA") + and _ver(v["pandas"]) >= (2, 2) + and v["scipy"] not in ("", "NA") + and _ver(v["scipy"]) < (1, 10) + ): + v["scipy"] = random.choice([s for s in scipy_pool if s and (1, 10) <= _ver(s) < (1, 15)]) + + # --- awkward / numpy 2.x support --- + + # awkward <2.6 uses numpy.AxisError, which numpy 2.0 removed (test_io.py then fails + # at collection). awkward is picked from the numpy 1.x pool before source builds blank + # numpy to latest, so a numpy 1.x pick can still end up here with numpy 2.x. + # MAINT: 2026-08-04 confirmed with awkward 2.0.10 + numpy 2.4 + if not np_is_1x and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 6): + v["awkward"] = random.choice(AWKWARD_VERSIONS["2.x"]) + + # --- awkward / Python version availability --- + + # awkward <2.7 has no py3.13 builds, <2.8 none for py3.14, and <2.13 none for py3.15 + # (awkward is noarch, but each release pins one awkward-cpp, and awkward-cpp 56 is the + # first with py315 builds) + if _ver(pyver) >= (3, 15) and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 13): + v["awkward"] = random.choice([a for a in AWKWARD_VERSIONS["2.x"] if _ver(a) >= (2, 13)]) + elif _ver(pyver) >= (3, 14) and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 8): + v["awkward"] = random.choice(["2.8", "2.9", ""]) + elif _ver(pyver) >= (3, 13) and v["awkward"] == "2.6": + v["awkward"] = random.choice(["2.7", "2.8", "2.9", ""]) + + # --- awkward / numpy 2.5 deprecations --- + + # awkward <2.10 builds numpy.datetime64("NaT") with a generic unit at import time, + # which numpy 2.5 deprecates; `filterwarnings = error` in pyproject.toml then turns the + # import into a collection error. Unpinned numpy is latest, which is >=2.5. + # MAINT: 2026-08-04 confirmed with awkward 2.8.12 and 2.9.0 + numpy 2.5 + if ( + _ver(v["numpy"]) >= (2, 5) + and v["awkward"] not in ("", "NA") + and _ver(v["awkward"]) < (2, 10) + ): + v["awkward"] = random.choice([a for a in AWKWARD_VERSIONS["2.x"] if _ver(a) >= (2, 10)]) + + # --- networkx / numpy, scipy, pandas floors --- + + # conda-forge networkx 3.5 through 3.7 declare `constrains: numpy >=1.25, scipy >=1.11.2, + # pandas >=2.0`. conda enforces those against our pins, and since networkx is pinned too + # the solve fails outright instead of backing off. Older pins are safe because a pin like + # "3.4" can resolve down to 3.4.0, whose floors (numpy >=1.22, scipy >=1.9, pandas >=1.4) + # sit at or below everything in our pools. Unpinned networkx is latest, so it gets the + # 3.5+ floors. Compare scipy at the minor level: the "1.11" pin resolves to 1.11.4, which + # clears the 1.11.2 floor. + # MAINT: 2026-10-18 conda-forge networkx 3.6.1 and 3.7 constrain numpy >=1.25, scipy >=1.11.2 + if _ver(v["networkx"]) >= (3, 5) and ( + _ver(v["numpy"]) < (1, 25) or _ver(v["scipy"]) < (1, 11) or _ver(v["pandas"]) < (2, 0) + ): + # Only the numpy 1.x pools reach here, and those Pythons all have pre-3.5 networkx. + v["networkx"] = random.choice([n for n in NETWORKX_VERSIONS[pyver] if _ver(n) < (3, 5)]) + + # --- matplotlib / numpy floor (notebooks job) --- + + # matplotlib is installed unpinned for the notebooks task. matplotlib 3.11's runtime + # check requires numpy >=1.25. conda-forge's 3.11.1 builds declare that floor, so the + # solver backs off on its own, but the 3.11.0 builds still say only "numpy >=1.23": + # the solver pairs 3.11.0 with a numpy 1.24 pin and the notebooks die at + # `import matplotlib`. Pin below 3.11 whenever the numpy pick sits under that floor; + # empty means latest is fine. + # MAINT: 2026-09-16 conda-forge matplotlib-base 3.11.0 depends on "numpy >=1.23", + # 3.11.1 on "numpy >=1.25" + if v["numpy"] and _ver(v["numpy"]) < (1, 25): + v["matplotlib"] = "<3.11" + else: + v["matplotlib"] = "" + + # --- numba constraints --- + + # numba minimum by Python version (see NUMBA_MIN_PYTHON) + if v["numba"] not in ("", "NA"): + min_ver = _min_for_python(NUMBA_MIN_PYTHON, pyver) + if _ver(v["numba"]) < min_ver: + pool = [n for n in numba_pool if _ver(n) >= min_ver] + v["numba"] = random.choice(pool) if pool else "" + + # numba's numpy ceiling (see NUMBA_MAX_NUMPY). This covers every numba/numpy + # pairing rule at once: 0.57 against numpy >=1.25, the pre-0.62 pins against + # numpy 2.x, and each later pin against the numpy that outgrew it. Re-pick from + # the pins that clear both this numpy and the Python minimum, or skip numba. + ceiling = _numba_ceiling(v["numba"]) + if ceiling is not None and npver >= ceiling: + min_ver = _min_for_python(NUMBA_MIN_PYTHON, pyver) + pool = [n for n in numba_pool if n and NUMBA_MAX_NUMPY[n] > npver and _ver(n) >= min_ver] + v["numba"] = random.choice(pool) if pool else "NA" + + # --- sparse --- + # No rules needed: sparse caps neither Python nor numpy, and the workflow drops it + # whenever numba is skipped or NA. + + +# --------------------------------------------------------------------------- +# Main logic +# --------------------------------------------------------------------------- + + +def pick_versions(pyver, source_type): + """Pick random compatible dependency versions. + + Parameters + ---------- + pyver : str + Python version like "3.12" + source_type : str + One of "conda-forge", "wheel", "source", "upstream" + + Returns + ------- + dict + Package name -> version string ("" = latest, "NA" = skip) + """ + # Step 1: Pick numpy + numpy_pool = NUMPY_VERSIONS[pyver] + npver = random.choice(numpy_pool) + + # Upstream needs numpy 2 + if source_type == "upstream" and npver.startswith("1."): + npver = random.choice([v for v in numpy_pool if not v.startswith("1.")] or [""]) + + np_is_1x = npver.startswith("1.") if npver else False + + # Step 2: Pick numpy-dependent deps + if np_is_1x: + scipy_pool = SCIPY_VERSIONS["1.x"].get(pyver, SCIPY_VERSIONS["1.x"]["3.11"]) + pandas_pool = PANDAS_VERSIONS["1.x"].get(pyver, PANDAS_VERSIONS["1.x"]["3.11"]) + awkward_pool = AWKWARD_VERSIONS["1.x"].get(pyver, AWKWARD_VERSIONS["1.x"]["3.11"]) + numba_pool = NUMBA_VERSIONS["1.x"] + else: + scipy_pool = SCIPY_VERSIONS["2.x"] + pandas_pool = PANDAS_VERSIONS["2.x"] + awkward_pool = AWKWARD_VERSIONS["2.x"] + numba_pool = NUMBA_VERSIONS["2.x"] + + v = { + "numpy": npver, + "scipy": random.choice(scipy_pool), + "pandas": random.choice(pandas_pool), + "awkward": random.choice(awkward_pool), + "numba": random.choice(numba_pool), + "networkx": random.choice(NETWORKX_VERSIONS[pyver]), + "pyyaml": random.choice(PYYAML_VERSIONS[pyver]), + "sparse": random.choice(SPARSE_VERSIONS[pyver]), + } + + # Source builds have issues with some numpy/scipy/pandas versions; + # blank them before constraints so numba/etc constraints see the right numpy. + if source_type == "source": + v["numpy"] = "" + v["scipy"] = "" + v["pandas"] = "" + + # Step 3: Apply compatibility constraints + apply_constraints(v, pyver, scipy_pool, numba_pool) + + # Step 4: Pick psg version + v["psg"] = _pick_psg(v["numpy"], pyver, source_type) + + return v + + +def _pick_psg(npver, pyver, source_type): + """Pick python-suitesparse-graphblas version.""" + if source_type == "upstream": + return "" + + np_is_1x = npver.startswith("1.") if npver else False + eq = "=" if source_type == "conda-forge" else "==" + + if np_is_1x: + pool = PSG_VERSIONS_NP1.get(source_type, {}).get(pyver, []) + if not pool: + return "" + return f"{eq}{random.choice(pool)}" + pool = PSG_VERSIONS_NP2 + if (floors := _min_for_python(PSG_MIN_PYTHON, pyver)) is not None: + floor = _ver(floors[source_type]) + pool = [ver for ver in pool if _ver(ver) >= floor] + return random.choice([f"{eq}{ver}" for ver in pool] + [""]) + + +# --------------------------------------------------------------------------- +# Output formatting +# --------------------------------------------------------------------------- + +# Map internal names to the short variable names used in the workflow +_VAR_NAMES = { + "numpy": "npver", + "scipy": "spver", + "pandas": "pdver", + "awkward": "akver", + "networkx": "nxver", + "pyyaml": "yamlver", + "sparse": "sparsever", + "numba": "numbaver", + "psg": "psgver", +} + +_SUMMARY_NAMES = { + "numpy": "np", + "scipy": "sp", + "pandas": "pd", + "awkward": "ak", + "networkx": "nx", + "pyyaml": "yaml", + "sparse": "sparse", + "numba": "numba", + "psg": "psg", + "matplotlib": "mpl", +} + + +def format_output(v): + """Format version dict as bash-eval-safe key=value lines. + + Values use '=X.Y' prefix for conda install (e.g., npver='=2.2'). + Empty means latest, NA means skip. Latest is capped by LATEST_CEILINGS + (e.g., pdver='<3.1'). + psg already has its prefix baked in. + """ + lines = [] + for key in ("numpy", "scipy", "pandas", "awkward", "networkx", "pyyaml", "sparse", "numba"): + var = _VAR_NAMES[key] + val = v[key] + if val == "" and key in LATEST_CEILINGS: + lines.append(f"{var}='{LATEST_CEILINGS[key]}'") + elif val in ("NA", ""): + lines.append(f"{var}='{val}'") + else: + lines.append(f"{var}='={val}'") + + # psg already has = or == prefix + lines.append(f"psgver='{v['psg']}'") + # matplotlib carries a bare ceiling (e.g. "<3.11") or is empty for latest + lines.append(f"mplver='{v['matplotlib']}'") + return "\n".join(lines) + + +def format_summary(v): + """One-line summary for CI log.""" + parts = [] + for key in ( + "numpy", + "scipy", + "pandas", + "awkward", + "networkx", + "numba", + "pyyaml", + "sparse", + "psg", + "matplotlib", + ): + name = _SUMMARY_NAMES[key] + # psg carries its own "=" or "==" pin prefix; strip it so the summary reads + # "psg=10.5.0.0" like every other entry instead of "psg==10.5.0.0". + val = v[key].lstrip("=") if key == "psg" else v[key] + if val == "NA": + parts.append(f"{name}=NA") + elif val == "": + parts.append(f"{name}=latest{LATEST_CEILINGS.get(key, '')}") + else: + parts.append(f"{name}={val}") + return "versions: " + " ".join(parts) + + +# --------------------------------------------------------------------------- +# Validation (for testing the script itself) +# --------------------------------------------------------------------------- + + +def check_ceilings(): + """Check that every scipy and numba pin has a numpy-ceiling row. Returns errors. + + A pin with no row would silently skip the numpy-ceiling rule and be paired with a + numpy it refuses to import against. + """ + tables = [ + ( + "SCIPY_MAX_NUMPY", + "scipy", + SCIPY_MAX_NUMPY, + [SCIPY_VERSIONS["2.x"], *SCIPY_VERSIONS["1.x"].values()], + ), + ("NUMBA_MAX_NUMPY", "numba", NUMBA_MAX_NUMPY, list(NUMBA_VERSIONS.values())), + ] + errors = [] + for name, dep, table, pools in tables: + pins = {p for pool in pools for p in pool if p} + if missing := sorted(pins - table.keys(), key=_ver): + errors.append(f"{name} has no numpy ceiling for {dep} {', '.join(missing)}") + if extra := sorted(table.keys() - pins, key=_ver): + errors.append(f"{name} has rows for untested {dep} {', '.join(extra)}") + return errors + + +def check_python_pools(): + """Check that every per-Python table covers the same Pythons. Returns list of errors. + + Adding a Python to some tables but not others would otherwise surface as a KeyError + from somewhere deep in the picking, or (for the tables consulted only by a + constraint) as a silently wrong pick. + """ + tables = { + "NETWORKX_VERSIONS": NETWORKX_VERSIONS, + "PYYAML_VERSIONS": PYYAML_VERSIONS, + "SPARSE_VERSIONS": SPARSE_VERSIONS, + "NUMBA_MIN_PYTHON": NUMBA_MIN_PYTHON, + } + errors = [] + for name, table in tables.items(): + missing = sorted(NUMPY_VERSIONS.keys() - table.keys(), key=_ver) + extra = sorted(table.keys() - NUMPY_VERSIONS.keys(), key=_ver) + if missing: + errors.append(f"{name} has no entry for Python {', '.join(missing)}") + if extra: + errors.append(f"{name} has entries for unsupported Python {', '.join(extra)}") + return errors + + +def validate(v, pyver, source_type=None): + """Check that a version combination satisfies all constraints. Returns list of errors.""" + errors = [] + np_is_1x = v["numpy"].startswith("1.") if v["numpy"] else False + + # scipy >=1.15 requires numpy >=1.26.4 + if v["numpy"] in ("1.24", "1.25") and _ver(v["scipy"]) >= (1, 15): + errors.append(f"scipy {v['scipy']} requires numpy >=1.26.4, got {v['numpy']}") + + # scipy's declared numpy ceiling (see SCIPY_MAX_NUMPY) + ceiling = SCIPY_MAX_NUMPY.get(v["scipy"]) + if ceiling is not None and _numpy_ver(v, pyver) >= ceiling: + cap = ".".join(str(x) for x in ceiling) + errors.append(f"scipy {v['scipy']} requires numpy <{cap}, got {v['numpy'] or 'latest'}") + + # scipy Python availability + if _ver(pyver) >= (3, 15) and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 18): + errors.append(f"scipy {v['scipy']} has no py{pyver} build (needs >=1.18)") + elif _ver(pyver) >= (3, 14) and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 16): + errors.append(f"scipy {v['scipy']} has no py{pyver} build (needs >=1.16)") + elif _ver(pyver) >= (3, 13) and v["scipy"] == "1.13": + errors.append(f"scipy 1.13 has no py{pyver} build") + if v["scipy"] and _ver(pyver) < (floor := SCIPY_MIN_PYTHON.get(v["scipy"], (0,))): + need = ".".join(str(x) for x in floor) + errors.append(f"scipy {v['scipy']} requires Python >={need}, has no py{pyver} build") + + # pandas Python availability + if _ver(pyver) >= (3, 15) and v["pandas"] not in ("", "NA") and _ver(v["pandas"]) < (3, 0): + errors.append(f"pandas {v['pandas']} has no py{pyver} build") + elif _ver(pyver) >= (3, 14) and v["pandas"] == "2.2": + errors.append(f"pandas 2.2 has no py{pyver} build") + + # pandas 3.0 requirements (numba >=0.60 is unreachable; see `apply_constraints`) + if v["pandas"] == "3.0" and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 15): + errors.append(f"pandas 3.0 requires scipy >=1.14.1, got {v['scipy']}") + + # conda-forge pandas >=2.2 constrains scipy >=1.10.0 (fails the conda solve) + if ( + v["pandas"] not in ("", "NA") + and _ver(v["pandas"]) >= (2, 2) + and v["scipy"] not in ("", "NA") + and _ver(v["scipy"]) < (1, 10) + ): + errors.append(f"pandas {v['pandas']} constrains scipy >=1.10.0, got {v['scipy']}") + + # awkward Python availability + if _ver(pyver) >= (3, 15) and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 13): + errors.append(f"awkward {v['awkward']} has no py{pyver} build (needs >=2.13)") + elif _ver(pyver) >= (3, 14) and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 8): + errors.append(f"awkward {v['awkward']} has no py{pyver} build (needs >=2.8)") + elif _ver(pyver) >= (3, 13) and v["awkward"] == "2.6": + errors.append(f"awkward 2.6 has no py{pyver} build") + + # awkward <2.6 requires numpy 1.x (numpy.AxisError) + if not np_is_1x and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 6): + errors.append(f"awkward {v['awkward']} doesn't support numpy 2.x") + + # matplotlib 3.11 requires numpy >=1.25 at runtime (notebooks job) + if v["numpy"] and _ver(v["numpy"]) < (1, 25) and v.get("matplotlib", "") != "<3.11": + errors.append(f"matplotlib unpinned alongside numpy {v['numpy']} (needs <3.11)") + + # awkward <2.10 requires numpy <2.5 (generic-unit datetime64("NaT") at import) + if ( + _ver(v["numpy"]) >= (2, 5) + and v["awkward"] not in ("", "NA") + and _ver(v["awkward"]) < (2, 10) + ): + errors.append(f"awkward {v['awkward']} requires numpy <2.5, got {v['numpy'] or 'latest'}") + + # networkx >=3.5 constrains numpy >=1.25, scipy >=1.11.2 (the "1.11" pin clears it), + # and pandas >=2.0 + if _ver(v["networkx"]) >= (3, 5): + nx = v["networkx"] or "latest" + if _ver(v["numpy"]) < (1, 25): + errors.append(f"networkx {nx} requires numpy >=1.25, got {v['numpy']}") + if _ver(v["scipy"]) < (1, 11): + errors.append(f"networkx {nx} requires scipy >=1.11.2, got {v['scipy']}") + if _ver(v["pandas"]) < (2, 0): + errors.append(f"networkx {nx} requires pandas >=2.0, got {v['pandas']}") + + # numba Python minimums + numba_min = _min_for_python(NUMBA_MIN_PYTHON, pyver) + if v["numba"] not in ("", "NA") and _ver(v["numba"]) < numba_min: + errors.append(f"numba {v['numba']} doesn't support Python {pyver}") + + # numba's declared numpy ceiling (see NUMBA_MAX_NUMPY) + ceiling = _numba_ceiling(v["numba"]) + if ceiling is not None and _numpy_ver(v, pyver) >= ceiling: + cap = ".".join(str(x) for x in ceiling) + errors.append( + f"numba {v['numba'] or 'latest'} requires numpy <{cap}, " + f"got {v['numpy'] or 'latest'}" + ) + + # sparse: the per-Python pool is the whole rule, so nothing may leave it + if v["sparse"] not in SPARSE_VERSIONS[pyver]: + errors.append(f"sparse {v['sparse']} isn't in the py{pyver} pool") + + # psg build availability (the floor depends on the Python and where psg comes from) + floors = _min_for_python(PSG_MIN_PYTHON, pyver) + if floors is not None and source_type in floors and v.get("psg"): + psgver = v["psg"].lstrip("=") + floor = floors[source_type] + if _ver(psgver) < _ver(floor): + errors.append(f"psg {psgver} has no py{pyver} {source_type} build (needs >={floor})") + + return errors + + +def stress_test(n=10000): + """Run n random picks for each Python/source combo and validate all.""" + total = 0 + failures = 0 + for pyver in NUMPY_VERSIONS: + for source in ("conda-forge", "wheel", "source", "upstream"): + for _ in range(n): + v = pick_versions(pyver, source) + errs = validate(v, pyver, source) + total += 1 + if errs: + failures += 1 + print(f"FAIL py{pyver} {source}: {v}", file=sys.stderr) + for e in errs: + print(f" - {e}", file=sys.stderr) + print(f"Stress test: {total} combos, {failures} failures", file=sys.stderr) + return failures + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main(): + parser = argparse.ArgumentParser(description="Pick random CI dependency versions") + parser.add_argument("--python", help="Python version (e.g. 3.12)") + parser.add_argument( + "--source", + choices=["conda-forge", "wheel", "source", "upstream"], + help="Package source type", + ) + parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility") + parser.add_argument( + "--validate", action="store_true", help="Run stress test to validate all constraints" + ) + args = parser.parse_args() + + if args.seed is not None: + random.seed(args.seed) + + if pool_errors := check_python_pools() + check_ceilings(): + print("Error: version pools are inconsistent:", file=sys.stderr) + for e in pool_errors: + print(f" - {e}", file=sys.stderr) + sys.exit(1) + + if args.validate: + failures = stress_test() + sys.exit(1 if failures else 0) + + if not args.python or not args.source: + parser.error("--python and --source are required (unless --validate)") + + pyver = args.python + if pyver not in NUMPY_VERSIONS: + print(f"Error: unsupported Python version {pyver}", file=sys.stderr) + print(f"Supported: {', '.join(NUMPY_VERSIONS.keys())}", file=sys.stderr) + sys.exit(1) + + v = pick_versions(pyver, args.source) + + # Print summary to stderr (visible in CI logs) + print(format_summary(v), file=sys.stderr) + + # Print bash-eval-safe key=value to stdout + print(format_output(v)) + + +if __name__ == "__main__": + main() diff --git a/scripts/create_pickle.py b/scripts/create_pickle.py index 10fe58630..f1baac753 100755 --- a/scripts/create_pickle.py +++ b/scripts/create_pickle.py @@ -4,6 +4,7 @@ Note that the exact binary of the pickle files may differ depending on which Python version is used to create them. """ + import argparse import pickle from pathlib import Path diff --git a/scripts/find_aiisms.sh b/scripts/find_aiisms.sh new file mode 100755 index 000000000..12ebb407c --- /dev/null +++ b/scripts/find_aiisms.sh @@ -0,0 +1,103 @@ +#!/usr/bin/env bash +# Find AI-ism candidates in branch-changed files: prose tells left by AI agents. +# A helper for the gbdoc workflow (.claude/skills/gbdoc/SKILL.md). Use, adjust, +# copy/paste, etc. as necessary; this finds candidates, judgment does the rest. +# +# Scans text files changed on the current branch (vs a base ref) for: +# 1. The dash family: em-dash, en-dash, and "--" used as a separator. +# 2. Stray unicode: bytes outside printable ASCII (smart quotes, ellipsis, +# arrows, ...). +# +# Output is `file:line:content` so hits are jump-to-able. It greps the changed +# files directly, so PRE-EXISTING hits in those files show up too. Only touch +# prose the branch actually changed; leave intentional dashes/unicode alone. +# RST underlines and "# ---" banners are already skipped, "--flag" CLI usage is +# not matched, and some files (e.g. README.md) use unicode on purpose. +# +# Uses plain `grep -E` so it works with BSD grep (macOS) and GNU grep alike; +# no PCRE / `-P` needed. The em-dash and en-dash in `dash_re` below are literal +# on purpose; they are the characters being hunted. +# +# Usage: +# scripts/find_aiisms.sh [base-ref] [path-filter] +# scripts/find_aiisms.sh # vs main, whole repo +# scripts/find_aiisms.sh main graphblas/core/ # vs main, scoped subtree +# scripts/find_aiisms.sh --strict # exit 1 if anything is found +set -euo pipefail + +strict=0 +if [ "${1:-}" = "--strict" ]; then + strict=1 + shift +fi +base="${1:-main}" +path_filter="${2:-}" + +if ! git rev-parse --verify --quiet "$base" >/dev/null; then + echo "Unknown base ref: ${base}" >&2 + exit 2 +fi + +# Collect text-y files changed on the branch (drop deleted files). +files=() +while IFS= read -r f; do + files+=("$f") +done < <( + if [ -n "$path_filter" ]; then + git diff "${base}...HEAD" --name-only --diff-filter=d -- "$path_filter" + else + git diff "${base}...HEAD" --name-only --diff-filter=d + fi | grep -E '\.(py|rst|md|toml|ya?ml|sh|cfg)$' || true +) + +if [ "${#files[@]}" -eq 0 ]; then + echo "No changed text files vs ${base}." + exit 0 +fi + +# Dash family: em-dash, en-dash, " -- " separator, or "word--" jammed form. +# Pure-punctuation lines (RST underlines, "# ---" banners) and "--flag" CLI +# usage are naturally skipped: none have an alnum-adjacent or space-bounded +# "--", nor a real em/en-dash character. +dash_re='—|–| -- |[[:alnum:]]--' +# Stray bytes: anything outside printable ASCII, allowing tab and CR. Run under +# LC_ALL=C so grep matches byte-wise and the check stays locale-independent. +_ws="$(printf '\t\r')" +unicode_re="[^${_ws} -~]" + +scan() { # $1 = label, $2 = regex, $3 = "C" to grep under LC_ALL=C + local label="$1" re="$2" loc="$3" f m hits="" + for f in "${files[@]}"; do + [ -f "$f" ] || continue + if [ "$loc" = "C" ]; then + m="$(LC_ALL=C grep -nE "$re" "$f" 2>/dev/null || true)" + else + m="$(grep -nE "$re" "$f" 2>/dev/null || true)" + fi + if [ -n "$m" ]; then + hits="${hits}$(printf '%s\n' "$m" | sed "s|^|${f}:|")"$'\n' + fi + done + echo "=== ${label} ===" + if [ -n "$hits" ]; then + printf '%s' "$hits" + echo "($(printf '%s' "$hits" | grep -c . || true) candidate line(s))" + return 1 + fi + echo "(none)" + return 0 +} + +rc=0 +scan "dash family (em/en-dash, '--' as separator)" "$dash_re" "" || rc=1 +echo +scan "stray unicode (non-ASCII bytes)" "$unicode_re" "C" || rc=1 +echo +echo "Note: results include any pre-existing hits in changed files. Only touch" +echo "prose the branch changed; leave intentional dashes/unicode alone." +echo "See .claude/skills/gbdoc/SKILL.md for the full workflow." + +if [ "$strict" = "1" ]; then + exit "$rc" +fi +exit 0 diff --git a/scripts/jit_diagnostics.py b/scripts/jit_diagnostics.py new file mode 100755 index 000000000..6959bfb2d --- /dev/null +++ b/scripts/jit_diagnostics.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python +"""Print JIT diagnostic information for debugging compiler configuration. + +Run with: python scripts/jit_diagnostics.py + +This prints the GraphBLAS JIT configuration, available compilers, sysconfig +values, and attempts a JIT compilation to verify everything works. +""" + +import os +import pathlib +import shutil +import sys +import sysconfig + + +def main(): + print("=" * 60) + print("JIT Diagnostics") + print("=" * 60) + + print("\n--- Platform ---") + print(f"sys.platform: {sys.platform}") + print(f"Python: {sys.version}") + + print("\n--- Environment ---") + conda = os.environ.get("CONDA_PREFIX", "") + print(f"CONDA_PREFIX: {conda or 'NOT SET'}") + print(f"GITHUB_ACTIONS: {os.environ.get('GITHUB_ACTIONS', 'NOT SET')}") + + print("\n--- sysconfig ---") + for key in ["CC", "CXX", "CFLAGS", "LDFLAGS", "LIBS"]: + val = sysconfig.get_config_var(key) + if val and len(str(val)) > 100: + val = str(val)[:100] + "..." + print(f" {key}: {val}") + print(f" include: {sysconfig.get_path('include')}") + + print("\n--- Compiler search ---") + candidates = ["cc", "gcc", "clang"] + if sys.platform == "linux": + candidates.append("x86_64-conda-linux-gnu-cc") + elif sys.platform == "darwin": + candidates.extend( + [ + "x86_64-apple-darwin13.4.0-clang", + "arm64-apple-darwin20.0.0-clang", + ] + ) + for name in candidates: + which = shutil.which(name) + conda_path = pathlib.Path(conda) / "bin" / name if conda else None + conda_exists = conda_path.exists() if conda_path else False + print(f" {name}: which={which}, conda={'yes' if conda_exists else 'no'}") + + # Import graphblas + try: + import graphblas as gb + + gb.init("suitesparse") + except Exception as e: + print(f"\nERROR: Could not initialize graphblas: {e}") + return 1 + + from graphblas.core.ss import _IS_SSGB7 + + if _IS_SSGB7: + print("\nSuiteSparse:GraphBLAS 7.x: JIT not available") + return 0 + + print("\n--- GraphBLAS JIT settings after graphblas init ---") + for key in [ + "jit_c_control", + "jit_c_compiler_name", + "jit_c_compiler_flags", + "jit_c_linker_flags", + "jit_c_libraries", + "jit_c_cmake_libs", + "jit_cache_path", + ]: + val = gb.ss.config[key] + if isinstance(val, str) and len(val) > 120: + val = val[:120] + "..." + print(f" {key}: {val}") + + # Check if default compiler exists + jit_cc = gb.ss.config["jit_c_compiler_name"] + print("\n--- Compiler path analysis ---") + print(f" Default compiler: {jit_cc}") + print(f" Exists? {pathlib.Path(jit_cc).exists()}") + cc_basename = pathlib.Path(jit_cc).name + if conda: + local_cc = pathlib.Path(conda) / "bin" / cc_basename + print(f" Conda equivalent: {local_cc}") + print(f" Exists? {local_cc.exists()}") + fallback_cc = pathlib.Path(conda) / "bin" / "cc" + print(f" Fallback (cc): {fallback_cc}") + print(f" Exists? {fallback_cc.exists()}") + + # Check for problematic flags + flags = gb.ss.config["jit_c_compiler_flags"] + import re + + if isysroot := re.search(r"-isysroot\s+(\S+)", flags): + path = isysroot.group(1) + print(f"\n -isysroot: {path}") + print(f" Exists? {pathlib.Path(path).exists()}") + fdebug = re.findall(r"-fdebug-prefix-map=(\S+)", flags) + for d in fdebug: + print(f" -fdebug-prefix-map: {d}") + + # Try the fix + print("\n--- Attempting gb.ss.fix_jit_config ---") + result = gb.ss.fix_jit_config() + # True=JIT working, False=probe failed, None=no conda env + desc = "working" if result is True else "probe failed" if result is False else "no conda" + print(f" Result: {result} ({desc})") + if result is True: + print(f" Compiler: {gb.ss.config['jit_c_compiler_name']}") + flags_after = gb.ss.config["jit_c_compiler_flags"] + if len(flags_after) > 120: + flags_after = flags_after[:120] + "..." + print(f" Flags: {flags_after}") + + # Try JIT compilation + print("\n--- JIT compilation test ---") + if gb.ss.config["jit_c_control"] == "off": + print(" SKIPPED (JIT is off)") + else: + try: + from graphblas import dtypes + + prev_burble = gb.ss.config["burble"] + gb.ss.config["burble"] = True + dtype = dtypes.ss.register_new( + "jit_diag_test", + "typedef struct { int val ; } jit_diag_test ;", + ) + gb.ss.config["burble"] = prev_burble + print(f" SUCCESS: registered type '{dtype.name}'") + except Exception as e: + print(f" FAILED: {type(e).__name__}: {e}") + try: + err_log = gb.ss.config["jit_error_log"] + if err_log: + print(f" JIT error log: {err_log[:500]}") + except Exception: + pass + + # Print final JIT state + print("\n--- Final JIT state ---") + print(f" jit_c_control: {gb.ss.config['jit_c_control']}") + print(f" jit_c_compiler_name: {gb.ss.config['jit_c_compiler_name']}") + + print(f"\n{'=' * 60}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())