From 22d42f615187a3a2c89b6400a898c5b72ef396a8 Mon Sep 17 00:00:00 2001 From: Jim Kitchen <2807270+jim22k@users.noreply.github.com> Date: Wed, 26 Feb 2025 10:02:04 -0600 Subject: [PATCH 01/52] Add 9.4.5.0 to tests (#557) --- .github/workflows/test_and_build.yml | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index bfc17834b..af7525928 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -217,24 +217,24 @@ jobs: 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", ""]))') + psgver=$(python -c 'import random ; print(random.choice(["=9.3.1.0", "=9.4.5.0", ""]))') psg=python-suitesparse-graphblas${psgver} else - psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", ""]))') + psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", "==9.4.5.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", ""]))') + psgver=$(python -c 'import random ; print(random.choice(["=9.3.1.0", "=9.4.5.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", ""]))') + psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", "==9.4.5.0", ""]))') fi fi # python-suitsparse-graphblas support is the same for Python 3.10 and 3.11 @@ -242,21 +242,21 @@ jobs: 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", ""]))') + psgver=$(python -c 'import random ; print(random.choice(["=9.3.1.0", "=9.4.5.0", ""]))') fi 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", ""]))') + psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", "==9.4.5.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", ""]))') + psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", "==9.4.5.0", ""]))') fi fi @@ -357,7 +357,7 @@ jobs: ${{ matrix.slowtask == 'pytest_bizarro' && 'black' || '' }} \ ${{ matrix.slowtask == 'notebooks' && 'matplotlib nbconvert jupyter "ipython>=7" drawsvg' || '' }} \ ${{ steps.sourcetype.outputs.selected == 'upstream' && 'cython' || '' }} \ - ${{ steps.sourcetype.outputs.selected != 'wheel' && '"graphblas>=7.4,<9.4"' || '' }} \ + ${{ steps.sourcetype.outputs.selected != 'wheel' && '"graphblas>=7.4,<9.5"' || '' }} \ ${{ contains(steps.pyver.outputs.selected, 'pypy') && 'pypy' || '' }} \ ${{ matrix.os == 'windows-latest' && 'cmake' || 'm4' }} \ # ${{ matrix.os != 'windows-latest' && 'pytest-forked' || '' }} # to investigate crashes From fec13d7b33ba0aba6ca52d2b13151128fae9e541 Mon Sep 17 00:00:00 2001 From: Jim Kitchen <2807270+jim22k@users.noreply.github.com> Date: Sat, 21 Mar 2026 17:15:15 -0500 Subject: [PATCH 02/52] fix segfault, and make repr work with pandas 3 (#570) --- graphblas/core/formatting.py | 32 +++++++++++++++++++++----------- graphblas/core/ss/matrix.py | 23 ++++++++++++++++------- graphblas/core/ss/vector.py | 23 ++++++++++++++++------- graphblas/tests/test_matrix.py | 2 +- graphblas/tests/test_vector.py | 7 ++----- 5 files changed, 56 insertions(+), 31 deletions(-) diff --git a/graphblas/core/formatting.py b/graphblas/core/formatting.py index 0b6252101..5fe9b6972 100644 --- a/graphblas/core/formatting.py +++ b/graphblas/core/formatting.py @@ -119,7 +119,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 +167,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 +209,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 +252,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 +264,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 +315,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 diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index 509c56113..de6395cf6 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -4088,13 +4088,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( diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index fdde7eb92..4acef3fe5 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -1659,13 +1659,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( diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 24f0e73d7..b972ef260 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -2794,7 +2794,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]) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index db80cdf71..d9bf84495 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -2221,11 +2221,8 @@ 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: + with pytest.raises(SyntaxError): + # The dtype name is too long to embed in the blob; dtype= must be provided 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) From c726047fa1f905bc17c121bc08269d7ca8d04ea7 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 22 Apr 2026 15:07:06 -0500 Subject: [PATCH 03/52] Maintenance overhaul: Python 3.14 support, drop 3.10, add IndexBinaryOp (#571) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Maintenance overhaul: Python 3.14 support, drop 3.10, add IndexBinaryOp Python version support: - Add Python 3.14, drop 3.10 (minimum now 3.11) - Update numba guards to python_version<'3.15' in pyproject.toml - Update all tool target-versions (black, ruff, pylint, pyupgrade) New feature: - Add IndexBinaryOp: user-defined f(x,ix,jx,y,iy,jy,theta)->z operators that bind theta to produce BinaryOps for use in ewise operations (requires SuiteSparse:GraphBLAS >=9.4, guarded with hasattr) Packaging modernization: - PEP 639: license = "Apache-2.0" with license-files - Bump setuptools >=77, numpy >=1.24, numba >=0.57 - Remove suitesparse-graphblas <10 upper bound (support psg 10.x) - Deprecate fast_matrix_market (no longer maintained) CI updates: - Standardize all GitHub Actions to latest (checkout v6, setup-python v6, choose-random-action v4.1, pypi-publish v1.13) - Add psg 10.x versions (10.0.1, 10.1.1, 10.3.1) to CI test matrix - Fix graphblas C library constraint (remove <9.5 cap) - Fix numba+numpy2 : numba >=0.62 now tested with numpy 2.x - Split numba version overrides: 0.61+ for py3.13, 0.63+ for py3.14 - Update pre-commit hooks, flake8-bugbear, flake8-simplify - Remove stale warning filters, update dependency version ranges Multi-version SuiteSparse:GraphBLAS support: - Handle GPU_ID->GPU_IDS rename in SS:GraphBLAS 10.2+ with hasattr guards - Remove awkward v1 dead code, fix XXX error messages, update copyright - Clean up stale MAINT comments and documentation references * Oops, don't run test_indexbinary.py on older psg versions; fix psg constraint * Fix more random version edges * Fix JIT tests (hopefully), simplify CI dependency version selection JIT: - Fix JIT tests by auto-detecting conda compiler from $CONDA_PREFIX/bin/ and stripping build-time -isysroot and -fdebug-prefix-map flags - Replace ~70 lines of hardcoded per-platform JIT config with generic _fix_jit_config() that works on all platforms - Tests now pytest.skip() when JIT unavailable (previously passed silently) - Add scripts/jit_diagnostics.py for CI debugging CI (test_and_build.yml): - Restructure dependency version selection: pick numpy first, then choose other deps based on numpy 1.x vs 2.x, then apply constraints - Remove psg 9.4.5 from Python 3.14 (no conda builds) - Fix numba 0.63 + numpy 2.4 incompatibility (require numba 0.64) - Remove suitesparse extras <10 cap in pyproject.toml - Add psg 10.x versions to all test paths * Fix bizarro scalar failures for IndexBinaryOp, improve JIT diagnostics - Add `# pragma: is_grbscalar` to IndexBinaryOp theta scalar creation so bizarro scalars sed doesn't break GxB_BinaryOp_new_IndexOp (requires GrB_Scalar, not C scalar) - Add burble and error log capture to JIT diagnostic script to reveal the actual compiler command and error on macOS ARM (still failing) - Add more compiler fallback candidates (clang, gcc) in _fix_jit_config * Fix CI failures: macOS JIT sysroot, dep version constraints, warning filter JIT: - Replace -isysroot with local macOS SDK via `xcrun --show-sdk-path` instead of stripping it (conda cross-compiler needs a sysroot) - Move subprocess import to module level CI dependency constraints: - Add numpy 2.3+/scipy constraint (scipy <1.15.1 needs numpy <2.3) - Move numba/numpy 2.4 constraint to be Python-version-independent - Move upstream numpy override before dep selection to prevent stale scipy/pandas/awkward versions being chosen for numpy 1.x then overridden to numpy 2.x Warning filter: - Restore np.find_common_type deprecation filter (needed for numpy 1.25 with older scipy/networkx) * Fix JIT tests: probe compilation before running, don't replace working compilers - Only replace the JIT compiler path if the default doesn't exist; some psg builds use /usr/bin/cc which already works - Add a probe compilation after fixing config: if JIT compilation fails for any reason (wrong flags, missing libs), turn JIT off and let tests skip cleanly instead of failing - Move upstream numpy override before dep selection to prevent stale scipy/pandas being chosen for numpy 1.x then overridden to 2.x - Add scipy/numpy 2.3+ constraint (scipy <1.15.1 needs numpy <2.3) - Make numba/numpy 2.4 constraint Python-version-independent - Restore np.find_common_type deprecation warning filter - Add macos-15-intel and ubuntu-24.04-arm to CI test matrix * Clean. Better — single try/yield/finally, clear three-way branching with comments explaining each case Fix JIT: three-way probe result, sysconfig won't override failed probe - _fix_jit_config returns True/False/None: working, probe-failed, no-conda - Fixture uses single try/yield/finally with clear three-way branch - When probe fails (False), JIT stays off — sysconfig won't retry - Sysconfig path only runs when there's no conda env (None) - Add pandas 3.0 requires numba >=0.60 constraint - Add macos-15-intel and ubuntu-24.04-arm to CI matrix - Stress testing in CI * Fix pandas 3.0 requires scipy >=1.14.1 constraint * Remove stress testing in CI; warn on auto-detected fast_matrix_market - Deprecation warning now fires on engine="auto" when fast_matrix_market is installed (previously only warned on explicit engine="fmm") * More more dependency tweak (pandas 3.0); more stress testing * Stress test pass. Remove stress testing * Another update to random versions in CI * Add (fix) IndexBinaryOp JIT support and move CI version selection to Python IndexBinaryOp ss namespace: - Create core/ss/indexbinary.py with JIT register_new() for C-string compilation via GxB_IndexBinaryOp_new, matching the indexunary pattern - Update indexbinary/ss.py to import register_new (was empty stub) - Add test_jit_indexbinary to test_ssjit.py (theta binding, mixed types, nested names, duplicate registration) - Verify indexbinary.ss.register_new exists in test_dir_and_module CI dependency version selection: - Replace ~190 lines of cascading bash if/elif with a Python script (scripts/ci_pick_versions.py) that declares version pools and constraints, with a built-in stress test (160k combos, 0 failures) - Workflow calls eval "$(python scripts/ci_pick_versions.py ...)" - Also fixes scipy >=1.15 + numpy <1.26 incompatibility --- .flake8 | 2 +- .github/workflows/debug.yml | 4 +- .github/workflows/imports.yml | 12 +- .github/workflows/lint.yml | 6 +- .github/workflows/publish_pypi.yml | 10 +- .github/workflows/test_and_build.yml | 221 ++-------- .pre-commit-config.yaml | 99 +++-- README.md | 5 +- dev-requirements.txt | 2 +- docs/conf.py | 2 +- docs/env.yml | 2 +- docs/getting_started/index.rst | 1 - docs/user_guide/io.rst | 2 - graphblas/__init__.py | 1 + graphblas/core/base.py | 20 +- graphblas/core/expr.py | 4 + graphblas/core/infix.py | 22 +- graphblas/core/mask.py | 2 + graphblas/core/matrix.py | 3 +- graphblas/core/operator/__init__.py | 1 + graphblas/core/operator/indexbinary.py | 507 +++++++++++++++++++++++ graphblas/core/operator/utils.py | 16 +- graphblas/core/recorder.py | 2 +- graphblas/core/scalar.py | 4 + graphblas/core/ss/__init__.py | 2 +- graphblas/core/ss/context.py | 12 +- graphblas/core/ss/descriptor.py | 5 +- graphblas/core/ss/indexbinary.py | 134 +++++++ graphblas/core/ss/matrix.py | 2 +- graphblas/core/ss/vector.py | 2 +- graphblas/core/vector.py | 2 +- graphblas/indexbinary/__init__.py | 33 ++ graphblas/indexbinary/ss.py | 6 + graphblas/io/_awkward.py | 17 +- graphblas/io/_matrixmarket.py | 34 ++ graphblas/monoid/numpy.py | 2 +- graphblas/ss/_core.py | 6 +- graphblas/tests/test_core.py | 2 +- graphblas/tests/test_formatting.py | 1 - graphblas/tests/test_indexbinary.py | 288 +++++++++++++ graphblas/tests/test_infix.py | 84 ++-- graphblas/tests/test_io.py | 2 +- graphblas/tests/test_matrix.py | 8 +- graphblas/tests/test_numpyops.py | 8 +- graphblas/tests/test_op.py | 4 +- graphblas/tests/test_scalar.py | 2 +- graphblas/tests/test_ss_utils.py | 4 +- graphblas/tests/test_ssjit.py | 264 +++++++----- graphblas/tests/test_vector.py | 1 - graphblas/viz.py | 4 +- pyproject.toml | 80 ++-- scripts/check_versions.sh | 28 +- scripts/ci_pick_versions.py | 532 +++++++++++++++++++++++++ scripts/create_pickle.py | 1 + scripts/jit_diagnostics.py | 163 ++++++++ 55 files changed, 2197 insertions(+), 486 deletions(-) create mode 100644 graphblas/core/operator/indexbinary.py create mode 100644 graphblas/core/ss/indexbinary.py create mode 100644 graphblas/indexbinary/__init__.py create mode 100644 graphblas/indexbinary/ss.py create mode 100644 graphblas/tests/test_indexbinary.py create mode 100755 scripts/ci_pick_versions.py create mode 100755 scripts/jit_diagnostics.py 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/workflows/debug.yml b/.github/workflows/debug.yml index 6c2b202b1..45401c704 100644 --- a/.github/workflows/debug.yml +++ b/.github/workflows/debug.yml @@ -15,7 +15,7 @@ jobs: strategy: fail-fast: false matrix: - pyver: [3.10] + pyver: [3.11] testopts: - "--blocking" # - "--non-blocking --record --runslow" @@ -26,7 +26,7 @@ jobs: # - "conda-forge" steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 persist-credentials: false diff --git a/.github/workflows/imports.yml b/.github/workflows/imports.yml index e24d0d4db..4be4926c7 100644 --- a/.github/workflows/imports.yml +++ b/.github/workflows/imports.yml @@ -14,7 +14,7 @@ jobs: pyver: ${{ steps.pyver.outputs.selected }} steps: - name: RNG for os - uses: ddradar/choose-random-action@v3.0.0 + uses: ddradar/choose-random-action@v4.1.0 id: os with: contents: | @@ -26,14 +26,14 @@ jobs: 1 1 - name: RNG for Python version - uses: ddradar/choose-random-action@v3.0.0 + uses: ddradar/choose-random-action@v4.1.0 id: pyver with: contents: | - 3.10 3.11 3.12 3.13 + 3.14 weights: | 1 1 @@ -45,14 +45,14 @@ jobs: # 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"] # os: ["ubuntu-latest", "macos-latest", "windows-latest"] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 with: fetch-depth: 0 persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v6 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..8bceeeadd 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@v6 with: fetch-depth: 0 persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - uses: pre-commit/action@v3.0.1 diff --git a/.github/workflows/publish_pypi.yml b/.github/workflows/publish_pypi.yml index 32926c5c8..ad9284f66 100644 --- a/.github/workflows/publish_pypi.yml +++ b/.github/workflows/publish_pypi.yml @@ -8,20 +8,19 @@ on: 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@v6 with: fetch-depth: 0 persist-credentials: false - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - name: Install build dependencies run: | python -m pip install --upgrade pip @@ -36,7 +35,8 @@ jobs: - name: Check with twine run: python -m twine check --strict dist/* - name: Publish to PyPI - uses: pypa/gh-action-pypi-publish@v1.12.4 + uses: pypa/gh-action-pypi-publish@v1.13.0 + if: startsWith(github.ref, 'refs/tags/') && github.repository == 'python-graphblas/python-graphblas' with: user: __token__ password: ${{ secrets.PYPI_TOKEN }} diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index af7525928..d31e8c534 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -50,7 +50,7 @@ jobs: backend: ${{ steps.backend.outputs.selected }} steps: - name: RNG for mapnumpy - uses: ddradar/choose-random-action@v3.0.0 + uses: ddradar/choose-random-action@v4.1.0 id: mapnumpy with: contents: | @@ -64,7 +64,7 @@ jobs: 1 1 - name: RNG for backend - uses: ddradar/choose-random-action@v3.0.0 + uses: ddradar/choose-random-action@v4.1.0 id: backend with: contents: | @@ -93,35 +93,50 @@ 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. KMP_DUPLICATE_LIB_OK: ${{ contains(matrix.os, 'macos') && 'TRUE' || 'FALSE' }} steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 persist-credentials: false - name: RNG for Python version - uses: ddradar/choose-random-action@v3.0.0 + uses: ddradar/choose-random-action@v4.1.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 + 3.14 weights: | 1 1 1 1 - name: RNG for source of python-suitesparse-graphblas - uses: ddradar/choose-random-action@v3.0.0 + uses: ddradar/choose-random-action@v4.1.0 id: sourcetype with: # Weights must be natural numbers, so set weights to very large to skip one @@ -149,164 +164,25 @@ jobs: auto-activate-base: false - name: Update env 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). + # 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 "v='${{ steps.pyver.outputs.selected }}'.split('.');print(f'{v[0]}.{v[1]}')") + eval "$(python scripts/ci_pick_versions.py --python "$pyver" --source ${{ steps.sourcetype.outputs.selected }})" - # 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="" - fi - - # 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) 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", "=9.4.5.0", ""]))') - psg=python-suitesparse-graphblas${psgver} - else - psgver=$(python -c 'import random ; print(random.choice(["==9.3.1.0", "==9.4.5.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", "=9.4.5.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", "==9.4.5.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", "=9.4.5.0", ""]))') - fi + if [[ ${{ steps.sourcetype.outputs.selected}} == "conda-forge" && -n "$psgver" ]] ; 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", "==9.4.5.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", "==9.4.5.0", ""]))') - fi fi - # 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", ""]))') - fi - - # Handle NumPy 2 - if [[ $npver != =1.* ]] ; then - # Only pandas >=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 - 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 @@ -315,36 +191,23 @@ jobs: if [[ ${{ contains(steps.pyver.outputs.selected, 'pypy') }} ]]; 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}" @@ -353,11 +216,11 @@ jobs: $(command -v mamba || command -v conda) install -c nodefaults \ packaging 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' || '' }} \ ${{ steps.sourcetype.outputs.selected == 'upstream' && 'cython' || '' }} \ - ${{ steps.sourcetype.outputs.selected != 'wheel' && '"graphblas>=7.4,<9.5"' || '' }} \ + ${{ steps.sourcetype.outputs.selected != 'wheel' && '"graphblas>=7.4"' || '' }} \ ${{ contains(steps.pyver.outputs.selected, 'pypy') && 'pypy' || '' }} \ ${{ matrix.os == 'windows-latest' && 'cmake' || 'm4' }} \ # ${{ matrix.os != 'windows-latest' && 'pytest-forked' || '' }} # to investigate crashes @@ -387,10 +250,8 @@ jobs: (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 run: | A=${{ needs.rngs.outputs.mapnumpy == 'A' || '' }} ; B=${{ needs.rngs.outputs.mapnumpy == 'B' || '' }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 43e28b8fe..20b632c30 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -16,20 +16,23 @@ 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 - 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 +40,142 @@ repos: - id: name-tests-test args: ["--pytest-test-first"] - repo: https://github.com/abravalheri/validate-pyproject - rev: v0.23 + rev: v0.25 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.3.3 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: 8.0.1 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.21.2 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.3.1 hooks: - id: black - id: black-jupyter - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.9.6 + rev: v0.15.7 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.3.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.3.0 + - flake8-bugbear==25.11.29 + - flake8-simplify==0.30.0 - repo: https://github.com/codespell-project/codespell - rev: v2.4.1 + rev: v2.4.2 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.15.7 hooks: - - id: ruff + - id: ruff-check + # - id: ruff-format # Prefer black, but may temporarily uncomment this to see - 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" + 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" hooks: - id: shellcheck - repo: https://github.com/rbubley/mirrors-prettier - rev: v3.5.1 + rev: v3.8.1 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 - repo: https://github.com/rhysd/actionlint - rev: v1.7.7 + rev: v1.7.11 hooks: - id: actionlint - repo: https://github.com/python-jsonschema/check-jsonschema - rev: 0.31.1 + rev: 0.37.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 - hooks: - - id: zizmor + # - repo: https://github.com/woodruffw/zizmor-pre-commit + # rev: v1.14.2 + # hooks: + # - id: zizmor - 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:.*_.*" # 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 +187,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/README.md b/README.md index 1080314c7..7b157c979 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 diff --git a/dev-requirements.txt b/dev-requirements.txt index a281672ec..c6b586369 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 diff --git a/docs/conf.py b/docs/conf.py index 283f6d047..da8d0dc67 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 diff --git a/docs/env.yml b/docs/env.yml index 78a50afbe..caec812d6 100644 --- a/docs/env.yml +++ b/docs/env.yml @@ -3,7 +3,7 @@ channels: - conda-forge - nodefaults dependencies: - - python=3.10 + - python=3.11 - pip # python-graphblas dependencies - donfig diff --git a/docs/getting_started/index.rst b/docs/getting_started/index.rst index 3726131d2..0c02b6194 100644 --- a/docs/getting_started/index.rst +++ b/docs/getting_started/index.rst @@ -34,7 +34,6 @@ to work. - `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`` GraphBLAS Fundamentals ---------------------- diff --git a/docs/user_guide/io.rst b/docs/user_guide/io.rst index f27b40bd3..8ce80184e 100644 --- a/docs/user_guide/io.rst +++ b/docs/user_guide/io.rst @@ -146,5 +146,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/graphblas/__init__.py b/graphblas/__init__.py index 63110eeeb..86759f570 100644 --- a/graphblas/__init__.py +++ b/graphblas/__init__.py @@ -49,6 +49,7 @@ def get_config(): "core", "dtypes", "exceptions", + "indexbinary", "indexunary", "io", "monoid", diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 24a49ba1a..15f66bc2f 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -266,28 +266,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): diff --git a/graphblas/core/expr.py b/graphblas/core/expr.py index efec2db5f..47ff18a0b 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)}.") @@ -473,6 +475,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)}.") diff --git a/graphblas/core/infix.py b/graphblas/core/infix.py index 24c109639..73b3b91b5 100644 --- a/graphblas/core/infix.py +++ b/graphblas/core/infix.py @@ -70,6 +70,8 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): def is_grbscalar(self): return not self.is_cscalar + __hash__ = None + # Begin auto-generated code: Scalar __and__ = wrapdoc(Scalar.__and__)(property(automethods.__and__)) __array__ = wrapdoc(Scalar.__array__)(property(automethods.__array__)) @@ -134,10 +136,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 +163,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): diff --git a/graphblas/core/mask.py b/graphblas/core/mask.py index 3bda2188a..4a8412a60 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)}.") diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index bf20cc953..698bf7b4f 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -636,7 +636,7 @@ 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( @@ -3979,6 +3979,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/indexbinary.py b/graphblas/core/operator/indexbinary.py new file mode 100644 index 000000000..cb07e0eaa --- /dev/null +++ b/graphblas/core/operator/indexbinary.py @@ -0,0 +1,507 @@ +import inspect +from types import FunctionType + +from ... import _STANDARD_OPERATOR_NAMES, indexbinary +from ...dtypes import BOOL, INT8, UINT64, lookup_dtype +from ...exceptions import UdfParseError, check_status_carg +from .. import _has_numba, ffi, lib +from ..dtypes import _sample_values +from .base import OpBase, ParameterizedUdf, TypedOpBase, _deserialize_parameterized + +_has_idxbinop = hasattr(lib, "GxB_IndexBinaryOp_new") + +if _has_numba: + import numba + +ffi_new = ffi.new + + +class _BoundIndexBinaryOp(TypedOpBase): + """A BinaryOp created by binding a theta value to an IndexBinaryOp.""" + + __slots__ = () + opclass = "BinaryOp" + + +class TypedBuiltinIndexBinaryOp(TypedOpBase): + __slots__ = () + opclass = "IndexBinaryOp" + + def __call__(self, theta=None): + """Bind a theta value to create a BinaryOp that can be used in operations. + + Parameters + ---------- + theta : scalar, optional + The theta parameter to bind. Default is 0 (False). + + Returns + ------- + TypedOpBase + A BinaryOp created from this IndexBinaryOp with the given theta. + + """ + from ..scalar import Scalar + + if theta is None: + theta = False + if not isinstance(theta, Scalar): + theta = Scalar.from_value(theta, is_cscalar=False, name="") # pragma: is_grbscalar + elif theta._is_cscalar: + # fmt: off + val = theta.value + theta = Scalar.from_value(val, is_cscalar=False, name="") # pragma: is_grbscalar + # fmt: on + 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.__new__(_BoundIndexBinaryOp) + rv.parent = self.parent + rv.name = self.name + rv.type = self.type + rv.return_type = self.return_type + rv.gb_obj = new_binop[0] + rv.gb_name = f"{self.name}_bound" + rv._type2 = self._type2 + return rv + + @property + def thunk_type(self): + return self.type if self._type2 is None else self._type2 + + +class TypedUserIndexBinaryOp(TypedOpBase): + __slots__ = () + opclass = "IndexBinaryOp" + + 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) + + @property + def orig_func(self): + return self.parent.orig_func + + @property + def _numba_func(self): + return self.parent._numba_func + + thunk_type = TypedBuiltinIndexBinaryOp.thunk_type + + 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) + + def __reduce__(self): + name = f"indexbinary.{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 IndexBinaryOp.register_anonymous(func, name, parameterized=True) + if (rv := IndexBinaryOp._find(name)) is not None: + return rv + return IndexBinaryOp.register_new(name, func, parameterized=True) + + +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/column indices of ``x`` and ``iy, jy`` are + the row/column indices of ``y``. + + An IndexBinaryOp can be converted to a BinaryOp by binding a theta value, + which makes it usable in any operation that accepts a BinaryOp (eWiseMult, + eWiseAdd, mxm, etc.). + + IndexBinaryOps are located in the ``graphblas.indexbinary`` namespace. + + There are no built-in IndexBinaryOps; all are user-defined. + """ + + __slots__ = "orig_func", "_is_udt", "_numba_func" + _module = indexbinary + _modname = "indexbinary" + _custom_dtype = 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 + indexbinary_udf = numba.njit(func) + 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 + + # Numba is unable to handle BOOL correctly right now + input_type = INT8 if type_ == BOOL else type_ + return_type = INT8 if ret_type == BOOL else ret_type + + # Build wrapper: 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)(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): + if not _has_idxbinop: + raise RuntimeError( + "IndexBinaryOp requires SuiteSparse:GraphBLAS 9.4+ " + "(python-suitesparse-graphblas 9.3.1+)" + ) + if dtype2 is None: + dtype2 = dtype + dtypes = (dtype, dtype2) + if dtypes in self._udt_types: + return self._udt_ops[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, + dtype2.numba_type, + ) + numba_func.compile(sig) + ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) + indexbinary_wrapper, wrapper_sig = _get_udt_wrapper_indexbinary( + numba_func, ret_type, dtype, dtype2 + ) + + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(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, + dtype2._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, + ) + 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): + """Register an IndexBinaryOp without registering it in 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 able + to be compiled with ``numba.njit``. + ``func`` takes seven input parameters--``(x, ix, jx, y, iy, jy, theta)``-- + where ``x`` and ``y`` are element values, ``ix, jx`` and ``iy, jy`` + are their row/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, which means + 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 and save it to ``graphblas.indexbinary`` namespace. + + Parameters + ---------- + name : str + The name of the operator. This will show up as ``gb.indexbinary.{name}``. + func : FunctionType + The function to compile. For all current backends, this must be able + to be compiled with ``numba.njit``. + ``func`` takes seven input parameters--``(x, ix, jx, y, iy, jy, theta)``-- + where ``x`` and ``y`` are element values, ``ix, jx`` and ``iy, jy`` + are their row/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 + If True, delay compilation until the operator is 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}, + ) + 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): DataType} + self._udt_ops = {} # {(dtype, dtype2): TypedUserIndexBinaryOp} + + 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"indexbinary.{self.name}") in _STANDARD_OPERATOR_NAMES: + return name + return (self._deserialize, (self.name, self.orig_func)) + + def __call__(self, theta=None, dtype=None): + """Bind a theta value to create a BinaryOp. + + Parameters + ---------- + theta : scalar, optional + The theta parameter to bind. Default is 0 (False). + dtype : dtype, optional + The dtype to use. If not provided, it will be inferred from theta. + + Returns + ------- + TypedOpBase + A BinaryOp created 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 not isinstance(theta, Scalar): + theta = Scalar.from_value(theta, is_cscalar=False, name="") # pragma: is_grbscalar + elif theta._is_cscalar: + # fmt: off + val = theta.value + theta = Scalar.from_value(val, is_cscalar=False, name="") # pragma: is_grbscalar + # fmt: on + if dtype is None: + dtype = theta.dtype + else: + dtype = _lookup_dtype(dtype) + typed_op = self[dtype] + return typed_op(theta) + + +def _get_udt_wrapper_indexbinary(numba_func, return_type, dtype, dtype2): + """Build a wrapper function for UDT IndexBinaryOp: z = f(x, ix, jx, y, iy, jy, theta).""" + nt = numba.types + ztype = INT8 if return_type == BOOL else return_type + xtype = INT8 if dtype == BOOL else dtype + ytype = INT8 if dtype2 == BOOL else dtype2 + + wrapper_sig = nt.void( + nt.CPointer(ztype.numba_type), + nt.CPointer(xtype.numba_type), + UINT64.numba_type, + UINT64.numba_type, + nt.CPointer(ytype.numba_type), + UINT64.numba_type, + UINT64.numba_type, + nt.CPointer(ytype.numba_type), + ) + + d = {"numba": numba, "numba_func": numba_func} + xderef = "bool(x_ptr[0])" if dtype == BOOL else "x_ptr[0]" + yderef = "bool(y_ptr[0])" if dtype2 == BOOL else "y_ptr[0]" + tderef = "bool(t_ptr[0])" if dtype2 == BOOL else "t_ptr[0]" + call = f"numba_func({xderef}, ix, jx, {yderef}, iy, jy, {tderef})" + if return_type == BOOL: + call = f"bool({call})" + text = f"def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):\n z_ptr[0] = {call}\n" + exec(text, d) # noqa: S102 + return d["wrapper"], wrapper_sig diff --git a/graphblas/core/operator/utils.py b/graphblas/core/operator/utils.py index 1442a9b5e..edd06b878 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -1,6 +1,17 @@ from types import BuiltinFunctionType, FunctionType, ModuleType -from ... import backend, binary, config, indexunary, monoid, op, select, semiring, unary +from ... import ( + backend, + binary, + config, + indexbinary, + indexunary, + monoid, + op, + select, + semiring, + unary, +) from ...dtypes import UINT64, lookup_dtype, unify from ..expr import InfixExprBase from .base import ( @@ -14,6 +25,7 @@ find_opclass, ) from .binary import BinaryOp +from .indexbinary import IndexBinaryOp from .indexunary import IndexUnaryOp from .monoid import Monoid from .select import SelectOp @@ -261,6 +273,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..c8148e8e2 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -136,6 +136,8 @@ def __eq__(self, other): """ return self.isequal(other) + __hash__ = None + def __ne__(self, other): return not self.isequal(other) @@ -1001,6 +1003,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 +1078,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 diff --git a/graphblas/core/ss/__init__.py b/graphblas/core/ss/__init__.py index 10a6fed94..79859eb53 100644 --- a/graphblas/core/ss/__init__.py +++ b/graphblas/core/ss/__init__.py @@ -1,5 +1,5 @@ import suitesparse_graphblas as _ssgb -(version_major, version_minor, version_bug) = map(int, _ssgb.__version__.split(".")[:3]) +version_major, version_minor, version_bug = map(int, _ssgb.__version__.split(".")[:3]) _IS_SSGB7 = version_major == 7 diff --git a/graphblas/core/ss/context.py b/graphblas/core/ss/context.py index f93d1ec1c..67c7a6c20 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, @@ -37,13 +40,14 @@ def __init__(self, engage=True, *, stack=True, nthreads=None, chunk=None, gpu_id 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: @@ -66,7 +70,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..20fc126cf 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": { diff --git a/graphblas/core/ss/indexbinary.py b/graphblas/core/ss/indexbinary.py new file mode 100644 index 000000000..89a38b3fc --- /dev/null +++ b/graphblas/core/ss/indexbinary.py @@ -0,0 +1,134 @@ +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, + TypedUserIndexBinaryOp, + _has_idxbinop, +) + +ffi_new = ffi.new + + +class TypedJitIndexBinaryOp(TypedOpBase): + __slots__ = "_jit_c_definition" + opclass = "IndexBinaryOp" + + def __init__(self, parent, name, type_, return_type, gb_obj, jit_c_definition, dtype2=None): + super().__init__(parent, name, type_, return_type, gb_obj, name, dtype2=dtype2) + self._jit_c_definition = jit_c_definition + + @property + def jit_c_definition(self): + return self._jit_c_definition + + thunk_type = TypedUserIndexBinaryOp.thunk_type + + 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. + + 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) ; }"`` + 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, theta_type) in rv.types + or rv._udt_types is not None + and (x_type, theta_type) in rv._udt_types + ): + raise TypeError( + f"IndexBinaryOp gb.indexbinary.{name} already defined for " + f"({x_type}, {theta_type}) input types" + ) + else: + rv = IndexBinaryOp(name, is_udt=True) + gb_obj = ffi_new("GxB_IndexBinaryOp*") + 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[]", funcname.encode()), + ffi_new("char[]", jit_c_definition.encode()), + ), + "IndexBinaryOp", + gb_obj[0], + ) + op = TypedJitIndexBinaryOp( + rv, funcname, x_type, ret_type, gb_obj[0], jit_c_definition, dtype2=theta_type + ) + rv._add(op, is_jit=True) + setattr(module, funcname, rv) + return rv diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index de6395cf6..e45c7014b 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -376,7 +376,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): diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index 4acef3fe5..20b9ce297 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -262,7 +262,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): diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 8bac4198e..8c73ecc48 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -528,7 +528,7 @@ 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( diff --git a/graphblas/indexbinary/__init__.py b/graphblas/indexbinary/__init__.py new file mode 100644 index 000000000..34f33fd4c --- /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/io/_awkward.py b/graphblas/io/_awkward.py index b30984251..fd4a9bab8 100644 --- a/graphblas/io/_awkward.py +++ b/graphblas/io/_awkward.py @@ -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..eff8c9031 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 @@ -40,9 +42,25 @@ 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() + 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( @@ -104,9 +122,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/monoid/numpy.py b/graphblas/monoid/numpy.py index b9ff2b502..aae8385df 100644 --- a/graphblas/monoid/numpy.py +++ b/graphblas/monoid/numpy.py @@ -95,7 +95,7 @@ ): # 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 + # The behavioral check above is kept as a safety net; it's correct for all numba/numpy combos _monoid_identities["fmax"].update( { "BOOL": False, diff --git a/graphblas/ss/_core.py b/graphblas/ss/_core.py index b42ea72b4..9c1d043bc 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -215,10 +215,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, diff --git a/graphblas/tests/test_core.py b/graphblas/tests/test_core.py index 3586eb4a8..a15da6f54 100644 --- a/graphblas/tests/test_core.py +++ b/graphblas/tests/test_core.py @@ -67,7 +67,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__") diff --git a/graphblas/tests/test_formatting.py b/graphblas/tests/test_formatting.py index faadc983b..720ab79e5 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) diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py new file mode 100644 index 000000000..80aa8d917 --- /dev/null +++ b/graphblas/tests/test_indexbinary.py @@ -0,0 +1,288 @@ +import pickle + +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 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 op is not None + 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 result is not None + assert hasattr(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 is not None + 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_index_aware(): + """Test that indices are correctly passed to the 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(): + """Test that theta=0 works correctly.""" + + 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(): + """Test passing a graphblas Scalar as theta.""" + + 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 + + +def test_pickle_registered(): + def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + indexbinary.register_new("pickle_test_op", add_theta) + 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 + + delattr(indexbinary, "pickle_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_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_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" + + +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..a1a59fd69 100644 --- a/graphblas/tests/test_infix.py +++ b/graphblas/tests/test_infix.py @@ -454,34 +454,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 @@ -564,34 +564,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 +716,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 diff --git a/graphblas/tests/test_io.py b/graphblas/tests/test_io.py index 7e786f0da..9443bbcfa 100644 --- a/graphblas/tests/test_io.py +++ b/graphblas/tests/test_io.py @@ -272,7 +272,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 diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index b972ef260..b5d00337d 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -3429,7 +3429,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] @@ -3848,11 +3848,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) diff --git a/graphblas/tests/test_numpyops.py b/graphblas/tests/test_numpyops.py index 999c6d5e0..534312c0b 100644 --- a/graphblas/tests/test_numpyops.py +++ b/graphblas/tests/test_numpyops.py @@ -40,7 +40,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,7 +48,7 @@ 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)) @@ -116,7 +116,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 +201,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..8a072a56b 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -869,13 +869,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 diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index e93511914..dbfb301a3 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -292,7 +292,7 @@ def test_neg(): ( 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 diff --git a/graphblas/tests/test_ss_utils.py b/graphblas/tests/test_ss_utils.py index 2df7ab939..40774186c 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -256,8 +256,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: diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 4cea0b563..ed6f429e7 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1,7 +1,7 @@ import os import pathlib -import platform -import sys +import re +import subprocess import sysconfig import numpy as np @@ -9,8 +9,9 @@ from numpy.testing import assert_array_equal import graphblas as gb -from graphblas import backend, binary, dtypes, indexunary, select, unary +from graphblas import backend, binary, dtypes, indexbinary, indexunary, select, unary from graphblas.core import _supports_udfs as supports_udfs +from graphblas.core.operator.indexbinary import _has_idxbinop from graphblas.core.ss import _IS_SSGB7 from .conftest import autocompute, burble @@ -26,108 +27,122 @@ pytest.skip("not suitesparse backend", allow_module_level=True) +def _fix_jit_config(): + """Fix the GraphBLAS JIT configuration for the current conda environment. + + The graphblas C library bakes in build-time compiler paths from conda-build, + which don't exist in the user's environment. This function: + 1. Replaces the compiler path with the equivalent from $CONDA_PREFIX/bin/ + 2. Replaces -isysroot with the local macOS SDK (via xcrun), or strips it on Linux + 3. Strips -fdebug-prefix-map flags referencing build paths + + Returns + ------- + True: JIT configured and verified working + False: JIT configuration attempted but compilation failed (don't retry) + None: No conda environment; caller should try a different approach + + Only modifies jit_c_compiler_name, jit_c_compiler_flags, and jit_c_control. + Linker flags and libraries are left at their defaults (already have correct + $CONDA_PREFIX paths substituted by the graphblas C library). + """ + conda_prefix = os.environ.get("CONDA_PREFIX", "") + if not conda_prefix: + return None # No conda env; caller should try sysconfig instead + + # Check if the default compiler already works (e.g., /usr/bin/cc on macOS). + # Only replace it if the baked-in path doesn't exist. + jit_cc = gb.ss.config["jit_c_compiler_name"] + if pathlib.Path(jit_cc).exists(): + # Default compiler exists — don't replace it, just fix flags + pass + else: + # Replace build-time path with local conda equivalent. + cc_basename = pathlib.Path(jit_cc).name + bin_dir = pathlib.Path(conda_prefix) / "bin" + for candidate in [cc_basename, "cc", "clang", "gcc"]: + local_cc = bin_dir / candidate + if local_cc.exists(): + break + else: + return False + gb.ss.config["jit_c_compiler_name"] = str(local_cc) + + # Fix compiler flags: fix build-time-only paths that don't exist in the + # user's environment + flags = gb.ss.config["jit_c_compiler_flags"] + # -isysroot : macOS SDK path from conda-build (e.g., /opt/conda-sdks/MacOSX10.13.sdk). + # The conda cross-compiler needs an explicit sysroot. Replace with the local SDK if available. + 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 available (linux, or macOS without Xcode CLT) + flags = re.sub(r"-isysroot\s+\S+", "", flags) + # -fdebug-prefix-map==: debug path remapping from conda-build + flags = re.sub(r"-fdebug-prefix-map=\S+", "", flags) + gb.ss.config["jit_c_compiler_flags"] = flags + + gb.ss.config["jit_c_control"] = "on" + + # Verify the JIT actually works by attempting a trivial compilation. + # If it fails (e.g., missing libraries, wrong flags), turn JIT off. + try: + from graphblas import dtypes as _dtypes + + _dtypes.ss.register_new("_jit_probe", "typedef struct { int _probe ; } _jit_probe ;") + except Exception: + gb.ss.config["jit_c_control"] = "off" + return False + else: + return True + + @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, JIT is broken, turn off. + - 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() + if result is True: + pass # Conda JIT configured and verified + elif result is False: + # Probe failed — JIT doesn't work with this psg build. + # Don't try sysconfig; if the conda compiler can't compile + # GraphBLAS JIT kernels, Python's sysconfig compiler won't either. + gb.ss.config["jit_c_control"] = "off" + else: + # No conda env (result is None). 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 + if libs: + 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 ( - 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" - yield - return - - # 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 try: yield finally: @@ -148,7 +163,7 @@ def test_jit_udt(): ) return if gb.ss.config["jit_c_control"] == "off": - return + pytest.skip("JIT not available (no C compiler configured)") with burble(): dtype = dtypes.ss.register_new( "myquaternion", "typedef struct { float x [4][4] ; int color ; } myquaternion ;" @@ -196,7 +211,7 @@ def test_jit_unary(v): unary.ss.register_new("square", cdef, "FP32", "FP32") return if gb.ss.config["jit_c_control"] == "off": - return + pytest.skip("JIT not available (no C compiler configured)") with burble(): square = unary.ss.register_new("square", cdef, "FP32", "FP32") assert not hasattr(unary, "square") @@ -234,7 +249,7 @@ def test_jit_binary(v): binary.ss.register_new("absdiff", cdef, "FP64", "FP64", "FP64") return if gb.ss.config["jit_c_control"] == "off": - return + pytest.skip("JIT not available (no C compiler configured)") with burble(): absdiff = binary.ss.register_new( "absdiff", @@ -302,7 +317,7 @@ def test_jit_indexunary(v): indexunary.ss.register_new("diffy", cdef, "FP64", "FP64", "FP64") return if gb.ss.config["jit_c_control"] == "off": - return + pytest.skip("JIT not available (no C compiler configured)") with burble(): diffy = indexunary.ss.register_new("diffy", cdef, "FP64", "FP64", "FP64") assert not hasattr(indexunary, "diffy") @@ -356,6 +371,69 @@ def test_jit_indexunary(v): assert ("FP32", "FP64") not in diffy +@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) ; }" + ) + if gb.ss.config["jit_c_control"] == "off": + pytest.skip("JIT not available (no C compiler configured)") + 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} + assert "FP64" in add_theta + assert add_theta["FP64"].return_type == dtypes.FP64 + assert add_theta["FP64"].jit_c_definition == cdef + # 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) + # 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="IndexBinaryOp gb.indexbinary.ss.add_theta already defined for .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) + (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", "FP32") in add_theta + assert ("FP32", "FP64") not in add_theta + + def test_jit_select(v): cdef = ( # Why does this one insist on `const` for `x` argument? @@ -367,7 +445,7 @@ def test_jit_select(v): select.ss.register_new("woot", cdef, "INT32", "INT32") return if gb.ss.config["jit_c_control"] == "off": - return + pytest.skip("JIT not available (no C compiler configured)") with burble(): woot = select.ss.register_new("woot", cdef, "INT32", "INT32") assert not hasattr(select, "woot") diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index d9bf84495..65460fa75 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -27,7 +27,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 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..ef2d11b64 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,15 @@ 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 :: Only", "Intended Audience :: Developers", "Intended Audience :: Other Audience", @@ -55,14 +55,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.15'", # make optional where numba is not supported ] [project.urls] @@ -72,9 +72,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 +83,10 @@ 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.15'", # 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 >=0.14; python_version<'3.13'", # make optional, b/c sparse needs numba + # 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 @@ -102,11 +102,11 @@ test = [ ] 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.15'", # 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.15'", # make optional, b/c datashade needs numba ] [tool.setuptools] @@ -122,6 +122,7 @@ packages = [ "graphblas.core.operator", "graphblas.core.ss", "graphblas.dtypes", + "graphblas.indexbinary", "graphblas.indexunary", "graphblas.io", "graphblas.monoid", @@ -140,7 +141,7 @@ dirty_template = "{tag}+{ccount}.g{sha}.dirty" [tool.black] line-length = 100 -target-version = ["py310", "py311", "py312", "py313"] +target-version = ["py311", "py312", "py313", "py314"] [tool.isort] sections = ["FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER"] @@ -154,7 +155,7 @@ 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) @@ -167,34 +168,19 @@ filterwarnings = [ # 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:", ] [tool.coverage.run] @@ -222,7 +208,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 @@ -313,6 +299,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", # Pattern passed to `match=` contains metacharacters but is neither escaped nor raw "PERF401", # Use a list comprehension to create a transformed list (Note: poorly implemented atm) # Intentionally ignored @@ -328,6 +315,7 @@ ignore = [ "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 "PLR0124", # Name compared with itself, consider replacing `x == x` (Note: too strict) "PLR0911", # Too many return statements "PLR0912", # Too many branches @@ -340,6 +328,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", # This suppression comment is invalid because... (Note: `fmt` is 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,7 +337,6 @@ 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) @@ -374,6 +362,7 @@ 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/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,6 +371,9 @@ ignore = [ "NPY002", # numba doesn't support rng generator yet ] "graphblas/core/utils.py" = ["PLE0302"] # `__set__` is used as a property +"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" = [ @@ -398,7 +390,7 @@ ignore = [ "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 +403,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 +415,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/check_versions.sh b/scripts/check_versions.sh index 5aa88e045..afcf3e0ac 100755 --- a/scripts/check_versions.sh +++ b/scripts/check_versions.sh @@ -3,15 +3,19 @@ # 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. +conda search 'flake8-bugbear[channel=conda-forge]>=25.11.29' +conda search 'flake8-simplify[channel=conda-forge]>=0.30.0' +conda search 'numpy[channel=conda-forge]>=2.4' +conda search 'pandas[channel=conda-forge]>=3.0' +conda search 'scipy[channel=conda-forge]>=1.17' +conda search 'networkx[channel=conda-forge]>=3.6' +conda search 'awkward[channel=conda-forge]>=2.9' +conda search 'sparse[channel=conda-forge]>=0.15' +# fast_matrix_market is deprecated (no longer maintained; last supported Python is 3.12) +conda search 'numba[channel=conda-forge]>=0.64' +conda search 'pyyaml[channel=conda-forge]>=6.0' +conda search 'python-suitesparse-graphblas[channel=conda-forge]>=10.3.1' +# 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..349791396 --- /dev/null +++ b/scripts/ci_pick_versions.py @@ -0,0 +1,532 @@ +#!/usr/bin/env python +"""Pick random dependency versions for CI testing. + +Randomly selects compatible dependency versions for python-graphblas CI. +Replaces the bash-based version selection 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 +# --------------------------------------------------------------------------- + +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", ""], + "3.13": ["2.1", "2.2", "2.3", "2.4", ""], + "3.14": ["2.3", "2.4", ""], +} + +# 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", ""], +} + +PANDAS_VERSIONS = { + "1.x": { + "3.11": ["1.5", "2.0", "2.1", "2.2", "2.3", ""], + "3.12": ["2.1", "2.2", "2.3", ""], + }, + "2.x": ["2.2", "2.3", "3.0", ""], +} + +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", ""], +} + +NUMBA_VERSIONS = { + "1.x": ["0.57", "0.58", "0.59", "0.60", "0.61", ""], + "2.x": ["0.62", "0.63", "0.64", ""], +} + +# Deps that only depend on Python version (not numpy) +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.13": ["3.4", "3.5", "3.6", ""], + "3.14": ["3.6", ""], +} + +PYYAML_VERSIONS = { + "3.11": ["5.4", "6.0", ""], + "3.12": ["6.0", ""], + "3.13": ["6.0", ""], + "3.14": ["6.0", ""], +} + +SPARSE_VERSIONS = { + "3.11": ["0.14", "0.15", ""], + "3.12": ["0.14", "0.15", ""], + "3.13": "NA", + "3.14": "NA", +} + +# PSG versions by numpy branch and source type. +# conda-forge uses "=" prefix, wheel/source use "==". +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"], + }, + "wheel": { + "3.11": ["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"], + }, + "source": { + "3.11": [ + "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", + ], + "3.12": ["8.2.0.1", "8.2.1.0"], + }, +} + +PSG_VERSIONS_NP2 = { + "no_py314": ["9.3.1.0", "9.4.5.0", "10.0.1.1", "10.1.1.0", "10.3.1.0"], + "py314_only": ["10.0.1.1", "10.1.1.0", "10.3.1.0"], +} + + +# --------------------------------------------------------------------------- +# 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(".")) + + +# --------------------------------------------------------------------------- +# 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. + """ + # --- 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 <1.13 doesn't support numpy 2.x (safety net) + np_is_1x = v["numpy"].startswith("1.") if v["numpy"] else False + if not np_is_1x and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 13): + v["scipy"] = random.choice([s for s in scipy_pool if _ver(s) >= (1, 13)]) + + # scipy <1.15.1 requires numpy <2.3; scipy <1.16 requires numpy <2.5 + if _ver(v["numpy"]) >= (2, 3): + if v["scipy"] in ("1.13", "1.14"): + v["scipy"] = random.choice(["1.16", "1.17", ""]) + elif v["scipy"] == "1.15": + v["scipy"] = random.choice(["1.15", "1.16", "1.17", ""]) + + # numpy 1.26 + scipy 1.9 conflict + if v["numpy"] == "1.26" and v["scipy"] == "1.9": + v["scipy"] = random.choice(["1.10", "1.11", ""]) + + # --- scipy / Python version availability --- + + # scipy <1.14 has no py3.13 builds; scipy <1.16 has no py3.14 builds + if pyver == "3.14" and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 16): + v["scipy"] = random.choice(["1.16", "1.17", ""]) + elif pyver == "3.13" and v["scipy"] == "1.13": + v["scipy"] = random.choice(["1.14", "1.15", "1.16", "1.17", ""]) + + # --- pandas constraints --- + + # pandas <2.3 has no py3.14 builds + if pyver == "3.14" and v["pandas"] == "2.2": + v["pandas"] = random.choice(["2.3", "3.0", ""]) + + # pandas 3.0 requires numba >=0.60 and scipy >=1.14.1 + if v["pandas"] == "3.0": + if v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 60): + v["numba"] = "0.60" + if v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 15): + v["scipy"] = random.choice(["1.15", "1.16", "1.17", ""]) + + # --- awkward / Python version availability --- + + # awkward <2.7 has no py3.13 builds; awkward <2.8 has no py3.14 builds + if pyver == "3.14" and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 8): + v["awkward"] = random.choice(["2.8", "2.9", ""]) + elif pyver == "3.13" and v["awkward"] == "2.6": + v["awkward"] = random.choice(["2.7", "2.8", "2.9", ""]) + + # --- numba constraints --- + + # numba minimum by Python version: 0.59 for 3.12, 0.61 for 3.13, 0.63 for 3.14 + numba_min = {"3.11": (0, 57), "3.12": (0, 59), "3.13": (0, 61), "3.14": (0, 63)} + if v["numba"] not in ("", "NA"): + min_ver = numba_min[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 <0.64 requires numpy <2.4 + if _ver(v["numpy"]) >= (2, 4) and v["numba"] in ("0.62", "0.63"): + v["numba"] = "0.64" + + # numba <0.62 doesn't support numpy 2.x + if not np_is_1x and v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 62): + v["numba"] = "NA" + + # --- sparse --- + + # sparse doesn't support Python 3.13+ + if pyver in ("3.13", "3.14"): + v["sparse"] = "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": _pick_scalar_or_list(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_scalar_or_list(pool): + """Handle pools that are either "NA" (string) or a list of choices.""" + return pool if isinstance(pool, str) else random.choice(pool) + + +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["py314_only"] if pyver == "3.14" else PSG_VERSIONS_NP2["no_py314"] + 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", +} + + +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. + 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 in ("NA", ""): + lines.append(f"{var}='{val}'") + else: + lines.append(f"{var}='={val}'") + + # psg already has = or == prefix + lines.append(f"psgver='{v['psg']}'") + 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", + ): + name = _SUMMARY_NAMES[key] + val = v[key] + if val == "NA": + parts.append(f"{name}=NA") + elif val == "": + parts.append(f"{name}=latest") + else: + parts.append(f"{name}={val}") + return "versions: " + " ".join(parts) + + +# --------------------------------------------------------------------------- +# Validation (for testing the script itself) +# --------------------------------------------------------------------------- + + +def validate(v, pyver): + """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 <1.13 requires numpy 1.x + if not np_is_1x and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 13): + errors.append(f"scipy {v['scipy']} doesn't support numpy 2.x") + + # scipy <1.15.1 requires numpy <2.3 + if _ver(v["numpy"]) >= (2, 3) and v["scipy"] in ("1.13", "1.14"): + errors.append(f"scipy {v['scipy']} requires numpy <2.3, got {v['numpy']}") + + # numpy 1.26 + scipy 1.9 + if v["numpy"] == "1.26" and v["scipy"] == "1.9": + errors.append("numpy 1.26 + scipy 1.9 conflict") + + # scipy Python availability + if pyver == "3.14" and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 16): + errors.append(f"scipy {v['scipy']} has no py3.14 build") + if pyver == "3.13" and v["scipy"] == "1.13": + errors.append("scipy 1.13 has no py3.13 build") + + # pandas Python availability + if pyver == "3.14" and v["pandas"] == "2.2": + errors.append("pandas 2.2 has no py3.14 build") + + # pandas 3.0 requirements + if v["pandas"] == "3.0": + if v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 60): + errors.append(f"pandas 3.0 requires numba >=0.60, got {v['numba']}") + if 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']}") + + # awkward Python availability + if pyver == "3.14" and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 8): + errors.append(f"awkward {v['awkward']} has no py3.14 build") + if pyver == "3.13" and v["awkward"] == "2.6": + errors.append("awkward 2.6 has no py3.13 build") + + # numba Python minimums + numba_min = {"3.11": (0, 57), "3.12": (0, 59), "3.13": (0, 61), "3.14": (0, 63)} + if v["numba"] not in ("", "NA") and _ver(v["numba"]) < numba_min[pyver]: + errors.append(f"numba {v['numba']} doesn't support Python {pyver}") + + # numba <0.64 requires numpy <2.4 + if v["numba"] in ("0.62", "0.63") and _ver(v["numpy"]) >= (2, 4): + errors.append(f"numba {v['numba']} requires numpy <2.4, got {v['numpy']}") + + # numba <0.62 requires numpy 1.x + if not np_is_1x and v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 62): + errors.append(f"numba {v['numba']} doesn't support numpy 2.x") + + # sparse Python availability + if pyver in ("3.13", "3.14") and v["sparse"] != "NA": + errors.append(f"sparse doesn't support Python {pyver}") + + 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) + 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 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/jit_diagnostics.py b/scripts/jit_diagnostics.py new file mode 100755 index 000000000..b1166171f --- /dev/null +++ b/scripts/jit_diagnostics.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python +"""Print JIT diagnostic information for debugging compiler configuration. + +Run with: python scripts/jit_diagnostics.py +Or during CI: python -m scripts.jit_diagnostics (from repo root) + +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 defaults (from compiled C library) ---") + 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 _fix_jit_config ---") + from graphblas.tests.test_ssjit import _fix_jit_config + + result = _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()) From 6f1eb02439b033b36f1474340327eaf7ecd1b6e7 Mon Sep 17 00:00:00 2001 From: Jim Kitchen <2807270+jim22k@users.noreply.github.com> Date: Thu, 23 Jul 2026 08:49:03 -0500 Subject: [PATCH 04/52] Adds IndexBinaryOp and fixes for JIT (#576) * Adds IndexBinaryOp and fixes for JIT * updates --- .github/workflows/test_and_build.yml | 21 +- .pre-commit-config.yaml | 42 +- dev-requirements.txt | 5 +- docs/_static/custom.css | 8 +- docs/conf.py | 2 +- docs/env.yml | 8 +- docs/getting_started/index.rst | 8 +- docs/index.rst | 81 +- docs/user_guide/fundamentals.rst | 4 +- docs/user_guide/index.rst | 1 + docs/user_guide/io.rst | 2 +- docs/user_guide/operations.rst | 6 +- docs/user_guide/operators.rst | 173 ++-- docs/user_guide/udf.rst | 193 +++- docs/user_guide/udt.rst | 278 +++++ environment.yml | 6 +- graphblas/core/dtypes.py | 230 ++++- graphblas/core/infixmethods.py | 12 +- graphblas/core/operator/agg.py | 63 +- graphblas/core/operator/base.py | 533 +++++++++- graphblas/core/operator/binary.py | 396 ++++--- graphblas/core/operator/indexbinary.py | 269 +++-- graphblas/core/operator/indexunary.py | 89 +- graphblas/core/operator/monoid.py | 147 ++- graphblas/core/operator/select.py | 82 +- graphblas/core/operator/semiring.py | 93 +- graphblas/core/operator/udt_utils.py | 994 ++++++++++++++++++ graphblas/core/operator/unary.py | 104 +- graphblas/core/operator/utils.py | 9 + graphblas/core/scalar.py | 11 +- graphblas/core/ss/_utils.py | 47 + graphblas/core/ss/jit_config.py | 293 ++++++ graphblas/core/ss/matrix.py | 30 +- graphblas/core/ss/vector.py | 32 +- graphblas/core/utils.py | 8 + graphblas/exceptions.py | 28 +- graphblas/indexbinary/__init__.py | 4 +- graphblas/monoid/numpy.py | 6 +- graphblas/ss/__init__.py | 17 +- graphblas/ss/_core.py | 6 +- graphblas/tests/conftest.py | 58 +- graphblas/tests/test_dtype.py | 193 ++++ graphblas/tests/test_formatting.py | 11 + graphblas/tests/test_indexbinary.py | 443 +++++++- graphblas/tests/test_matrix.py | 4 +- graphblas/tests/test_op.py | 1303 ++++++++++++++++++++++++ graphblas/tests/test_pickle.py | 178 ++++ graphblas/tests/test_scalar.py | 33 +- graphblas/tests/test_ssjit.py | 1107 ++++++++++++++++++-- graphblas/tests/test_vector.py | 2 +- pyproject.toml | 60 +- scripts/ci_pick_versions.py | 26 +- scripts/find_aiisms.sh | 103 ++ scripts/jit_diagnostics.py | 9 +- 54 files changed, 6886 insertions(+), 985 deletions(-) create mode 100644 docs/user_guide/udt.rst create mode 100644 graphblas/core/operator/udt_utils.py create mode 100644 graphblas/core/ss/_utils.py create mode 100644 graphblas/core/ss/jit_config.py create mode 100755 scripts/find_aiisms.sh diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index d31e8c534..8e435cd4a 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -169,9 +169,10 @@ jobs: pyver=$(python -c "v='${{ steps.pyver.outputs.selected }}'.split('.');print(f'{v[0]}.{v[1]}')") eval "$(python scripts/ci_pick_versions.py --python "$pyver" --source ${{ steps.sourcetype.outputs.selected }})" - # Set up psg conda package variable (conda-forge installs via conda, others via pip) + # 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}} == "conda-forge" && -n "$psgver" ]] ; then + if [[ ${{ steps.sourcetype.outputs.selected}} == "conda-forge" ]] ; then psg=python-suitesparse-graphblas${psgver} fi @@ -214,7 +215,7 @@ jobs: set -x # echo on $(command -v mamba || command -v conda) install -c nodefaults \ - packaging pytest coverage pytest-randomly cffi donfig tomli c-compiler make \ + 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} ${psg} \ ${{ matrix.slowtask == 'pytest_bizarro' && 'black' || '' }} \ @@ -226,24 +227,26 @@ jobs: # ${{ matrix.os != 'windows-latest' && 'pytest-forked' || '' }} # to investigate crashes - name: Build extension module run: | + # ``python -m pip`` ensures the conda env's interpreter drives the + # install, not a system ``pip`` on PATH. if [[ ${{ steps.sourcetype.outputs.selected }} == "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" + # python -m 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 # 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 + # 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 [[ ${{ 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 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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 20b632c30..ac60ef29a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -24,7 +24,7 @@ repos: - id: check-illegal-windows-names - id: check-merge-conflict - id: check-ast - - id: check-json # No json files yet + - id: check-json # no JSON files yet; enabled for the future - id: check-toml - id: check-yaml - id: check-executables-have-shebangs @@ -53,6 +53,9 @@ repos: # 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 + # 9.0 is still in alpha (9.0.0a3) and a pre-release pin defeats + # pre-commit.ci's autoupdate "avoid prereleases" logic; stay on the last + # stable 8.x until 9.0 GAs. rev: 8.0.1 hooks: - id: isort @@ -73,7 +76,7 @@ repos: - id: black - id: black-jupyter - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.7 + rev: v0.15.13 hooks: - id: ruff-check args: [--fix-only, --show-fixes] @@ -98,10 +101,10 @@ repos: - tomli; python_version<'3.11' files: ^(graphblas|docs)/ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.7 + rev: v0.15.13 hooks: - id: ruff-check - # - id: ruff-format # Prefer black, but may temporarily uncomment this to see + # - id: ruff-format # we prefer black; uncomment to compare with ruff's formatter - repo: https://github.com/sphinx-contrib/sphinx-lint rev: v1.0.2 hooks: @@ -119,7 +122,7 @@ repos: hooks: - id: shellcheck - repo: https://github.com/rbubley/mirrors-prettier - rev: v3.8.1 + rev: v3.8.3 hooks: - id: prettier args: [--prose-wrap=preserve] @@ -138,12 +141,13 @@ repos: rev: v0.9.3 hooks: - id: taplo-format + args: ["--option", "column_width=100"] - repo: https://github.com/rhysd/actionlint - rev: v1.7.11 + rev: v1.7.12 hooks: - id: actionlint - repo: https://github.com/python-jsonschema/check-jsonschema - rev: 0.37.0 + rev: 0.37.2 hooks: - id: check-dependabot - id: check-github-workflows @@ -156,6 +160,11 @@ repos: # rev: v1.14.2 # hooks: # - id: zizmor + - repo: https://github.com/pysentry/pysentry-pre-commit + rev: v0.4.6 + hooks: + - id: pysentry + args: ["--format=human", "--fail-on=low", "--sources=pypa,osv,pypi"] - repo: meta hooks: - id: check-hooks-apply @@ -176,6 +185,25 @@ repos: 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 + )$ # 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 diff --git a/dev-requirements.txt b/dev-requirements.txt index c6b586369..a986ece58 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -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/conf.py b/docs/conf.py index da8d0dc67..b6b269114 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -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 caec812d6..e47eea420 100644 --- a/docs/env.yml +++ b/docs/env.yml @@ -3,7 +3,7 @@ channels: - conda-forge - nodefaults dependencies: - - python=3.11 + - 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 0c02b6194..cbd0bb6a6 100644 --- a/docs/getting_started/index.rst +++ b/docs/getting_started/index.rst @@ -30,10 +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 + - `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 `_ 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, 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..6a66c295c 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``. @@ -187,26 +187,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 +221,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``. + +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[float](0.5) # bind theta + C << A.ewise_mult(B, binop) + + # For mxm, wrap in a Semiring; the additive monoid must accept the + # bound op's return type. + 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 ----------- diff --git a/docs/user_guide/udf.rst b/docs/user_guide/udf.rst index e7b984b44..670cbc6e3 100644 --- a/docs/user_guide/udf.rst +++ b/docs/user_guide/udf.rst @@ -1,55 +1,190 @@ -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. + +Lambdas are auto-registered as anonymous wherever a UnaryOp is expected: + +.. code-block:: python -.. csv-table:: w - :class: matrix - :header: 0,1,2,3,4,5 + v.apply(lambda x: x % 5 - 2).new() - 1,3,,3,9,15 +Operator signatures +------------------- +The function signature depends on the operator type: -Similar methods exist for BinaryOp and IndexUnaryOp. User-defined Monoids and Semirings are -constructed out of previously defined and built-in UnaryOps and BinaryOps. +================== ==================================================== ===================================== +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) +================== ==================================================== ===================================== -Auto-registration of Lambdas ----------------------------- +``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). -As a convenience, any lambda expression used in place of a UnaryOp will be automatically -compiled as registered anonymously. +Parameterized UDFs +------------------ -Example lambda usage: +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: + +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[float](0.5) # theta = 0.5 + A = Matrix.from_coo([0, 1], [0, 1], [1.0, 2.0]) + B = Matrix.from_coo([0, 1], [0, 1], [3.0, 4.0]) + C = A.ewise_mult(B, bound).new() + # C[0,0] = (1+3)*0.5 = 2; C[1,1] = (2+4)*0.5 = 3 + +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 resulting Semiring is monomorphic in the bound IBO's input/output types +(SuiteSparse builds exactly one ``GrB_Semiring`` for that type pair, rather +than the type-polymorphic table the standard semirings carry). To reuse the +same IBO at a different type, bind theta again under that type and build a +new Semiring. 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 -.. csv-table:: - :class: matrix - :header: 0,1,2,3,4,5 + 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. - -1,0,,1,1,2 +The compile happens on first lookup (``unary.heavy_op[int]``). diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst new file mode 100644 index 000000000..18af2cf3e --- /dev/null +++ b/docs/user_guide/udt.rst @@ -0,0 +1,278 @@ + +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. + +**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. +- Strings, Python objects, and nested UDTs are not. +- 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 + +The lift is field-by-field for record UDTs and element-by-element for array +UDTs. Field types that don't support the operation raise ``KeyError`` on the +first lookup. ``binary.min``, ``binary.max``, and ``binary.floordiv`` reject +UDTs with any complex leaf (no ordering, no integer modulus); use ``plus``, +``minus``, ``times``, or ``truediv`` for complex arithmetic, or register a +custom binary op. + +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. + +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 the UDT name and 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. 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 # 'Edge', or None if not JIT-able + udt.jit_c_definition # 'typedef struct { ... } Edge ;' + + op = binary.plus[udt] + op.jit_c_name # 'plus_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 (or the UDT name, if you supplied one) is a C reserved word + (``class``, ``return``, ...) or a stdlib macro/typedef pulled in by + ``GraphBLAS.h`` (``NULL``, ``FILE``, ``M_PI``, ``complex``, ...). 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). +- 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). + +Anonymous UDTs (no ``name=`` argument) still take the JIT path: a synthetic +``_gbudt_NNN`` C name is minted so SuiteSparse always has a registerable +identifier. Pass ``name=`` only for readable JIT cache filenames and +introspection. + +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``) is emitted with the cause and the remediation. The warning +fires once per ``(op, dtype)`` rather than once per process, so distinct +fallback causes each surface a warning. Silence it 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. + +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 auto-fixes this at import. 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), and ``jit_c_control`` is +bumped from the SS default ``'run'`` (run cached kernels only; no compile, +no load from disk) to ``'on'`` (compile, load, and run). When the default +config is already valid, only the mode bump applies. + +Call the helper manually to re-fix or verify:: + + gb.ss.fix_jit_config() # repair compiler path (full 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/core/dtypes.py b/graphblas/core/dtypes.py index 2d4178b14..8eb09725b 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -1,3 +1,4 @@ +import dataclasses import warnings from ast import literal_eval @@ -15,7 +16,20 @@ 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 +38,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 +80,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 +104,85 @@ 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`` may differ from ``datatype.name``: if the user + didn't supply a name, or supplied one that isn't a valid C identifier, + ``_udt_c_typedef`` synthesizes a fresh ``_gbudt_NNN`` 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_utils 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_utils import _has_jit_set + + 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 + 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 + + +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 +193,56 @@ 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``. + + 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``. Must be a valid C identifier (and not a C reserved word) + to be usable on the SuiteSparse JIT path; otherwise JIT is skipped + silently and ops fall back to the Numba cfunc path. + + 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: @@ -103,12 +260,20 @@ def register_anonymous(dtype, name=None): else: raise 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") @@ -131,10 +296,7 @@ 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: @@ -159,6 +321,8 @@ def register_anonymous(dtype, name=None): 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 @@ -440,13 +604,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/infixmethods.py b/graphblas/core/infixmethods.py index 71201ad9c..67248223a 100644 --- a/graphblas/core/infixmethods.py +++ b/graphblas/core/infixmethods.py @@ -65,7 +65,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 +75,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 +92,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 +111,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 +129,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/operator/agg.py b/graphblas/core/operator/agg.py index 6b463a8a6..7afeb9e46 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 @@ -88,14 +89,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 +138,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 +190,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..bf685ba5c 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -1,10 +1,12 @@ +import re from functools import lru_cache from operator import getitem 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 ...exceptions import UdfParseError, check_status_carg +from .. import _has_numba, _supports_udfs, ffi, lib from ..expr import InfixExprBase from ..utils import output_type @@ -93,6 +95,20 @@ def _hasop(module, name): ) +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-05-24: still hits on Numba 0.65. Re-test periodically and + drop the INT8 routing when upstream is fixed. + """ + return INT8 if dtype == BOOL else dtype + + class OpPath: def __init__(self, parent, name): self._parent = parent @@ -165,67 +181,347 @@ 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): + 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``. + """ + wrapper = numba.cfunc(wrapper_sig, nopython=True)(wrapper) + 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. + """ + try: + return lookup_dtype(numba_ret_type) + except (ValueError, TypeError): + pass + if isinstance(numba_ret_type, numba.core.types.BaseTuple): + from .udt_utils import _iter_record_leaves + + n = len(numba_ret_type.types) + + def _leaves(d): + """Yield ``(python_path, c_path, leaf_dtype)`` for 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): + continue + 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." + ) + 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 _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, 1)`` line + to add to the wrapper body (empty for non-record 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), + ) + # Array UDT: pass the raw pointer. The UDF carrays it if it wants. + return "", f"{var}_ptr", nt.CPointer(dtype.numba_type) + if dtype == BOOL: + # Numba can't compile bool ptrs (numba/numba#5395); expose them + # as int8 and cast on deref. + # MAINT 2026-05-24: still hits on Numba 0.65; 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, otherwise ``"direct"``. + ``write_info`` is the tuple-unpack field tuple for ``"record_fields"``, + or a ``(BL, BR, zname)`` 3-tuple for ``"direct"``. + """ + nt = numba.types ztype = INT8 if return_type == BOOL else return_type - xtype = INT8 if dtype == BOOL else dtype + 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`` is the list of Python access paths to each leaf + # field (``"['a']"`` for flat, ``"['outer']['inner_a']"`` for + # nested). The wrapper iterates these to write the flat tuple + # ``_result`` back leaf-by-leaf. + from .udt_utils import _iter_record_leaves + + leaf_python_paths = tuple(py 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", + leaf_python_paths, + ) + 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]"), + ) + # Array UDT: write via z_ptr[0]. The UDF receives the carray and + # writes in place. + return "", ret_ptr_type, "direct", ("", "[0]", "z_ptr[0]") + 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. + """ + if zkind == "record_fields": + field_assigns = "".join( + f" z[0]{path} = _result[{i}]\n" for i, path in enumerate(zinfo) + ) + return ( + f"{signature_line}\n" + f"{body_setup}" + f" _result = {call_expr}\n" + f"{field_assigns}" + ) + BL, BR, zname = zinfo + return f"{signature_line}\n{body_setup} {zname} = {BL}{call_expr}{BR}\n" + + 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 - wrapper_args = [nt.CPointer(ztype.numba_type), nt.CPointer(xtype.numba_type)] + 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] if include_indexes: wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) + ysetup, yderef_expr, yarg = "", "", "" if dtype2 is not None: - ytype = INT8 if dtype2 == BOOL else dtype2 - wrapper_args.append(nt.CPointer(ytype.numba_type)) + ysetup, yderef, yptr_type = _input_operand(dtype2, "y") + wrapper_args.append(yptr_type) + yarg = ", y_ptr" + yderef_expr = f", {yderef}" 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 = ")" + 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})" - 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]" + text = _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr) + from .udt_utils import _compile_codegen - 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]" + 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 - if include_indexes: - rcidx = ", row, col" + def _get_udt_wrapper_indexbinary( + numba_func, return_type, dtype, dtype2, *, numba_ret_type=None + ): + """Build a Numba cfunc wrapper for IndexBinaryOp UDFs on UDTs. + + Signature: ``f(x, ix, jx, y, iy, jy, theta) -> z``. ``dtype2`` is the + shared type of ``y`` and ``theta``. + """ + 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(dtype2, "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, + ) - 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" + 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) + from .udt_utils import _compile_codegen + + wrapper = _compile_codegen( + text, + func_name="wrapper", + source_label=f"", + extra_ns={"numba_func": numba_func}, ) - exec(text, d) # pylint: disable=exec-used - return d["wrapper"], wrapper_sig + return 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``. +_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 +533,18 @@ class TypedOpBase: "gb_obj", "gb_name", "_type2", + "_jit_c_info", + "_owns_gb_obj_inst", "__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. + # Specific instances can override via ``_owns_gb_obj_inst`` (set by + # the constructor); ``SelectOp._from_indexunary`` aliases an existing + # ``GrB_IndexUnaryOp`` and must clear ownership to avoid a double free. + _owns_gb_obj = False def __init__(self, parent, name, type_, return_type, gb_obj, gb_name, dtype2=None): self.parent = parent @@ -248,6 +554,24 @@ 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 + # Per-instance ownership override; defaults to the class attribute. + # ``SelectOp._from_indexunary`` flips this to ``False`` on aliasing + # TypedUserSelectOps so only the IndexUnaryOp frees the handle. + self._owns_gb_obj_inst = type(self)._owns_gb_obj + + @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 +592,52 @@ 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, so + # gate on the per-instance owns flag (defaults to the class + # attribute; the alias case overrides to False). Mirrors the + # ``Matrix.__del__`` / ``Vector.__del__`` pattern. + if not getattr(self, "_owns_gb_obj_inst", False): + 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 +647,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 +662,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] != "_") @@ -362,7 +760,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 @@ -500,6 +898,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: diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 3ee089fe4..b3fac1cb9 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -32,46 +32,140 @@ ParameterizedUdf, TypedOpBase, _call_op, - _deserialize_parameterized, _hasop, ) +# Imported unconditionally (plain dict, no numba): ``_compile_udt`` consults it +# even when numba is absent. +from .udt_utils import BUILTIN_UDT_BINARY_OPS as _BUILTIN_UDT_BINARY_OPS + if _has_numba: import numba - from .base import _get_udt_wrapper + from .base 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``. Returns + ``(wrapper_func, wrapper_sig)``. + """ + from .udt_utils import _check_udt_pair, _get_udt_info, _iter_record_leaves + + 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) + + 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 + + nt = numba.types + cmp_op = "==" if is_eq else "!=" + join = " and " if is_eq else " or " + + if udt_info[0] == "array": + # Array UDT: unroll element-by-element so the scalar broadcast + # case can drop ``numba.carray`` on that side entirely. + # ``_get_udt_info`` already flattens the multi-dim shape. + # The wrapper sees the UDT 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. + _base, N = udt_info[1] + base_numba = numba.from_dtype(_base) + x_ptr_type = ( + nt.CPointer(numba.from_dtype(dtype.np_type)) + if x_is_scalar + else nt.CPointer(base_numba) + ) + y_ptr_type = ( + nt.CPointer(numba.from_dtype(dtype2.np_type)) + if y_is_scalar + else nt.CPointer(base_numba) + ) + wrapper_sig = nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type) + if x_is_scalar: + terms = [f"(x_ptr[0] {cmp_op} y[{i}])" for i in range(N)] + arrays = f" y = numba.carray(y_ptr, {N})\n" + elif y_is_scalar: + terms = [f"(x[{i}] {cmp_op} y_ptr[0])" for i in range(N)] + arrays = f" x = numba.carray(x_ptr, {N})\n" + else: + terms = [f"(x[{i}] {cmp_op} y[{i}])" for i in range(N)] + arrays = f" x = numba.carray(x_ptr, {N})\n y = numba.carray(y_ptr, {N})\n" + body = join.join(terms) + src = f"def wrapper(z_ptr, x_ptr, y_ptr):\n{arrays} z_ptr[0] = {body}\n" + wrapper = _compile_codegen( + src, + func_name="wrapper", + source_label=f"", + ) + return wrapper, wrapper_sig + + # Record UDT (possibly nested): chain leaf comparisons. For an + # array-valued sub-field, unroll its elements so the scalar side + # never sees a numpy fancy-broadcast (Numba's record access can't + # express that as a single expression). + np_type = udt_dtype.np_type + terms = [] + for py_path, _c, leaf_dtype in _iter_record_leaves(np_type): + if leaf_dtype.subdtype is not None: + _base, shape = leaf_dtype.subdtype + sub_N = int(reduce(mul, shape)) + if x_is_scalar: + subterms = [f"(x_ptr[0] {cmp_op} y[0]{py_path}[{i}])" for i in range(sub_N)] + elif y_is_scalar: + subterms = [f"(x[0]{py_path}[{i}] {cmp_op} y_ptr[0])" for i in range(sub_N)] + else: + reducer = ".all()" if is_eq else ".any()" + terms.append(f"(x[0]{py_path} {cmp_op} y[0]{py_path}){reducer}") + continue + terms.append("(" + join.join(subterms) + ")") + else: + 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}" + 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): @@ -129,6 +223,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 +303,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) @@ -297,12 +378,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,6 +388,16 @@ def _pair_dtype(op, dtype, dtype2): return op[INT64] +if _has_numba: + from .udt_utils import ( + _compile_codegen, + _has_jit_set, + _maybe_warn_jit_skipped, + compile_udt_binary_wrapper, + set_jit_c_on_op, + ) + + class BinaryOp(OpBase): """Takes two inputs and returns one output, possibly of a different data type. @@ -463,11 +550,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( @@ -525,109 +609,75 @@ def _compile_udt(self, dtype, dtype2): if dtypes in self._udt_types: return self._udt_ops[dtypes] - 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() + # 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 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}") - 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] = 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), + if self.name in ("eq", "ne") and not self._anonymous 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 self.name in _BUILTIN_UDT_BINARY_OPS: + # 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. Only + # same-type pairs apply: mixed UDT+scalar ops don't have a + # meaningful single-type C definition. ``_has_jit_set`` gates access + # to ``lib.GrB_BinaryOp_set_String``, which is absent on SS < 8. + is_jittable = self.name in _BUILTIN_UDT_BINARY_OPS or self.name in ("eq", "ne") + 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_utils 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) return op @classmethod @@ -733,7 +783,12 @@ 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: binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) @@ -915,29 +970,40 @@ 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) 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 are auto-generated from per-field / + # per-element scalar ops. + if _has_numba: + for op_name in _BUILTIN_UDT_BINARY_OPS: + binop = getattr(binary, op_name, None) + if binop is not None: + binop._udt_types = {} + binop._udt_ops = {} + binop._custom_dtype = _udt_dtype cls._initialized = True def __init__( @@ -963,15 +1029,6 @@ def __init__( 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)) - __call__ = TypedBuiltinBinaryOp.__call__ is_commutative = TypedBuiltinBinaryOp.is_commutative commutes_to = ParameterizedBinaryOp.commutes_to @@ -983,3 +1040,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 index cb07e0eaa..a19f5ffec 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -6,21 +6,100 @@ from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, ffi, lib from ..dtypes import _sample_values -from .base import OpBase, ParameterizedUdf, TypedOpBase, _deserialize_parameterized +from .base import OpBase, ParameterizedUdf, TypedOpBase _has_idxbinop = hasattr(lib, "GxB_IndexBinaryOp_new") if _has_numba: import numba + from .base 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): + # ``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 infer from the value. A raw user-defined-type + value (array/tuple) can't be inferred, so raise a clear error that names the + value and the working alternatives. + """ + from ..scalar import Scalar + + 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}; pass an explicit dtype " + "(e.g., op(value, dtype=...)) or a typed Scalar. Required for user-defined types." + ) from exc + raise + + class _BoundIndexBinaryOp(TypedOpBase): - """A BinaryOp created by binding a theta value to an IndexBinaryOp.""" + """A BinaryOp produced by binding a theta value to an IndexBinaryOp. - __slots__ = () + 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",) 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 + # ``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, self.type, self._theta)) class TypedBuiltinIndexBinaryOp(TypedOpBase): @@ -28,29 +107,38 @@ class TypedBuiltinIndexBinaryOp(TypedOpBase): opclass = "IndexBinaryOp" def __call__(self, theta=None): - """Bind a theta value to create a BinaryOp that can be used in operations. + """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. Default is 0 (False). + The theta parameter to bind. Defaults to 0 (False). Returns ------- TypedOpBase - A BinaryOp created from this IndexBinaryOp with the given theta. + 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): - theta = Scalar.from_value(theta, is_cscalar=False, name="") # pragma: is_grbscalar + # 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 - val = theta.value - theta = Scalar.from_value(val, is_cscalar=False, name="") # pragma: is_grbscalar + # 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 new_binop = ffi_new("GrB_BinaryOp*") check_status_carg( @@ -58,14 +146,16 @@ def __call__(self, theta=None): "BinaryOp", new_binop[0], ) - rv = _BoundIndexBinaryOp.__new__(_BoundIndexBinaryOp) - rv.parent = self.parent - rv.name = self.name - rv.type = self.type - rv.return_type = self.return_type - rv.gb_obj = new_binop[0] - rv.gb_name = f"{self.name}_bound" - rv._type2 = self._type2 + 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 return rv @property @@ -76,6 +166,7 @@ def thunk_type(self): class TypedUserIndexBinaryOp(TypedOpBase): __slots__ = () opclass = "IndexBinaryOp" + _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) @@ -112,35 +203,22 @@ def _call(self, *args, **kwargs): idxbinop._parameterized_info = (self, args, kwargs) return IndexBinaryOp.register_anonymous(idxbinop, self.name, is_udt=self._is_udt) - def __reduce__(self): - name = f"indexbinary.{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 IndexBinaryOp.register_anonymous(func, name, parameterized=True) - if (rv := IndexBinaryOp._find(name)) is not None: - return rv - return IndexBinaryOp.register_new(name, func, parameterized=True) - class IndexBinaryOp(OpBase): - """Takes two inputs with their indices, plus a thunk, and returns one output. + """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/column indices of ``x`` and ``iy, jy`` are - the row/column indices of ``y``. - - An IndexBinaryOp can be converted to a BinaryOp by binding a theta value, - which makes it usable in any operation that accepts a BinaryOp (eWiseMult, - eWiseAdd, mxm, etc.). + where ``ix, jx`` are the row and column indices of ``x``, and ``iy, jy`` + are the row and column indices of ``y``. - IndexBinaryOps are located in the ``graphblas.indexbinary`` namespace. + Binding a theta value (``ibo[dtype](theta)``) produces a BinaryOp usable + directly in ``ewise_mult`` and ``ewise_add``. 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. - There are no built-in IndexBinaryOps; all are user-defined. + IndexBinaryOps live in the ``graphblas.indexbinary`` namespace. There are + no built-ins; all IndexBinaryOps are user-defined. """ __slots__ = "orig_func", "_is_udt", "_numba_func" @@ -204,12 +282,11 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): elif type_ == BOOL and ret_type.name == "INT64" and return_types.get(INT8) == INT8: ret_type = INT8 - # Numba is unable to handle BOOL correctly right now - 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: z = f(x, ix, jx, y, iy, jy, theta) - # C signature: void(z*, x*, ix, jx, y*, iy, jy, theta*) + # 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), @@ -282,7 +359,9 @@ def indexbinary_wrapper( def _compile_udt(self, dtype, dtype2): if not _has_idxbinop: - raise RuntimeError( + # 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+)" ) @@ -304,10 +383,13 @@ def _compile_udt(self, dtype, dtype2): UINT64.numba_type, dtype2.numba_type, ) - numba_func.compile(sig) - ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) + _compile_udf_for_udt( + numba_func, sig, op_kind="indexbinary", 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) indexbinary_wrapper, wrapper_sig = _get_udt_wrapper_indexbinary( - numba_func, ret_type, dtype, dtype2 + numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexbinary_wrapper) @@ -340,22 +422,22 @@ def _compile_udt(self, dtype, dtype2): @classmethod def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): - """Register an IndexBinaryOp without registering it in the ``indexbinary`` namespace. + """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 able - to be compiled with ``numba.njit``. - ``func`` takes seven input parameters--``(x, ix, jx, y, iy, jy, theta)``-- - where ``x`` and ``y`` are element values, ``ix, jx`` and ``iy, jy`` - are their row/column indices (int64), and ``theta`` is a scalar parameter. + 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, which means + 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. @@ -372,24 +454,24 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals @classmethod def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): - """Register a new IndexBinaryOp and save it to ``graphblas.indexbinary`` namespace. + """Register a new IndexBinaryOp under the ``graphblas.indexbinary`` namespace. Parameters ---------- name : str - The name of the operator. This will show up as ``gb.indexbinary.{name}``. + The name of the operator. Available afterwards as ``gb.indexbinary.{name}``. func : FunctionType - The function to compile. For all current backends, this must be able - to be compiled with ``numba.njit``. - ``func`` takes seven input parameters--``(x, ix, jx, y, iy, jy, theta)``-- - where ``x`` and ``y`` are element values, ``ix, jx`` and ``iy, jy`` - are their row/column indices (int64), and ``theta`` is a scalar parameter. + 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 - If True, delay compilation until the operator is used. + When True, defer compilation until the operator is first used. Examples -------- @@ -401,7 +483,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: idxbinop = ParameterizedIndexBinaryOp(name, func, is_udt=is_udt) @@ -420,7 +502,7 @@ def _initialize(cls): if cls._initialized: return super()._initialize(include_in_ops=False) - # No built-in IndexBinaryOps to register + # No built-in IndexBinaryOps to register. cls._initialized = True def __init__(self, name, func=None, *, anonymous=False, is_udt=False, numba_func=None): @@ -432,29 +514,20 @@ def __init__(self, name, func=None, *, anonymous=False, is_udt=False, numba_func self._udt_types = {} # {(dtype, dtype2): DataType} self._udt_ops = {} # {(dtype, dtype2): TypedUserIndexBinaryOp} - 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"indexbinary.{self.name}") in _STANDARD_OPERATOR_NAMES: - return name - return (self._deserialize, (self.name, self.orig_func)) - def __call__(self, theta=None, dtype=None): - """Bind a theta value to create a BinaryOp. + """Bind a theta value to produce a BinaryOp. Parameters ---------- theta : scalar, optional - The theta parameter to bind. Default is 0 (False). + The theta parameter to bind. Defaults to 0 (False). dtype : dtype, optional - The dtype to use. If not provided, it will be inferred from theta. + The dtype to use. Inferred from ``theta`` when not provided. Returns ------- TypedOpBase - A BinaryOp created from this IndexBinaryOp with the given theta. + A BinaryOp built from this IndexBinaryOp with the given theta. """ from ...dtypes import lookup_dtype as _lookup_dtype @@ -463,11 +536,14 @@ def __call__(self, theta=None, dtype=None): if theta is None: theta = False if not isinstance(theta, Scalar): - theta = Scalar.from_value(theta, is_cscalar=False, name="") # pragma: is_grbscalar + # An explicit dtype supplies a UDT; with no dtype the value is + # inferred (and a raw UDT value raises a clear error). + theta = _theta_to_scalar(theta, _lookup_dtype(dtype) if dtype is not None else None) elif theta._is_cscalar: # fmt: off - val = theta.value - theta = Scalar.from_value(val, is_cscalar=False, name="") # pragma: is_grbscalar + # 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 if dtype is None: dtype = theta.dtype @@ -477,31 +553,4 @@ def __call__(self, theta=None, dtype=None): return typed_op(theta) -def _get_udt_wrapper_indexbinary(numba_func, return_type, dtype, dtype2): - """Build a wrapper function for UDT IndexBinaryOp: z = f(x, ix, jx, y, iy, jy, theta).""" - nt = numba.types - ztype = INT8 if return_type == BOOL else return_type - xtype = INT8 if dtype == BOOL else dtype - ytype = INT8 if dtype2 == BOOL else dtype2 - - wrapper_sig = nt.void( - nt.CPointer(ztype.numba_type), - nt.CPointer(xtype.numba_type), - UINT64.numba_type, - UINT64.numba_type, - nt.CPointer(ytype.numba_type), - UINT64.numba_type, - UINT64.numba_type, - nt.CPointer(ytype.numba_type), - ) - - d = {"numba": numba, "numba_func": numba_func} - xderef = "bool(x_ptr[0])" if dtype == BOOL else "x_ptr[0]" - yderef = "bool(y_ptr[0])" if dtype2 == BOOL else "y_ptr[0]" - tderef = "bool(t_ptr[0])" if dtype2 == BOOL else "t_ptr[0]" - call = f"numba_func({xderef}, ix, jx, {yderef}, iy, jy, {tderef})" - if return_type == BOOL: - call = f"bool({call})" - text = f"def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):\n z_ptr[0] = {call}\n" - exec(text, d) # noqa: S102 - return d["wrapper"], wrapper_sig +ParameterizedIndexBinaryOp._op_class = IndexBinaryOp diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 6fdacbcc1..c7af8ed78 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 .base 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. @@ -152,11 +143,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( @@ -220,32 +208,17 @@ def _compile_udt(self, dtype, dtype2): 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 +231,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 +281,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 +321,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) @@ -428,15 +401,7 @@ def __init__( 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)) - __call__ = TypedBuiltinIndexUnaryOp.__call__ + + +ParameterizedIndexUnaryOp._op_class = IndexUnaryOp diff --git a/graphblas/core/operator/monoid.py b/graphblas/core/operator/monoid.py index e3f218a90..eb3aed6c1 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,79 @@ 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 _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 + if np.issubdtype(scalar_dtype, np.integer): + 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) + if val is None: + 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) + raise KeyError(f"monoid.{monoid_name} does not work with {top_dtype}") + + ffi_new = ffi.new @@ -67,10 +137,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,15 +208,18 @@ 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): @@ -221,7 +299,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 +308,25 @@ 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) + identity_val = self._identity + if identity_val is None and self.name in _BUILTIN_UDT_MONOIDS: + # Auto-generate the identity for built-in monoids on UDTs. + 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 +334,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 @@ -348,11 +439,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 +532,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..32a2509a5 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -1,17 +1,13 @@ 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 .base import _get_udt_wrapper -ffi_new = ffi.new + from .base import _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper class TypedBuiltinSelectOp(TypedOpBase): @@ -29,6 +25,7 @@ def __call__(self, val, thunk=None): class TypedUserSelectOp(TypedOpBase): __slots__ = () opclass = "SelectOp" + _owns_gb_obj = True # underlying object is a GrB_IndexUnaryOp 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 +58,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 +97,10 @@ def _from_indexunary(cls, iop): t.return_type, t.gb_obj, ) + # Aliases the IndexUnaryOp's allocation. The IndexUnaryOp + # owns the free; clearing here prevents a double free when + # both ends are GC'd. + op._owns_gb_obj_inst = False else: op = cls._typed_class( obj, @@ -143,31 +128,15 @@ def _compile_udt(self, dtype, dtype2): # 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], + return _finalize_udt_op( + self, dtype, dtype2, BOOL, select_wrapper, wrapper_sig, TypedUserSelectOp ) - op = TypedUserSelectOp( - self, - self.name, - dtype, - BOOL, - new_select[0], - dtype2=dtype2, - ) - 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 +150,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 +201,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 +291,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)) - __call__ = TypedBuiltinSelectOp.__call__ + + +ParameterizedSelectOp._op_class = SelectOp diff --git a/graphblas/core/operator/semiring.py b/graphblas/core/operator/semiring.py index a8d18f1bf..32e20d0bf 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) @@ -207,8 +222,23 @@ 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)}") + # 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`` 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. + from .indexbinary import _BoundIndexBinaryOp + + if isinstance(binaryop, _BoundIndexBinaryOp): + return cls._build_from_bound_indexbinary(name, monoid, binaryop, anonymous=anonymous) if type(binaryop) is not BinaryOp: - raise TypeError(f"binaryop must be a BinaryOp, not {type(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 +273,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 +326,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 +335,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 +349,11 @@ 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). A bound IndexBinaryOp (``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 +378,10 @@ 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. A bound IndexBinaryOp + (``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/udt_utils.py b/graphblas/core/operator/udt_utils.py new file mode 100644 index 000000000..fc33ea8c0 --- /dev/null +++ b/graphblas/core/operator/udt_utils.py @@ -0,0 +1,994 @@ +"""Shared utilities for auto-generating UDT operator implementations. + +Provides Numba wrapper generators and JIT C code generators for element-wise +operations on record UDTs (struct types) and array UDTs (e.g., FP64[3]). + +Compilation is lazy: functions here are only called from ``_compile_udt`` +methods when an operator is first used with a particular UDT. +""" + +import ast +import itertools +import linecache +from functools import reduce +from operator import mul + +import numpy as np + +from ... import backend +from .. import _has_numba, ffi, lib + +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") + namespace = {"min": min, "max": max, "abs": abs} + if _has_numba: + namespace["numba"] = numba + if extra_ns: + namespace.update(extra_ns) + exec(code, namespace) + return namespace[func_name] + + +BUILTIN_UDT_BINARY_OPS = { + "plus": "+", + "minus": "-", + "times": "*", + "truediv": "/", + "floordiv": "//", + "min": "min", + "max": "max", +} + +BUILTIN_UDT_UNARY_OPS = { + "ainv": "-", + "abs": "abs", +} + +# 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 and are skipped via _op_supports_field_dtype. + np.dtype(np.complex64): "GxB_FC32_t", + np.dtype(np.complex128): "GxB_FC64_t", +} + +# Ops whose JIT C codegen would produce uncompilable kernels on complex fields +# (no ordering for min/max, no integer-mod for floordiv). +_OPS_NOT_FOR_COMPLEX = frozenset({"min", "max", "floordiv"}) + +# C operator equivalents for JIT code generation +_C_INFIX_OPS = {"+": "+", "-": "-", "*": "*", "/": "/", "//": "/"} + +# 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 UDT type names or record field names in the +# JIT C source. Covers two failure modes from the JIT include chain +# (``GraphBLAS.h`` transitively pulls in ````, ````, +# ````, ````, ````, ````): +# +# 1. **Macro names**: the preprocessor expands them inside any declarator +# or expression they appear in, so a field declared ``double M_PI ;`` +# becomes ``double 3.14159... ;`` and won't compile. The C-standard +# separate namespace for struct members doesn't help here. +# 2. **Typedef names** at the *outer* type-name position: emitting +# ``typedef struct { ... } FILE ;`` collides with the existing +# ``FILE`` typedef in scope (most compilers reject as a redefinition). +# Struct member names of the same spelling are fine (members live in +# their own namespace), but our gate runs the same check at both +# positions for simplicity. +# +# Either way the JIT compile fails, SuiteSparse swallows the failure +# silently, and the op runs through the slower Numba cfunc path. Block +# these eagerly so the user sees the "without JIT" warning instead. +_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", + } +) + + +def _is_valid_c_identifier(name): + """True if ``name`` is a valid C identifier and not a reserved word.""" + return isinstance(name, str) and name.isidentifier() and name not in _C_RESERVED + + +# Process-global counter for synthesizing C names when the user-supplied +# Python name isn't a valid C identifier (or is a C reserved word, 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) + + +def _pick_c_type_name(python_name): + """Return a valid C identifier for use as ``GxB_JIT_C_NAME`` on a UDT. + + When ``python_name`` is already a valid C identifier (and not reserved), + return it unchanged so introspection and JIT cache filenames remain + readable. Otherwise mint a fresh ``_gbudt_NNN`` name. This lets UDTs + that were anonymous in the Python sense (no ``name=`` passed, name with + dots / special chars, name that collides with a C keyword) still take + the JIT path; users never see the synthetic name unless they read + ``DataType.jit_c_name``. + """ + if _is_valid_c_identifier(python_name): + return python_name + return f"_gbudt_{next(_synthetic_udt_counter)}" + + +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 _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`` (one side is a scalar broadcast), no check fires. Without this + gate the generated wrapper code raises a cryptic Numba ``TypingError`` + on the first field access. + """ + 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" and detail_x != detail_y: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must share field names; got {list(detail_x)} vs " + f"{list(detail_y)}." + ) + if kind_x == "array" and detail_x != detail_y: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"array UDTs must share base dtype and flat size; " + f"got {detail_x} vs {detail_y}." + ) + + +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 (the codegen treats them opaquely; arithmetic + on an array-typed sub-field isn't supported by either the cfunc or the + JIT path today). + """ + 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 _udt_c_typedef(python_name, np_type): + """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 (may differ from + ``python_name`` if that wasn't a valid C identifier; see + ``_pick_c_type_name``). 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 = _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) + 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. + """ + if np_type.names is None: + return np_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 + + +def _op_supports_field_dtypes(op_name, np_type): + """Return True if the JIT codegen for ``op_name`` can handle every field + type in ``np_type``. + + Currently the only restriction is complex fields: ``min`` / ``max`` / + ``floordiv`` have no defined semantics on C99 ``_Complex`` (no ordering, + no integer mod), so we skip JIT and let SuiteSparse use the Numba cfunc + instead. The cfunc itself errors on these combinations, so this gate + only moves the failure earlier and gives a clearer message. + + Recurses into nested record fields so a complex field nested inside an + outer record is still detected. + """ + if op_name not in _OPS_NOT_FOR_COMPLEX: + return True + if np_type.names is not None: + return not any(leaf.kind == "c" for _, _, leaf in _iter_record_leaves(np_type)) + if np_type.subdtype is not None: + return np_type.subdtype[0].kind != "c" + 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): + """Python-source builder; sibling of :func:`_c_expr_binary` for JIT C.""" + if py_op in _FUNC_BINARY_OPS: + return f"{py_op}({x_expr}, {y_expr})" + return f"{x_expr} {py_op} {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: + return f"{py_op}({operand})" + return f"{py_op}{operand}" + + def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scalar=False): + """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 (in base.py) 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. + """ + if arity == 2: + parts = [] + for path in leaf_paths: + x_expr = "x" if x_is_scalar else f"x{path}" + y_expr = "y" if y_is_scalar else f"y{path}" + parts.append(_expr_binary(py_op, x_expr, y_expr)) + sig = "x, y" + else: + parts = [_expr_unary(py_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})" + src = f"def _op({sig}):\n return {ret}\n" + op_func = _compile_codegen( + src, + func_name="_op", + source_label=f"", + ) + return numba.njit(op_func) + + def _make_array_wrapper( + size, + base_numba_type, + arity, + py_op, + *, + x_scalar_type=None, + y_scalar_type=None, + ): + """Build a cfunc-ready wrapper for an array UDT (element-by-element). + + When ``x_scalar_type`` or ``y_scalar_type`` is set, that side is a plain + scalar pointer (broadcast to all elements). + + Returns (wrapper_func, wrapper_sig). + """ + nt = numba.types + if arity == 2: + x_ref = "x_ptr[0]" if x_scalar_type else "x[{i}]" + y_ref = "y_ptr[0]" if y_scalar_type else "y[{i}]" + assigns = "\n".join( + f" z[{i}] = {_expr_binary(py_op, x_ref.format(i=i), y_ref.format(i=i))}" + for i in range(size) + ) + params = "z_ptr, x_ptr, y_ptr" + arrays = f" z = numba.carray(z_ptr, {size})\n" + if not x_scalar_type: + arrays += f" x = numba.carray(x_ptr, {size})\n" + if not y_scalar_type: + arrays += f" y = numba.carray(y_ptr, {size})\n" + x_numba = nt.CPointer(x_scalar_type) if x_scalar_type else nt.CPointer(base_numba_type) + y_numba = nt.CPointer(y_scalar_type) if y_scalar_type else nt.CPointer(base_numba_type) + sig = nt.void(nt.CPointer(base_numba_type), x_numba, y_numba) + else: + assigns = "\n".join( + f" z[{i}] = {_expr_unary(py_op, f'x[{i}]')}" for i in range(size) + ) + params = "z_ptr, x_ptr" + arrays = f" z = numba.carray(z_ptr, {size})\n x = numba.carray(x_ptr, {size})\n" + sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(base_numba_type)) + src = f"def _op({params}):\n{arrays}{assigns}\n" + op_func = _compile_codegen( + src, + func_name="_op", + source_label=f"", + ) + 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): + """Compile a built-in element-wise binary op for a UDT. + + Handles two cases: + + - Both sides are the same UDT: 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). + + 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 .base import _get_udt_wrapper, _resolve_udt_return_type + + info_x = _get_udt_info(dtype) + info_y = _get_udt_info(dtype2) + _check_udt_pair(op_name, dtype, dtype2, info_x, info_y) + + # Pick the UDT side. Both sides may be UDTs; the pre-check above + # ensures they share a shape in that case. + if info_x is not None: + udt_dtype = dtype + udt_info = info_x + elif info_y is not None: + udt_dtype = dtype2 + udt_info = info_y + else: + 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])." + ) + + # ``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 not _op_supports_field_dtypes(op_name, udt_dtype.np_type): + raise 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." + ) + + 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 + kind, detail = udt_info + + if kind == "record": + from .base 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']"``. + leaf_paths = [py for py, _c, _d in _iter_record_leaves(udt_dtype.np_type)] + func = _make_record_func( + leaf_paths, + 2, + py_op, + x_is_scalar=x_is_scalar, + y_is_scalar=y_is_scalar, + ) + sig = (dtype.numba_type, dtype2.numba_type) + _compile_udf_for_udt( + func, sig, op_kind="binary", op_name=op_name, dtypes=(dtype, dtype2) + ) + numba_ret_type = func.overloads[sig].signature.return_type + ret_type = _resolve_udt_return_type(numba_ret_type, udt_dtype) + wrapper, wrapper_sig = _get_udt_wrapper( + func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type + ) + else: + base_dtype, size = detail + ret_type = udt_dtype + wrapper, wrapper_sig = _make_array_wrapper( + size, + numba.from_dtype(base_dtype), + 2, + py_op, + x_scalar_type=numba.from_dtype(dtype.np_type) if x_is_scalar else None, + y_scalar_type=numba.from_dtype(dtype2.np_type) if y_is_scalar else None, + ) + return wrapper, wrapper_sig, ret_type + + def compile_udt_unary_wrapper(op_name, py_op, dtype): + """Compile a built-in element-wise unary op for a UDT. + + Returns (wrapper_func, wrapper_sig, ret_type). + Raises KeyError if the dtype is not supported. + """ + from .base 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 .base 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) + wrapper, wrapper_sig = _get_udt_wrapper( + func, ret_type, dtype, numba_ret_type=numba_ret_type + ) + else: + base_dtype, size = detail + ret_type = dtype + wrapper, wrapper_sig = _make_array_wrapper(size, numba.from_dtype(base_dtype), 1, py_op) + return wrapper, wrapper_sig, ret_type + + +# JIT C code generators below. + + +def _c_expr_binary(py_op, lhs, rhs, field_dtype=None): + """Return a C expression for a binary op: e.g., ``(x->a) + (y->a)``. + + ``field_dtype`` is the numpy dtype of the *result* element. It is only + consulted for ``floordiv`` (``//``), which needs Python ``//`` semantics + rather than C ``/`` (trunc toward zero for ints, true division for + floats). Other ops are type-agnostic at the C level. + """ + if py_op == "min": + # Match Python ``min(a, b) = b if b < a else a`` so NaN propagates + # from the first operand (cfunc / numba follows the same rule). + # The naive ``(a < b ? a : b)`` would silently swallow NaN to the + # right-hand side and disagree with the cfunc path. + return f"(({rhs}) < ({lhs}) ? ({rhs}) : ({lhs}))" + if py_op == "max": + return f"(({rhs}) > ({lhs}) ? ({rhs}) : ({lhs}))" + if py_op == "//": + return _c_floordiv_expr(lhs, rhs, field_dtype) + c_op = _C_INFIX_OPS.get(py_op, py_op) + return f"({lhs}) {c_op} ({rhs})" + + +def _c_floordiv_expr(lhs, rhs, field_dtype): + """Return a C expression for Python-semantics floor division. + + Python ``//`` is floor (rounds toward negative infinity); C ``/`` is + trunc toward zero for ints and true division for floats. The two only + agree for non-negative integer operands; for everything else the JIT + path silently disagreed with the Numba cfunc path before this helper. + + Float fields use ``floor()`` / ``floorf()`` from ````, which is + available in the JIT kernel via SuiteSparse's include chain + (``GraphBLAS.h`` -> ````). Signed integer fields use the + standard trunc-to-floor adjustment. Unsigned integers don't need + adjusting because both operands are non-negative. + """ + if field_dtype is None: + # Caller didn't pass dtype info. The C ``/`` semantics match Python + # ``//`` for non-negative integer operands only. + return f"({lhs}) / ({rhs})" + kind = field_dtype.kind + if kind == "f": + if field_dtype.itemsize == 4: + return f"floorf((float)({lhs}) / (float)({rhs}))" + return f"floor((double)({lhs}) / (double)({rhs}))" + if kind in ("u", "b"): + return f"({lhs}) / ({rhs})" + # Signed integer: trunc-toward-zero is one greater than floor when the + # signs of ``a`` and ``b`` differ and the division has a non-zero + # remainder; subtract 1 in that case. + return f"(({lhs}) / ({rhs}) - ((({lhs}) % ({rhs}) != 0) && ((({lhs}) < 0) != (({rhs}) < 0))))" + + +def _c_expr_unary(py_op, operand, field_dtype=None): + """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 is not None: + 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}") + + +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), + or if the op isn't defined on the dtype's field types (e.g. ``min`` on a + complex field). + """ + np_type = dtype.np_type + if not _op_supports_field_dtypes(op_name, np_type): + return None + # 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 + # synthetic ``_gbudt_NNN`` name (for UDTs whose Python name isn't a valid + # C identifier) is recorded; using it keeps op codegen consistent. + pinned_name = dtype.jit_c_name + typedef_info = _udt_c_typedef(pinned_name or dtype.name, dtype.np_type) + 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 expression builder so + # type-sensitive ops (currently floordiv) can emit correct C. + assigns = " ".join( + f"z->{c} = {_c_expr_binary(py_op, 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) + if arity == 2: + assigns = " ".join( + f"z->v[{i}] = {_c_expr_binary(py_op, f'x->v[{i}]', f'y->v[{i}]', base_dtype)} ;" + for i in range(size) + ) + else: + assigns = " ".join( + f"z->v[{i}] = {_c_expr_unary(py_op, f'x->v[{i}]', base_dtype)} ;" + for i in range(size) + ) + return c_name, f"void {c_name} ({params}) {{ {assigns} }}" + + +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 unroll + their elements at codegen time. + + 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 + pinned_name = dtype.jit_c_name + # ``_udt_c_typedef`` returns None when the numpy layout doesn't match a + # C compiler's layout for the same struct. + typedef_info = _udt_c_typedef(pinned_name or dtype.name, np_type) + 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: unroll all elements. + _base, shape = np_type.subdtype + size = reduce(mul, shape) + terms = [f"((x->v[{i}]) {op} (y->v[{i}]))" for i in range(size)] + elif np_type.names is not None: + terms = [f"((x->{c}) {op} (y->{c}))" for _py, c, _d in _iter_record_leaves(np_type)] + else: + return None + body = join.join(terms) if terms else ("1" if is_eq else "0") + params = f"_Bool *z, const {type_name} *x, const {type_name} *y" + return c_name, f"void {c_name} ({params}) {{ *z = {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. + """ + set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) + set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) + + +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 JIT isn't available or + the dtype can't be expressed in C; callers cache the returned strings + on the op for introspection (see ``TypedUserBinaryOp.jit_c_source``). + """ + if not _has_jit_set: + return None + result = _make_jit_c_comparison_definition(op_name, dtype, is_eq=is_eq) + if result is None: + return None + _set_jit_c_strings(gb_obj, *result, set_string_func) + 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 JIT setters are available and 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``). + """ + if not _has_jit_set: + return None + result = _make_jit_c_definition(op_name, py_op, dtype, arity) + if result is None: + return None + _set_jit_c_strings(gb_obj, *result, set_string_func) + return result diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index 26e0ca61c..a88cba8a6 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 .base 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,6 +119,15 @@ def _one(x): return 1 # pragma: no cover (numba) +if _has_numba: + from .udt_utils import BUILTIN_UDT_UNARY_OPS as _BUILTIN_UDT_UNARY_OPS + from .udt_utils import ( + _maybe_warn_jit_skipped, + 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. @@ -199,11 +200,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( @@ -252,26 +250,41 @@ def unary_wrapper(z, x): def _compile_udt(self, dtype, dtype2): if dtype in self._udt_types: return self._udt_ops[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 + ) + # ``set_jit_c_on_op`` is a no-op (returns None) when + # ``_has_jit_set`` is False (SS < 8). + 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): @@ -357,7 +370,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) @@ -442,6 +455,13 @@ def _initialize(cls): 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 +483,7 @@ 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 diff --git a/graphblas/core/operator/utils.py b/graphblas/core/operator/utils.py index edd06b878..6f2df5535 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -40,6 +40,15 @@ 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() except Exception: # pragma: no cover (debug) # Exceptions here can often get ignored by Python import traceback diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index c8148e8e2..0a4acbbb8 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -251,6 +251,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 @@ -369,8 +374,10 @@ def value(self, 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) + # 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 the GrB_Scalar storage. + self.gb_obj[0 : self.dtype.np_type.itemsize] = arr.tobytes() else: self.gb_obj[0] = val self._empty = False diff --git a/graphblas/core/ss/_utils.py b/graphblas/core/ss/_utils.py new file mode 100644 index 000000000..6ab0c68f7 --- /dev/null +++ b/graphblas/core/ss/_utils.py @@ -0,0 +1,47 @@ +import itertools + +from ...exceptions import _error_code_lookup +from .. import ffi, lib +from ..dtypes import _string_to_dtype + + +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 diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py new file mode 100644 index 000000000..f77b55651 --- /dev/null +++ b/graphblas/core/ss/jit_config.py @@ -0,0 +1,293 @@ +"""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. +""" + +import os +import pathlib +import platform +import re +import subprocess + +from .. import lib # noqa: F401 (sets up cffi) + +# 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()``. + """ + cfg = _ss_config() + 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/``, and strips build-time-only flags (``-isysroot``, + ``-fdebug-prefix-map``). + + 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, try to JIT-register a trivial UDT to verify + the compiler actually works. SuiteSparse auto-flips ``jit_c_control`` + from ``'on'`` to ``'load'`` on a failed compile; the probe absorbs + that first failure so user-visible ops afterwards see a stable + ``'load'`` (cache-only) state and punt to generic cleanly. + + 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 (typically ``'load'``). + 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() + if conda_prefix := os.environ.get("CONDA_PREFIX", ""): + rv = _fix_from_conda(cfg, conda_prefix) + elif use_sysconfig: + rv = _fix_from_sysconfig(cfg) + else: + return None + if rv is None: + return None + # 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() + if not probe: + return True + return _probe_jit(cfg) + + +def _fix_from_conda(cfg, conda_prefix): + """Conda-aware fix; swap the compiler path with one from $CONDA_PREFIX/bin/.""" + jit_cc = cfg["jit_c_compiler_name"] + if not pathlib.Path(jit_cc).exists(): + cc_basename = pathlib.Path(jit_cc).name + bin_dir = pathlib.Path(conda_prefix) / "bin" + for candidate in [cc_basename, "cc", "clang", "gcc"]: + local_cc = bin_dir / candidate + if local_cc.exists(): + cfg["jit_c_compiler_name"] = str(local_cc) + break + else: + return None # nothing usable + _fix_compiler_flags(cfg) + cfg["jit_c_control"] = "on" + 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 + cfg["jit_c_control"] = "on" + 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, + ) + + +def _probe_jit(cfg): + """Probe a trivial JIT compile to verify the config works. + + On failure, SuiteSparse will have flipped ``jit_c_control`` from + ``'on'`` to ``'load'`` (its built-in response to a failed compile); + we leave that state alone. Both ``'load'`` and ``'off'`` cause + downstream ops to punt to the generic kernel, but ``'load'`` preserves + any pre-compiled kernels in the cache. + """ + from ... import dtypes as _dtypes + + # The probe dtype is installed at ``dtypes.ss._jit_probe`` on first + # success; reuse it on a second probe so a repeat call doesn't appear to + # fail. Without the hasattr short-circuit, ``register_new`` would raise + # ``ValueError("name unavailable")`` on the second call. + probe_name = "_jit_probe" + if hasattr(_dtypes.ss, probe_name): + return True + try: + _dtypes.ss.register_new(probe_name, "typedef struct { int _probe ; } _jit_probe ;") + except Exception: + # ``register_new`` can raise ``JitError`` (bad path / flags / arch / + # SDK), ``RuntimeError`` (SS<8 has no JIT), or one of its input + # validation ``ValueError``s. The probe's contract is "did this + # work?", so absorb every failure mode here. + return False + return True + + +def _auto_fix_jit_at_import(): + """Run :func:`fix_jit_config` 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 + delegates the failure-prone work to ``_probe_jit``, which catches + everything internally. + + The probe is the load-bearing piece: without it, SS would surface + ``JitError`` on the first user-triggered JIT compile (a failed + compile is only converted to a silent ``'load'`` fallback on + subsequent calls). + """ + cfg = _ss_config() + if "jit_c_control" not in cfg: + return + if jit_compiler_is_usable(): + if cfg["jit_c_control"] in ("run", "load"): + cfg["jit_c_control"] = "on" + else: + fix_jit_config(use_sysconfig=True, probe=False) + if jit_compiler_is_usable() and cfg["jit_c_control"] in ("run", "load"): + cfg["jit_c_control"] = "on" + if cfg.get("jit_c_control") == "on": + _probe_jit(cfg) + + +# 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 most likely cause (bogus compiler path, ``jit_c_control`` off, or + UDT not C-expressible) is named in the message along with the remediation. + """ + key = (op_name, dtype_name) + if key in _warned_no_jit_for: + return + _warned_no_jit_for.add(key) + import warnings as _warnings + + cfg = _ss_config() + if not jit_compiler_is_usable(): + cause = ( + "the JIT compiler path is not usable " + f"({cfg.get('jit_c_compiler_name', '')!r}); " + "call ``gb.ss.fix_jit_config()`` to repair it" + ) + elif cfg.get("jit_c_control") != "on": + cause = ( + f"jit_c_control is {cfg.get('jit_c_control')!r} (must be 'on' to compile); " + "set ``gb.ss.config['jit_c_control'] = 'on'`` to enable compilation" + ) + else: + # Compiler is usable and mode is OK; the UDT itself isn't C-expressible. + # The introspection properties (``DataType.jit_c_definition`` and + # ``TypedUserBinaryOp.jit_c_source``) show what was generated, or + # ``None`` when codegen was skipped. + loc = f" (op={op_name!r}, dtype={dtype_name!r})" if op_name else "" + cause = ( + "this UDT is not expressible as a C struct" + f"{loc} (a field name is a C reserved word or stdlib macro, a " + "field type isn't in the numpy-to-C map, a field is array-typed, " + "or the record has a packed layout). The op still works via the " + "Numba cfunc path; only the JIT speedup is lost" + ) + from ...exceptions import NoJITWarning + + _warnings.warn( + f"UDT operator running without JIT compilation: {cause}. " + f"Operations will use the Numba function-pointer fallback " + f"(typically 2-3x slower for elementwise ops, since SuiteSparse " + f"can't inline the kernel into its eWise and reduce templates). " + f"This warning fires once per (op, dtype) per process; silence with " + f"``warnings.filterwarnings('ignore', category=gb.exceptions.NoJITWarning)`` " + f"or by message match.", + NoJITWarning, + stacklevel=3, + ) diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index e45c7014b..0580b7581 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -10,7 +10,6 @@ from ...exceptions import _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 ( @@ -25,6 +24,7 @@ values_to_numpy_buffer, wrapdoc, ) +from ._utils import _resolve_serialized_dtype from .config import BaseConfig from .descriptor import get_descriptor @@ -4147,33 +4147,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) diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index 20b9ce297..6c79a338b 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -10,7 +10,6 @@ from ...exceptions import _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 ( @@ -21,6 +20,7 @@ values_to_numpy_buffer, wrapdoc, ) +from ._utils import _resolve_serialized_dtype from .config import BaseConfig from .descriptor import get_descriptor from .matrix import _concat_mn, njit @@ -1641,7 +1641,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. @@ -1718,33 +1718,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/utils.py b/graphblas/core/utils.py index e9a29b3a9..49a8bd79c 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -361,6 +361,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/exceptions.py b/graphblas/exceptions.py index 05cac988a..b78023db3 100644 --- a/graphblas/exceptions.py +++ b/graphblas/exceptions.py @@ -91,7 +91,33 @@ 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 couldn't because the JIT compiler is unusable, ``jit_c_control`` + is off, or the UDT isn't expressible as a C struct. Fires once per + ``(op, dtype)`` pair per process, so a user who registers several + UDTs gets one warning per pair regardless of cause. + + 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 = { diff --git a/graphblas/indexbinary/__init__.py b/graphblas/indexbinary/__init__.py index 34f33fd4c..60c73fa96 100644 --- a/graphblas/indexbinary/__init__.py +++ b/graphblas/indexbinary/__init__.py @@ -1,5 +1,5 @@ -# All items are dynamically added by classes in operator.py -# This module acts as a container of IndexBinaryOp instances +# All items are dynamically added by classes in operator.py. +# This module acts as a container of IndexBinaryOp instances. _delayed = {} diff --git a/graphblas/monoid/numpy.py b/graphblas/monoid/numpy.py index aae8385df..ab0a6bdfa 100644 --- a/graphblas/monoid/numpy.py +++ b/graphblas/monoid/numpy.py @@ -93,9 +93,11 @@ 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 - # The behavioral check above is kept as a safety net; it's correct for all numba/numpy combos + # 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, diff --git a/graphblas/ss/__init__.py b/graphblas/ss/__init__.py index 1f059771b..394dcc391 100644 --- a/graphblas/ss/__init__.py +++ b/graphblas/ss/__init__.py @@ -1,7 +1,22 @@ 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 + + # Auto-fix the JIT config 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). When the compiler is already usable, the + # call only bumps ``jit_c_control`` from SS's default ``run`` to ``on`` so + # first-time UDT auto-lift actually JIT-compiles. Skipped on vanilla: the + # auto-fix 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 9c1d043bc..91e1496e9 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -128,14 +128,14 @@ 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** + **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 diff --git a/graphblas/tests/conftest.py b/graphblas/tests/conftest.py index 964325e0d..4270fec17 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,44 @@ import graphblas as gb from graphblas.core import _supports_udfs as supports_udfs +# 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 +58,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 +106,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() diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index ecbca707f..1866806a6 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -8,7 +8,9 @@ import graphblas as gb from graphblas import core, dtypes +from graphblas.core import _supports_udfs as supports_udfs # noqa: F401 from graphblas.core import lib +from graphblas.core.operator.udt_utils import _has_jit_set # noqa: F401 from graphblas.core.utils import _NP2 from graphblas.dtypes import lookup_dtype @@ -222,6 +224,175 @@ def test_record_dtype_from_dict(): 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_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 + + +@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"): + plus_udt = gb.binary.plus[udt] + 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 _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 == "_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 == "_PinUDT_first" + assert "_PinUDT_first" in udt.jit_c_definition + + def test_dtype_to_from_string(): types = [dtypes.BOOL, dtypes.FP64] for c in string.ascii_letters: @@ -248,6 +419,28 @@ def test_dtype_to_from_string(): assert dtype == dtype2 +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 720ab79e5..a6522dcef 100644 --- a/graphblas/tests/test_formatting.py +++ b/graphblas/tests/test_formatting.py @@ -81,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") diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 80aa8d917..493ac5ea9 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -1,5 +1,6 @@ import pickle +import numpy as np import pytest import graphblas as gb @@ -19,7 +20,6 @@ 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 op is not None assert "add_with_theta" in op.name assert int in op.types or dtypes.INT64 in op.types @@ -29,8 +29,8 @@ 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 result is not None 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]) @@ -50,7 +50,7 @@ def lazy_op(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert "lazy_op" in dir(indexbinary) op = indexbinary.lazy_op - assert op is not None + assert op.name == "lazy_op" delattr(indexbinary, "lazy_op") @@ -84,7 +84,7 @@ def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) def test_index_aware(): - """Test that indices are correctly passed to the function.""" + """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 @@ -136,7 +136,7 @@ def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) def test_default_theta(): - """Test that theta=0 works correctly.""" + """``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 @@ -162,7 +162,7 @@ def is_close(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) def test_scalar_theta(): - """Test passing a graphblas Scalar as 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 @@ -192,20 +192,224 @@ def inner(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert C.to_coo()[2][0] == 26 # (3+5)*2 + 10 = 26 -def test_pickle_registered(): - def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) - return x + y + theta +# 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 - indexbinary.register_new("pickle_test_op", add_theta) - 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 +def _ibo_return_x(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + - delattr(indexbinary, "pickle_test_op") +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_bind_raw_udt_theta_without_dtype_errors(): + """A raw UDT theta with no dtype can't be inferred; the error is clear and actionable.""" + indexbinary.register_new("raw_no_dtype_op", _ibo_return_x, is_udt=True) + try: + op = indexbinary.raw_no_dtype_op + with pytest.raises(TypeError, match="Cannot infer a dtype for theta"): + op(np.array([1, 2, 3], dtype=np.int64)) + with pytest.raises(TypeError, match="Cannot infer a dtype for theta"): + op((1, 2, 3)) + finally: + delattr(indexbinary, "raw_no_dtype_op") + + +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 resource + import sys + + 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(): @@ -278,6 +482,211 @@ def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert bound.opclass == "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 test_dir_and_module(): assert "register_new" in dir(indexbinary) assert "register_anonymous" in dir(indexbinary) diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index b5d00337d..4befea2cf 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -1886,7 +1886,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) @@ -2837,7 +2837,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): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 8a072a56b..300753bc8 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1358,6 +1358,1195 @@ def _this_or_that(val, idx, _, thunk): # pragma: no cover (numba) assert result.isequal(w) +@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) + + +@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) + + +@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) + + 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_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) + + 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_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] + + +@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] + + +@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, 0.1), (0, 0.1), (0, 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() + assert result.isequal(_record_expected(record_udt, 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_utils 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 + arr_dtype = np.dtype((np.float64, (7,))) + arr_udt = dtypes.register_anonymous(arr_dtype, "Vec7") + lib.GrB_Type_get_String(arr_udt._carg, buf, lib.GxB_JIT_C_DEFINITION) + defn = ffi.string(buf).decode() + assert "double v [7]" in defn + assert "Vec7" 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_utils 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}_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_JitOpTest" in defn + assert "jp" in defn + + # Array UDT JIT definitions + arr_dtype = np.dtype((np.float64, (5,))) + arr_udt = dtypes.register_anonymous(arr_dtype, "Vec5Jit") + 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_Vec5Jit" in defn + assert "v[0]" in defn + assert "v[4]" 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_Vec5Jit" 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_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 +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 +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") +@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 +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") +@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)) @@ -1468,6 +2657,76 @@ def test_is_idempotent(): 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: @@ -1477,3 +2736,47 @@ def test_ops_have_ss(): for mod in modules: with pytest.raises(AttributeError): mod.ss + + +@pytest.mark.skipif("not supports_udfs") +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_utils 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) diff --git a/graphblas/tests/test_pickle.py b/graphblas/tests/test_pickle.py index 724f43d76..10431ed06 100644 --- a/graphblas/tests/test_pickle.py +++ b/graphblas/tests/test_pickle.py @@ -1,3 +1,4 @@ +import multiprocessing import pickle from pathlib import Path @@ -338,3 +339,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 dbfb301a3..da9205b53 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -198,6 +198,21 @@ def test_isclose(): assert Scalar.from_value(None, dtype="FP64").isclose(Scalar.from_value(None, dtype="FP32")) +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_nvals(s): assert s.nvals == 1 s.clear() @@ -300,16 +315,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 diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index ed6f429e7..1d6165d42 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1,7 +1,5 @@ +import contextlib import os -import pathlib -import re -import subprocess import sysconfig import numpy as np @@ -9,7 +7,18 @@ from numpy.testing import assert_array_equal import graphblas as gb -from graphblas import backend, binary, dtypes, indexbinary, 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.ss import _IS_SSGB7 @@ -27,79 +36,17 @@ pytest.skip("not suitesparse backend", allow_module_level=True) -def _fix_jit_config(): - """Fix the GraphBLAS JIT configuration for the current conda environment. +# 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) - The graphblas C library bakes in build-time compiler paths from conda-build, - which don't exist in the user's environment. This function: - 1. Replaces the compiler path with the equivalent from $CONDA_PREFIX/bin/ - 2. Replaces -isysroot with the local macOS SDK (via xcrun), or strips it on Linux - 3. Strips -fdebug-prefix-map flags referencing build paths - - Returns - ------- - True: JIT configured and verified working - False: JIT configuration attempted but compilation failed (don't retry) - None: No conda environment; caller should try a different approach - - Only modifies jit_c_compiler_name, jit_c_compiler_flags, and jit_c_control. - Linker flags and libraries are left at their defaults (already have correct - $CONDA_PREFIX paths substituted by the graphblas C library). - """ - conda_prefix = os.environ.get("CONDA_PREFIX", "") - if not conda_prefix: - return None # No conda env; caller should try sysconfig instead - - # Check if the default compiler already works (e.g., /usr/bin/cc on macOS). - # Only replace it if the baked-in path doesn't exist. - jit_cc = gb.ss.config["jit_c_compiler_name"] - if pathlib.Path(jit_cc).exists(): - # Default compiler exists — don't replace it, just fix flags - pass - else: - # Replace build-time path with local conda equivalent. - cc_basename = pathlib.Path(jit_cc).name - bin_dir = pathlib.Path(conda_prefix) / "bin" - for candidate in [cc_basename, "cc", "clang", "gcc"]: - local_cc = bin_dir / candidate - if local_cc.exists(): - break - else: - return False - gb.ss.config["jit_c_compiler_name"] = str(local_cc) - - # Fix compiler flags: fix build-time-only paths that don't exist in the - # user's environment - flags = gb.ss.config["jit_c_compiler_flags"] - # -isysroot : macOS SDK path from conda-build (e.g., /opt/conda-sdks/MacOSX10.13.sdk). - # The conda cross-compiler needs an explicit sysroot. Replace with the local SDK if available. - 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 available (linux, or macOS without Xcode CLT) - flags = re.sub(r"-isysroot\s+\S+", "", flags) - # -fdebug-prefix-map==: debug path remapping from conda-build - flags = re.sub(r"-fdebug-prefix-map=\S+", "", flags) - gb.ss.config["jit_c_compiler_flags"] = flags - - gb.ss.config["jit_c_control"] = "on" - - # Verify the JIT actually works by attempting a trivial compilation. - # If it fails (e.g., missing libraries, wrong flags), turn JIT off. - try: - from graphblas import dtypes as _dtypes - - _dtypes.ss.register_new("_jit_probe", "typedef struct { int _probe ; } _jit_probe ;") - except Exception: - gb.ss.config["jit_c_control"] = "off" - return False - else: - return True +# Capture the post-import JIT state before the autouse ``_setup_jit`` fixture +# runs. ``_auto_fix_jit_at_import`` probes once; ``jit_c_control`` is ``'on'`` +# iff the env can actually JIT-compile. Tests that assert on the import-time +# state read this constant; the fixture may transiently mutate the live config. +_JIT_WORKS_AT_IMPORT = ( + not _IS_SSGB7 and backend == "suitesparse" and gb.ss.config["jit_c_control"] == "on" +) @pytest.fixture(scope="module", autouse=True) @@ -109,7 +56,9 @@ def _setup_jit(): Strategy: 1. _fix_jit_config(): fix conda-baked compiler paths and probe. - Returns True: JIT works, proceed. - - Returns False: probe failed, JIT is broken, turn off. + - Returns False: probe failed. SuiteSparse will have left + ``jit_c_control = 'load'`` (compile disabled, cache loading + still allowed); we leave that state untouched. - Returns None: no conda env, try sysconfig instead. 2. Sysconfig fallback: for non-conda installs (pure pip). """ @@ -121,15 +70,12 @@ def _setup_jit(): prev = gb.ss.config["jit_c_control"] result = _fix_jit_config() - if result is True: - pass # Conda JIT configured and verified - elif result is False: - # Probe failed — JIT doesn't work with this psg build. - # Don't try sysconfig; if the conda compiler can't compile - # GraphBLAS JIT kernels, Python's sysconfig compiler won't either. - gb.ss.config["jit_c_control"] = "off" - else: - # No conda env (result is None). Try sysconfig for non-conda installs. + # ``True``: conda JIT configured and verified. ``False``: probe failed + # (``_probe_jit`` leaves ``jit_c_control`` at ``'load'``; 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") @@ -140,8 +86,6 @@ def _setup_jit(): gb.ss.config["jit_c_compiler_flags"] = f"{cflags} -I{include}" if libs: gb.ss.config["jit_c_libraries"] = libs - else: - gb.ss.config["jit_c_control"] = "off" try: yield @@ -149,6 +93,94 @@ def _setup_jit(): gb.ss.config["jit_c_control"] = prev +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 + + +@pytest.mark.skipif("_IS_SSGB7") +def test_auto_fix_jit_at_import_left_compiler_usable(): + """After ``import graphblas.ss``, a probe-confirmed compiler implies + ``jit_c_control == 'on'``. + """ + 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 import-time 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 (``'load'`` + # or ``'run'``) 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_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 + # tries to register a one-shot ``_jit_probe`` UDT and previously + # interpreted 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]) @@ -162,8 +194,7 @@ def test_jit_udt(): "myquaternion", "typedef struct { float x [4][4] ; int color ; } myquaternion ;" ) return - if gb.ss.config["jit_c_control"] == "off": - pytest.skip("JIT not available (no C compiler configured)") + _require_jit_on() with burble(): dtype = dtypes.ss.register_new( "myquaternion", "typedef struct { float x [4][4] ; int color ; } myquaternion ;" @@ -210,15 +241,14 @@ def test_jit_unary(v): 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": - pytest.skip("JIT not available (no C compiler configured)") + _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") @@ -242,14 +272,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": - pytest.skip("JIT not available (no C compiler configured)") + _require_jit_on() with burble(): absdiff = binary.ss.register_new( "absdiff", @@ -264,7 +294,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") @@ -307,6 +337,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) " @@ -316,8 +347,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": - pytest.skip("JIT not available (no C compiler configured)") + _require_jit_on() with burble(): diffy = indexunary.ss.register_new("diffy", cdef, "FP64", "FP64", "FP64") assert not hasattr(indexunary, "diffy") @@ -328,7 +358,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") @@ -371,6 +401,7 @@ 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 = ( @@ -378,8 +409,7 @@ def test_jit_indexbinary(v): "double *y, GrB_Index iy, GrB_Index jy, double *theta) " "{ (*z) = (*x) + (*y) + (*theta) ; }" ) - if gb.ss.config["jit_c_control"] == "off": - pytest.skip("JIT not available (no C compiler configured)") + _require_jit_on() with burble(): add_theta = indexbinary.ss.register_new("add_theta", cdef, "FP64", "FP64", "FP64", "FP64") assert not hasattr(indexbinary, "add_theta") @@ -434,9 +464,11 @@ def test_jit_indexbinary(v): assert ("FP32", "FP64") not in add_theta +@pytest.mark.slow def test_jit_select(v): cdef = ( - # Why does this one insist on `const` for `x` argument? + # SelectOps don't write to their input array, so SuiteSparse requires + # the x argument to be ``const``. "void woot (bool *z, const int32_t *x, GrB_Index i, GrB_Index j, int32_t *y) " "{ (*z) = ((*x) + i + j == (*y)) ; }" ) @@ -444,8 +476,7 @@ def test_jit_select(v): with pytest.raises(RuntimeError, match="JIT was added"): select.ss.register_new("woot", cdef, "INT32", "INT32") return - if gb.ss.config["jit_c_control"] == "off": - pytest.skip("JIT not available (no C compiler configured)") + _require_jit_on() with burble(): woot = select.ss.register_new("woot", cdef, "INT32", "INT32") assert not hasattr(select, "woot") @@ -456,7 +487,7 @@ def test_jit_select(v): assert woot.types == {(dtypes.INT32, dtypes.INT32): dtypes.BOOL} assert "INT32" in woot assert woot["INT32"].return_type == dtypes.BOOL - # 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="woot does not work with .INT64, INT64. types"): v << woot(v, 1) v = v.dup("INT32") @@ -505,6 +536,872 @@ def test_jit_select(v): assert ("INT32", "INT64") not in woot +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("_IS_SSGB7") +def test_udt_jit_c_source_introspection(): + """``jit_c_source`` and ``jit_c_name`` should expose what SS sees. + + Covers record UDTs, array UDTs, and ops where no JIT definition was set + (built-in scalar ops, mixed UDT+scalar binary ops). Also verifies the + matching ``jit_c_definition`` / ``jit_c_name`` properties on the dtype. + + Uses field names unique to this test to avoid colliding with other + test UDTs. ``register_anonymous`` shares one DataType per ``np.dtype``, + so two tests with the same dtype share state (including cached JIT C + info) and would race when run in the same session. + """ + record_dtype = np.dtype([("introsp_a", np.int64), ("introsp_b", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_IntrospectUDT") + + # dtype-level introspection + assert udt.jit_c_name == "_IntrospectUDT" + assert "typedef struct" in udt.jit_c_definition + assert "int64_t introsp_a" in udt.jit_c_definition + assert "double introsp_b" in udt.jit_c_definition + + # Builtin scalar dtype: no JIT C definition + assert dtypes.INT64.jit_c_name is None + assert dtypes.INT64.jit_c_definition is None + + # Auto-lifted record UDT binary op + plus_udt = binary.plus[udt] + assert plus_udt.jit_c_name == "plus__IntrospectUDT" + src = plus_udt.jit_c_source + assert src is not None + assert "z->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__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 "z->v[0] = (x->v[0]) * (y->v[0])" in src_arr + assert "z->v[10] = (x->v[10]) * (y->v[10])" 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") +@pytest.mark.skipif("_IS_SSGB7") +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 == "_PinnedOrig" + + op = binary.plus[udt] + # Op signature must reference the pinned struct name, not the renamed one. + assert op.jit_c_name == "plus__PinnedOrig" + assert "_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_propagates_nan(): + """``binary.min``/``max`` on a float UDT must propagate NaN like Python/numba. + + Regression: the JIT codegen used to emit ``(a < b ? a : b)``, which + silently swallows NaN to the right-hand side. Python ``min(a, b)`` (and + numba's ``min``) returns ``a`` when neither comparison is true (NaN + involved), so ``min(NaN, 1.0) == NaN`` and ``min(1.0, NaN) == 1.0``. + The fix swaps the ternary to ``(b < a ? b : a)``. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + _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.float64)]), "_NanJitMM" + ) + 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) + u[i] = (1.0 + i, nan if i % 2 else 3.0 + i) + + w = v.ewise_mult(u, binary.min).new() + assert np.isnan(w[0].new().value[0]) # min(NaN, 1.0) -> NaN + assert w[1].new().value[1] == 3.0 # min(3.0, NaN) -> 3.0 (NaN swallowed) + assert w[2].new().value[1] == 4.0 # min(4.0, 5.0) -> 4.0 (normal case) + + w = v.ewise_mult(u, binary.max).new() + assert np.isnan(w[0].new().value[0]) # max(NaN, 1.0) -> NaN + assert w[1].new().value[1] == 3.0 # max(3.0, NaN) -> 3.0 (NaN swallowed) + assert w[2].new().value[1] == 5.0 # max(4.0, 5.0) -> 5.0 (normal case) + + +@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 _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + _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 _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + _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.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 _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + _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") +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 _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, 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_skips_jit(): + """Mixed UDT+scalar arithmetic and eq/ne stay on the cfunc path. + + The JIT codegen helpers (``set_jit_c_on_op`` and + ``set_jit_c_comparison_on_op``) only run when ``dtype == dtype2``. A + mixed pair would need a per-side type in the C signature plus one C + name per (op, left_type, right_type) triple; for now we keep the + cfunc fallback for broadcast and skip JIT. 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 _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + 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 pairs do not JIT (cfunc handles broadcast correctly; this is + # the documented limit, not an oversight). + assert binary.plus[record, dtypes.INT64].jit_c_source is None + assert binary.plus[dtypes.INT64, record].jit_c_source is None + assert binary.eq[record, dtypes.INT64].jit_c_source is None + assert binary.ne[record, dtypes.FP64].jit_c_source is None + + def test_context_importable(): if _IS_SSGB7: with pytest.raises(ImportError, match="Context was added"): diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 65460fa75..8c55baffd 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1600,7 +1600,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): diff --git a/pyproject.toml b/pyproject.toml index ef2d11b64..18c969706 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,12 +94,7 @@ 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.15'", # make optional where numba is not supported @@ -181,6 +176,25 @@ filterwarnings = [ # 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] @@ -196,11 +210,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" @@ -254,7 +264,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) @@ -299,7 +309,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", # Pattern passed to `match=` contains metacharacters but is neither escaped nor raw + "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 @@ -309,13 +319,13 @@ 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 + "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 @@ -328,7 +338,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", # This suppression comment is invalid because... (Note: `fmt` is for black, not ruff) + "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,7 +358,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) @@ -362,6 +372,7 @@ 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_utils.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 @@ -371,21 +382,10 @@ ignore = [ "NPY002", # numba doesn't support rng generator yet ] "graphblas/core/utils.py" = ["PLE0302"] # `__set__` is used as a property -"scripts/ci_pick_versions.py" = [ - "S311", -] # random.choice is intentional for CI version selection +"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__`) diff --git a/scripts/ci_pick_versions.py b/scripts/ci_pick_versions.py index 349791396..427a69be5 100755 --- a/scripts/ci_pick_versions.py +++ b/scripts/ci_pick_versions.py @@ -1,8 +1,8 @@ #!/usr/bin/env python -"""Pick random dependency versions for CI testing. +"""Pick random, compatible dependency versions for python-graphblas CI. -Randomly selects compatible dependency versions for python-graphblas CI. -Replaces the bash-based version selection in test_and_build.yml. +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)" @@ -85,8 +85,10 @@ "3.14": "NA", } -# PSG versions by numpy branch and source type. -# conda-forge uses "=" prefix, wheel/source use "==". +# 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": [ @@ -106,20 +108,6 @@ "3.11": ["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"], }, - "source": { - "3.11": [ - "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", - ], - "3.12": ["8.2.0.1", "8.2.1.0"], - }, } PSG_VERSIONS_NP2 = { 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 index b1166171f..3c94f19bd 100755 --- a/scripts/jit_diagnostics.py +++ b/scripts/jit_diagnostics.py @@ -2,7 +2,6 @@ """Print JIT diagnostic information for debugging compiler configuration. Run with: python scripts/jit_diagnostics.py -Or during CI: python -m scripts.jit_diagnostics (from repo root) This prints the GraphBLAS JIT configuration, available compilers, sysconfig values, and attempts a JIT compilation to verify everything works. @@ -66,7 +65,7 @@ def main(): from graphblas.core.ss import _IS_SSGB7 if _IS_SSGB7: - print("\nSuiteSparse:GraphBLAS 7.x — JIT not available") + print("\nSuiteSparse:GraphBLAS 7.x: JIT not available") return 0 print("\n--- GraphBLAS JIT defaults (from compiled C library) ---") @@ -111,10 +110,8 @@ def main(): print(f" -fdebug-prefix-map: {d}") # Try the fix - print("\n--- Attempting _fix_jit_config ---") - from graphblas.tests.test_ssjit import _fix_jit_config - - result = _fix_jit_config() + 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})") From fb687dd61be352b4c7f2c45a9e9359b73e0da193 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 5 Aug 2026 12:12:44 -0500 Subject: [PATCH 05/52] Compare the pypy check against true in the dependency-selection step (#580) * Ask scipy's mmread for a sparse array scipy 1.18 deprecated the sparse-matrix return of `mmread` and 1.20 flips the default, so the call first warns and then changes type under us. We read only `.shape`, `.row`, `.col`, and `.data`, which sparse matrices and sparse arrays both provide. Pass `spmatrix=False` when scipy accepts it, detected from the signature rather than pinned to a version, and only when the caller has not already chosen a value. `test_mmread_mmwrite` is the existing coverage for this path, but it skips whenever scipy's own example files are absent, which happens on some conda-forge builds. The new test writes its own Matrix Market input so it runs there too. * Compare the pypy check against true in the dependency-selection step `if [[ ${{ contains(steps.pyver.outputs.selected, 'pypy') }} ]]` substitutes to `if [[ false ]]` on CPython jobs. In bash that tests whether the string "false" is non-empty, which is always true, so the pypy branch was taken on every job and cleared awkward, numpy, and sparse for CPython runs as well. Compare against `true` explicitly, matching the form the surrounding condition in the same step already uses. * Skip pysentry on pre-commit.ci pre-commit.ci runs hooks in a sandbox without network access, and pysentry downloads the PyPA/OSV advisory databases at run time, so every pre-commit.ci run failed with a download error. Skipping it there keeps the audit in local runs, where the network exists. * Keep the CI awkward pick compatible with the numpy pick * Skip the IndexBinaryOp memory test where the resource module is missing * Keep the CI networkx pick compatible with the numpy, scipy, and pandas picks * Guard the unary JIT string setter behind _has_jit_set * Adapt tests to SuiteSparse builds older than 9 * Gate the remaining JIT introspection tests on SuiteSparse 9 --- .github/workflows/test_and_build.yml | 5 +- .pre-commit-config.yaml | 4 +- graphblas/core/operator/binary.py | 2 +- graphblas/core/operator/unary.py | 15 +++--- graphblas/io/_matrixmarket.py | 15 ++++++ graphblas/tests/test_dtype.py | 2 + graphblas/tests/test_indexbinary.py | 4 +- graphblas/tests/test_io.py | 37 ++++++++++++++ graphblas/tests/test_op.py | 16 +++++++ graphblas/tests/test_ssjit.py | 34 ++++++++----- graphblas/tests/test_vector.py | 10 +++- scripts/check_versions.sh | 2 +- scripts/ci_pick_versions.py | 72 +++++++++++++++++++++++++++- 13 files changed, 191 insertions(+), 27 deletions(-) diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index 8e435cd4a..9ece8cfed 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -189,7 +189,10 @@ jobs: 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 npver="" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ac60ef29a..3582142d2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,7 +10,9 @@ 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 can only run locally. + skip: [pylint, no-commit-to-branch, pysentry] fail_fast: false default_language_version: python: python3 diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index b3fac1cb9..97e0a9e70 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -653,7 +653,7 @@ def _compile_udt(self, dtype, dtype2): # Set JIT C definition for auto-generated built-in UDT ops. Only # same-type pairs apply: mixed UDT+scalar ops don't have a # meaningful single-type C definition. ``_has_jit_set`` gates access - # to ``lib.GrB_BinaryOp_set_String``, which is absent on SS < 8. + # to ``lib.GrB_BinaryOp_set_String``, which is absent on SS < 9. is_jittable = self.name in _BUILTIN_UDT_BINARY_OPS or self.name in ("eq", "ne") if _has_numba and _has_jit_set and dtype == dtype2 and is_jittable: if self.name in _BUILTIN_UDT_BINARY_OPS: diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index a88cba8a6..22822921a 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -122,6 +122,7 @@ def _one(x): if _has_numba: from .udt_utils import BUILTIN_UDT_UNARY_OPS as _BUILTIN_UDT_UNARY_OPS from .udt_utils import ( + _has_jit_set, _maybe_warn_jit_skipped, compile_udt_unary_wrapper, set_jit_c_on_op, @@ -264,12 +265,14 @@ def _compile_udt(self, dtype, dtype2): op = _finalize_udt_op( self, dtype, None, ret_type, unary_wrapper, wrapper_sig, TypedUserUnaryOp ) - # ``set_jit_c_on_op`` is a no-op (returns None) when - # ``_has_jit_set`` is False (SS < 8). - 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) + # ``_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}") diff --git a/graphblas/io/_matrixmarket.py b/graphblas/io/_matrixmarket.py index eff8c9031..cf425e2e9 100644 --- a/graphblas/io/_matrixmarket.py +++ b/graphblas/io/_matrixmarket.py @@ -42,6 +42,7 @@ 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. " @@ -67,10 +68,24 @@ def mmread(source, engine="auto", *, dup_op=None, name=None, **kwargs): "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 diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 1866806a6..3a5026b6a 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -419,6 +419,8 @@ 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 diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 493ac5ea9..c56bb86f9 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -354,9 +354,11 @@ def test_bound_ibo_memory_bounded(): growth should stay below a few MB after GC drops them. """ import gc - import resource 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) diff --git a/graphblas/tests/test_io.py b/graphblas/tests/test_io.py index 9443bbcfa..c332820dd 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 @@ -304,6 +305,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]))) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 300753bc8..ef5227f07 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -28,6 +28,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, @@ -1973,6 +1974,15 @@ def test_udt_op_compilation_is_lazy(): @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) @@ -2069,6 +2079,10 @@ def test_udt_eq_packed_mixed_width(): @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( @@ -2303,6 +2317,8 @@ def test_udt_eq_ne_scalar_broadcast_array_2d(): @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( diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 1d6165d42..fc174c768 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -21,6 +21,7 @@ ) from graphblas.core import _supports_udfs as supports_udfs from graphblas.core.operator.indexbinary import _has_idxbinop +from graphblas.core.operator.udt_utils import _has_jit_set from graphblas.core.ss import _IS_SSGB7 from .conftest import autocompute, burble @@ -537,7 +538,9 @@ def test_jit_select(v): @pytest.mark.skipif("not supports_udfs") -@pytest.mark.skipif("_IS_SSGB7") +# The JIT string setters (GrB_Type_set_String with GxB_JIT_C_NAME) arrived in +# SS 9, so every jit_c_* property is None on 7.x and 8.x alike. +@pytest.mark.skipif("not _has_jit_set") def test_udt_jit_c_source_introspection(): """``jit_c_source`` and ``jit_c_name`` should expose what SS sees. @@ -606,7 +609,8 @@ def test_udt_jit_c_source_introspection(): @pytest.mark.skipif("not supports_udfs") -@pytest.mark.skipif("_IS_SSGB7") +# 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 @@ -958,8 +962,8 @@ def test_anonymous_udt_with_no_name_still_jits(): 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 _IS_SSGB7: - pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + 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) @@ -1001,8 +1005,8 @@ def test_complex_field_udt_jits_arithmetic(): without extra includes. ``abs`` uses ``cabs`` / ``cabsf`` which match Numba's ``abs(complex)`` behavior (real magnitude in the real part). """ - if _IS_SSGB7: - pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + 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)]) @@ -1113,8 +1117,8 @@ def test_nested_record_udt_jits(): the nested record fields (Numba can't ``setitem`` a tuple to a record field, but leaf scalar writes work). """ - if _IS_SSGB7: - pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + 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 @@ -1167,6 +1171,10 @@ def test_nested_record_udt_jits(): @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. @@ -1306,8 +1314,8 @@ def test_eq_ne_udt_jit_matches_cfunc(shape): record, complex, array UDT), and that NaN propagation matches IEEE (``eq(NaN, NaN) == False``) on variants flagged ``nan_at_zero``. """ - if _IS_SSGB7: - pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + 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") @@ -1385,8 +1393,10 @@ def test_udt_scalar_broadcast_skips_jit(): routing so a future change doesn't accidentally regress it (either direction is fine, but it should be a conscious decision). """ - if _IS_SSGB7: - pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + 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", diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 8c55baffd..52c382742 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -12,6 +12,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, @@ -2220,8 +2221,13 @@ def test_udt(): if suitesparse: vv = Vector.ss.deserialize(v.ss.serialize(), dtype=long_udt) assert v.isequal(vv, check_dtype=True) - with pytest.raises(SyntaxError): - # The dtype name is too long to embed in the blob; dtype= must be provided + # 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) diff --git a/scripts/check_versions.sh b/scripts/check_versions.sh index afcf3e0ac..3977bd8a0 100755 --- a/scripts/check_versions.sh +++ b/scripts/check_versions.sh @@ -12,7 +12,7 @@ conda search 'numpy[channel=conda-forge]>=2.4' conda search 'pandas[channel=conda-forge]>=3.0' conda search 'scipy[channel=conda-forge]>=1.17' conda search 'networkx[channel=conda-forge]>=3.6' -conda search 'awkward[channel=conda-forge]>=2.9' +conda search 'awkward[channel=conda-forge]>=2.11' conda search 'sparse[channel=conda-forge]>=0.15' # fast_matrix_market is deprecated (no longer maintained; last supported Python is 3.12) conda search 'numba[channel=conda-forge]>=0.64' diff --git a/scripts/ci_pick_versions.py b/scripts/ci_pick_versions.py index 427a69be5..68a435165 100755 --- a/scripts/ci_pick_versions.py +++ b/scripts/ci_pick_versions.py @@ -50,12 +50,15 @@ "2.x": ["2.2", "2.3", "3.0", ""], } +# The "2.x" pool starts at 2.6 because awkward <2.6 still uses numpy.AxisError, +# which numpy 2.0 removed. 2.10 and 2.11 are the only pins that 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.x": ["2.6", "2.7", "2.8", "2.9", "2.10", "2.11", ""], } NUMBA_VERSIONS = { @@ -138,7 +141,11 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): 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. + 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. """ # --- scipy / numpy constraints --- @@ -184,6 +191,15 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): if v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 15): v["scipy"] = random.choice(["1.15", "1.16", "1.17", ""]) + # --- 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; awkward <2.8 has no py3.14 builds @@ -192,6 +208,35 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): elif 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 and 3.6 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-08-04 conda-forge networkx 3.6.1 constrains 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)]) + # --- numba constraints --- # numba minimum by Python version: 0.59 for 3.12, 0.61 for 3.13, 0.63 for 3.14 @@ -433,6 +478,29 @@ def validate(v, pyver): if pyver == "3.13" and v["awkward"] == "2.6": errors.append("awkward 2.6 has no py3.13 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") + + # 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 = {"3.11": (0, 57), "3.12": (0, 59), "3.13": (0, 61), "3.14": (0, 63)} if v["numba"] not in ("", "NA") and _ver(v["numba"]) < numba_min[pyver]: From bc2e3e70b96d3a820a8144be959ee25d6af7efc2 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 5 Aug 2026 12:15:51 -0500 Subject: [PATCH 06/52] Move the docs build to ubuntu-24.04 and miniforge3 (#581) * Move the docs build to ubuntu-24.04 and miniforge3 Read the Docs no longer offers the ubuntu-20.04 build image, so the docs build fails before it reaches sphinx. `build.os` now accepts ubuntu-22.04, ubuntu-24.04, ubuntu-26.04, and ubuntu-lts-latest. `mambaforge-4.10` is still accepted but is marked deprecated, and the mambaforge series is superseded by miniforge3, whose conda solves with libmamba anyway. The comment the old pin carried (mamba makes builds faster) no longer describes a difference. --- .readthedocs.yml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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 From cba3b44ac2d2df8d6c01138f7da6ee8c191fd2fe Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 5 Aug 2026 12:17:20 -0500 Subject: [PATCH 07/52] Register the markers scipy's test modules declare (#582) * Register the markers scipy's test modules declare `test_mmread_mmwrite` imports `scipy.io.tests.test_mmio` for its example matrices. That module declares two markers this project does not define: `thread_unsafe` at module scope, and `fail_slow(10)` on one test. Importing the module evaluates both. We run with `--strict-markers`, so each unregistered name fails collection with "'' not found in `markers` configuration option". Measured on pytest 9.1.1 and on 8.4.2; the route is the same on both, because the strict flag is what decides it. Without `--strict-markers` the same import instead raises PytestUnknownMarkWarning, which `filterwarnings = error` would turn into a failure as well, so the entries are needed either way. This is why CI fails where a local run may not: some conda-forge builds omit `scipy/io/tests` entirely, which the test already notes and skips on. Against a scipy that does ship them, main fails both parametrizations of `test_mmread_mmwrite`. Register both as unused-but-declared, and note at the import site why the entries exist. --- graphblas/tests/test_io.py | 3 +++ pyproject.toml | 14 +++++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/graphblas/tests/test_io.py b/graphblas/tests/test_io.py index c332820dd..a8d80e9ca 100644 --- a/graphblas/tests/test_io.py +++ b/graphblas/tests/test_io.py @@ -166,6 +166,9 @@ 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 diff --git a/pyproject.toml b/pyproject.toml index 18c969706..9fe272c98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -156,7 +156,19 @@ addopts = [ "--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 From d9b1a66203e14323f13a2c5b3bff38241a119caa Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 5 Aug 2026 12:21:15 -0500 Subject: [PATCH 08/52] Zero a UDT scalar's alignment padding instead of leaving it uninitialized (#584) * Zero a UDT scalar's alignment padding instead of leaving it uninitialized Setting a record UDT scalar builds a numpy buffer with `np.empty` and fills it with `arr[:] = val`, which writes the fields and skips the padding between them. `tobytes()` then copies the whole buffer into the GrB_Scalar, so uninitialized process memory becomes part of the stored value and reads back out of `Scalar.value`. For `[("a", int8), ("b", float64)]` under C alignment that is 7 of the 16 bytes. Allocate with `np.zeros` instead. The regression test dirties a same-sized buffer before allocating, because numpy serves small allocations from a cache of recently freed blocks and an otherwise-fresh page reads back as zeros whether or not the bug is present. Reverting the fix makes it fail on 8 runs out of 8. * Zero UDT padding explicitly instead of relying on numpy's copy strategy --- graphblas/core/scalar.py | 74 +++++++++++++++++++++++++++++++--- graphblas/tests/test_scalar.py | 35 ++++++++++++++++ 2 files changed, 103 insertions(+), 6 deletions(-) diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 0a4acbbb8..8dd712ba8 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -1,4 +1,5 @@ import itertools +from functools import lru_cache import numpy as np @@ -367,17 +368,18 @@ def value(self, val): val = val.value # raise below if wrong type (as determined by cffi) if self.dtype._is_udt: np_type = self.dtype.np_type + # Zeros, not empty: when numpy writes only the fields, a + # record's alignment padding is whatever the allocation + # happened to contain. ``_udt_bytes`` zeros the padding in the + # other direction, when numpy copies whole items instead. if np_type.subdtype is None: - arr = np.empty(1, dtype=np_type) + arr = np.zeros(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 = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) arr[:] = val - # 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 the GrB_Scalar storage. - self.gb_obj[0 : self.dtype.np_type.itemsize] = arr.tobytes() + self.gb_obj[0 : np_type.itemsize] = _udt_bytes(np_type, arr) else: self.gb_obj[0] = val self._empty = False @@ -1156,6 +1158,66 @@ def _dict_to_record(np_type, d): return tuple(rv) +@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) + + _MATERIALIZE = Scalar.from_value(lib.GrB_MATERIALIZE, is_cscalar=True, name="GrB_MATERIALIZE") _COMPLETE = Scalar.from_value(lib.GrB_COMPLETE, is_cscalar=True, name="GrB_COMPLETE") diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index da9205b53..aa11d1fe5 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -213,6 +213,41 @@ def test_isclose_udt_raises(): 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() + + def test_nvals(s): assert s.nvals == 1 s.clear() From ec6692967d87bf75e36a5b99e9a00b991abba7f0 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 5 Aug 2026 13:03:12 -0500 Subject: [PATCH 09/52] Read a nested array UDT scalar at its full extent (#585) * Read a nested array UDT scalar at its full extent numpy keeps subarray dtypes layered instead of flattening them, so a UDT holding a 6-long array of 5-long FP64 arrays reports `subdtype` as a 40-byte inner dtype with shape `(6,)` rather than FP64 with shape `(6, 5)`. Viewing all 240 bytes of the scalar as that inner dtype raised ValueError: Changing the dtype to a subarray type is only supported if the total itemsize is unchanged so reading `.value` on such a scalar failed outright. Walk the `subdtype` chain to the base element type, accumulating the shape, and view once at the bottom. The test asserts the values rather than only that the read succeeds, because a wrong collapse yields a correctly-shaped but transposed array. * Skip the nested array UDT roundtrip on SuiteSparse before 9 --- graphblas/core/scalar.py | 11 ++++++++++- graphblas/tests/test_scalar.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 8dd712ba8..5b5e56299 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -356,7 +356,16 @@ def value(self): 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]) + # Collapse nested subarray dtypes before viewing. numpy keeps them + # layered, so a 6-long array of 5-long FP64 arrays has + # ``subdtype == (dtype((' Date: Wed, 5 Aug 2026 13:05:00 -0500 Subject: [PATCH 10/52] Fix two numpy monoid identities and guard the table against drift (#586) * Fix monoid.numpy.fmax identity for INT16/INT32/INT64 (mapnumpy=False) With config["mapnumpy"]=False, monoid.numpy.fmax for INT16/INT32/INT64 used int8's minimum (-128) as its identity instead of the type's own minimum. Because SuiteSparse seeds a non-empty max-style reduction with the monoid identity, an fmax reduction over values below -128 returned -128 instead of the true maximum. Example (mapnumpy=False): fmax-reduce of [-1000, -2000, -3000] over INT16 returned -128; it now returns -1000. The identities now match the builtin maximum monoid byte-for-byte. INT8 (-128), the unsigned types (0), and the entire fmin block are unchanged; _fmin_is_float (imported by semiring/numpy.py) is untouched. Pre-existing since a22be711 (Oct 2022); not a regression, and only reachable with mapnumpy=False (with mapnumpy=True fmax resolves to the builtin max monoid). Add a regression test that reproduces the bug end-to-end, guarded to skip (not error) on the numba #8478 combo where the integer identities are absent. * Fix monoid.numpy.logical_or[BOOL] identity and add a table-wide guard test monoid.numpy.logical_or[BOOL] used True as its identity, but False is OR's identity (True is its annihilator). With config["mapnumpy"]=False the numpy monoid registers a fresh UDF whose .identity was therefore wrong; with mapnumpy=True it resolves to the builtin lor (identity False), so the flag silently flipped a public .identity value. SuiteSparse does not seed OR reductions with the identity, so no reduction returned a wrong result; the defect is wrong .identity metadata and mapnumpy non-transparency. Add test_numpy_monoid_identity_matches_builtin: for every op in _numpy_to_graphblas, the numpy monoid's identity must equal the mapped builtin monoid's identity, compared cast-to-dtype (so benign compact literals like bitwise_and's -1 == band's 255 do not trip it). It reads _monoid_identities directly (which register_new consumes verbatim on the mapnumpy=False path), so it mutates no operator cache. This one sweep would have caught both this and the fmax bug in the previous commit. Also harden the fmax regression test teardown to save/restore the numpy monoid and binary cache entries rather than pop them, so it cannot leave a UDF binary op cached (which desynced monoid / binary / op.numpy under --mapnumpy --runslow). * Guard the unmapped numpy monoid identities against silent drift test_numpy_monoid_identity_matches_builtin pins only the 11 numpy monoids in _numpy_to_graphblas, each against its builtin. The 6 unmapped monoids (gcd, hypot, logaddexp, logaddexp2, maximum, minimum) have no builtin counterpart, so nothing tested their identities. maximum's identity table is byte-identical to fmax's and minimum's to fmin's, and both fmax and fmin are covered by that test, so test_numpy_monoid_unmapped_identity_consistency asserts that equality over the shared dtype keys: a transitive guard that fails if maximum/minimum ever drift from fmax/fmin. It also pins logaddexp/logaddexp2 to -inf. gcd/hypot get a presence-only check, since 0 is their conventional identity but numpy gcd/hypot are not sign-preserving (a strict f(e,x)==x check would falsely fail on negatives). Reads _monoid_identities directly (no operator-cache mutation, no mapnumpy toggle). No source change; verified to discriminate (a forced maximum != fmax makes it fail). * Clear the binary and op numpy caches in the fmax identity test --- graphblas/monoid/numpy.py | 8 +- graphblas/tests/test_numpyops.py | 133 +++++++++++++++++++++++++++++++ 2 files changed, 137 insertions(+), 4 deletions(-) diff --git a/graphblas/monoid/numpy.py b/graphblas/monoid/numpy.py index ab0a6bdfa..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, @@ -103,11 +103,11 @@ "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/tests/test_numpyops.py b/graphblas/tests/test_numpyops.py index 534312c0b..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 @@ -52,6 +53,138 @@ def func(x): # pragma: no cover (numba) 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)) From 10a06bbac0a3106d84b5568d97268e31b16ea908 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 6 Aug 2026 10:38:54 -0500 Subject: [PATCH 11/52] Keep the CI matplotlib pick compatible with the numpy pick (#621) The overnight matrix produced the stack's first two post-restack failures, both the same one line: the notebooks task on macos-latest and windows-latest picked numpy 1.24 while installing matplotlib unpinned, and matplotlib 3.11.0's runtime check requires numpy >=1.25. Its conda-forge metadata does not carry that floor, so the solver pairs them happily and the notebooks die at import. Same class and same fix shape as the networkx and awkward pick constraints: the version picker now emits mplver, a bare ceiling of <3.11 whenever the numpy pick sits under 1.25 and empty otherwise, the notebooks install line consumes it, and validate() checks the pairing. The --validate stress test passes 160000 combos with zero failures, and both emission cases are confirmed by running the picker. --- .github/workflows/test_and_build.yml | 2 +- scripts/ci_pick_versions.py | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index 9ece8cfed..735b33889 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -222,7 +222,7 @@ jobs: pyyaml${yamlver} ${sparse} pandas${pdver} scipy${spver} numpy${npver} ${awkward} \ 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"' || '' }} \ ${{ contains(steps.pyver.outputs.selected, 'pypy') && 'pypy' || '' }} \ diff --git a/scripts/ci_pick_versions.py b/scripts/ci_pick_versions.py index 68a435165..5b8fe0f7e 100755 --- a/scripts/ci_pick_versions.py +++ b/scripts/ci_pick_versions.py @@ -237,6 +237,19 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): # 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.0's + # runtime check requires numpy >=1.25, but its conda-forge metadata does not, + # so the solver happily pairs matplotlib 3.11 with a numpy 1.24 pin and the + # notebooks then die at `import matplotlib`. Pin below 3.11 whenever the + # numpy pick sits under that floor; empty means latest is fine. + # MAINT: 2026-08-05 confirmed with matplotlib 3.11.0 + numpy 1.24.4 + if v["numpy"] and _ver(v["numpy"]) < (1, 25): + v["matplotlib"] = "<3.11" + else: + v["matplotlib"] = "" + # --- numba constraints --- # numba minimum by Python version: 0.59 for 3.12, 0.61 for 3.13, 0.63 for 3.14 @@ -380,6 +393,7 @@ def _pick_psg(npver, pyver, source_type): "sparse": "sparse", "numba": "numba", "psg": "psg", + "matplotlib": "mpl", } @@ -401,6 +415,8 @@ def format_output(v): # 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) @@ -417,6 +433,7 @@ def format_summary(v): "pyyaml", "sparse", "psg", + "matplotlib", ): name = _SUMMARY_NAMES[key] val = v[key] @@ -482,6 +499,10 @@ def validate(v, pyver): 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) From 978c17bbfa21e71bf969e587aa841ebb1385e710 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 6 Aug 2026 10:41:05 -0500 Subject: [PATCH 12/52] Keep the CI numba pick solvable against numpy 1.25 and 1.26 (#622) * Keep the CI numba pick solvable against numpy 1.25 and 1.26 The overnight matrix hit a conda solver failure on the bottom rung: the picker drew numpy 1.25 with numba 0.57, and conda-forge numba 0.57 requires numpy <1.25, so the environment could not solve and the job died before running anything. The picker's constraints covered numba's numpy 2.x boundaries but not this 1.x-era ceiling. Bump numba 0.57 to 0.58 whenever the numpy pick is 1.25 or 1.26, following the existing 0.62/0.63-to-0.64 pattern, and teach validate() the pairing. The boundary is verified by dry-run solves against conda-forge: 0.57+1.25 conflicts; 0.58+1.25, 0.58+1.26, and 0.59+1.26 all solve. The --validate stress test passes 160000 combos, and 200 fresh draws produce zero bad pairings. * Keep the CI scipy pick solvable against pandas 2.2 and 2.3 The matrix hit another conda solver failure: the picker drew pandas 2.2 with scipy 1.9, reachable only on the py3.11 numpy 1.x path, and the environment could not solve, so the job died before running anything. conda-forge pandas >=2.2 carries run_constrained "scipy >=1.10.0" (the latest 2.2.3 build and all 2.3 builds; pandas 2.0 and 2.1 declare no scipy constraint), and scipy 1.9.3 py311 builds exist even on win-64, so the pairing fails purely on the pandas constraint. Verified with conda search --info against conda-forge. Re-pick scipy from the pool's 1.10 through 1.14 pins whenever a pandas 2.2 or 2.3 pin lands next to a scipy pin below 1.10, and teach validate() the pairing. The replacement stays pinned below 1.15 because scipy 1.9 only survives to this point alongside numpy 1.24 or 1.25, whose scipy ceiling is 1.15. Unpinned pandas is left alone: the solver backs it off below 2.2 on its own, which keeps scipy 1.9 coverage. The --validate stress test passes 160000 combos, 4800 fresh draws produce zero bad pairings, and forcing the old bad pair through apply_constraints now always lands scipy in 1.10 through 1.14 with no validate() errors. --- scripts/ci_pick_versions.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/scripts/ci_pick_versions.py b/scripts/ci_pick_versions.py index 5b8fe0f7e..738ef99a8 100755 --- a/scripts/ci_pick_versions.py +++ b/scripts/ci_pick_versions.py @@ -191,6 +191,21 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): if v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 15): v["scipy"] = random.choice(["1.15", "1.16", "1.17", ""]) + # 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 @@ -264,6 +279,13 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): if _ver(v["numpy"]) >= (2, 4) and v["numba"] in ("0.62", "0.63"): v["numba"] = "0.64" + # conda-forge numba 0.57 requires numpy <1.25, so that pairing fails the + # conda solve outright. Verified by dry-run solves: 0.57+1.25 conflicts; + # 0.58+1.25, 0.58+1.26, and 0.59+1.26 all solve. + # MAINT: 2026-08-05 confirmed against conda-forge with mamba dry-runs + if v["numba"] == "0.57" and np_is_1x and _ver(v["numpy"]) >= (1, 25): + v["numba"] = "0.58" + # numba <0.62 doesn't support numpy 2.x if not np_is_1x and v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 62): v["numba"] = "NA" @@ -489,6 +511,15 @@ def validate(v, pyver): if 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 pyver == "3.14" and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 8): errors.append(f"awkward {v['awkward']} has no py3.14 build") @@ -503,6 +534,10 @@ def validate(v, pyver): 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)") + # conda-forge numba 0.57 requires numpy <1.25 (fails the conda solve) + if v["numba"] == "0.57" and v["numpy"] and _ver(v["numpy"]) >= (1, 25) and np_is_1x: + errors.append(f"numba 0.57 requires numpy <1.25, got {v['numpy']}") + # awkward <2.10 requires numpy <2.5 (generic-unit datetime64("NaT") at import) if ( _ver(v["numpy"]) >= (2, 5) From 26123a72394d5d46da9be0a341a232ecc7a9b72f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 6 Aug 2026 21:48:46 -0500 Subject: [PATCH 13/52] Skip pickle3 load in test_udt on Linux aarch64 with numpy 1.x (#624) pickle3.pkl and pickle3-vanilla.pkl were 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 the flags back into a char and raises ValueError("incorrect value for flags variable (overflow)") when the value does not round-trip; on Linux aarch64 plain char is unsigned per AAPCS64, so (char)-112 == 144 and the load fails deterministically. numpy 2.0 widened descriptor flags to npy_uint64 and normalizes signed pickled values, so it loads the file everywhere. Regenerating the pickles cannot fix this while numpy 1.x remains in the CI pools: a file written with positive flags fails the same round-trip check in the reverse direction on signed-char platforms. So skip only the file-load portion of test_udt on the affected combination; the in-process roundtrip asserts above it still run on all platforms. Verified with python -m pytest graphblas/tests/test_pickle.py::test_udt on macOS x86 (passes, skip does not trigger) and by exercising the skip condition against simulated version/platform/machine combinations. --- graphblas/tests/test_pickle.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/graphblas/tests/test_pickle.py b/graphblas/tests/test_pickle.py index 10431ed06..3659b666c 100644 --- a/graphblas/tests/test_pickle.py +++ b/graphblas/tests/test_pickle.py @@ -1,5 +1,7 @@ import multiprocessing import pickle +import platform +import sys from pathlib import Path import numpy as np @@ -313,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) From 2e03b5490fedad7a5f31f1385c1829f76c53a24b Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 7 Aug 2026 00:09:32 -0500 Subject: [PATCH 14/52] Pick CI's random matrix values with bash instead of a marketplace action (#626) The rngs gate jobs and every build_and_test job resolved ddradar/choose-random-action before starting, and GitHub's action resolution service has been failing that lookup with repeated 5xx errors ("Internal Server Error occurred while resolving ddradar/choose-random-action@v4.1.0", also plain Service Unavailable). A failure in a rngs job cancels everything gated behind it: the whole 11-job matrix in test_and_build.yml and the test_imports job in imports.yml, which triggers on every pull request. Ten of the last ten red Test runs died exactly there before running anything. All six RNG steps (four in test_and_build.yml, two in imports.yml) drew uniformly from small lists, which bash can do by itself: index an array with RANDOM modulo its length. RANDOM is a bash builtin, so it works in the ubuntu default shell, git-bash on Windows, and the bash 3.2 that macOS runners provide, and the jobs no longer download anything to make their picks. Weighting a choice is now expressed by repeating it in the list, and skipping one by removing it, which replaces the action's natural-number weights mechanism (the old comment about setting weights very large to skip an entry is gone with it). Verified: both workflows parse as YAML, the snippet produces uniform draws under bash 3.2.57 (macOS's /bin/bash, with -e -u -o pipefail), and no uses: of the action remain in any workflow. --- .github/workflows/imports.yml | 34 ++++------- .github/workflows/test_and_build.yml | 84 ++++++++++------------------ 2 files changed, 43 insertions(+), 75 deletions(-) diff --git a/.github/workflows/imports.yml b/.github/workflows/imports.yml index 4be4926c7..7f1596c0a 100644 --- a/.github/workflows/imports.yml +++ b/.github/workflows/imports.yml @@ -13,32 +13,22 @@ 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@v4.1.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@v4.1.0 id: pyver - with: - contents: | - 3.11 - 3.12 - 3.13 - 3.14 - weights: | - 1 - 1 - 1 - 1 + run: | + choices=(3.11 3.12 3.13 3.14) + selected=${choices[$((RANDOM % ${#choices[@]}))]} + echo "selected: $selected (from: ${choices[*]})" + echo "selected=$selected" >> "$GITHUB_OUTPUT" test_imports: needs: rngs runs-on: ${{ needs.rngs.outputs.os }} diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index 735b33889..69fab3770 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -49,34 +49,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@v4.1.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@v4.1.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 }} @@ -118,39 +109,26 @@ jobs: 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@v4.1.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.11 - 3.12 - 3.13 - 3.14 - 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) + 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@v4.1.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 id: setup_conda From 9682283e04f2fb6ece6bc2e0a2ab4c038bb14992 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 11 Sep 2026 11:46:45 -0500 Subject: [PATCH 15/52] Adapt to SuiteSparse:GraphBLAS 10.5.0 and read its version from the library (#628) 10.5.0 stopped checking, after a build, whether every value happened to be the same: "GrB_Matrix_build and GrB_Vector_build: no longer do a post-iso check; they leave the matrix in non-iso format even if all the entries are the same." Four tests read that check as if it were python-graphblas's own contract and went red. Each one wanted an iso object rather than an inference, so each now asks for one with ``ss.build_scalar``, which means the same thing on every SuiteSparse version and produces the same storage formats the tests were selecting for. ``from_coo`` is left alone: adding a scan of the values to recover the old behavior would re-add the cost upstream just removed. The version gates had a second, quieter problem. ``graphblas.core.ss`` derived the C library version by parsing ``suitesparse_graphblas.__version__``, which is the version of the Python wrapper. The two normally agree, but a development build between releases reports the previous one: psg 10.4.1.0 against a 10.5.0 library here, so every gate saw 10.4.1 and any 10.5 branch would have been silently skipped. Read the library's own GxB_IMPLEMENTATION_* constants instead, falling back to the old parse where a build does not expose them. test_openmp_enabled keeps its hard assertion and gains a message saying what is wrong and that the fault is in the installed build. It caught a real one: conda-forge's first graphblas 10.5.0 build on osx-arm64 shipped without OpenMP. The build 1 rebuild (2026-08-26) restored OpenMP and the assertion passes again. --- graphblas/core/ss/__init__.py | 14 +++++++- graphblas/tests/test_formatting.py | 7 +++- graphblas/tests/test_io.py | 14 ++++++-- graphblas/tests/test_ss_utils.py | 58 +++++++++++++++++++++++++----- 4 files changed, 80 insertions(+), 13 deletions(-) diff --git a/graphblas/core/ss/__init__.py b/graphblas/core/ss/__init__.py index 79859eb53..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/tests/test_formatting.py b/graphblas/tests/test_formatting.py index a6522dcef..2b412ac1d 100644 --- a/graphblas/tests/test_formatting.py +++ b/graphblas/tests/test_formatting.py @@ -5032,7 +5032,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_io.py b/graphblas/tests/test_io.py index a8d80e9ca..962595aaa 100644 --- a/graphblas/tests/test_io.py +++ b/graphblas/tests/test_io.py @@ -430,18 +430,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) diff --git a/graphblas/tests/test_ss_utils.py b/graphblas/tests/test_ss_utils.py index 40774186c..2c9e01a1a 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -10,6 +10,24 @@ pytest.skip("gb.ss and A.ss only available with suitesparse backend", allow_module_level=True) +# ``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): v0 = Vector(int, 5) @@ -19,9 +37,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 +95,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" @@ -200,8 +229,19 @@ def test_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(): From bb063229965100ae57221c1154dfe24f7afef948 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 17 Sep 2026 21:47:22 -0500 Subject: [PATCH 16/52] Cover SuiteSparse:GraphBLAS 10.4 and 10.5 in the CI psg pool (#629) * Adapt to SuiteSparse:GraphBLAS 10.5.0 and read its version from the library 10.5.0 stopped checking, after a build, whether every value happened to be the same: "GrB_Matrix_build and GrB_Vector_build: no longer do a post-iso check; they leave the matrix in non-iso format even if all the entries are the same." Four tests read that check as if it were python-graphblas's own contract and went red. Each one wanted an iso object rather than an inference, so each now asks for one with ``ss.build_scalar``, which means the same thing on every SuiteSparse version and produces the same storage formats the tests were selecting for. ``from_coo`` is left alone: adding a scan of the values to recover the old behavior would re-add the cost upstream just removed. The version gates had a second, quieter problem. ``graphblas.core.ss`` derived the C library version by parsing ``suitesparse_graphblas.__version__``, which is the version of the Python wrapper. The two normally agree, but a development build between releases reports the previous one: psg 10.4.1.0 against a 10.5.0 library here, so every gate saw 10.4.1 and any 10.5 branch would have been silently skipped. Read the library's own GxB_IMPLEMENTATION_* constants instead, falling back to the old parse where a build does not expose them. test_openmp_enabled keeps its hard assertion and gains a message saying what is wrong and that the fault is in the installed build. It caught a real one: conda-forge's first graphblas 10.5.0 build on osx-arm64 shipped without OpenMP. The build 1 rebuild (2026-08-26) restored OpenMP and the assertion passes again. * Cover SuiteSparse:GraphBLAS 10.4 and 10.5 in the CI psg pool The pool stopped at psg 10.3.1.0, so CI installed nothing newer except on the draw that picks "latest", which does not pin and so never produces a reproducible combination. Three releases have landed since, and each carries a SuiteSparse:GraphBLAS this project has to work against: 10.4.0 moved a field in the matrix struct, and 10.5.0 dropped the post-build iso check that four tests used to rely on. psg 10.5.0.0 is on PyPI and, since 2026-08-26, on conda-forge, so the conda-forge, wheel and source draws all cover it. `scripts/check_versions.sh` tracks the newest version we test, so its psg floor moves with the pool. Only the psg pools change here. The other pools are revised further up this branch stack, so raising them belongs there rather than in a commit that would conflict with every rebase. * Improvements to scripts/ci_pick_versions.py (AI-maintained) * Audit CI dependency pools against PyPI and conda-forge Add scripts/audit_{pypi,conda}_versions.py, which check the pools and tables in ci_pick_versions.py against what is actually published. The --validate stress test proves the picker agrees with the constraints it documents; it cannot notice when a constraint drifts from reality, since that happens when a package releases, not when we commit. A new weekly workflow runs both audits. In the picker, replace the hand-written scipy/numba compatibility rules with SCIPY_MAX_NUMPY and NUMBA_MAX_NUMPY ceiling tables, add SCIPY_MIN_PYTHON, and gate the tables behind check_ceilings() and check_python_pools() so a half-updated table fails the run instead of producing a silently wrong pick. Bump the pools to numpy 2.5, scipy 1.18, awkward 2.13, numba 0.67 and sparse 0.19, and offer sparse on py3.13/3.14 now that numba supports them. Also bump pre-commit hooks, update check_versions.sh to match, and ignore PLR0917, CPY and ISC004 (all newly stable in ruff 0.16). Fixes found while reviewing the above: - test_and_build.yml discarded the picker's exit status. `eval "$(cmd)"` reports eval's status, and that step's shell has no -e, so a picker failure left every *ver variable unset and quietly installed the latest of everything. Assign first, then fail the step. - pyproject: bound the sparse py3.14 marker below 3.15. sparse depends on numba, so it has to drop out wherever numba does. - _pick_scipy filtered ceilings and requires-python but not the one numpy floor scipy imposes, leaving the order dependence it exists to remove. - audit_pypi: report a missing scipy release instead of raising KeyError, honor a `<=` numpy cap rather than skipping the check, and keep the notes section inside the network guard. --- .github/workflows/audit_versions.yml | 80 +++++ .github/workflows/test_and_build.yml | 10 +- .pre-commit-config.yaml | 38 ++- pyproject.toml | 12 +- scripts/audit_conda_versions.py | 394 ++++++++++++++++++++++++ scripts/audit_pypi_versions.py | 424 ++++++++++++++++++++++++++ scripts/check_versions.sh | 21 +- scripts/ci_pick_versions.py | 434 +++++++++++++++++++-------- 8 files changed, 1263 insertions(+), 150 deletions(-) create mode 100644 .github/workflows/audit_versions.yml create mode 100755 scripts/audit_conda_versions.py create mode 100755 scripts/audit_pypi_versions.py diff --git a/.github/workflows/audit_versions.yml b/.github/workflows/audit_versions.yml new file mode 100644 index 000000000..646e96f59 --- /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@v6 + with: + persist-credentials: false + - uses: actions/setup-python@v6 + 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@v6 + with: + persist-credentials: false + - uses: conda-incubator/setup-miniconda@v3 + 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/test_and_build.yml b/.github/workflows/test_and_build.yml index 69fab3770..574bcd3d8 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -145,7 +145,15 @@ jobs: # 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 "v='${{ steps.pyver.outputs.selected }}'.split('.');print(f'{v[0]}.{v[1]}')") - eval "$(python scripts/ci_pick_versions.py --python "$pyver" --source ${{ steps.sourcetype.outputs.selected }})" + # 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 ${{ steps.sourcetype.outputs.selected }}) ; then + echo "::error::ci_pick_versions.py failed; refusing to install unpinned dependencies" + exit 1 + fi + eval "$picks" # 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". diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3582142d2..3b36390ad 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -42,23 +42,20 @@ repos: - id: name-tests-test args: ["--pytest-test-first"] - repo: https://github.com/abravalheri/validate-pyproject - rev: v0.25 + rev: "0.26" 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.3 + 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 - # 9.0 is still in alpha (9.0.0a3) and a pre-release pin defeats - # pre-commit.ci's autoupdate "avoid prereleases" logic; stay on the last - # stable 8.x until 9.0 GAs. - rev: 8.0.1 + rev: 9.0.1 hooks: - id: isort # Let's keep `pyupgrade` even though `ruff --fix` probably does most of it @@ -73,12 +70,12 @@ repos: - id: auto-walrus args: [--line-length, "100"] - repo: https://github.com/psf/black-pre-commit-mirror - rev: 26.3.1 + rev: 26.5.1 hooks: - id: black - id: black-jupyter - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.13 + rev: v0.16.8 hooks: - id: ruff-check args: [--fix-only, --show-fixes] @@ -92,10 +89,10 @@ repos: additional_dependencies: &flake8_dependencies # These versions need updated manually - flake8==7.3.0 - - flake8-bugbear==25.11.29 + - flake8-bugbear==26.9.9 - flake8-simplify==0.30.0 - repo: https://github.com/codespell-project/codespell - rev: v2.4.2 + rev: v2.4.3 hooks: - id: codespell types_or: [python, markdown, rst, toml, yaml] @@ -103,7 +100,7 @@ repos: - tomli; python_version<'3.11' files: ^(graphblas|docs)/ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.13 + rev: v0.16.8 hooks: - id: ruff-check # - id: ruff-format # we prefer black; uncomment to compare with ruff's formatter @@ -115,16 +112,18 @@ repos: # `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 + # 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.11.0.1" + rev: "v0.11.0.1-1" hooks: - id: shellcheck - repo: https://github.com/rbubley/mirrors-prettier - rev: v3.8.3 + rev: v3.9.7 hooks: - id: prettier args: [--prose-wrap=preserve] @@ -149,7 +148,7 @@ repos: hooks: - id: actionlint - repo: https://github.com/python-jsonschema/check-jsonschema - rev: 0.37.2 + rev: 0.38.0 hooks: - id: check-dependabot - id: check-github-workflows @@ -163,7 +162,7 @@ repos: # hooks: # - id: zizmor - repo: https://github.com/pysentry/pysentry-pre-commit - rev: v0.4.6 + rev: v0.5.0 hooks: - id: pysentry args: ["--format=human", "--fail-on=low", "--sources=pypa,osv,pypi"] @@ -206,6 +205,15 @@ repos: | 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 diff --git a/pyproject.toml b/pyproject.toml index 9fe272c98..cdf8867c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,7 +85,14 @@ io = [ "python-graphblas[networkx,scipy]", "python-graphblas[numba]; python_version<'3.15'", # make optional where numba is not supported "awkward >=2.0", - "sparse >=0.14; python_version<'3.13'", # make optional, b/c sparse needs numba + # sparse needs numba, so it stays optional. numba gained py3.13 support in 0.61 and + # py3.14 in 0.63; sparse 0.16 is the first release that postdates py3.13, and 0.18 the + # first that postdates py3.14. + "sparse >=0.14; python_version<'3.13'", + "sparse >=0.16; python_version=='3.13'", + # Upper guard for the same reason as numba above: sparse depends on numba, so it has + # to drop out wherever numba does, not just trail it. + "sparse >=0.18; python_version>='3.14' and python_version<'3.15'", # fast-matrix-market is no longer maintained; last supported Python is 3.12 ] viz = ["python-graphblas[networkx,scipy]", "matplotlib >=3.6"] @@ -342,6 +349,7 @@ ignore = [ "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) @@ -362,6 +370,7 @@ ignore = [ # 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) @@ -394,6 +403,7 @@ 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 diff --git a/scripts/audit_conda_versions.py b/scripts/audit_conda_versions.py new file mode 100755 index 000000000..0d2fe4c16 --- /dev/null +++ b/scripts/audit_conda_versions.py @@ -0,0 +1,394 @@ +#!/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.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 3977bd8a0..d9959556d 100755 --- a/scripts/check_versions.sh +++ b/scripts/check_versions.sh @@ -6,16 +6,23 @@ # # 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. -conda search 'flake8-bugbear[channel=conda-forge]>=25.11.29' +# +# 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.9.9' conda search 'flake8-simplify[channel=conda-forge]>=0.30.0' -conda search 'numpy[channel=conda-forge]>=2.4' +conda search 'numpy[channel=conda-forge]>=2.5' conda search 'pandas[channel=conda-forge]>=3.0' -conda search 'scipy[channel=conda-forge]>=1.17' +conda search 'scipy[channel=conda-forge]>=1.18' conda search 'networkx[channel=conda-forge]>=3.6' -conda search 'awkward[channel=conda-forge]>=2.11' -conda search 'sparse[channel=conda-forge]>=0.15' +conda search 'awkward[channel=conda-forge]>=2.13' +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.64' +conda search 'numba[channel=conda-forge]>=0.67' conda search 'pyyaml[channel=conda-forge]>=6.0' -conda search 'python-suitesparse-graphblas[channel=conda-forge]>=10.3.1' +# 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.0' # conda search 'python[channel=conda-forge]>=3.11 *pypy*' diff --git a/scripts/ci_pick_versions.py b/scripts/ci_pick_versions.py index 738ef99a8..2fcfafc02 100755 --- a/scripts/ci_pick_versions.py +++ b/scripts/ci_pick_versions.py @@ -23,13 +23,21 @@ # "" 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.15, 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. 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", ""], - "3.13": ["2.1", "2.2", "2.3", "2.4", ""], - "3.14": ["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", ""], } # Deps that depend on numpy version (1.x vs 2.x path). @@ -39,7 +47,32 @@ "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", ""], + "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.5; -# see the awkward constraints in `apply_constraints`. +# which numpy 2.0 removed. Only >=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.x": ["2.6", "2.7", "2.8", "2.9", "2.10", "2.11", "2.12", "2.13", ""], } NUMBA_VERSIONS = { "1.x": ["0.57", "0.58", "0.59", "0.60", "0.61", ""], - "2.x": ["0.62", "0.63", "0.64", ""], + "2.x": ["0.62", "0.63", "0.64", "0.65", "0.66", "0.67", ""], +} + +# 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)} + +# The numpy ceiling each numba pin declares. numba caps `numpy =3.9), so hold +# them to the wheel floor rather than test a build psg never claimed works. +# MAINT: 2026-09-11 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"}, } @@ -131,6 +217,62 @@ def _ver(s): 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 # --------------------------------------------------------------------------- @@ -147,6 +289,10 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): 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 @@ -154,42 +300,34 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): 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 <1.13 doesn't support numpy 2.x (safety net) - np_is_1x = v["numpy"].startswith("1.") if v["numpy"] else False - if not np_is_1x and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 13): - v["scipy"] = random.choice([s for s in scipy_pool if _ver(s) >= (1, 13)]) - - # scipy <1.15.1 requires numpy <2.3; scipy <1.16 requires numpy <2.5 - if _ver(v["numpy"]) >= (2, 3): - if v["scipy"] in ("1.13", "1.14"): - v["scipy"] = random.choice(["1.16", "1.17", ""]) - elif v["scipy"] == "1.15": - v["scipy"] = random.choice(["1.15", "1.16", "1.17", ""]) - - # numpy 1.26 + scipy 1.9 conflict - if v["numpy"] == "1.26" and v["scipy"] == "1.9": - v["scipy"] = random.choice(["1.10", "1.11", ""]) + # 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; scipy <1.16 has no py3.14 builds - if pyver == "3.14" and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 16): - v["scipy"] = random.choice(["1.16", "1.17", ""]) - elif pyver == "3.13" and v["scipy"] == "1.13": - v["scipy"] = random.choice(["1.14", "1.15", "1.16", "1.17", ""]) + if _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 - if pyver == "3.14" and v["pandas"] == "2.2": + if _ver(pyver) >= (3, 14) and v["pandas"] == "2.2": v["pandas"] = random.choice(["2.3", "3.0", ""]) - # pandas 3.0 requires numba >=0.60 and scipy >=1.14.1 - if v["pandas"] == "3.0": - if v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 60): - v["numba"] = "0.60" - if v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 15): - v["scipy"] = random.choice(["1.15", "1.16", "1.17", ""]) + # 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 @@ -218,9 +356,9 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): # --- awkward / Python version availability --- # awkward <2.7 has no py3.13 builds; awkward <2.8 has no py3.14 builds - if pyver == "3.14" and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 8): + if _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 pyver == "3.13" and v["awkward"] == "2.6": + 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 --- @@ -254,12 +392,14 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): # --- matplotlib / numpy floor (notebooks job) --- - # matplotlib is installed unpinned for the notebooks task. matplotlib 3.11.0's - # runtime check requires numpy >=1.25, but its conda-forge metadata does not, - # so the solver happily pairs matplotlib 3.11 with a numpy 1.24 pin and the - # notebooks then die at `import matplotlib`. Pin below 3.11 whenever the - # numpy pick sits under that floor; empty means latest is fine. - # MAINT: 2026-08-05 confirmed with matplotlib 3.11.0 + numpy 1.24.4 + # 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: @@ -267,34 +407,22 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): # --- numba constraints --- - # numba minimum by Python version: 0.59 for 3.12, 0.61 for 3.13, 0.63 for 3.14 - numba_min = {"3.11": (0, 57), "3.12": (0, 59), "3.13": (0, 61), "3.14": (0, 63)} + # numba minimum by Python version (see NUMBA_MIN_PYTHON) if v["numba"] not in ("", "NA"): - min_ver = numba_min[pyver] + 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 <0.64 requires numpy <2.4 - if _ver(v["numpy"]) >= (2, 4) and v["numba"] in ("0.62", "0.63"): - v["numba"] = "0.64" - - # conda-forge numba 0.57 requires numpy <1.25, so that pairing fails the - # conda solve outright. Verified by dry-run solves: 0.57+1.25 conflicts; - # 0.58+1.25, 0.58+1.26, and 0.59+1.26 all solve. - # MAINT: 2026-08-05 confirmed against conda-forge with mamba dry-runs - if v["numba"] == "0.57" and np_is_1x and _ver(v["numpy"]) >= (1, 25): - v["numba"] = "0.58" - - # numba <0.62 doesn't support numpy 2.x - if not np_is_1x and v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 62): - v["numba"] = "NA" - - # --- sparse --- - - # sparse doesn't support Python 3.13+ - if pyver in ("3.13", "3.14"): - v["sparse"] = "NA" + # 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" # --------------------------------------------------------------------------- @@ -347,7 +475,7 @@ def pick_versions(pyver, source_type): "numba": random.choice(numba_pool), "networkx": random.choice(NETWORKX_VERSIONS[pyver]), "pyyaml": random.choice(PYYAML_VERSIONS[pyver]), - "sparse": _pick_scalar_or_list(SPARSE_VERSIONS[pyver]), + "sparse": random.choice(SPARSE_VERSIONS[pyver]), } # Source builds have issues with some numpy/scipy/pandas versions; @@ -366,11 +494,6 @@ def pick_versions(pyver, source_type): return v -def _pick_scalar_or_list(pool): - """Handle pools that are either "NA" (string) or a list of choices.""" - return pool if isinstance(pool, str) else random.choice(pool) - - def _pick_psg(npver, pyver, source_type): """Pick python-suitesparse-graphblas version.""" if source_type == "upstream": @@ -384,7 +507,10 @@ def _pick_psg(npver, pyver, source_type): if not pool: return "" return f"{eq}{random.choice(pool)}" - pool = PSG_VERSIONS_NP2["py314_only"] if pyver == "3.14" else PSG_VERSIONS_NP2["no_py314"] + 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] + [""]) @@ -458,7 +584,9 @@ def format_summary(v): "matplotlib", ): name = _SUMMARY_NAMES[key] - val = v[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 == "": @@ -473,7 +601,56 @@ def format_summary(v): # --------------------------------------------------------------------------- -def validate(v, pyver): +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 @@ -482,34 +659,28 @@ def validate(v, pyver): 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 <1.13 requires numpy 1.x - if not np_is_1x and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 13): - errors.append(f"scipy {v['scipy']} doesn't support numpy 2.x") - - # scipy <1.15.1 requires numpy <2.3 - if _ver(v["numpy"]) >= (2, 3) and v["scipy"] in ("1.13", "1.14"): - errors.append(f"scipy {v['scipy']} requires numpy <2.3, got {v['numpy']}") - - # numpy 1.26 + scipy 1.9 - if v["numpy"] == "1.26" and v["scipy"] == "1.9": - errors.append("numpy 1.26 + scipy 1.9 conflict") + # 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 pyver == "3.14" and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 16): - errors.append(f"scipy {v['scipy']} has no py3.14 build") - if pyver == "3.13" and v["scipy"] == "1.13": - errors.append("scipy 1.13 has no py3.13 build") + if _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") + 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 pyver == "3.14" and v["pandas"] == "2.2": - errors.append("pandas 2.2 has no py3.14 build") + if _ver(pyver) >= (3, 14) and v["pandas"] == "2.2": + errors.append(f"pandas 2.2 has no py{pyver} build") - # pandas 3.0 requirements - if v["pandas"] == "3.0": - if v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 60): - errors.append(f"pandas 3.0 requires numba >=0.60, got {v['numba']}") - if 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']}") + # 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 ( @@ -521,10 +692,10 @@ def validate(v, pyver): errors.append(f"pandas {v['pandas']} constrains scipy >=1.10.0, got {v['scipy']}") # awkward Python availability - if pyver == "3.14" and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 8): - errors.append(f"awkward {v['awkward']} has no py3.14 build") - if pyver == "3.13" and v["awkward"] == "2.6": - errors.append("awkward 2.6 has no py3.13 build") + if _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") + 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): @@ -534,10 +705,6 @@ def validate(v, pyver): 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)") - # conda-forge numba 0.57 requires numpy <1.25 (fails the conda solve) - if v["numba"] == "0.57" and v["numpy"] and _ver(v["numpy"]) >= (1, 25) and np_is_1x: - errors.append(f"numba 0.57 requires numpy <1.25, got {v['numpy']}") - # awkward <2.10 requires numpy <2.5 (generic-unit datetime64("NaT") at import) if ( _ver(v["numpy"]) >= (2, 5) @@ -558,21 +725,30 @@ def validate(v, pyver): errors.append(f"networkx {nx} requires pandas >=2.0, got {v['pandas']}") # numba Python minimums - numba_min = {"3.11": (0, 57), "3.12": (0, 59), "3.13": (0, 61), "3.14": (0, 63)} - if v["numba"] not in ("", "NA") and _ver(v["numba"]) < numba_min[pyver]: + 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 <0.64 requires numpy <2.4 - if v["numba"] in ("0.62", "0.63") and _ver(v["numpy"]) >= (2, 4): - errors.append(f"numba {v['numba']} requires numpy <2.4, got {v['numpy']}") - - # numba <0.62 requires numpy 1.x - if not np_is_1x and v["numba"] not in ("", "NA") and _ver(v["numba"]) < (0, 62): - errors.append(f"numba {v['numba']} doesn't support numpy 2.x") - - # sparse Python availability - if pyver in ("3.13", "3.14") and v["sparse"] != "NA": - errors.append(f"sparse 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 @@ -585,7 +761,7 @@ def stress_test(n=10000): for source in ("conda-forge", "wheel", "source", "upstream"): for _ in range(n): v = pick_versions(pyver, source) - errs = validate(v, pyver) + errs = validate(v, pyver, source) total += 1 if errs: failures += 1 @@ -618,6 +794,12 @@ def main(): 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) From 6a967a735cc4015d4796c9bbb775539b71d5f9d3 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 18 Sep 2026 12:34:22 -0500 Subject: [PATCH 17/52] Hold the owner of a borrowed GraphBLAS handle in TypedOpBase (#587) * Hold the owner of a borrowed GraphBLAS handle in TypedOpBase `SelectOp._from_indexunary` reuses the IndexUnaryOp's `GrB_IndexUnaryOp` rather than allocating a second one, and marked the borrowing typed op `_owns_gb_obj_inst = False` so that only one side would free it. Nothing kept the owner alive, though: `register_anonymous` drops the IndexUnaryOp on the way out, so its typed op is collected and frees the handle while the SelectOp is still pointing at it. The result is a dangling pointer rather than a NULL one. `__del__` frees through a synthesized cell, `ffi.new(f"{c_type_name}*", gb_obj)`, so `GrB_*_free` clears that temporary and leaves the op's own `gb_obj` holding the old address. Reproduced at 6f1eb02: def _ne_thunk(x, i, j, thunk): return x != thunk sel = SelectOp.register_anonymous(_ne_thunk) gc.collect() Vector.from_coo([0, 1, 2], [1, 5, 9]).select(sel, 5).new() -> UninitializedObject Whether it raises depends on whether anything happens to retain the IndexUnaryOp. With `x > thunk` it survives, because a traceback caught inside the per-dtype compile loop keeps a frame that references it. Replace the per-instance opt-out with `_gb_obj_owner`, naming the owning typed op. That suppresses our free and keeps the handle alive for as long as we can reach it; disclaiming ownership without naming an owner is what left the handle dangling. The opt-out being replaced was added in 6f1eb02 to stop a double free, so the sequence is leak, double free, opt-out, dangling handle, and now an owned reference. * Note the borrowed-handle exception on TypedUserSelectOp._owns_gb_obj Read on its own, `_owns_gb_obj = True` says every TypedUserSelectOp frees its handle. The ones `_from_indexunary` builds do not: they borrow the IndexUnaryOp's `GrB_IndexUnaryOp` and set `_gb_obj_owner`, which suppresses the free. Worth a comment because the mistake is silent. Freeing from both sides is a use-after-free, and a freed handle usually still reads as valid: with the ownership removed, select kept returning correct results even after 30 fresh allocations churned the allocator. No test would catch it. --- graphblas/core/operator/base.py | 28 +++++++++++++++------------- graphblas/core/operator/select.py | 17 ++++++++++++----- graphblas/tests/test_op.py | 21 +++++++++++++++++++++ 3 files changed, 48 insertions(+), 18 deletions(-) diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index bf685ba5c..a893ba3ad 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -534,16 +534,13 @@ class TypedOpBase: "gb_name", "_type2", "_jit_c_info", - "_owns_gb_obj_inst", + "_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. - # Specific instances can override via ``_owns_gb_obj_inst`` (set by - # the constructor); ``SelectOp._from_indexunary`` aliases an existing - # ``GrB_IndexUnaryOp`` and must clear ownership to avoid a double free. _owns_gb_obj = False def __init__(self, parent, name, type_, return_type, gb_obj, gb_name, dtype2=None): @@ -558,10 +555,12 @@ def __init__(self, parent, name, type_, return_type, gb_obj, gb_name, dtype2=Non # for this typed op; ``None`` for built-in ops and for UDT ops with # no JIT path. self._jit_c_info = None - # Per-instance ownership override; defaults to the class attribute. - # ``SelectOp._from_indexunary`` flips this to ``False`` on aliasing - # TypedUserSelectOps so only the IndexUnaryOp frees the handle. - self._owns_gb_obj_inst = type(self)._owns_gb_obj + # 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): @@ -594,11 +593,14 @@ def __reduce__(self): def __del__(self): # Free the SuiteSparse handle we allocated. Built-in typed ops alias - # SuiteSparse's permanent built-in singletons and must never free, so - # gate on the per-instance owns flag (defaults to the class - # attribute; the alias case overrides to False). Mirrors the - # ``Matrix.__del__`` / ``Vector.__del__`` pattern. - if not getattr(self, "_owns_gb_obj_inst", False): + # 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: diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 32a2509a5..3cc506d84 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -25,7 +25,11 @@ def __call__(self, val, thunk=None): class TypedUserSelectOp(TypedOpBase): __slots__ = () opclass = "SelectOp" - _owns_gb_obj = True # underlying object is a GrB_IndexUnaryOp + # 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) @@ -97,10 +101,13 @@ def _from_indexunary(cls, iop): t.return_type, t.gb_obj, ) - # Aliases the IndexUnaryOp's allocation. The IndexUnaryOp - # owns the free; clearing here prevents a double free when - # both ends are GC'd. - op._owns_gb_obj_inst = False + # 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, diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index ef5227f07..e439fc149 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1359,6 +1359,27 @@ def _this_or_that(val, idx, _, thunk): # pragma: no cover (numba) assert result.isequal(w) +@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)) + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_tuple_return_binaryop(record_udt): From f093e36515ddeb05bd9e936fe313dcfc17db91bd Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 18 Sep 2026 12:39:11 -0500 Subject: [PATCH 18/52] Compile UDFs under Numba's numpy error model (#588) A UDF that divides by zero silently corrupted the result. Numba's default error model raises ZeroDivisionError inside the cfunc, where Numba prints the traceback and returns without writing, so GraphBLAS keeps whatever the output element already held. Measured at 6f1eb02: def _idiv(x, y): return x // y op = BinaryOp.register_anonymous(_idiv, "_probe_idiv") op(Vector.from_coo([0, 1], [10, 20]) & Vector.from_coo([0, 1], [2, 0])).new() -> [5, 5] # numpy gives [5, 0]; the 5 is the previous element The float case is the same shape: `[0.5, 0.5]` where numpy gives `[0.5, inf]`. The only user-visible signal is an "Exception ignored" traceback on stderr. Pass `error_model="numpy"` at every `numba.njit` and `numba.cfunc` site for user-defined ops. Setting it on the cfunc alone is not enough: `_build` calls `.compile(sig)` on the Dispatcher first, a Dispatcher keeps one compilation per signature, and whichever compile happens first fixes the model. Reverting only the Dispatcher argument, leaving it on the cfunc, reproduces `[5, 5]`. This changes what a dividing UDF returns, but the behaviour it replaces is an unwritten element rather than an error, so there is no correct result being taken away. --- graphblas/core/operator/base.py | 2 +- graphblas/core/operator/binary.py | 15 +++++++-- graphblas/core/operator/indexbinary.py | 12 +++++-- graphblas/core/operator/indexunary.py | 8 +++-- graphblas/core/operator/udt_utils.py | 2 +- graphblas/core/operator/unary.py | 10 ++++-- graphblas/tests/test_op.py | 46 ++++++++++++++++++++++++++ 7 files changed, 82 insertions(+), 13 deletions(-) diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index a893ba3ad..287e525d5 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -189,7 +189,7 @@ def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, t from ``typed_user_cls.opclass``. ``dtype2`` is ``None`` for unary ops; the rest pass both. Returns the cached ``TypedUser*Op``. """ - wrapper = numba.cfunc(wrapper_sig, nopython=True)(wrapper) + wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) 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}*") diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 97e0a9e70..5329b4415 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -525,7 +525,14 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): 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. + 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 @@ -581,7 +588,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( @@ -981,7 +990,7 @@ def _initialize(cls): (binary.any, _second), ]: binop.orig_func = func - binop._numba_func = numba.njit(func) if _has_numba else 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.second._numba_func diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index a19f5ffec..de2ca5f71 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -246,7 +246,9 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if name is None: name = getattr(func, "__name__", "") success = False - indexbinary_udf = numba.njit(func) + # 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 ) @@ -333,7 +335,9 @@ def indexbinary_wrapper( ): # 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)(indexbinary_wrapper) + 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( @@ -392,7 +396,9 @@ def _compile_udt(self, dtype, dtype2): numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) - indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexbinary_wrapper) + 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( diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index c7af8ed78..5820e66c5 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -116,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 ) @@ -176,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( diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index fc33ea8c0..3e2a60d45 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -558,7 +558,7 @@ def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scala func_name="_op", source_label=f"", ) - return numba.njit(op_func) + return numba.njit(op_func, error_model="numpy") def _make_array_wrapper( size, diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index 22822921a..64f68eb48 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -176,7 +176,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 @@ -231,7 +233,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( @@ -453,7 +457,7 @@ 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 = {} diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index e439fc149..de1c01b88 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1359,6 +1359,52 @@ def _this_or_that(val, idx, _, thunk): # pragma: no cover (numba) assert result.isequal(w) +@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"]) + + 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) + + @pytest.mark.skipif("not supports_udfs") def test_select_op_outlives_source_indexunary(): """A SelectOp keeps alive the IndexUnaryOp whose GraphBLAS handle it borrows. From 3d9d7c54fb07ab91d0f0af4de793e749f26ec8fe Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 24 Sep 2026 15:05:45 -0500 Subject: [PATCH 19/52] Pass and receive numpy views in array-UDT UDFs (#589) * Pass and receive numpy views in array-UDT UDFs A UDF over an array UDT wrote past the end of its output element. The cfunc wrapper passed the raw element pointer and stored the return value with `z_ptr[0] = ...`, so what landed in the buffer was Numba's NestedArray descriptor (data pointer, shape, strides) rather than the element payload. The descriptor is the wider of the two, so the store ran off the end of the element into memory SuiteSparse owns. Measured against the code this replaces, driving the wrapper's cfunc through ctypes over a guarded buffer: (6,) np itemsize 48 | numba value type size 56 | overrun 8 bytes (2, 3) np itemsize 48 | numba value type size 72 | overrun 24 bytes z[:6] = [100.0, 101.0, 102.0, 103.0, 104.0, 105.0] z[6:] = [5e-324, -1.0, -1.0, ...] # -1.0 is the guard byte pattern indices written past the element: [6] Both operands and the output are now bound with `numba.carray(ptr, shape)` in the UDT's declared shape, so the UDF receives a numpy view it can index (`x[i, j]`, `x.shape`) and the wrapper slice-assigns the result back with `z[:] = ...`. Array-typed record leaves slice-assign for a second reason: Numba's record-field setitem copies the destination's extent regardless of the source's, so a short source was read past its end. The return-type resolution has to move with the wrapper. Once the UDF receives a view, a UDF that builds its result (`x + y`) instead of returning an operand types as a plain Numba `Array`, which `lookup_dtype` does not recognize; `_resolve_udt_return_type` gains an `Array` branch matching it back to an input array UDT by base element type and rank. Splitting the two apart would leave array-UDT UDFs that build a result broken outright. An `Array` type carries `ndim` but not its extents, so a return whose shape does not fit the UDT is not a type error. It is now memory-safe, but still silent: the slice-assign's ValueError is raised inside a cfunc, which Numba prints and swallows, leaving the element as SuiteSparse found it. * Pin the array-UDT any/first/second wrapper to one element's bytes Binary ops outside _BUILTIN_UDT_BINARY_OPS (any, first, second) compile through the generic _numba_func branch of BinaryOp._compile_udt. The old wrapper there loaded and stored an array-UDT 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, regardless of the 32-byte payload. SuiteSparse's generic reduce keeps a UDT accumulator in a stack array sized to the element, so each fold overflowed it by 24 bytes and clobbered a spilled GrB_Type pointer: segfault on some builds, SIGBUS or a silently wrong answer on others. The carray-based wrapper rework on this branch writes exactly itemsize bytes; this test keeps it that way. The test builds the wrapper for binary.any the same way _compile_udt does, compiles it with numba.cfunc, and calls it via ctypes on heap buffers with slack, so a regression trips an assert instead of corrupting a stack frame. The z guard sentinel (0xAB) must differ from the sentinel in y's trailing slack (0xCD): the descriptor load/store is a byte-preserving copy of the source element plus its trailing bytes, so with one shared sentinel the overflow would rewrite z's guard with identical values and go undetected. A public-path any-reduce smoke runs after the byte-level checks, so a regression fails the assert before reaching the code that can crash the process. Verified both directions: passes here (host numba 0.65.1 and a conda env with numba 0.66 plus python-suitesparse-graphblas 10.0.1.1), and on an export of main the assert trips with exactly 24 overflow bytes while the payload check still passes, demonstrating the blind spot a same-sentinel check would have. * Resolve array-UDT UDF returns to the operand they came from A UDF that returns an array-UDT operand as-is types as that operand's NestedArray, which went through lookup_dtype. That registry keys on the Numba type, which a flat FP64[2, 3] and a layered FP64[3][2] share, so it named whichever UDT was registered last. binary.second, monoid.any, and assignment back into the operand's vector then failed with DomainMismatch. Match a NestedArray return to its operand by identity first, and raise if two distinct operands share its type. A NestedArray that matches no operand, such as a record's array field, is no longer matched by rank alone to an operand of a different length. The "matches no input array UDT" error now describes every array operand. Fold the two wrapper byte-guard tests into one, over a 1-D and a 2-D element. A user UDF and binary.any build the same wrapper, and only distinct sentinels see the old overrun reliably: it copied y's trailing bytes, so the dropped test caught it only because the memory after its buffer happened not to match the guard. * Use unique dtype to avoid clobbering types in lookup_dtype * Flatten nested subarray dtypes so each UDT layout has one DataType numpy keeps np.dtype((FP64[3], (2,))) distinct from FP64[2, 3], but nothing else does: same bytes, same expanded array, same Numba type. As separate DataTypes the nested one broke auto-lifted ops, monoids, and JIT, could not mix with its flat twin, and made lookup_dtype's Numba-type keys resolve to whichever was registered last. - register_anonymous and dtypes.ss.register_new flatten nested subarrays, including inside records - Drop UDT Numba types as registry keys; lookup_dtype converts a Record or NestedArray to its numpy dtype instead - Simplify _resolve_udt_return_type: known extents resolve directly, and only an array built in the UDF is matched to an operand by rank - View, rather than cast, values whose dtype differs only in nesting; numpy's cast between the two scrambles values - Fix stale field-type rules in the UDT guide and note that records with array fields skip the JIT * Build the nested-subarray Vector only where build can type dup_op test_nested_subarray_record_values_load_intact failed on every CI job that ran the suitesparse-vanilla backend without Numba: KeyError: "any does not work with ({'li_k': INT64, 'li_pts': FP64[2, 3]}, {'li_k': INT64, 'li_pts': FP64[2, 3]}) types" On that backend, Vector.build defaults dup_op to binary.any, and a UDT gets binary.any only by compiling it with Numba. The values_to_numpy_buffer and Scalar checks need neither, so they still run everywhere; only the Vector half is gated. * Keep the array-UDT test types under SuiteSparse 8's 128-byte cap SuiteSparse built without variable-length arrays (MSVC, so the Windows builds) rejects a UDT larger than 128 bytes with GrB_INVALID_VALUE before 9.0. The Windows job that drew psg 8.0.2 failed two tests on it: test_nested_subarrays_register_flat (FP32[3, 2, 7], 168 bytes) and test_udt_array_udf_known_extents_never_match_by_rank (FP64[21], 168 bytes). test_udt_array_udf_returning_operand_keeps_its_type registers 152- and 160-byte types as well; it is marked slow, and that job did not run slow tests. None of the three is about size, so keep their shapes and shrink the element type (INT16 or FP32) rather than skip them on SuiteSparse < 9. --- docs/user_guide/udt.rst | 32 ++- graphblas/core/dtypes.py | 70 +++++- graphblas/core/operator/base.py | 134 ++++++++++-- graphblas/core/operator/udt_utils.py | 8 +- graphblas/core/scalar.py | 11 +- graphblas/core/ss/dtypes.py | 5 +- graphblas/core/utils.py | 11 +- graphblas/tests/test_dtype.py | 117 +++++++++- graphblas/tests/test_op.py | 312 +++++++++++++++++++++++++++ graphblas/tests/test_scalar.py | 15 +- 10 files changed, 663 insertions(+), 52 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 18af2cf3e..d43331257 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -32,7 +32,9 @@ A UDT is any ``numpy.dtype`` you register with python-graphblas. There are three 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. +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``: @@ -68,7 +70,11 @@ Field type rules: - Numeric scalar types (``int``, ``float``, ``bool``, ``complex``, the corresponding numpy scalar types) are supported. -- Strings, Python objects, and nested UDTs are not. +- 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}``. +- Strings and Python objects are not supported. - Dataclass annotation strings (e.g., ``"int"``) resolve through ``lookup_dtype``. @@ -173,6 +179,26 @@ For nested record UDTs the tuple is *flat over the leaves*. Given 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. +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 of the same shape: + +.. 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] + 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. @@ -216,6 +242,8 @@ JIT is skipped when: 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``). diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 8eb09725b..73b905306 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -10,6 +10,7 @@ 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") @@ -200,6 +201,10 @@ def register_anonymous(dtype, name=None): 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 @@ -259,6 +264,7 @@ 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. The # Python-side ``rv.name`` updates here, but the SuiteSparse-side @@ -318,9 +324,6 @@ def register_anonymous(dtype, name=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 @@ -532,6 +535,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 @@ -585,6 +594,61 @@ 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 _dtype_to_string(dtype): """Convert a dtype to a string that can be safely evaluated to recreate the dtype. diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 287e525d5..28c2291f8 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -272,7 +272,9 @@ def _resolve_udt_return_type(numba_ret_type, *dtypes): 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. + 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) @@ -329,17 +331,80 @@ def _leaves(d): 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, 1)`` line - to add to the wrapper body (empty for non-record cases). + - ``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. """ @@ -351,8 +416,12 @@ def _input_operand(dtype, var): f"{var}[0]", nt.CPointer(dtype.numba_type), ) - # Array UDT: pass the raw pointer. The UDF carrays it if it wants. - return "", f"{var}_ptr", 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. @@ -364,9 +433,11 @@ 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, otherwise ``"direct"``. - ``write_info`` is the tuple-unpack field tuple for ``"record_fields"``, - or a ``(BL, BR, zname)`` 3-tuple for ``"direct"``. + 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 @@ -378,18 +449,21 @@ def _output_handler(return_type, numba_ret_type): and return_type._is_udt and return_type.np_type.names is not None ): - # ``write_info`` is the list of Python access paths to each leaf - # field (``"['a']"`` for flat, ``"['outer']['inner_a']"`` for - # nested). The wrapper iterates these to write the flat tuple - # ``_result`` back leaf-by-leaf. + # ``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. from .udt_utils import _iter_record_leaves - leaf_python_paths = tuple(py for py, _c, _d in _iter_record_leaves(return_type.np_type)) + 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", - leaf_python_paths, + leaves, ) if return_type._is_udt: if return_type.np_type.subdtype is None: @@ -399,9 +473,13 @@ def _output_handler(return_type, numba_ret_type): "direct", ("", "", "z[0]"), ) - # Array UDT: write via z_ptr[0]. The UDF receives the carray and - # writes in place. - return "", ret_ptr_type, "direct", ("", "[0]", "z_ptr[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]") @@ -410,11 +488,29 @@ 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. + 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. """ + 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} = _result[{i}]\n" for i, path in enumerate(zinfo) + f" z[0]{path}{'[:]' if is_array else ''} = _result[{i}]\n" + for i, (path, is_array) in enumerate(zinfo) ) return ( f"{signature_line}\n" diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index 3e2a60d45..b55a17f13 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -349,9 +349,11 @@ def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): ``("['outer']['inner_a']", "outer.inner_a", float64_dtype)``. Only record-in-record nesting is recognised. Array-typed fields are - yielded as a single leaf (the codegen treats them opaquely; arithmetic - on an array-typed sub-field isn't supported by either the cfunc or the - JIT path today). + yielded as a single leaf, which the Numba path handles: the generated + expression operates on the whole sub-array and the wrapper slice-assigns + it back. The JIT path never sees such a record, because a subarray dtype + has no entry in ``NP_TO_C_TYPES`` and ``_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] diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 5b5e56299..9755acb3c 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -8,6 +8,7 @@ from ..exceptions import EmptyObject, check_status from . import _has_numba, _supports_udfs, automethods, ffi, lib, utils from .base import BaseExpression, BaseType, call +from .dtypes import _view_if_same_layout from .expr import AmbiguousAssignOrExtract from .operator import get_typed_op from .utils import _Pointer, output_type, wrapdoc @@ -356,15 +357,7 @@ def value(self): rv = np.array(ffi.buffer(scalar.gb_obj[0 : np_type.itemsize])) if np_type.subdtype is None: return rv.view(np_type)[0] - # Collapse nested subarray dtypes before viewing. numpy keeps them - # layered, so a 6-long array of 5-long FP64 arrays has - # ``subdtype == (dtype((' 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): diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 3a5026b6a..2e3e17d87 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -8,10 +8,10 @@ import graphblas as gb from graphblas import core, dtypes -from graphblas.core import _supports_udfs as supports_udfs # noqa: F401 +from graphblas.core import _supports_udfs as supports_udfs from graphblas.core import lib -from graphblas.core.operator.udt_utils import _has_jit_set # noqa: F401 -from graphblas.core.utils import _NP2 +from graphblas.core.operator.udt_utils import _has_jit_set +from graphblas.core.utils import _NP2, values_to_numpy_buffer from graphblas.dtypes import lookup_dtype suitesparse = gb.backend == "suitesparse" @@ -215,10 +215,119 @@ 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}" diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index de1c01b88..9bc3f2cdc 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1926,6 +1926,318 @@ def test_udt_auto_monoid(): assert any_res in ((7, 8.0), (11, 12.0)) +@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.base 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_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 base as _base + + # 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 = _base.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(_base, "lookup_dtype", refuse_nested) + with pytest.raises(UdfParseError, match="matches no input array UDT"): + _base._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(): diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 49e7221b4..b8c00ffec 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -259,18 +259,17 @@ def test_udt_scalar_padding_is_zeroed(): reason="SuiteSparse < 9 rejects a 240-byte UDT on builds without VLA support", ) def test_udt_scalar_nested_array_roundtrip(): - """A UDT whose element is itself an array dtype reads back at full extent. + """A UDT declared as an array of arrays reads back at full extent. - numpy keeps subarray dtypes layered rather than flattening them, 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". + 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.itemsize == 240 - assert udt.np_type.subdtype[0].subdtype is not None # genuinely nested + 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 From f33f8c641f1f877b7c3cbd4173d4602f579523e4 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sat, 26 Sep 2026 00:23:26 -0500 Subject: [PATCH 20/52] Stop `gb.ss` import from turning the SuiteSparse JIT on (#630) * Stop `gb.ss` import from turning the SuiteSparse JIT on Importing ``graphblas.ss`` raised ``jit_c_control`` from SuiteSparse's default ``'run'`` to ``'on'``. That submodule is reached by attribute access, so ``print(gb.ss.about["library_version"])``, the line a user writes for a bug report, silently changed the numerical machinery of every later call in the process. ``'on'`` is also what makes SuiteSparse load kernels from its on-disk cache, and that cache is keyed on the version triple alone. Two builds sharing a triple but differing in a JIT-visible header therefore share kernels: locally, kernels compiled by a 10.4.0 dated "July 31, 2026" were loaded by a 10.4.0 dated "Aug 7, 2026", whose GB_Matrix_content.h had moved a field. A four-element complement-masked assign returned GrB_OUT_OF_MEMORY, and a full suite run showed 25 unrelated-looking failures. Reported upstream with a standalone C reproducer. The import now repairs the conda-baked compiler path and nothing else. The bump to ``'on'`` moves to ``_enable_jit_for_udt``, called when a UDT registers a type or arms an op with C source: acts that plainly involve compiling C, so enabling the compiler is not a surprise. An explicit ``'off'`` or ``'pause'`` is honored; only SuiteSparse's own non-compiling defaults are raised. ``fix_jit_config()`` is unchanged, since enabling compilation is what its caller asked for. Two details the hooks make necessary: ``_probe_jit`` registers a UDT to verify the compiler, and registering a UDT now asks to enable the JIT, which would probe again. The inner probe registered the probe type and the outer one then died redeclaring it through cffi, so a working environment got ``False`` from the first ``fix_jit_config()`` and ``True`` from every call after. ``_probing_jit`` makes the re-entrant request answer yes without starting a second probe. ``_maybe_warn_no_jit`` blamed ``jit_c_control`` for a UDT that is not expressible as a C struct, and now names the one cause its call site can produce. It used to re-derive the cause from the live config, which is not the state codegen read. Codegen reads the dtype; it never looks at the control or the compiler path, so neither can be why it declined. On a host whose JIT probe fails the control reads back ``'run'`` and ``_enable_jit_for_udt`` has cached its ``False``, which is enough for the old message to blame the control again. That is the ubuntu-latest and macos-latest conda runners, so it was a failing test as well as a message that sent the reader to a setting which cannot fix it. Dropping the environment branches costs no diagnostic: a UDT that does have a C form is armed and never reaches this warning, so an unusable compiler was already silent here. ``_enable_jit_for_udt`` caches whether this process can compile, which is settled once, but it was also returning before re-raising the control, which is separate state that anything may have moved since. Restoring a saved ``'run'`` after an earlier op had enabled the JIT therefore disabled compilation for the rest of the process, silently: the next op is armed with C source SuiteSparse will never build, and nothing warns, because codegen succeeded. The cached path now raises ``'run'`` and ``'load'`` the same way the uncached one does, and still honors an explicit ``'off'`` or ``'pause'``. Before this commit the question could not arise: the import forced ``'on'``, so a saved value was always ``'on'`` too. Its test drives a stand-in config so the assertion holds on every host, and pins the live round trip separately: a library built without the JIT clamps every write down to ``'run'``, and would never read back the ``'on'`` the test is about. One prose correction while in the file: ``_probe_jit`` claimed SuiteSparse flips the control to ``'load'`` after a failed compile, as though that were the only outcome. ``git grep -n "GB_jit_control = "`` in SuiteSparse finds eight places that lower it. ``GB_jitifyer_set_control`` is the one that comes first: it takes the minimum of the requested value and ``'run'`` on a library built without the JIT, so ``'on'`` cannot be set there however good the compiler is. On the load path, a compile failure gives ``'load'``, a load failure ``'run'``, and a failed hash insert ``'pause'``. The docstring says the value is not ours to predict, so the contract here is the bool. The new subprocess tests read the control through the C API rather than ``gb.ss.config``, because reaching for ``gb.ss`` is the thing under test. * Say where the JIT is enabled now in the UDT user guide The "JIT compiler auto-fix" section still said the import bumps jit_c_control from 'run' to 'on'. After the previous commit the import repairs the compiler path and nothing else, and the control is raised the first time a UDT registers a type or arms an operator with C source. Also note that fix_jit_config() sets 'on' as well as repairing and probing. * Keep an explicit JIT opt-out from switching the JIT off for good _enable_jit_for_udt cached False when the first UDT op found jit_c_control at 'off' or 'pause', though an opt-out says nothing about whether the process can compile. A saved 'run' restored afterwards was then never raised, so every later UDT op ran without the JIT, silently. An opt-out is now honored without being cached, and the function returns whether a kernel armed now will be compiled. 'load' is now honored the same way instead of being raised. SuiteSparse sets it after a compile fails so that the failure is not repeated, and raising it at the next arming ran the failing compiler again at every later UDT op, raising JitError each time on SuiteSparse 9.4 and later. It also left 'off' as the only way to stop a host with cached kernels and no compiler from trying to compile. Also say in fix_jit_config's docstring that it sets 'on', and which SuiteSparse versions raise the JitError its probe absorbs. In the user guide, say that NoJITWarning covers only a UDT with no C form and which calls raise the control. Share one helper between the two subprocess tests, and word as hypotheticals the test docstrings for failures that never reached main. --- docs/user_guide/udt.rst | 36 ++-- graphblas/core/operator/udt_utils.py | 7 + graphblas/core/ss/dtypes.py | 6 + graphblas/core/ss/jit_config.py | 211 ++++++++++++++-------- graphblas/exceptions.py | 7 +- graphblas/ss/__init__.py | 14 +- graphblas/tests/test_dtype.py | 39 +++- graphblas/tests/test_ssjit.py | 254 ++++++++++++++++++++++++--- 8 files changed, 454 insertions(+), 120 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index d43331257..d37653250 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -256,12 +256,13 @@ Anonymous UDTs (no ``name=`` argument) still take the JIT path: a synthetic identifier. Pass ``name=`` only for readable JIT cache filenames and introspection. -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``) is emitted with the cause and the remediation. The warning -fires once per ``(op, dtype)`` rather than once per process, so distinct -fallback causes each surface a warning. Silence it by category or by -message:: +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 @@ -281,16 +282,25 @@ 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 auto-fixes this at import. 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), and ``jit_c_control`` is -bumped from the SS default ``'run'`` (run cached kernels only; no compile, -no load from disk) to ``'on'`` (compile, load, and run). When the default -config is already valid, only the mode bump applies. +python-graphblas repairs the compiler path when ``graphblas.ss`` is +imported, which the first access to ``gb.ss`` does. 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`` is raised from the SS default ``'run'`` (run cached kernels +only; no compile, no load from disk) to ``'on'`` (compile, load, and run) when +a UDT operator gets C source for SuiteSparse to compile, such as the first +``binary.plus[udt]``, or when ``gb.dtypes.ss.register_new`` registers a type +from a C typedef. Any other setting is kept, including ``'off'``, ``'pause'``, +and ``'load'`` (load and run cached kernels, never compile). SuiteSparse +itself sets ``'load'`` after a compile fails, so keeping it means a failing +compile is tried once, not again at every later UDT operation. Call the helper manually to re-fix or verify:: - gb.ss.fix_jit_config() # repair compiler path (full probe) + 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 diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index b55a17f13..ede25b467 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -956,7 +956,14 @@ def _set_jit_c_strings(gb_obj, c_name, c_defn, set_string_func): 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() set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) diff --git a/graphblas/core/ss/dtypes.py b/graphblas/core/ss/dtypes.py index 1b0d93506..eeca52d0b 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -40,6 +40,12 @@ def register_new(name, jit_c_definition, *, np_type=None): 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 _enable_jit_for_udt + + _enable_jit_for_udt() + 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()) diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py index f77b55651..8cdc19e25 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -62,8 +62,8 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): 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/``, and strips build-time-only flags (``-isysroot``, - ``-fdebug-prefix-map``). + ``$CONDA_PREFIX/bin/``, strips build-time-only flags (``-isysroot``, + ``-fdebug-prefix-map``), and sets ``jit_c_control`` to ``'on'``. Parameters ---------- @@ -73,10 +73,10 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): to restrict the repair to a conda environment only. probe : bool, default True After fixing the config, try to JIT-register a trivial UDT to verify - the compiler actually works. SuiteSparse auto-flips ``jit_c_control`` - from ``'on'`` to ``'load'`` on a failed compile; the probe absorbs - that first failure so user-visible ops afterwards see a stable - ``'load'`` (cache-only) state and punt to generic cleanly. + 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 ------- @@ -84,30 +84,47 @@ def fix_jit_config(*, use_sysconfig=True, probe=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 (typically ``'load'``). + 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() - if conda_prefix := os.environ.get("CONDA_PREFIX", ""): - rv = _fix_from_conda(cfg, conda_prefix) - elif use_sysconfig: - rv = _fix_from_sysconfig(cfg) - else: - return None + rv = _repair_jit_compiler(cfg, use_sysconfig=use_sysconfig) if rv is None: return None + # Enabling compilation is what the caller asked for. The import-time + # repair deliberately stops short of this; see ``_auto_fix_jit_at_import``. + 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 $CONDA_PREFIX/bin/.""" jit_cc = cfg["jit_c_compiler_name"] @@ -122,7 +139,6 @@ def _fix_from_conda(cfg, conda_prefix): else: return None # nothing usable _fix_compiler_flags(cfg) - cfg["jit_c_control"] = "on" return True @@ -139,7 +155,6 @@ def _fix_from_sysconfig(cfg): cfg["jit_c_compiler_flags"] = f"{cflags} -I{include}" if libs := sysconfig.get_config_var("LIBS"): cfg["jit_c_libraries"] = libs - cfg["jit_c_control"] = "on" return True @@ -182,11 +197,13 @@ def _strip_mismatched_arch(flags): def _probe_jit(cfg): """Probe a trivial JIT compile to verify the config works. - On failure, SuiteSparse will have flipped ``jit_c_control`` from - ``'on'`` to ``'load'`` (its built-in response to a failed compile); - we leave that state alone. Both ``'load'`` and ``'off'`` cause - downstream ops to punt to the generic kernel, but ``'load'`` preserves - any pre-compiled kernels in the cache. + ``jit_c_control`` is left wherever SuiteSparse puts it, and that 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. """ from ... import dtypes as _dtypes @@ -197,6 +214,8 @@ def _probe_jit(cfg): probe_name = "_jit_probe" if hasattr(_dtypes.ss, probe_name): return True + global _probing_jit + _probing_jit = True try: _dtypes.ss.register_new(probe_name, "typedef struct { int _probe ; } _jit_probe ;") except Exception: @@ -205,34 +224,94 @@ def _probe_jit(cfg): # validation ``ValueError``s. The probe's contract is "did this # work?", so absorb every failure mode here. return False + finally: + _probing_jit = False return True def _auto_fix_jit_at_import(): - """Run :func:`fix_jit_config` 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 - delegates the failure-prone work to ``_probe_jit``, which catches - everything internally. - - The probe is the load-bearing piece: without it, SS would surface - ``JitError`` on the first user-triggered JIT compile (a failed - compile is only converted to a silent ``'load'`` fallback on - subsequent calls). + """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`` at whatever SuiteSparse set. 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. + Compilation is enabled later, by :func:`_enable_jit_for_udt`, when a UDT + actually needs a kernel built. """ cfg = _ss_config() if "jit_c_control" not in cfg: return - if jit_compiler_is_usable(): - if cfg["jit_c_control"] in ("run", "load"): - cfg["jit_c_control"] = "on" - else: - fix_jit_config(use_sysconfig=True, probe=False) - if jit_compiler_is_usable() and cfg["jit_c_control"] in ("run", "load"): + if not jit_compiler_is_usable(): + _repair_jit_compiler(cfg) + + +# Tri-state: ``None`` until a UDT first asks for a JIT kernel 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 + +# True only while ``_probe_jit`` is registering its own UDT. That registration +# goes through ``dtypes.ss.register_new``, which asks to enable the JIT, which +# would probe again; the second probe would register the same name a second +# time and the first one would then die on cffi's "multiple declarations". +_probing_jit = False + + +def _enable_jit_for_udt(): + """Enable JIT compilation the first time a UDT needs a kernel built. + + SuiteSparse defaults ``jit_c_control`` to ``'run'``, 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). + + This is where that bump belongs, rather than at import: registering a + type from its C typedef 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. An explicit ``'off'``, ``'pause'``, + or ``'load'`` is honored, 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 + if _probing_jit: + # Re-entered from the probe's own registration. The probe is the thing + # deciding this answer, so say yes and let it finish rather than + # starting a second one inside it. + return True + cfg = _ss_config() + control = cfg.get("jit_c_control") + if control in (None, "off", "pause", "load"): + # ``None`` is a library with no JIT at all. The others are honored as + # set and not cached, so a saved ``'run'`` restored afterwards is + # still raised. + return False + if _jit_enabled_for_udt is None: + if not jit_compiler_is_usable(): + _repair_jit_compiler(cfg) + if _jit_enabled_for_udt := jit_compiler_is_usable(): cfg["jit_c_control"] = "on" - if cfg.get("jit_c_control") == "on": - _probe_jit(cfg) + # 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 # Keyed by ``(op_name, dtype_name)`` so each distinct pair warns once. @@ -244,8 +323,15 @@ def _auto_fix_jit_at_import(): 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 most likely cause (bogus compiler path, ``jit_c_control`` off, or - UDT not C-expressible) is named in the message along with the remediation. + The only caller is ``udt_utils._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: @@ -253,41 +339,20 @@ def _maybe_warn_no_jit(*, op_name="", dtype_name=""): _warned_no_jit_for.add(key) import warnings as _warnings - cfg = _ss_config() - if not jit_compiler_is_usable(): - cause = ( - "the JIT compiler path is not usable " - f"({cfg.get('jit_c_compiler_name', '')!r}); " - "call ``gb.ss.fix_jit_config()`` to repair it" - ) - elif cfg.get("jit_c_control") != "on": - cause = ( - f"jit_c_control is {cfg.get('jit_c_control')!r} (must be 'on' to compile); " - "set ``gb.ss.config['jit_c_control'] = 'on'`` to enable compilation" - ) - else: - # Compiler is usable and mode is OK; the UDT itself isn't C-expressible. - # The introspection properties (``DataType.jit_c_definition`` and - # ``TypedUserBinaryOp.jit_c_source``) show what was generated, or - # ``None`` when codegen was skipped. - loc = f" (op={op_name!r}, dtype={dtype_name!r})" if op_name else "" - cause = ( - "this UDT is not expressible as a C struct" - f"{loc} (a field name is a C reserved word or stdlib macro, a " - "field type isn't in the numpy-to-C map, a field is array-typed, " - "or the record has a packed layout). The op still works via the " - "Numba cfunc path; only the JIT speedup is lost" - ) 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: {cause}. " - f"Operations will use the Numba function-pointer fallback " - f"(typically 2-3x slower for elementwise ops, since SuiteSparse " - f"can't inline the kernel into its eWise and reduce templates). " - f"This warning fires once per (op, dtype) per process; silence with " - f"``warnings.filterwarnings('ignore', category=gb.exceptions.NoJITWarning)`` " - f"or by message match.", + 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/exceptions.py b/graphblas/exceptions.py index b78023db3..df674b55e 100644 --- a/graphblas/exceptions.py +++ b/graphblas/exceptions.py @@ -109,10 +109,9 @@ class NoJITWarning(UserWarning): 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 couldn't because the JIT compiler is unusable, ``jit_c_control`` - is off, or the UDT isn't expressible as a C struct. Fires once per - ``(op, dtype)`` pair per process, so a user who registers several - UDTs gets one warning per pair regardless of cause. + 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: diff --git a/graphblas/ss/__init__.py b/graphblas/ss/__init__.py index 394dcc391..33a342373 100644 --- a/graphblas/ss/__init__.py +++ b/graphblas/ss/__init__.py @@ -11,12 +11,14 @@ 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 - # Auto-fix the JIT config 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). When the compiler is already usable, the - # call only bumps ``jit_c_control`` from SS's default ``run`` to ``on`` so - # first-time UDT auto-lift actually JIT-compiles. Skipped on vanilla: the - # auto-fix reads ``gb.ss.config``, which needs GxB callables vanilla strips. + # 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 SuiteSparse set it, 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/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 2e3e17d87..a9550ff51 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -461,8 +461,12 @@ def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): # 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"): + 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 SuiteSparse 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. @@ -478,6 +482,39 @@ def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): 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, SuiteSparse leaves ``jit_c_control`` + at ``'run'`` and ``_enable_jit_for_udt`` caches its ``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``. diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index fc174c768..34e4767d6 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1,6 +1,9 @@ import contextlib import os +import subprocess +import sys import sysconfig +from pathlib import Path import numpy as np import pytest @@ -23,6 +26,7 @@ from graphblas.core.operator.indexbinary import _has_idxbinop from graphblas.core.operator.udt_utils import _has_jit_set from graphblas.core.ss import _IS_SSGB7 +from graphblas.core.ss.jit_config import _enable_jit_for_udt from .conftest import autocompute, burble @@ -41,13 +45,12 @@ # 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) -# Capture the post-import JIT state before the autouse ``_setup_jit`` fixture -# runs. ``_auto_fix_jit_at_import`` probes once; ``jit_c_control`` is ``'on'`` -# iff the env can actually JIT-compile. Tests that assert on the import-time -# state read this constant; the fixture may transiently mutate the live config. -_JIT_WORKS_AT_IMPORT = ( - not _IS_SSGB7 and backend == "suitesparse" and gb.ss.config["jit_c_control"] == "on" -) +# ``_auto_fix_jit_at_import`` repairs the compiler path but leaves +# ``jit_c_control`` alone, so compilation is enabled on demand by +# ``_enable_jit_for_udt``. Ask it once here, before the autouse ``_setup_jit`` +# fixture mutates the live config: it answers whether this environment can +# JIT-compile at all, and caches that answer for the rest of the process. +_JIT_WORKS_AT_IMPORT = not _IS_SSGB7 and backend == "suitesparse" and _enable_jit_for_udt() @pytest.fixture(scope="module", autouse=True) @@ -57,9 +60,8 @@ def _setup_jit(): Strategy: 1. _fix_jit_config(): fix conda-baked compiler paths and probe. - Returns True: JIT works, proceed. - - Returns False: probe failed. SuiteSparse will have left - ``jit_c_control = 'load'`` (compile disabled, cache loading - still allowed); we leave that state untouched. + - 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). """ @@ -71,10 +73,9 @@ def _setup_jit(): prev = gb.ss.config["jit_c_control"] result = _fix_jit_config() - # ``True``: conda JIT configured and verified. ``False``: probe failed - # (``_probe_jit`` leaves ``jit_c_control`` at ``'load'``; don't try - # sysconfig, since if the conda compiler can't build GraphBLAS JIT - # kernels, Python's sysconfig compiler won't either). Both done. + # ``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") @@ -111,26 +112,233 @@ def _jit_mode(mode): gb.ss.config["jit_c_control"] = prev +def _run_in_fresh_process(script): + """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. + """ + repo_root = Path(__file__).resolve().parents[2] + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + env=dict(os.environ, PYTHONPATH=str(repo_root)), + ) + 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") -def test_auto_fix_jit_at_import_left_compiler_usable(): - """After ``import graphblas.ss``, a probe-confirmed compiler implies - ``jit_c_control == 'on'``. +def test_ss_import_leaves_jit_c_control_alone(): + """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. """ + lines, report = _run_in_fresh_process(_SS_IMPORT_PROBE) + assert lines["after"] == lines["before"], report + + +# 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 registers a UDT, and registering a UDT asks to enable the JIT, + which would start a second probe inside the first. The inner probe would + register the probe type and the outer one would then die redeclaring it, + so a working environment would answer ``False`` the first time and + ``True`` every time after. + """ + lines, report = _run_in_fresh_process(_FIX_JIT_TWICE_PROBE) + assert lines["first"] == lines["second"], report + + +@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 import-time 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 (``'load'`` - # or ``'run'``) varies by SuiteSparse version. + # 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) + # 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 + + +@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 a saved ``'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'`` and ``'load'``. + """ + 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) + for start in ("off", "pause", "load"): + 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 + 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(): + """The same raise, through SuiteSparse's own config rather than a stand-in.""" + 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 + gb.ss.config["jit_c_control"] = "run" + assert _enable_jit_for_udt() is True + assert gb.ss.config["jit_c_control"] == "on" + finally: + gb.ss.config["jit_c_control"] = prev_control + jit_config._jit_enabled_for_udt = prev_enabled + + +@pytest.mark.skipif("_IS_SSGB7") +@pytest.mark.skipif("not supports_udfs") +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"] = "run" + # 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``.""" From 5aa3d4ceb57df1f384fedae202c320a882065a53 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 1 Oct 2026 12:10:29 -0500 Subject: [PATCH 21/52] Refresh the sparse pool, fix CI on SuiteSparse 8, and fix ss.register_new after #630 (#631) * Refresh the sparse pool and lift its stale Python 3.13 exclusion The pools refresh below this commit left sparse behind: the pool stopped at 0.15 while 0.16 through 0.19.2 exist, and py3.13/3.14 hardcoded "NA" under a comment saying sparse does not support them. That was true in the 0.15 era, when numba could not run on those Pythons; it is not true now. Constraints read off conda-forge metadata rather than assumed: every sparse release here is noarch with numba >=0.49 and numpy >=1.17 and no numpy ceiling, so only the Python floors matter. 0.16/0.17 declare python >=3.10 and 0.18/0.19 declare >=3.11. On py3.13/3.14 the pool starts at 0.16 because 0.14/0.15 predate those Pythons; on py3.11/3.12 the old entries stay. No new pairing rules are needed: the workflow already skips sparse whenever numba is skipped or NA, which covers sparse's only runtime dependency. The validate() Python check moves from "any sparse on py3.13+" to the 0.16 floor, and the stress test passes 160000 combos. Verified beyond metadata on this machine (py3.14, macOS arm64): sparse 0.16.0 and 0.19.2 pass the pydata-sparse io tests, and 0.18.0 passed the full suite. **Edit: rebased to include the changes on main, which largely supercedes this** * Test and allow every sparse release on Python 3.13 and 3.14 The CI pools started sparse at 0.16 on py3.13 and 0.18 on py3.14, and the io extra in pyproject.toml had the same floors. Each was the first sparse release to come out after that Python, on the assumption that older ones might not work there. They do. sparse is noarch, and every release from 0.14 on declares numba >=0.49, numpy >=1.17 and a Python floor of 3.11 or lower (>=3.8 for 0.14/0.15, >=3.10 for 0.16/0.17, >=3.11 for 0.18/0.19), with no ceiling on any of them. Each pin, as conda-forge resolves it (0.14.0, 0.15.5, 0.16.0, 0.17.0, 0.18.0 and 0.19.2), solves on both Pythons and passes the pydata-sparse tests with the oldest and newest numpy/numba pair each Python draws: numpy 2.1 with numba 0.62 and numpy 2.5 with numba 0.67 on py3.13, numpy 2.3 with numba 0.63 and numpy 2.5 with numba 0.67 on py3.14. So every Python now draws from the same pool, 0.14 through 0.19 plus latest, and the io extra requires sparse >=0.14 wherever numba installs. No pairing rules are needed, since the workflow already drops sparse whenever it skips numba. The stress test passes 160000 combinations, and both version audits pass. * Skip the failed-compile JIT test on SuiteSparse 8 test_failed_compile_is_not_retried_at_the_next_arming, from #630, arms binary.plus on a UDT with C source for SuiteSparse to compile, then makes that compile fail. An op is armed with C source only when the JIT string setters exist, and those arrived in SuiteSparse 9, so on 8.x jit_c_source is None and the test failed at its first assertion. It gets that far only where the JIT probe passes. In CI that happens only on some macOS runners that draw psg 8 (numpy 1.x); main's CI failed there after #630 merged. The test now skips when _has_jit_set is false, like the other tests that need an armed op. On 8.x the stand-in config tests still check that _enable_jit_for_udt leaves 'load' alone. Locally on psg 9.3.1.0, the only release in the CI pool that takes its pre-9.4 branch, it runs and passes. * Enable the JIT when registering from C source, and say why it failed gb.unary.ss.register_new and its binary, indexunary, select, and indexbinary siblings create the op from a C definition and no function pointer, so SuiteSparse has to compile the definition, or load it from its cache, before the op exists. Under its default jit_c_control of 'run' it does neither, and the registration fails with a bare NullPointer. Before #630, importing graphblas.ss raised the control to 'on', and the docstrings of these functions send users to gb.ss.config, so touching it first was enough; a process that never touched gb.ss already failed. #630 stopped the import from changing the control and moved the raise to _enable_jit_for_udt, which UDT op arming and gb.dtypes.ss.register_new call but these registrars did not. Since then nothing raises it for them, and each fails in a fresh process, on SuiteSparse 8.0.2 and 10.5 alike. The four op registrars and gb.dtypes.ss.register_new now create through one context manager, _compiling_c_source, which raises 'run' first and still honors an explicit 'off', 'pause', or 'load'. When creation fails anyway, SuiteSparse's error (NullPointer, InvalidValue, or from 9.4 JitError) has no message, so the same exception is raised again with the reason, read off jit_c_control before and after: the JIT is off or paused; it was left at 'load' and the kernel is not cached; this compile failed and SuiteSparse set 'load'; the compiler probe failed in this very call; or the control is still 'run' because no compiler was found or the library has no JIT. gb.dtypes.ss.register_new also accepted a name of GxB_MAX_NAME_LEN characters. That limit counts the NUL, so SuiteSparse cut the name by one, it no longer matched the typedef, and the compile failed. It is now rejected up front. The user guide's list of what raises the control includes the op registrars. New tests register one op of each kind in a fresh process, resetting the control to 'run' before each so none can rely on another, check each reason's message, and check the name limit. * Fix small errors in docstrings, messages, and docs Messages: Matrix, Vector, and Scalar __sizeof__ printed a literal {backend} (no f prefix); the Matrix diagonal size-mismatch message ran two sentences together, and the nrows/ncols error two words; the compression level error named only "lz4hc" where "zstd" also takes a level; the UDT repr size warning ended its first sentence twice. Docstrings: isequal documented check_dtype as check_dtypes; from_csc documented col_indices and said ncols is computed where it takes row_indices and computes nrows; from_dicts documented d for nested_dicts; Vector.diag documented a dtype it does not take; Vector.from_scalar documented nrows and ncols for size; serialize gave the zstd default level as 19 and a 1-9 range for both methods (zstd defaults to 1 and takes 1-19), and the descriptor docs gave the lz4hc default as 1 (it is 9); operator registration pointed at a graphblas.ops namespace, which is graphblas.op; several parameter lines numpydoc cannot parse ("name: type", "name :type"); two choice lists missing a closing quote; .mtz.gz for .mtx.gz in mmread; "`first` any `any`" in a comment. Docs: "the the" in the operators guide, the README converting the scipy matrix back instead of the GraphBLAS one, and two misspellings in a CSS comment. --- README.md | 2 +- docs/_static/matrix.css | 4 +- docs/user_guide/operators.rst | 2 +- docs/user_guide/udt.rst | 19 ++-- graphblas/core/dtypes.py | 2 +- graphblas/core/matrix.py | 16 +-- graphblas/core/operator/base.py | 4 +- graphblas/core/operator/binary.py | 2 +- graphblas/core/operator/monoid.py | 10 +- graphblas/core/operator/semiring.py | 2 +- graphblas/core/operator/unary.py | 2 +- graphblas/core/scalar.py | 4 +- graphblas/core/ss/binary.py | 29 ++--- graphblas/core/ss/descriptor.py | 4 +- graphblas/core/ss/dtypes.py | 24 ++-- graphblas/core/ss/indexbinary.py | 31 +++--- graphblas/core/ss/indexunary.py | 29 ++--- graphblas/core/ss/jit_config.py | 78 +++++++++++-- graphblas/core/ss/matrix.py | 6 +- graphblas/core/ss/unary.py | 27 +++-- graphblas/core/ss/vector.py | 6 +- graphblas/core/utils.py | 2 +- graphblas/core/vector.py | 16 +-- graphblas/io/_awkward.py | 2 +- graphblas/io/_matrixmarket.py | 2 +- graphblas/ss/_core.py | 2 +- graphblas/tests/test_ssjit.py | 164 +++++++++++++++++++++++++++- pyproject.toml | 11 +- scripts/ci_pick_versions.py | 22 ++-- 29 files changed, 386 insertions(+), 138 deletions(-) diff --git a/README.md b/README.md index 7b157c979..81df4fe64 100644 --- a/README.md +++ b/README.md @@ -267,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/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/user_guide/operators.rst b/docs/user_guide/operators.rst index 6a66c295c..59a38c791 100644 --- a/docs/user_guide/operators.rst +++ b/docs/user_guide/operators.rst @@ -331,7 +331,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/udt.rst b/docs/user_guide/udt.rst index d37653250..15cbedbb4 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -290,13 +290,18 @@ import changes nothing else, so reading ``gb.ss.about`` does not change what later operations compute. ``jit_c_control`` is raised from the SS default ``'run'`` (run cached kernels -only; no compile, no load from disk) to ``'on'`` (compile, load, and run) when -a UDT operator gets C source for SuiteSparse to compile, such as the first -``binary.plus[udt]``, or when ``gb.dtypes.ss.register_new`` registers a type -from a C typedef. Any other setting is kept, including ``'off'``, ``'pause'``, -and ``'load'`` (load and run cached kernels, never compile). SuiteSparse -itself sets ``'load'`` after a compile fails, so keeping it means a failing -compile is tried once, not again at every later UDT operation. +only; no compile, no load from disk) to ``'on'`` (compile, load, and run) by +each of these, since each asks SuiteSparse to compile C source: + +- a UDT operator that gets C source, such as the first ``binary.plus[udt]``; +- ``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 setting other than ``'run'`` is kept, including ``'off'``, ``'pause'``, and +``'load'`` (load and run cached kernels, never compile). 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:: diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 73b905306..93b3d4d8a 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -308,7 +308,7 @@ def register_anonymous(dtype, name=None): 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, ) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 698bf7b4f..4625e4258 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -368,7 +368,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,7 +377,7 @@ 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 @@ -532,7 +532,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`` @@ -1131,15 +1131,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``. @@ -1529,7 +1529,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 @@ -2974,7 +2974,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}." ) diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 28c2291f8..485857cfb 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -833,7 +833,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) @@ -903,7 +903,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 diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 5329b4415..2cdab6eff 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -402,7 +402,7 @@ 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__ = ( diff --git a/graphblas/core/operator/monoid.py b/graphblas/core/operator/monoid.py index eb3aed6c1..0b02d4226 100644 --- a/graphblas/core/operator/monoid.py +++ b/graphblas/core/operator/monoid.py @@ -226,7 +226,7 @@ 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" @@ -347,10 +347,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 @@ -381,10 +381,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 diff --git a/graphblas/core/operator/semiring.py b/graphblas/core/operator/semiring.py index 32e20d0bf..b333259e3 100644 --- a/graphblas/core/operator/semiring.py +++ b/graphblas/core/operator/semiring.py @@ -175,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" diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index 64f68eb48..fc9c92661 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -133,7 +133,7 @@ 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" diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 9755acb3c..ce191f374 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -182,7 +182,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). @@ -191,7 +191,7 @@ 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 diff --git a/graphblas/core/ss/binary.py b/graphblas/core/ss/binary.py index d53608818..87c213d32 100644 --- a/graphblas/core/ss/binary.py +++ b/graphblas/core/ss/binary.py @@ -106,20 +106,23 @@ 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 + 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[]", funcname.encode()), + ffi_new("char[]", jit_c_definition.encode()), + ), + "BinaryOp", + gb_obj[0], + ) op = TypedJitBinaryOp( rv, funcname, left_type, ret_type, gb_obj[0], jit_c_definition, dtype2=right_type ) diff --git a/graphblas/core/ss/descriptor.py b/graphblas/core/ss/descriptor.py index 20fc126cf..c8a2fcade 100644 --- a/graphblas/core/ss/descriptor.py +++ b/graphblas/core/ss/descriptor.py @@ -153,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. @@ -212,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 eeca52d0b..12f74fbce 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -31,9 +31,13 @@ def register_new(name, jit_c_definition, *, np_type=None): 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): 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`") @@ -42,15 +46,17 @@ def register_new(name, jit_c_definition, *, np_type=None): # 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 _enable_jit_for_udt - - _enable_jit_for_udt() + 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*") diff --git a/graphblas/core/ss/indexbinary.py b/graphblas/core/ss/indexbinary.py index 89a38b3fc..85b3253e0 100644 --- a/graphblas/core/ss/indexbinary.py +++ b/graphblas/core/ss/indexbinary.py @@ -111,21 +111,24 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): ) else: rv = IndexBinaryOp(name, is_udt=True) + from .jit_config import _compiling_c_source + gb_obj = ffi_new("GxB_IndexBinaryOp*") - 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[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), - ), - "IndexBinaryOp", - gb_obj[0], - ) + 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[]", funcname.encode()), + ffi_new("char[]", jit_c_definition.encode()), + ), + "IndexBinaryOp", + gb_obj[0], + ) op = TypedJitIndexBinaryOp( rv, funcname, x_type, ret_type, gb_obj[0], jit_c_definition, dtype2=theta_type ) diff --git a/graphblas/core/ss/indexunary.py b/graphblas/core/ss/indexunary.py index b60837acf..fa394eb4f 100644 --- a/graphblas/core/ss/indexunary.py +++ b/graphblas/core/ss/indexunary.py @@ -105,20 +105,23 @@ 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 + 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[]", funcname.encode()), + ffi_new("char[]", jit_c_definition.encode()), + ), + "IndexUnaryOp", + gb_obj[0], + ) op = TypedJitIndexUnaryOp( rv, funcname, input_type, ret_type, gb_obj[0], jit_c_definition, dtype2=thunk_type ) diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py index 8cdc19e25..8f3aa4a9a 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -20,6 +20,7 @@ warning at import or at first JIT use. """ +import contextlib import os import pathlib import platform @@ -241,7 +242,7 @@ def _auto_fix_jit_at_import(): submodule, and that must not change what any later operation computes or which kernels SuiteSparse is willing to load from its on-disk cache. Compilation is enabled later, by :func:`_enable_jit_for_udt`, when a UDT - actually needs a kernel built. + 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: @@ -250,7 +251,7 @@ def _auto_fix_jit_at_import(): _repair_jit_compiler(cfg) -# Tri-state: ``None`` until a UDT first asks for a JIT kernel while the +# Tri-state: ``None`` until 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 @@ -263,17 +264,19 @@ def _auto_fix_jit_at_import(): def _enable_jit_for_udt(): - """Enable JIT compilation the first time a UDT needs a kernel built. + """Enable JIT compilation the first time C source needs compiling. SuiteSparse defaults ``jit_c_control`` to ``'run'``, 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). + 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 from its C typedef 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. + 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. An explicit ``'off'``, ``'pause'``, or ``'load'`` is honored, and ``'load'`` is also where SuiteSparse leaves @@ -314,6 +317,67 @@ def _enable_jit_for_udt(): 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": + 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. diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index 0580b7581..2c9a4d99e 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -4059,10 +4059,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 diff --git a/graphblas/core/ss/unary.py b/graphblas/core/ss/unary.py index 0b7ced3c8..e0be45db9 100644 --- a/graphblas/core/ss/unary.py +++ b/graphblas/core/ss/unary.py @@ -90,19 +90,22 @@ 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 + 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[]", funcname.encode()), + ffi_new("char[]", jit_c_definition.encode()), + ), + "UnaryOp", + gb_obj[0], + ) op = TypedJitUnaryOp(rv, funcname, input_type, ret_type, gb_obj[0], jit_c_definition) rv._add(op, is_jit=True) setattr(module, funcname, rv) diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index 6c79a338b..80932da73 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -1630,10 +1630,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 diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index aef2f76a3..72f7219bd 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -153,7 +153,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: diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 8c73ecc48..5d8c8fef9 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -335,7 +335,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,7 +344,7 @@ 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 @@ -469,7 +469,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`` @@ -609,8 +609,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. @@ -800,10 +798,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. @@ -1316,7 +1312,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 diff --git a/graphblas/io/_awkward.py b/graphblas/io/_awkward.py index fd4a9bab8..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: diff --git a/graphblas/io/_matrixmarket.py b/graphblas/io/_matrixmarket.py index cf425e2e9..4d7fcca22 100644 --- a/graphblas/io/_matrixmarket.py +++ b/graphblas/io/_matrixmarket.py @@ -21,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``, diff --git a/graphblas/ss/_core.py b/graphblas/ss/_core.py index 91e1496e9..5e851b7cc 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -136,7 +136,7 @@ class GlobalConfig(BaseConfig): 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} + 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 diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 34e4767d6..eeeeac206 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -3,6 +3,7 @@ import subprocess import sys import sysconfig +import uuid from pathlib import Path import numpy as np @@ -195,6 +196,149 @@ def test_fix_jit_config_first_call_agrees_with_the_second(): 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 ``'run'``, 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"] = "run" + 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 SuiteSparse's 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 SuiteSparse's default in place. + monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", False) + gb.ss.config["jit_c_control"] = "run" + 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"] = "run" + 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_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.""" @@ -286,8 +430,11 @@ def test_enable_jit_for_udt_re_raises_the_live_control(): jit_config._jit_enabled_for_udt = prev_enabled -@pytest.mark.skipif("_IS_SSGB7") @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. @@ -444,6 +591,21 @@ def test_jit_udt(): assert dtype.np_type == np_type +@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: diff --git a/pyproject.toml b/pyproject.toml index cdf8867c8..9bd988a10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,14 +85,9 @@ io = [ "python-graphblas[networkx,scipy]", "python-graphblas[numba]; python_version<'3.15'", # make optional where numba is not supported "awkward >=2.0", - # sparse needs numba, so it stays optional. numba gained py3.13 support in 0.61 and - # py3.14 in 0.63; sparse 0.16 is the first release that postdates py3.13, and 0.18 the - # first that postdates py3.14. - "sparse >=0.14; python_version<'3.13'", - "sparse >=0.16; python_version=='3.13'", - # Upper guard for the same reason as numba above: sparse depends on numba, so it has - # to drop out wherever numba does, not just trail it. - "sparse >=0.18; python_version>='3.14' and python_version<'3.15'", + # 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.15'", # fast-matrix-market is no longer maintained; last supported Python is 3.12 ] viz = ["python-graphblas[networkx,scipy]", "matplotlib >=3.6"] diff --git a/scripts/ci_pick_versions.py b/scripts/ci_pick_versions.py index 2fcfafc02..ebd003d96 100755 --- a/scripts/ci_pick_versions.py +++ b/scripts/ci_pick_versions.py @@ -139,16 +139,20 @@ "3.14": ["6.0", ""], } -# sparse is noarch, so there are no per-Python builds to check; the split is by what -# each release plausibly saw. The 0.14/0.15 series began well before py3.13 (0.15.0 is -# Jan 2024), and 0.16/0.17 (Apr/May 2025) predate py3.14. 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-16 sparse 0.19.2 is latest; 0.18 verified against py3.14 locally +# 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.16", "0.17", "0.18", "0.19", ""], - "3.14": ["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", ""], } # PSG versions to pair with numpy 1.x (only reachable on py3.11/py3.12, since @@ -424,6 +428,10 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): 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 From 03cf0bcc66fcf4695957221302bb62c7608ddf22 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 2 Oct 2026 21:56:46 -0700 Subject: [PATCH 22/52] Reject UDT UDF returns that cannot fill the element (#590) * Reject UDT UDF returns that cannot fill the element A UDF over an array UDT whose return does not fit the element is not a type error. Numba's `Array` type carries `ndim` but not its extents, so returning `x[:2]` from a 9-element UDT types identically to returning `x`. The wrapper's `z[:] = ...` then raises inside the cfunc, where Numba prints the traceback and returns. GraphBLAS is handed no error, and the element keeps whatever was already in the buffer, so the caller gets a wrong answer and no exception. A record UDF that under-fills an array-typed leaf fails the same way and abandons the write part-way: every leaf after the short one, scalar leaves included, keeps what SuiteSparse had there. Extents can only be recovered by running the function, so `_get_udt_wrapper` and the IndexBinaryOp wrapper now call it once on stand-in operands and check that the result fits the UDT. Fitting is not the same as matching. Slice-assign broadcasts, so a `(1,)` return fills a `(6,)` element and a `(1, 3)` return fills both rows of a `(2, 3)` one; both work today and must keep working, and an equality check would reject them. `_fits_by_broadcast` applies numpy's assignment rule instead, and `test_udt_broadcast_matches_numba_slice_assign` pins it against Numba's own slice-assign in both directions, including the two ranks where broadcasting alone gives the wrong answer. The operands are ones rather than zeros for the same reason `dtypes._sample_values` avoids zeros: a UDF that divides by an operand would raise on a zero-filled probe, and a raising probe is treated as "cannot check" rather than as a failure. The check is best-effort, and this commit should not claim more. A UDF that raises on the probe values is not checked at all, so one that also returns a shape that does not fit still reaches the wrapper and still fails silently, exactly as it did before. The wrapper's slice-assign remains the only backstop that always runs. What the check buys is that the common case, a UDF that is simply wrong about the shape, becomes an error at registration instead of a wrong answer at apply time. Because the probe reports the shape for the values it was run on, the diagnostic says so. A UDF whose output shape depends on its operands' values rather than their types can otherwise be rejected on a shape it would never return for real data. This is a judgement call rather than a straight fix, and it is cheap to decline. Typing an op for a UDT now executes the user's function, where before it only compiled it. `OpBase.__contains__` is implemented as a typed lookup, so `udt in some_op` runs user code as a side effect of a membership test. Measured: registering an op runs no probe, since compilation stays lazy, but a single `udt in some_op` runs exactly one. That probe costs no extra compile, because it calls the function with an `ndarray` and the wrapper's `numba.carray` view needs that same specialization anyway. A/B with the check stubbed out: 93.4 ms median disabled against 92.1 ms enabled, two Numba specializations either way. Only array UDTs and records that actually have an array leaf are probed, so an ordinary record UDF is still typed without running user code. The array-UDT paragraph in `docs/user_guide/udt.rst`, added one commit earlier, described the rule as building "a new array of the same shape". It now states the rule the check enforces, and says the shape is learned by running the UDF. If this is declined, the array-UDT view change it sits on still stands. A return that does not fit stays memory-safe, because the slice-assign cannot overrun the element, and goes back to being silent. * Drop the flat-vs-layered ambiguity case from the shape-error test test_udt_array_udf_shape_errors expected a flat FP64[3, 4] and a layered FP64[4] x 3 to be separate UDTs, so that a UDF building an array from the pair was ambiguous. "Flatten nested subarray dtypes so each UDT layout has one DataType" (#589) made them the same DataType, so nothing is ambiguous and the test failed with "DID NOT RAISE UdfParseError" on every CI job with Numba. Keep the same-type check that block ended with, now on udt9. The flattening itself is tested where it is introduced. * Run a UDT UDF for its shape only when Numba cannot tell it The shape check ran the UDF on sample values for every array return and every record return with an array leaf. An operand, or operand field, returned as-is keeps its extents in Numba's `NestedArray` type, so check those from the type and run the UDF only for an array it builds (a plain `Array`). Returning an operand no longer runs user code when the op is typed, and fields returned into the wrong slots are caught even if the UDF would raise on the sample values. The two checkers become one, `_check_udf_fills_output`. Also reject a record leaf the UDF can return as `None` (typed `Optional`, or always `None`), scalar leaves included. On that path the wrapper's assignment raises inside the cfunc and the element comes back holding another element's values. The probe docstring's reason for using ones was wrong: UDFs compile with `error_model="numpy"`, so dividing by zero does not raise. Drop two test cases that repeated #589's resolver tests, and add a test pinning when the UDF is run. * Ignore the UDT repr size warning in the shape-check probe test Below SuiteSparse 9, register_anonymous stores a UDT's numpy repr in a type name capped at 128 characters, and warns when the repr is longer. _ProbeWhenRec's is 142, so the test failed on psg 8 draws that ran slow tests on the suitesparse backend. It is about the shape check, not serialization. * Reject record returns that cannot be written to their fields Two kinds of record UDF return got past the shape check. A tuple or list of the wrong length for an array field failed silently, like a wrong-shape array. Numba copies a sequence into a field element by element, only into a 1-D field and only at its exact length. It does not broadcast, so even a 1-tuple for a 3-element field raises inside the cfunc, and the element keeps whatever was in the buffer. A tuple's length is in its type; a list's is not, so a list is run on sample values like a built array. An array, tuple or list returned for a scalar field failed only at the wrapper compile, with Numba's long TypingError. It is now rejected from its type with a UdfParseError, as a field that can be None already is. Also test the probe's IndexUnaryOp path, which nothing covered. If the row and column were dropped from the probe call, the check silently did nothing. Pin that a scalar or tuple for an array field is checked without running the UDF, and that every leading 1 is dropped in the broadcast rule. udt.rst names the error and documents the record-field rules. --- docs/user_guide/udt.rst | 18 +- graphblas/core/operator/base.py | 201 +++++++++++++- graphblas/tests/test_indexbinary.py | 26 ++ graphblas/tests/test_op.py | 403 ++++++++++++++++++++++++++++ 4 files changed, 646 insertions(+), 2 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 15cbedbb4..ae2cef9e2 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -179,10 +179,17 @@ For nested record UDTs the tuple is *flat over the leaves*. Given 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 of the same shape: +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 @@ -199,6 +206,15 @@ the UDF may return one of its operands or build a new array of the same shape: 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. diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 485857cfb..e4c41b95e 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -3,10 +3,13 @@ from operator import getitem from types import BuiltinFunctionType, ModuleType +import numpy as np + from ... import _STANDARD_OPERATOR_NAMES, backend, op from ...dtypes import BOOL, INT8, UINT64, _supports_complex, lookup_dtype from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, _supports_udfs, ffi, lib +from ..dtypes import _sample_values from ..expr import InfixExprBase from ..utils import output_type @@ -494,7 +497,9 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): 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. + 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 ( @@ -521,6 +526,194 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): 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: + 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": + from .udt_utils import _iter_record_leaves + + # 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): + return # a wrong field count, not a wrong shape + 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 ): @@ -536,15 +729,19 @@ def _get_udt_wrapper( 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}):" @@ -586,6 +783,8 @@ def _get_udt_wrapper_indexbinary( UINT64.numba_type, tptr_type, ) + arg_dtypes = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, dtype2] + _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}" diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index c56bb86f9..19b65d720 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -300,6 +300,32 @@ def test_bind_raw_array_udt_theta(): 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_bind_raw_udt_theta_without_dtype_errors(): """A raw UDT theta with no dtype can't be inferred; the error is clear and actionable.""" indexbinary.register_new("raw_no_dtype_op", _ibo_return_x, is_udt=True) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 9bc3f2cdc..36152dfe0 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2052,6 +2052,409 @@ def _add_corner(x, y): ) +@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.base 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 base as _base + + probed = [] + real_probe = _base._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(_base, "_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(): From b07a40570fa9b66e1b7ef12044acdc72471d7b04 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 20:18:05 -0500 Subject: [PATCH 23/52] Support SuiteSparse:GraphBLAS 10.5.2 (#658) * Support SuiteSparse:GraphBLAS 10.5.2 Behavior fix. SuiteSparse:GraphBLAS 10.5.2 turned GxB_NONBLOCKING_GPU and GxB_BLOCKING_GPU from GrB_Mode values of their own (7099 and 7098) into #defines of GrB_NONBLOCKING and GrB_BLOCKING, and python-suitesparse-graphblas built against it no longer has them. Before, graphblas.ss (gb.ss.about, gb.ss.config, gb.ss.concat, and so on, and `from graphblas import *`) failed to import with that build: AttributeError: cffi library 'suitesparse_graphblas._graphblas' has no function, constant or global variable named 'GxB_NONBLOCKING_GPU' After, it imports, and gb.ss.about["mode"] reads "blocking" or "nonblocking" as before. The GPU modes go into About._modes only when the library has them, with setdefault, so a library that exposes them as aliases cannot turn "nonblocking" into "nonblocking_gpu". test_about now checks the mode against the one graphblas was initialized with. 10.5.2 also checks that imported values are long enough, which older versions got wrong (they compared the byte size of the values with the number of entries). test_ss_import_export imported a 14 x 7 bitmap with values for 7 x 7, which 10.5.2 rejects with InvalidValue and older versions filled from past the end of the values; it now gives values for the whole bitmap. * Cap an unpinned pandas below 3.1 in CI pandas 3.1.0 came out on 2026-10-22, on PyPI so far, and three test_formatting.py tests fail with it, on main and here: * test_sparse_vector_repr and test_sparse_matrix_repr: the pandas-based repr renderer appends its "..." row with ``df.loc["..."] = ...`` to integer columns, and pandas 3.1 warns ("Pandas4Warning: Setting an item of incompatible dtype is deprecated"), which ``filterwarnings = error`` turns into a failure. * test_autocompute: under pandas' default display.max_columns=0, pandas 3.1 fits one column fewer on each side of a wide repr's "..." than 3.0 did, and the test pins 3.0's output. conda-forge has pandas 3.1.0 now, so every CI draw of the newest pandas hits these. #605 renders reprs without pandas, which makes both moot, so rather than change the renderer it replaces, CI stays below pandas 3.1 until #605 is on main. The pools keep "" for the newest release. A new LATEST_CEILINGS table caps it when the pick is written out, so an unpinned pandas becomes ``pdver='<3.1'`` and the log's summary reads ``pd=latest<3.1``. Pins and rules are unchanged; an unpinned draw now gets the newest 3.0.x. A MAINT note says when to remove the cap, and check_versions.sh points to it. Only test_and_build.yml installs pandas to run the tests. imports.yml (pip, ``.[default]``) and debug.yml (conda) also install the newest pandas, but they only import graphblas, which pandas 3.1 does not affect. Checked: ``--validate`` passes (160000 combos). Dry-run solves with the cap land on pandas 3.0.6 for Python 3.15 on linux-64 and for 3.13 with numpy 2.2 on win-64, and on 2.3.3 for 3.11 with numpy 1.24 on osx-arm64. Locally, the three tests fail with pandas 3.1.0 and pass with 3.0.6. --- graphblas/ss/_core.py | 8 ++++++-- graphblas/tests/test_matrix.py | 3 +++ graphblas/tests/test_ss_utils.py | 2 ++ scripts/check_versions.sh | 1 + scripts/ci_pick_versions.py | 17 ++++++++++++++--- 5 files changed, 26 insertions(+), 5 deletions(-) diff --git a/graphblas/ss/_core.py b/graphblas/ss/_core.py index 5e851b7cc..f036b3bb7 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -259,9 +259,13 @@ 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, } diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 4befea2cf..ea8a29f26 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -2084,6 +2084,9 @@ 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) diff --git a/graphblas/tests/test_ss_utils.py b/graphblas/tests/test_ss_utils.py index 2c9e01a1a..b1469b7a0 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -223,6 +223,8 @@ 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) diff --git a/scripts/check_versions.sh b/scripts/check_versions.sh index d9959556d..f702f2f41 100755 --- a/scripts/check_versions.sh +++ b/scripts/check_versions.sh @@ -13,6 +13,7 @@ conda search 'flake8-bugbear[channel=conda-forge]>=26.9.9' conda search 'flake8-simplify[channel=conda-forge]>=0.30.0' 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.6' diff --git a/scripts/ci_pick_versions.py b/scripts/ci_pick_versions.py index ebd003d96..8c7d4786c 100755 --- a/scripts/ci_pick_versions.py +++ b/scripts/ci_pick_versions.py @@ -83,6 +83,14 @@ "2.x": ["2.2", "2.3", "3.0", ""], } +# Caps on an unpinned ("" = latest) pick, as conda specs. Pins in the pools are not +# affected; a capped package's newest release stays out of CI until the cap goes. +# MAINT: 2026-10-22 pandas 3.1.0 fits fewer columns in a repr's auto-width +# (display.max_columns=0), and warns (Pandas4Warning) when the pandas-based repr renderer +# appends its "..." row to a numeric column; both fail test_formatting.py. Remove the cap +# once #605 (the repr renderer that does not use pandas) is on main. +LATEST_CEILINGS = {"pandas": "<3.1"} + # The "2.x" pool starts at 2.6 because awkward <2.6 still uses numpy.AxisError, # which numpy 2.0 removed. Only >=2.10 pins survive numpy >=2.5; see the awkward # constraints in `apply_constraints`. @@ -557,14 +565,17 @@ 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. + 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 in ("NA", ""): + 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}'") @@ -598,7 +609,7 @@ def format_summary(v): if val == "NA": parts.append(f"{name}=NA") elif val == "": - parts.append(f"{name}=latest") + parts.append(f"{name}=latest{LATEST_CEILINGS.get(key, '')}") else: parts.append(f"{name}={val}") return "versions: " + " ".join(parts) From 932ca0a02aabb683f236968654e8cf6b3e2577a2 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 9 Nov 2026 09:43:39 -0600 Subject: [PATCH 24/52] Type UDT ops like built-in ops; fix UDT literals, casts and comparisons (#591) * Reject record UDT pairs whose leaf counts differ `_check_udt_pair` rejects two record operands that disagree on shape, but it compared only top-level field names. Two records can share those and still nest differently: `[("a", f8), ("b", f8)]` against `[("a", [("n1", f8), ("n2", f8)]), ("b", f8)]` both report fields `["a", "b"]` while contributing two leaves and three. The codegen walks one operand's leaf paths and applies them to both, so applying the scalar path `["a"]` to the nested side asks Numba to add a float to a record. Measured before this change, that pair gets through the checks and comes back from Numba's typing pass as: UdfParseError: binary.plus does not work with (_NestFlat, _NestDeep): No implementation of function Function() found for signature: >>> add(float64, Record(nst_n1[type=float64;offset=0], nst_n2[type=float64;offset=8];16;False)) That reports a compile failure for what is the same shape disagreement the three sibling checks in the same function already report as a KeyError. Comparing leaf counts makes it the fourth of those checks. The exception class changes, and that is visible in both directions: UdfParseError derives from GraphblasException, not from KeyError, so code catching KeyError to mean "no such op" now catches this pair, and code catching UdfParseError no longer does. `OpBase.__contains__` catches both, so `udt in binary.plus` is unaffected either way. * Ignore the UDT repr size warning in the nesting mismatch test * Compare record leaf paths and field shapes, not leaf counts Equal leaf counts do not make two records pair up. The codegen reads both operands at the left one's leaf paths, so records nested under different fields, or with differently named sub-records, still reached Numba's typing: a UdfParseError for plus, a raw TypingError for eq. The check now requires the same leaf paths in the same order, as it already does for top-level names. That also rejects reordered sub-record fields. They work here, but they give wrong answers once #593 pairs leaf dtypes by position and #595 branches on them. Array fields whose shapes do not broadcast were worse: (3,) with (4,) raised inside the cfunc, unseen, and left garbage in the result. Those are now KeyErrors, as is a scalar field on the left of an array one, which never compiled. The arithmetic ops also need the broadcast to fit the left field. A (1,) field left of a (3,) one was caught only by the UDF return check, worded as a UDF the user did not write. eq and ne reduce each field to one bool and keep accepting that pair. Fields that broadcast, such as (3,) with () or (1, 3), still combine; across 64 field-shape pairs under plus, min and eq, everything that worked before still does. The user guide now states the pairing rules. * Type UDT ops like built-in ops, cast UDT stores, broadcast array UDTs Lifted ops (plus, minus, times, truediv, floordiv, min, max, eq, ne, ainv, abs) type each UDT element or record leaf as the built-in op types that dtype, so results never narrow. A mixed pair gives an operand's type if it holds the result, else the promoted layout. Records pair on field names, order, nesting and field shapes; array UDTs broadcast as numpy arrays do. A Python number beside a UDT is weak per leaf (NEP 50), the same on numpy 1 and 2. A literal that must become a UDT element (tuple, array, user-op literal, ewise_union default) is converted only if no value changes. apply and ewise ops with a Monoid on a mixed pair use its BinaryOp. Storing into an object of another UDT casts each leaf as GraphBLAS casts built-ins; incompatible layouts and UDT/built-in pairs raise DomainMismatch. Assigning or setting an element from another UDT used to read raw bytes, past the end of a smaller element. ewise_add converts operands to the result type, which GraphBLAS cannot cast. Array codegen uses loops, in Numba and in C: typing an op on FP64[64,64] took ~112 s and now takes ~0.02 s. Also fixes several errors, eq on 2-D array fields, array-field truediv on small ints, a UDF registered under a lifted op's name, and packed vs aligned records. Docs and tests updated. * Cast UDT stores without temporaries, type UDT literals, fold in #592-#595 Store casts: an existing object casts in one pass, and element-wise lifted ops (apply, ewise_mult, ewise_union) write the target type directly. Any other store into another UDT type raises DomainMismatch and suggests obj << expr.new(). ewise_add on mixed-type Vectors/Matrices raises; infix + takes the union with each operand's zero (-0.0 in float fields). Literals beside a UDT: Python numbers and tuples/lists/dicts of them are weak per field (NEP 50), numpy values and Scalars are strong. A value that must become a UDT element (user op literal, select/IndexUnaryOp thunk, ewise_union default) must fit by type; one that fits only by value (0 for a bool field, 2.0 for an int field) still converts, with a DeprecationWarning. Lossy values raise. From #592-#595: min/max (fmin semantics, fixed -0.0 tie), floordiv, integer and complex division, and complex multiplication give the same bits on the Numba and C JIT paths. Array fields in records compute element by element. Also: extracting an element into another UDT type casts it instead of copying raw bytes (a segfault on main); Monoid literals are typed as their BinaryOp's; outer works on mixed types; the user guide's UDT rules are two tables. * Type Python numbers beside built-in dtypes as numpy 2 does (NEP 50) A Python number beside a Vector, Matrix or Scalar of a built-in dtype was typed through numpy, with a 0-d array of value 0 standing in for it (dtypes.unify). numpy 1 judged that array by its value, so the collection's dtype won and a large literal wrapped (int8_vec * 1000 multiplied by -24, int8_vec < 300 compared with 44); numpy 2 (NEP 50) made 0-d arrays strong, so every int was INT64 and every float FP64 (int8_vec + 1 was INT64, uint64_vec + 1 FP64, losing precision). The same script gave other dtypes, and on numpy 1 wrong answers, depending on the installed numpy. A Python bool, int, float or complex operand is now weak, by the table UDT leaves already use: it takes the other operand's dtype when that dtype's kind holds it, else INT64, FP64 or FC64, on numpy 1 and 2 alike. So int8_vec + 1 is INT8 and wraps as numpy does, int8_vec * 0.5 is FP64, fp32_vec * 0.5 is FP32, int8_vec + 300 raises OverflowError, and fp32_vec + 1e300 (or + 2**200) is inf with a RuntimeWarning, as in numpy 2. A strong operand keeps the old widening: int8_vec + np.int64(1). This covers apply (infix ops, binary.plus(v, 1), Monoids), Scalar ewise ops, ewise_union defaults, value thunks of select and IndexUnaryOps, and UDFs. A comparison keeps no result type, so an int out of range compares exactly (int8_vec < 300 is all True), and so does an op whose result is another type anyway: int8_vec / 300 is FP64, as numpy divides integers in float64, for UDT fields too. A to_dense fill value that does not fit widens, an int or a float (fp32_vec.to_dense(fill_value=1e300) is float64), rather than raising or becoming inf. numpy scalars, Scalars and typed ops stay strong, and a thunk that is not compared with the values, such as a row index, keeps its own type. unify no longer stands a scalar in as a value-0 array. eq and ne beside a UDT now type a Python number or sequence the same way, leaf by leaf, with the same exact fallback: fp32_udt == 0.1 compares in float32, as fp32_vec == 0.1 and numpy do. * Compare select value thunks exactly; isequal compares arrays as np.array_equal SuiteSparse:GraphBLAS's select with a value op (valueeq, valuelt, ...) typed for another type than the input's reads the thunk's bytes as the input's type (8.0.2 through 10.6.0; apply with the same op is right; fix proposed upstream in DrTimothyAldenDavis/GraphBLAS#463). On an INT8 vector, select(valueeq, 1.0) kept nothing and select(valuelt, 300) compared with 44, on main too. select now compares in the input's own type with a thunk that keeps the same entries (x < 2.5 is x < 3 for integers). isequal on array UDTs compares as np.array_equal: elements must have the same shape, apart from leading axes of length 1, so an FP64[3] of [1, 1, 1] is no longer isequal to an FP64[1] of [1]. == still broadcasts, and a Scalar's ==, which returns a bool, now computes binary.eq, as == on a Vector does (it called isequal). A literal beside a UDT Scalar is compared as == types it, so [0, 0, 0] == 0.5 is False (0.5 was truncated to 0), and Scalar(FP32 0.1) == 0.1 compares in FP32, as numpy does. A float too large for FP32 in a comparison (fp32_vec == 1e300) compares with infinity, as numpy 2 does, and now warns at the caller's line, as arithmetic does. A Python int from 2**63 to 2**64 - 1 now goes into a UINT64 Scalar, Vector or Matrix element (Scalar.from_value(2**63, "UINT64") and uint64_vec[0] = 2**63 raised OverflowError, on main too). * Compare a Python int beyond a 64-bit element's range exactly Commit 2 typed an int beyond an integer element's range by a type that holds it, so a comparison no longer raised, but beside INT64 and UINT64 that comparison runs in FP64, which rounds both: int64_vec < 2**63 was False for 2**63 - 1, uint64_vec == 2**64 was True for 2**64 - 1, and Scalar(2**63 - 1) == 2**63 was True, as for UDT int64 and uint64 leaves (main raised OverflowError). Such an int now compares as infinity of its sign, which every integer element compares with as with the int, so infix comparisons, apply, ewise_union defaults and Scalar == are exact, as in numpy 2, and beside a UDT a literal with one equals no element. An int too large for a float compares too (int64_vec < 10**400 and select(valueeq, 10**400) raised OverflowError). A float too large for a float32 UDT leaf (fp32_udt + 1e300, or a store) now warns at the caller's line, as a built-in FP32 does, instead of from inside python-graphblas, where the default filter showed it once per process. Drops a test helper nothing used. * Give each UDT layout its own JIT C name; Scalar == without Numba; CI fixes SuiteSparse's JIT knows a type by its C name alone, and a UDT's C name was the name it was registered under, so two layouts registered under one name shared JIT kernels: making an iso Vector non-iso ran a kernel compiled for the other layout, which corrupted the heap or gave wrong values (on main too; the new broadcasting test registers four layouts as _BcastX, and crashed 9 Linux CI jobs). A name already given to another layout now gets a fresh _gbudt_NNN, and gb.ss.register_new refuses a name another type has. Scalar == on a UDT runs binary.eq, which needs Numba; without Numba it raised KeyError, where it compared in Python before. It now compares the values in numpy, leaf by leaf in the type binary.eq compares in, and raises binary.eq's KeyError for UDTs that do not pair. Tests: guard the cases that need Numba, skip the 32 KB UDT on Windows (MSVC builds of SuiteSparse refuse a UDT over 1024 bytes), give the JIT name tests array lengths no other test registers, and expect numpy's own int dtype (int32 on Windows with numpy 1). * Spell out complex * and / for UDTs; Numba 0.68 follows Python 3.14 From Python 3.14, CPython recovers infinities that the textbook complex formulas leave as nan (C99 Annex G), and Numba 0.68 follows it there: 1j * (inf+infj) is -inf+infj and (inf+nanj) / (1+1j) is inf-infj, where numpy and the C JIT kernel give nan+nanj. The UDT cfunc used Numba's operators, so it disagreed with the C JIT kernel on CI's macOS job with Python 3.14 and Numba 0.68. The cfunc now spells out the formulas the C kernel uses (the textbook product, and Smith's division as CPython divided before 3.14), so the two agree on any Python and Numba. The user guide no longer calls a complex quotient numpy's: numpy multiplies by the divisor's reciprocal, so the last bit can differ, and numpy's overflows where the quotient is finite. * Move UDT and UDF logic out of the general modules Code moves, plus import changes and a few small helpers, so scalar, vector, matrix and base read mostly as the built-in path. - core/udt.py (new): rules for UDT elements: layout, op typing, store casts, the ewise_add rule, equality without Numba, and element values as bytes. The type rules moved here from udt_utils, which now only generates code and is renamed udt_codegen. - core/literals.py (new): NEP 50 typing of literal operands, for built-in and UDT types (moved from scalar.py and udt_utils). - operator/udf.py (new): the Numba-only UDF-on-UDT block from operator/base.py, plus _bool_to_int8. - scalar._union_default replaces the ewise_union default block, which was copied six times across Scalar, Vector and Matrix. - New helpers pulled out of scalar.py's inline code: _scalar_ewise_add_converted, _udt_eq_without_numba, _element_bytes and _element_value. core/udt.py imports operator, base and scalar names inside functions, after the built-in early returns, so built-in calls pay nothing extra. * Pass CODECOV_TOKEN to the coverage upload Codecov takes tokenless uploads only from forks. Every upload from this PR's CI runs was refused with "Token required because branch is protected", so the PR had no coverage report, and the PRs stacked on it had no base to compare against. The repository now has the secret. * Cover how core/literals.py types literal operands Tests for the literal-typing paths CI never ran: - Beside built-in dtypes: a comparison op given as a string with an int out of range (ewise_union with "<"), an op string that names no BinaryOp (the op reports it), an int beyond INT64 beside a bool (OverflowError, as numpy's bools & 2**70), a user's BinaryOp with an int out of range, and a complex literal on a backend without complex types (vanilla only). - Beside UDTs: an infinite or NaN Decimal, literals not laid out like an element (ragged, too long or too short for a field, not broadcastable) or not numbers, which numpy then rejects, a number beside a bytes field, a nested record that is not the last field, and the conversions that lose a number (an imaginary part, a float32 overflow, a fraction in a record's array field) or keep a bytes field as it is. - Scalar.isequal with a ragged literal on an array UDT (ValueError). * Cover select's rewrite of value thunks of another type Tests for the branches of operator/select.py's _value_select_in_own_type that CI never ran: a complex thunk beside integers and floats (only valueeq and valuene, which have an answer), a NaN thunk of another float type, an infinite thunk beside integers, and an empty thunk Scalar, which GraphBLAS rejects. * Cover UDT store casts, ewise_add casts and fused stores Tests for the paths in core/udt.py, Matrix and udt_codegen that CI never ran: - store casts between records whose field shapes differ or whose bytes fields differ in width (DomainMismatch saying why), a Scalar expression assigned to many elements of another UDT type, one element extracted from a Matrix into a Scalar of another type, and Matrix.power with a semiring whose result is another type; - without Numba no UDT cast compiles, so every such store raises (monkeypatched, so it runs everywhere); - Matrix.ewise_add with operands of another type than the result's (DomainMismatch, and DimensionMismatch first for other shapes), and Scalar.ewise_add with a Monoid and with a non-lifted op; - a fused store with a built-in Scalar operand, and fused stores that cannot be made (first, a user's op); - (x | y) + z of non-BOOL vectors, which reports the infix error. * Cover UDF return-type errors and the Numba error summary Tests for operator/udf.py: a UDF on a UDT that returns None, or an array with no array UDT input, or a tuple beside a UDT that is neither a record nor an array, raises UdfParseError naming the type; built-in ops on such a UDT raise their KeyError. The one-line summary of a Numba error falls back to the first non-empty line, then to a fixed sentence. * Cover the JIT layout checks for records a C struct cannot match A record with explicit offsets can have a C struct's size with a field elsewhere, as can such a record nested in one whose own offsets match, and a nested record with an array field has no C typedef. Each skips the JIT and still computes with the cfunc. * Remove branches no caller can reach in the UDT code Behavior-neutral; each guard tests a state its callers never produce: - literals._check_literal_fits: every caller passes a UDT, never None. - select._value_select_in_own_type: an int too large for any float never reaches the float branch (beside a float type it is typed as that type, or float() raises first, as in numpy). - udt_codegen._c_aligned_version is only called with records. - udt_codegen._c_expr_unary: both callers pass the field's dtype. - udt_codegen._make_jit_c_comparison_definition: _udt_c_typedef has already returned None for a UDT that is neither array nor record. - udt_codegen.set_jit_c_on_op and set_jit_c_comparison_on_op: callers check _has_jit_set first, since the setter they pass does not exist without it. * Exclude Numba-compiled bodies and unreachable guards from coverage - numba: _complex_times and _complex_truediv run only inside compiled code, as do two UDFs in test_op.py. - safety: guards no public call reaches: a record leaf lookup_dtype refuses (a void field fails to register first), a Numba error in the UDF probe (the wrapper has the same operand types), a tuple whose length differs from the record _resolve_udt_return_type matched, a C expression for an op or dtype no lifted op gives, and the frame walk in _warn_from_caller running off the stack. - only on failure: the udt_op_path fixture's demotion check, and the re-raise for a large UDT off Windows. * Drop test conditions that cannot change and a needless Numba skip - test_dtype: the JIT name test is skipped unless _has_jit_set, which needs the suitesparse backend, so it always has gb.ss. - test_ssjit: the two complex JIT tests are skipped without Numba, so their own "numba is None" skips never ran. - test_op: test_compile_codegen_helper compiles plain Python, so its Numba skip was needless; it now runs without Numba too, which covers _compile_codegen's no-Numba namespace in CI's no-Numba jobs. * Fix select on complex values with a thunk of another type Behavior fix: wrong results before, on main too (not a #591 regression). SuiteSparse:GraphBLAS reads the thunk of a value select as the input's type when the op is typed for another type, and the workaround from 69173110 converted the thunk only for bool, integer and float inputs. On an FC32 vector of [1, 2+1j, 3, nan], select(valueeq, np.complex128(1)) kept nothing (now [0]), valueeq with np.float64(3) or np.int64(3) kept nothing (now [2]), and valuene with np.complex128(2+1j) kept everything (now [0, 2, 3]); apply with the same ops was right. (Main also got a Python 3 wrong there; #591's typing of a weak int as FC32 already fixed that.) A complex input now gets the thunk in its own type when that converts exactly; otherwise (or with a NaN part) the thunk equals no element, so valueeq keeps nothing and valuene keeps everything. "Exactly" is as for float inputs: the thunk in the input's type must equal the thunk, part by part, in Python's int and float comparison, so np.int64(2**53 + 1) equals no complex64 element, as it equals no float32 one, where apply compares in complex128 and matches 2**53. The ordered value ops still raise KeyError on complex values, which have no order. A real input with a complex thunk keeps its rule: an imaginary part of 0 converts, any other equals nothing. * Compare UDTs that are neither records nor arrays in eq and ne Behavior fix. A UDT of bytes, dates or the like holds one value per element. binary.eq and binary.ne on it raised TypeError ('NoneType' object is not subscriptable), and so did Vector and Matrix isequal and ==: on main too. A Scalar's == raised it as well, where main compared the values (a #591 regression, from 69173110's Scalar == through binary.eq), as did == with a literal (b"ab"), which the literal checks read as an array UDT. On the vanilla backend a complex literal is such a UDT, so int8_vec == 1j raised the same TypeError. eq and ne now compare the two values as numpy does, as a record field of that type compares (bytes by value, NaT equal to nothing), and only two elements of the one UDT compare: any other pair raises KeyError, from _check_udt_pair, with or without Numba. So int8_vec == 1j on vanilla is a KeyError, and so is a Scalar's == with another type (another bytes UDT, or a number), where main compared the values (True or False). A literal beside such a UDT is checked as one value. A Scalar's == on such a UDT compares in numpy, as without Numba, so it needs no Numba type for the element: a float16 UDT, which Numba cannot compile, compares as on main. Vector and Matrix == on float16 still need Numba and raise its NotImplementedError, as main raised there too. * Type each product of Matrix.power for its UDT operands Behavior fix of a #591 regression: main, which typed UDT ops otherwise, computed int8 results here. power types the semiring for A, and GraphBLAS casts a built-in operand of another type to it, but not a UDT. With a semiring whose result is another type, such as plus_truediv on an int8 array UDT (float64 results), power(2) worked, while from n = 3 on the products of float64 squares with A or with each other raised a bare GrB_DOMAIN_MISMATCH from GrB_mxm. Each product of UDT operands that are not of the types the semiring was typed for now takes the semiring typed for its own operands, as any mxm of them does. power(n) into a float64 array object now equals the power of a float64 copy of A, as for built-in INT8, for every n. Unchanged: A.power(n, plus_truediv).new() makes an object of A's int8 type, and an mxm cannot cast a UDT as it stores, so it raises DomainMismatch by #591's store rule, where main computed int8 results (built-in INT8 truncates as it stores). * Select with a float64 thunk too large for FP32 without a warning Behavior fix, small: f32_vec.select(valueeq, np.float64(1e300)) gave the right result but numpy warned "overflow encountered in cast" from inside python-graphblas (an error under -W error), as the select workaround from 69173110 converted the thunk to float32. A #591 regression: main had no such conversion (it let SuiteSparse read the thunk's bytes as float32, which was wrong in another way). numpy's own f32_array == np.float64(1e300) does not warn. The conversion now ignores the overflow: the thunk is infinity in float32, which compares with every element as 1e300 does, so the selected elements are unchanged. * Refuse values that a bytes, date or duration leaf of a UDT would not keep Behavior fix. A literal that must become an element of a UDT and does not fit it by type converts, with a DeprecationWarning, only when its values convert exactly (#591's MAINT rule). The check skipped every leaf that holds no number, so whatever was given for a bytes, str, date or duration leaf counted as exact. Under a user's op, (1.5, 1) for a record of a bytes field and an int8 field became (b"1.5", 1), and 1/3 became b"0.33", with a warning that said "Its values convert exactly"; np.timedelta64(5, "ms") for a timedelta of seconds became 0 s, and a structured literal's six bytes became four, under the same warning. - main: converted silently (1.5 became b"1.5"). - 0e1ad03c: a record field converted with that warning; a bare UDT of bytes or durations raised "does not fit" (the check failed on its layout). - the commit "Compare UDTs that are neither records nor arrays in eq and ne" let a bare UDT convert with the warning too. - now: such a leaf keeps a value only if it converts back to it, so a Python number, a numpy bool, a string cut short or a duration in a coarser unit raises ValueError "does not fit". bytes or str that fit still convert, as does np.timedelta64(5000, "ms") as 5 s (with the warning, since ms to s is not a safe cast). Only a literal whose type does not fit reaches this check. A numpy number that numpy casts safely to the leaf's type fits by type: np.uint8(5) beside a 3-byte bytes UDT is b"5" without a warning, by numpy's rule. bytes or str alone beside a bare bytes UDT has no type, so numpy still cuts b"abcdef" to b"abc" there. A record literal holding a numpy number, text or a date gets no type at all and so no check; the next commit types it. * Type every leaf of a record literal, numpy numbers, text and dates too Behavior fix, in #591's literal typing (older than its last commits; main had it too). A tuple or dict literal for a record UDT got no type at all when a leaf held anything but Python numbers: a numpy number, text or a date. Then nothing checked it, and numpy converted it into the record, losing value where a Python number alone raises. In an int8 record, under a user's op, (np.float64(1.5), 1) became (1, 1) while (1.5, 1) raised "does not fit"; under a lifted op, int8_rec + (np.float64(1.5), 1) computed with (1, 1), and int8_rec == (np.float64(1.5), 2) was True. A record of a bytes and an int8 field holding (b"cd", 1) had (b"cd", 1.5) == it True, and a user's op stored (b"cd", 1.5) as (b"cd", 1); a record with a date field the same. - main and 0e1ad03c: all of these as above, silently. - now every leaf is typed. A numpy number or bool is strong, as alone or beside a built-in type: the leaf takes its dtype, promoted with any Python numbers in it, and a numpy number may stand for a whole nested record. Text or a date in a text or date field takes the field's own type when it converts there and back unchanged, so ops still pair the field with itself (rec + (b"ab", 1) is (b"abab", 2), and (np.timedelta64(5, "m"), 1) adds 300 s to a field of seconds, as on main); otherwise it keeps its own type, which does not fit: text at its own length, a numpy date or duration in its unit. Other dates (datetime, str) take the field's type. So the user's op raises "does not fit" for np.float64(1.5) or 1.5 in an int8 field next to text or a date, for text longer than its field (it was cut short) and for a date or duration in a finer unit that loses value (5 ms in seconds); np.int64(2) converts with the DeprecationWarning (exact, but not a safe cast) and np.int16(2) fits an int16 field; the lifted op gives a float64 field (2.5, 3), and + with 5 ms in a field of seconds raises KeyError (main added 0 s); and the comparisons are False. In a comparison, an int beyond its integer leaf's range now becomes infinity in that leaf only, not in the whole literal, so a text leaf beside it still pairs: rec == (b"ab", 300) is False (it raised). A text date is compared as a date in the exactness check. A sequence for an array UDT keeps its rule, the array numpy makes of it. * Refuse UDT ops on float16 elements by name; compare them in numpy for Scalar == Behavior fix. Numba has a float16 type but no code for it, so every op compiled on a UDT with a float16 leaf (a field, an array's element, or a bare float16 UDT) failed with a bare NotImplementedError: float16. - A record with a float16 field, + (1.0,) or == (1.5,): NotImplementedError on main, at 0e1ad03c and before this commit; now KeyError. - Its Scalar == (1.5,): True on main, NotImplementedError at 0e1ad03c and before; now True again. - A bare float16 UDT, Vector == or isequal: TypeError on main and at 0e1ad03c, NotImplementedError before; now KeyError. - An int8 record + (np.float16(1.5), 1) or == (np.float16(1.5), 2): main and 0e1ad03c cut 1.5 to 1 and answered (2, 3) and True, both wrong; before, KeyError ("no common type") and NotImplementedError; now KeyError. Every UDT op's _compile_udt (unary, binary, indexunary, select, indexbinary) now refuses a float16 leaf up front with the KeyError an op gives for a type it does not take: "binary.eq does not work with (...): Numba cannot compute on float16, which ... holds." A Scalar's == on a UDT with a float16 leaf compares in numpy, as for a UDT that is neither a record nor an array. A store that would cast to or from a UDT with a float16 leaf (another record or array UDT, or dup(dtype=...)), which raised the same NotImplementedError on main, at 0e1ad03c and before, now raises the DomainMismatch of a store that cannot cast, saying why: "... Numba cannot compute on float16, which ... holds." Widening float16 to float32 was not done: the result type would change, and results stored back as float16 would round. * Type a record's array field as an array UDT types its literal Behavior fix. A sequence for an array UDT is typed by the rule #591 gives it: Python numbers alone are weak, each taking the element's type when its kind holds them, and a numpy array, or a sequence holding a numpy value, is the array numpy makes of it. The same sequence for an array field of a record was typed element by element, numpy values strong and Python numbers weak beside the field, and a numpy array was read as Python numbers. So, with an int8[3] array UDT and a record with an int8[3] field: - (np.int8(1), 300, 3): the array UDT gives int64[3] and computes [2, 302, 6]; the record raised OverflowError (main and 0e1ad03c too). - np.array([1, 300, 3]): the array UDT gives int64[3]; the record raised OverflowError (0e1ad03c), and main wrapped 300 to 44 silently. - (np.float64(1.5), 2, 3): main and 0e1ad03c cut 1.5 to 1; the commit "Type every leaf of a record literal, numpy numbers, text and dates too" already gave a float64 field. A record's array field now takes the array UDT's rule: an array, or a sequence holding a numpy number or bool, is the array numpy makes of it, broadcast to the field's shape. So the record gives an int64[3] field and [2, 302, 6], as the array UDT does. Python numbers alone, and a numpy number for the whole field, are typed as before. A comparison keeps a Python int beyond the element's range exact, as it does everywhere (int8_vec == 300 is False), also inside such an array: the int is infinity in its place before numpy makes the array. So rec == ((np.int8(1), 2**70, 3), 4) and int8_array_udt == (np.int8(1), 2**70, 3) are False and != True. The array UDT raised numpy's OverflowError there (older than this commit), and the record field, False before, would have raised it under the array rule alone. A user's op with such a literal still raises OverflowError. * Refuse == between text, dates and numbers with eq's KeyError Behavior fix. binary.eq and ne compared each pair of record leaves with Numba, which cannot compare bytes with a number, so such a pair raised UdfParseError carrying Numba's typing error: - a record with a bytes and a float64 field, eq with a record of two float64 fields: GrB_DOMAIN_MISMATCH on main, UdfParseError at 0e1ad03c and before this commit; - the same record == 1 or != (1.0, 1.0), and its Scalar == (1.0, 1.0): main converted the number into the bytes field (1 became b"1") and compared the converted value; UdfParseError at 0e1ad03c and before; - a nested record == (b"x", np.float64(1)), whose number stands for a nested record holding a bytes field: False on main and at 0e1ad03c, UdfParseError before this commit. Text, a date, a duration and a number now compare only with their own kind; any other pair of leaves raises the KeyError binary.eq gives for a pair it does not take: "binary.eq does not work with (...): elements of |S4 and float64 do not compare." Scalar == raises it with or without Numba. This follows the rules #591 already has: plus finds no common type for such leaves (a KeyError), and a UDT that is neither a record nor an array pairs only with its own type. Answering False instead (a number for a text leaf can never be equal) would turn the record pair, an error since main, into an answer, and would differ from main anyway where main compared b"1" with 1 as equal. * Test the text-date and duration-for-a-date branches of record literal typing Tests only. Two branches of the record literal typing had no row: - a date given as text shorter than numpy's own date string ("1970-01-01" for a field of seconds), which _converts_back compares as a date: with a number that converts exactly, the literal converts with the DeprecationWarning, and without the branch it raised "does not fit"; - a numpy duration given for a date field, which keeps its own type: numpy converts np.timedelta64(5, "s") to a date by value, so without the guard it took the date field's type and failed in numpy's own conversion; with it, it raises "does not fit". * Let a UDF return an operand of a UDT that is neither a record nor an array Behavior fix. A user's op on a UDT of bytes, str or the like that returns an operand as it is (lambda x, y: x) failed to compile: Numba types the element as a char sequence, which no DataType is registered for, so the return type was not found. - main: UdfParseError "UDT UDF returned an unsupported type [char x 8]". - 0e1ad03c: TypeError 'NoneType' object is not subscriptable (the literal beside it was read as an array UDT). - before this commit: the UdfParseError again. - now: the return takes the operand's UDT when its Numba type is that UDT's, so the op computes; [b"abc"] for x, the literal for y. * Keep text that a bare text UDT would cut short out of it, and compare text widths Behavior fix. Text given beside a UDT of bytes or str (not a record) had no type, so numpy cut it to the UDT's width: with an S4 element b"abcd", b"abcdef" became b"abcd". - v == b"abcdef": TypeError on main and at 0e1ad03c; True before this commit; now False (!= is True). - Scalar == and Scalar.isequal with b"abcdef": True on main (isequal True at 0e1ad03c too), now False. - a user's op with b"abcdef": main and the commits before this one cut it silently (0e1ad03c raised TypeError); now ValueError "does not fit". Such text is now typed as a record's text field is: the UDT itself when it converts into it and back unchanged (b"ab", "ab" into S4, so bare_s4 == b"ab" works as before), else its own type at its own length, which does not fit. For eq and ne, that literal is text of the same kind at another width, so this commit also loosens, there and only there, the pairing rule of "Compare UDTs that are neither records nor arrays in eq and ne": two UDTs of text of the same kind (bytes with bytes, str with str) at different widths now pair under eq and ne and compare as text, as numpy compares them and as a record's S4 and S2 fields already do. An S4 and an S6 Vector compare element by element (they raised KeyError before; TypeError on main and at 0e1ad03c). Bytes against str, a number, float16 or any other kind still raise the pairing KeyError, as does every other op on such a pair. Stores are unchanged: v[0] = b"abcdef" still stores b"abcd", as every literal store converts by numpy's cast (int8_vec[0] = 1.5 stores 1). * Keep a NaT duration NaT in another unit on numpy before 2.3 Behavior fix. numpy before 2.3 converts a NaT timedelta64 scalar into another unit as the integer NaT is stored as: NaT ms assigned into timedelta64[s] is 9223372036854774 s, and NaT minutes is 0 s. An array of durations (astype) and a date convert NaT correctly, and numpy 2.3 and later convert it correctly every way. On numpy 2.2 (numpy 1.26 converts NaT the same way), with a UDT of timedelta64[s] or a record with a timedelta64[s] field: - Scalar.from_value, v[0] = ..., Scalar.value = ..., Vector.from_coo, and the Scalar a literal becomes, given NaT ms or NaT minutes: stored 9223372036854774 s or 0 s, as on main; now NaT. - A NaT ms literal beside the bare UDT: ValueError "does not fit", since it converted to a number and so not exactly; now it converts with the DeprecationWarning, as on numpy 2.4. - (NaT ms, 1) beside the record: typed as timedelta64[ms], so binary.plus raised KeyError "no common type"; now it takes the field's type, as on numpy 2.4. The conversion into an element, and the checks of whether a literal converts exactly, now hand numpy a NaT of no unit in place of each NaT duration (alone, or in a tuple or list), which every numpy converts to NaT in any unit. numpy 2.4 gives the same results before and after. * Keep the JIT layout test's record short enough to name on SuiteSparse 8 SuiteSparse before 9 has no GrB_NAME setter, so registration names a UDT by its numpy repr, which must fit in 128 bytes, and warns when it does not; pytest makes the warning an error. test_record_layouts_c_cannot_declare_skip_jit registered its record with an array field with align=True, which gives the dict form of the repr: 143 bytes. Its float64 fields have the same offsets packed or aligned, so it is now registered packed (a 71-byte repr); a C struct still cannot declare it, and the test runs the same lines. Every other UDT the tests register with a repr over 128 bytes is in a test that ignores that warning. * Convert a NaT duration in its leaf's unit, not as a NaT of no unit "Keep a NaT duration NaT in another unit on numpy before 2.3" handed numpy a NaT of no unit, np.timedelta64("NaT"), in place of each NaT duration. numpy 2.5 deprecates that generic unit ("The 'generic' unit for NumPy timedelta is deprecated"), so on numpy 2.5 every store or literal of a NaT duration into a duration UDT or field warned with a DeprecationWarning, which the test suite's filterwarnings makes an error: test_udt_nat_duration_in_another_unit_stays_nat, test_udt_sequence_and_array_literals_type_like_numbers and test_udt_literal_converts_only_when_exact failed. Each NaT duration is now converted into the unit of the leaf it fills by converting it as an array (astype), which keeps NaT on every numpy, before numpy assigns it. The value is walked beside the dtype: a tuple of a record's length as its fields, a list (or a tuple for any other dtype) as a sequence of elements, and an array UDT by its element type. Nothing makes a generic unit. numpy 2.5.4: DeprecationWarning before, NaT with no warning now. numpy 1.26, 2.2 and 2.4: NaT in the target unit before and after. New rows: a nested record's duration field, and an array UDT of durations given a list, where Numba is absent (Numba cannot type an array of durations, so registering one fails with it). * Skip the JIT part of the one-name UDT layouts test without a C compiler test_udt_layouts_under_one_name_get_their_own_jit_c_names set jit_c_control to 'on' and compiled kernels for both layouts, so on a host whose JIT cannot compile (no C compiler, or a broken one) it failed with JitError instead of skipping. CI always installs a compiler, so CI never saw it. It now asks _enable_jit_for_udt, which probes with a real compile once per process, and skips the kernel part when the answer is no. The C name checks and the register_new check need no compiler, so they come first and run everywhere. * Expand an iso array UDT's value without an extra axis in exports Behavior fix. SuiteSparse exports the iso value of an array UDT with shape (1, *subarray shape), and normalize_values broadcast it to shape + values.shape, keeping that axis of length 1. With the suitesparse backend, for an iso Vector or Matrix of an array UDT such as INT64[4] (as v[:] = [0, 1, 2, 3] makes): - Vector.to_dense: TypeError "unsupported operand type(s) for +: 'int' and 'tuple'" on main and before this commit (it passed its size as an int); now an array of shape (size, 4). - Matrix.to_dense: shape (nrows, ncols, 1, 4); now (nrows, ncols, 4). - Matrix.to_csr, to_csc, to_dcsr and to_dcsc: values of shape (nvals, 1, 4); now (nvals, 4). - Matrix.to_dicts, which uses to_dcsr or to_dcsc: each value as [[0, 1, 2, 3]]; now [0, 1, 2, 3]. Vector.to_coo, Matrix.to_coo, the vanilla backend and records were already right. normalize_values now broadcasts to shape + values.shape[1:], and Vector.to_dense passes (size,). CI found it through #652's test_udt_mixed_jit_c_name_is_its_own (17 jobs, e.g. job 122633175811): SuiteSparse before 10.5 makes a vector built from equal values iso, so that test's constant inputs gave an iso result, and to_dense raised. With 10.5 the vectors are not iso. * Refuse a zero-size UDT, so a str literal raises TypeError on SuiteSparse 7 Behavior fix. register_anonymous passed a dtype's itemsize to GrB_Type_new without a check. SuiteSparse 8 and later refuse a size of 0, but SuiteSparse 7 takes it as 1. lookup_dtype auto-registers what it does not know, so on SuiteSparse 7 a str became the UDT " after): - A.ewise_mult(B, op()) on INT64: a BOOL op, so every value became True -> INT64, the values op computes. - u.ewise_mult(w, op(2)) on FP64: computed in INT64 (1.5 + 2 * 2.25 gave 5) -> FP64 (6.0). op(np.int8(2)): computed in INT8, where 300.0 became 127 and wrapped -> FP64. - A Semiring from op(0.5) with INT64, FP32 or INT8 matrices: KeyError -> computed in FP64, with or without is_udt. - The SuiteSparse user guide's argmax (an is_udt multiplier that makes the pair (k, A(i, k)), theta unused) on a matrix of pairs: iop(), iop(False), iop(0) and iop(0.0) were ops on BOOL, INT64 or FP64 values, which failed to compile for the pairs -> all four run, compiled for (pair, BOOL, BOOL), (pair, BOOL, INT64) and (pair, BOOL, FP64). - An is_udt op on INT64 data: op(0.5) or an FP32 Scalar theta made the op FP64 or FP32 and cast the inputs to it -> the inputs stay INT64, theta keeps its type, and Numba types the result (FP64 for x + y + theta). - An is_udt op on int32[3] data: op(np.array([10, 20, 30], np.int32)) raised "Cannot infer a dtype for theta" -> theta is an int32[3] value; an int64 array is an int64[3] one. op(5) gave GrB_DOMAIN_MISMATCH -> theta is an INT64. op((10, 20, 30)) still raises TypeError, now saying to give theta as a Scalar of its UDT; an op without is_udt raises KeyError for an array theta, as for a UDT Scalar. - A JIT op (indexbinary.ss.register_new) is registered for its x, y and theta types. op(theta) took the one registered at theta's type T for x, y and theta, and cast the inputs to T -> it takes the one for the inputs' types and theta's, without casting, as a JIT IndexUnaryOp takes its thunk: with an op registered at FP64, op(10.0) on FP64 data runs as before, op(10.0) on FP32 data ran in FP64 -> KeyError naming (FP32, FP32, FP64), and op(10) raised KeyError when bound -> when it runs. op.types is keyed (x, y, theta), so registering the same x and theta types for another y no longer raises "already defined", and op["FP64", "FP32"] is x and y, no longer x and theta: op["FP64", "FP64", "FP32"] is the op with an FP32 theta. - A bound op pickles with all its types: op["INT64", "FP64"](1) came back as op["INT64"](1), whose y is INT64. - op(theta).type: the bound type -> AttributeError, as for any untyped op; op(theta)[T] is op[T](theta). repr(op(2)): binary.f[INT64] -> indexbinary.f(2). op(theta)(A & B) works, as for other BinaryOps. - A Scalar theta is copied when given, so a later change to it does not reach the op, as when the op bound it at once. find_opclass now reads an opclass attribute on an OpBase subclass, so op(theta) is still a BinaryOp there. The user guide's examples drop the op[float](0.5) that hid the old typing. The comments on IndexUnaryOp's and SelectOp's UDT caches now say they are keyed by (dtype, thunk_type). * Refuse a record literal whose array field part does not broadcast to the field Behavior fix. Where a literal for a record UDT is typed (a lifted op such as +, ==, a user's op, an IndexUnaryOp thunk, an ewise_union default, and == on a Scalar), a part for an array field must now broadcast to the field's shape, or ValueError names both shapes. An array UDT's element takes the shape it broadcasts to with a literal, as numpy does, so INT8[3] plus np.array([[1, 300, 3]]) is INT64[1, 3]; that is unchanged. A record's field keeps its shape, so for the record the typing gave up, and the literal was converted into the field as numpy assigns it: the leading axis was dropped and 300 wrapped to 44, silently, even beside a user's op, which refuses a literal that does not fit. Stores (r[0] = ...) still convert as numpy assigns, as they do for an array UDT. With r a Vector of a record with an int8[3] field a and an int64 field b, holding ([1, 2, 3], 0) (before -> after): - r + (np.array([[1, 300, 3]]), 0): ([2, 46, 6], 0) -> ValueError "array([[ 1, 300, 3]]) has shape (1, 3), which does not fit field 'a' of R, of shape (3,): ...". - r.apply(user_op, right=(np.array([[1, 300, 3]]), 0)), with user_op returning y: ([1, 44, 3], 0) -> the same ValueError. - r == (np.array([[1, 300, 3]]), 0): False -> the same ValueError; with np.array([[1, 2, 3]], np.int8) it was True. - r + ([[1, 300, 3]], 0), and the other paths with it: numpy's "setting an array element with a sequence" ValueError -> this ValueError. - r + ([1, 2], 0): numpy's "could not broadcast input array from shape (2,) into shape (3,)" -> this ValueError. - r[0] = (np.array([[1, 300, 3]]), 0): ([1, 44, 3], 0), unchanged. * Cover more literals a UDT refuses and the Numba error of a lifted op Literal typing had arms no test reached: a comparison with a record literal of the wrong length (alone and nested beside an int beyond its leaf), a part holding numpy values that numpy cannot make an array of or that is not numbers, and text a bytes field cannot encode. An array of records now pins the UdfParseError a lifted op gives when Numba cannot compile it, and the NoJITWarning is checked to stay silent for a pair it warned about. Two arms of the monoid identity are excluded from coverage: the binary op compiles first, and it refuses what would reach them. * Give each type signature of a JIT IndexBinaryOp its own C name indexbinary.ss.register_new compiled every type signature of one op under the op's own name. SuiteSparse finds a compiled kernel in memory by its op's name and the types in the kernel's encoding, and an ewise or mxm kernel's encoding has x's, y's and z's types but not theta's. So two signatures that differ only in theta's type ran whichever kernel was compiled first. Now each signature's C function is compiled as name_N_x_M_y_theta, with N and M the lengths of the x and y parts (built-in types by name, UDTs by their C name): the C definition is the user's with the function renamed. A name SuiteSparse would cut at 127 characters is the op name's first 100 and a digest of the whole. The op's jit_c_name and jit_c_source show the C name and the source SuiteSparse compiles; jit_c_definition is still the definition given. JIT IndexUnaryOps and SelectOps are not affected: the thunk's type is in their kernels' encodings. Behavior change (before -> after): with add_theta registered at FP64 for x, y and theta, and then at an FP32 theta with (*z) = (*x) + (*y) + 2 * (double)(*theta), the FP32 signature on FP64 data with theta 10 gave [2, 2, 4, 0] (the FP64 kernel, reading the float as a double) -> [22, 22, 24, 20]. op.jit_c_name: None -> add_theta_4_FP64_4_FP64_FP32. * Name each op registered from C by its full name and definition The ss.register_new functions of unary, binary, indexunary (and select) gave SuiteSparse the op's last name as its C name, and indexbinary's the last name and types. SuiteSparse finds a compiled kernel in memory by its op's name, and checks an op's definition only when it loads a kernel from disk. So, in one process, an op registered again under its name with another definition (after delattr, as when a notebook cell is run again), or an op of the same last name in another namespace (ss.f and ss.nested.f), ran the kernel compiled first. Now every op registered from C source is compiled under a C name made of its last name, the types for an IndexBinaryOp (whose theta type is not in its kernels' encoding), and a 12-hex digest of these with the full dotted name and the definition: f_3a1c90e2d4b7, or add_theta_4_FP64_4_FP64_4_FP32_... for an IndexBinaryOp. A name SuiteSparse would cut at 127 characters is the last name's first 100 characters and a 16-hex digest. The function is renamed in the definition wherever it is named, except as a struct member (after "." or "->"). The same name and definition give the same C name, so a kernel is compiled once for them. jit_c_name and jit_c_source, None for these ops before (except IndexBinaryOp, from the previous commit), show the C name and the renamed source; jit_c_definition is still the definition given. The helper is ss._utils._jit_c_source. Behavior changes (before -> after), in one process: - binary.ss.register_new("f", x + y) and ("nested.f", x - y), both FP64: nested.f computed x + y -> x - y. - unary.ss.register_new("u", x * x), delattr(unary.ss, "u"), then register_new("u", x + 10): u computed x * x -> x + 10. * Refuse an empty Scalar theta for an IndexBinaryOp SuiteSparse binds theta without checking that it holds a value: it reads the Scalar's value array. Binding an empty Scalar crashed the interpreter, and op(s) with no type, which binds a copy of s, ran with theta 0. Both now raise EmptyObject, as apply does for an empty IndexUnaryOp thunk, in every bind: typed or not, is_udt or not, a JIT op's too. Behavior change (before -> after), with s = Scalar(int): - op[int](s) and op(s, dtype=int): segmentation fault -> EmptyObject. - op(s): an op that ran with theta 0 -> EmptyObject, at the bind. With s a C scalar (is_cscalar=True): segmentation fault -> EmptyObject. * Leave a dtypes.ss.register_new name free when the registration fails register_new recorded the name as a JIT C type name as soon as SuiteSparse made the type, before the numpy and Numba types were made. When one of those raised (an np_type with an array of bool, which Numba cannot type), the name stayed reserved, and a correct second call raised "name for dtype is unavailable". Main allowed the second call; the check against C type names came with "Give each UDT layout its own JIT C name". The name is now recorded after the type is registered. The type made by the failed call is never used, so no JIT kernel knows its name. Behavior change (before -> after), with d the typedef of pt_t: - register_new("pt_t", d, np_type=np.dtype((np.bool_, (2,)))) raises ValueError (unchanged), then register_new("pt_t", d, np_type=np.dtype([("x", "f8"), ("y", "f8")])): ValueError "'pt_t' name for dtype is unavailable" -> registers dtypes.ss.pt_t. * Type a Python date or duration in a record literal as numpy does A record literal's date or duration field takes the field's type only when its value converts into the field and back unchanged; otherwise it keeps its own type, so == compares it whole. A numpy date did this already, but a Python datetime.date, datetime.datetime or datetime.timedelta, or a date given as text, always took the field's type. So it was cut to the field's unit before == compared it: noon compared equal to its day in a field of days, and 5 s and 1 us equal to 5 s in a field of seconds. Main gave the same wrong answers (it gave them for numpy dates too). Now such a value is first made the numpy value it stands for (a datetime in microseconds, a timedelta in microseconds, a date as text in the unit it is written to), and is typed as that numpy value is. Behavior change (before -> after), with v a Vector of a record of an M8[D] and an m8[s] field holding (2020-01-01, 5 s), and s = v[0].new(): - v == (datetime(2020, 1, 1, 12), np.timedelta64(5, "s")), and with "2020-01-01T12:00" for the date: True -> False; the same for s ==. - v == (np.datetime64("2020-01-01"), timedelta(seconds=5, microseconds=1)): True -> False. - v.apply(user_op, right=(datetime(2020, 1, 1, 12), np.timedelta64(5, "s"))): stored the day, silently -> ValueError "does not fit", as for np.datetime64("2020-01-01T12:00"). - v + (datetime(2020, 1, 1, 12), ...): UdfParseError from Numba -> KeyError naming datetime64[D] and datetime64[us]. - A datetime at midnight, a date, and text such as "2020-01-01" still compare equal and fit the field. * Compare a UDT of dates or durations with one in another unit as numpy does A UDT that is a bare datetime64 or timedelta64 (not a record or an array) converted a date or duration literal into its own unit before == compared it, so noon compared equal to its day in a UDT of days. Main raised TypeError here for a Vector, and gave True for a Scalar; "Compare UDTs that are neither records nor arrays in eq and ne" made the Vector give True too. A numpy duration in another unit was a KeyError instead, even when it was the same duration. Now such a literal is typed as in a record's date field: the UDT itself when it converts into the UDT and back unchanged, else numpy's type for it (a datetime is in microseconds, a timedelta too, a date as text in the unit it is written to). Under eq and ne, a UDT of dates or durations pairs with one in another unit, as text of one kind pairs at another width, and they compare as numpy compares them. Behavior change (before -> after), with day a Vector of an M8[D] UDT holding 2020-01-01 and five one of an m8[s] UDT holding 5 s: - day == np.datetime64("2020-01-01T12:00"), datetime(2020, 1, 1, 12) or "2020-01-01T12:00": True -> False; != the opposite; the same for day[0].new() ==. - five == np.timedelta64(5000, "ms"): KeyError -> True; 5001 ms: KeyError -> False. five == timedelta(seconds=5, microseconds=1): True -> False. - binary.eq(five & a Vector of m8[ms] holding 5000 ms), and five.isequal of it: KeyError -> True. - day.apply(user_op, right=datetime(2020, 1, 1, 12)): stored the day, silently -> ValueError "does not fit". - five.apply(user_op, right=np.timedelta64(5000, "ms")): converted with a DeprecationWarning -> converts with none, as in a record's field. - day == timedelta(5): numpy's ValueError -> binary.eq's KeyError. * Type a literal beside a typed op in the op's own input type A literal operand beside a UDT became an element of that UDT even when the op was typed and takes that operand in another type: op[A, INT64] takes its thunk, or its second input, as an INT64. Apply and select pass a UDT thunk or operand to SuiteSparse as a pointer to its bytes, so the op read the bytes of the A element as an int64: silently wrong answers. Main did the same. A Scalar of the op's type gave the right answer. Now a typed op's own input type on the literal's side decides: its first input type for a left operand, else its second input or thunk type. GraphBLAS converts a value to a built-in input type, so the literal is typed as beside a built-in typed op (its own type: 2.5 for an INT64 thunk is converted to 2, as for an op typed on built-in types); a UDT input type takes only a literal that fits it. _literal_operand takes left=True from apply's left operand and from ewise_union's left default. Behavior change (before -> after), with v a Vector of a record A of two float64 fields holding (1.0, 2.0) and (3.0, 4.0), and user ops typed [A, INT64] (or [INT64, A]) that compare or add field a with the other: - v.apply(iu[A, INT64], 2): [False, False] -> [False, True], as with Scalar.from_value(2); the same for thunks 2.5 and np.int8(2). - v.select(se[A, INT64], 2), on a Vector or Matrix: empty -> index 1. - v.apply(f[A, INT64], right=2) and v.apply(g[INT64, A], left=2): about 4.6e18 for each element -> [3.0, 5.0]. - With B a record of two int64 fields: v.apply(h[A, B], right=(2, 3)): about 4.6e18 -> [3.0, 5.0]; v.apply(iu_rec[A, B], (2, 0)): [False, False] -> [False, True], and with (2.5, 0): [False, False] -> ValueError "does not fit". - Scalar(A).ewise_mult(2, f[A, INT64]): GraphBLAS's DomainMismatch -> 3.0. * Accumulate into a UDT Scalar as into a built-in Scalar The accumulator is typed on the object's dtype, and for a built-in Scalar GraphBLAS converts the value to that type before it accumulates: an INT8 Scalar of -1 with s += 0.5 (Scalar.__iadd__ is s(binary.plus) << 0.5) adds 0 and stays -1, as v(accum=binary.plus) << 0.5 does for a Vector of either kind. A UDT Scalar used the untyped op instead (_scalar_accum_op), so it computed -1 + 0.5 in FP64 and truncated the sum to 0. That was meant to match a Vector's v += 0.5, but that is v << v + 0.5, not an accumulate. Main gave the built-in answer. Now a UDT Scalar accumulates with the typed op, and the value is converted to the Scalar's type first, as storing it would convert it: a literal by Scalar.from_value, a Scalar or an element of another UDT by the store cast. A built-in Scalar into a UDT Scalar still raises DomainMismatch, as a store of it does, into a Vector too. Behavior change (before -> after), with s an INT8[3] UDT Scalar of [-1, -1, -1]: - s += 0.5 and s(accum=binary.plus) << 0.5: [0, 0, 0] -> [-1, -1, -1] (an INT8 Scalar gives -1; a Vector of the UDT gives [-1, -1, -1]). - s(accum=binary.plus) << (0.5, 1.5, 2.5): [0, 0, 1] -> [-1, 0, 1], as for a Vector of the UDT. - s(accum=binary.plus) << v[1], with v a Vector of an FP64[3] UDT holding [0.5, 1.5, 2.5]: [0, 0, 1] -> [-1, 0, 1], as with the same value in a Scalar of that UDT, which was cast first already. * Pair an array literal with an array UDT when their shapes broadcast udt.rst says a numpy array beside a UDT pairs as an array UDT of its own shape, and array UDTs pair when their shapes broadcast. A numpy array paired only when it broadcast to the UDT's shape, or had that shape apart from leading 1s. Any other array, even one whose shape broadcasts with the UDT's, was converted into the UDT, where numpy raised "could not broadcast", while the same values as a UDT operand worked. Main raised too (a KeyError for the UDT pair). Now such an array is typed in its own shape and pairs as an array UDT operand does; the result has the shape numpy's would. An array that broadcasts to the UDT's shape is typed in the UDT's shape, as before, which a literal that must become an element needs. Behavior change (before -> after), with v a Vector of an INT8[3] UDT holding [1, 2, 3] and u one of an INT8[1] UDT holding [1]: - v + np.array([[1, 2, 3], [4, 5, 6]]), and apply(binary.plus) with it on either side: ValueError "could not broadcast input array from shape (2,3) into shape (3,)" -> INT64[2, 3] [[2, 4, 6], [5, 7, 9]], as with the same values in an INT64[2, 3] UDT. - v == np.array([[1, 2, 3], [1, 2, 3]]), on a Vector or a Scalar: the same ValueError -> True. - u + np.array([1, 2, 3]), and u + (1, 2, 3): ValueError -> INT64[3] [2, 3, 4]. - v + [[1, 2, 3], [4, 5, 6]]: numpy's "setting an array element with a sequence" ValueError -> INT64[2, 3], as for the array. - v.ewise_union(v, binary.plus, a 2-by-3 array, 0): numpy's ValueError -> ValueError "does not fit". * Give IndexBinaryOp is_positional, so dtype in op and op[dtype] work OpBase.__getitem__ reads is_positional for a type the op does not have, and IndexBinaryOp had none, so asking an IndexBinaryOp (registered without is_udt) for a type it does not take raised AttributeError. Main did the same; IndexBinaryOp has not been released. It is now False, as on a bound IndexBinaryOp and every other op class. Behavior change (before -> after), with op an IndexBinaryOp of x + y + theta and R a record UDT: - R in op: AttributeError -> False, as for an IndexUnaryOp. - op[R] and op(Scalar.from_value((1.0, 2), R), dtype=R): AttributeError -> KeyError "... does not work with R". * Refuse a bound IndexBinaryOp in reduce with the usual TypeError reduce turns a BinaryOp that is part of a Monoid into that Monoid, and refuses any other BinaryOp with a TypeError naming it. It reads the op's monoid for that, and the error message its parent's; neither a bound IndexBinaryOp nor IndexBinaryOp had the attribute, so reduce raised AttributeError, for ibo[T](theta) on main and for the new untyped ibo(theta) too. Both now have monoid = None, as a BinaryOp of no Monoid has. IndexBinaryOp has not been released. Behavior change (before -> after), with v an INT64 Vector and ibo an IndexBinaryOp of x + y + theta: - v.reduce(ibo(1)) and v.reduce(ibo[int](1)): AttributeError "'_BoundIndexBinaryOp' object has no attribute 'monoid'" -> TypeError "Bad type for argument `op` in Vector.reduce(...) ... is not known to be part of a Monoid", as for binary.minus. - A.reduce_rowwise(ibo(1)): AttributeError -> NotImplementedException, as for binary.minus. --- .github/workflows/test_and_build.yml | 5 + .pre-commit-config.yaml | 1 - docs/user_guide/operators.rst | 8 +- docs/user_guide/types.rst | 27 + docs/user_guide/udf.rst | 52 +- docs/user_guide/udt.rst | 155 +- graphblas/core/base.py | 49 +- graphblas/core/dtypes.py | 110 +- graphblas/core/infixmethods.py | 11 + graphblas/core/literals.py | 835 ++++++ graphblas/core/matrix.py | 143 +- graphblas/core/operator/base.py | 649 +---- graphblas/core/operator/binary.py | 168 +- graphblas/core/operator/indexbinary.py | 256 +- graphblas/core/operator/indexunary.py | 9 +- graphblas/core/operator/monoid.py | 8 +- graphblas/core/operator/select.py | 140 +- graphblas/core/operator/semiring.py | 29 +- graphblas/core/operator/udf.py | 682 +++++ graphblas/core/operator/udt_codegen.py | 1536 +++++++++++ graphblas/core/operator/udt_utils.py | 1003 -------- graphblas/core/operator/unary.py | 30 +- graphblas/core/operator/utils.py | 23 +- graphblas/core/scalar.py | 271 +- graphblas/core/ss/_utils.py | 32 + graphblas/core/ss/binary.py | 14 +- graphblas/core/ss/dtypes.py | 12 +- graphblas/core/ss/indexbinary.py | 38 +- graphblas/core/ss/indexunary.py | 15 +- graphblas/core/ss/jit_config.py | 2 +- graphblas/core/ss/select.py | 7 +- graphblas/core/ss/unary.py | 14 +- graphblas/core/udt.py | 911 +++++++ graphblas/core/utils.py | 34 +- graphblas/core/vector.py | 120 +- graphblas/tests/test_dtype.py | 76 +- graphblas/tests/test_indexbinary.py | 293 ++- graphblas/tests/test_infix.py | 2 + graphblas/tests/test_matrix.py | 18 +- graphblas/tests/test_op.py | 3252 +++++++++++++++++++++++- graphblas/tests/test_scalar.py | 163 ++ graphblas/tests/test_ssjit.py | 344 ++- graphblas/tests/test_vector.py | 303 ++- pyproject.toml | 2 +- 44 files changed, 9680 insertions(+), 2172 deletions(-) create mode 100644 graphblas/core/literals.py create mode 100644 graphblas/core/operator/udf.py create mode 100644 graphblas/core/operator/udt_codegen.py delete mode 100644 graphblas/core/operator/udt_utils.py create mode 100644 graphblas/core/udt.py diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index 574bcd3d8..8dec76f13 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -358,6 +358,11 @@ jobs: coverage report --show-missing - name: codecov uses: codecov/codecov-action@v5 + 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' run: | diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3b36390ad..def06f9d5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -29,7 +29,6 @@ repos: - 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 diff --git a/docs/user_guide/operators.rst b/docs/user_guide/operators.rst index 59a38c791..fb827afe3 100644 --- a/docs/user_guide/operators.rst +++ b/docs/user_guide/operators.rst @@ -232,7 +232,8 @@ receives both values, the indices of *each* value, and a thunk parameter: 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``. +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: @@ -243,11 +244,10 @@ Example usage: gb.indexbinary.register_new("discounted_sum", discounted_sum) - binop = gb.indexbinary.discounted_sum[float](0.5) # bind theta + binop = gb.indexbinary.discounted_sum(0.5) # bind theta C << A.ewise_mult(B, binop) - # For mxm, wrap in a Semiring; the additive monoid must accept the - # bound op's return type. + # For mxm, wrap in a Semiring sr = gb.semiring.register_anonymous(gb.monoid.plus, binop) D << A.mxm(B, sr) diff --git a/docs/user_guide/types.rst b/docs/user_guide/types.rst index aaeba1c6b..36d93e8a4 100644 --- a/docs/user_guide/types.rst +++ b/docs/user_guide/types.rst @@ -30,6 +30,33 @@ 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. 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 ``v / 300`` raise, since its result is FP64 whatever the integer. 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 670cbc6e3..6255c1c56 100644 --- a/docs/user_guide/udf.rst +++ b/docs/user_guide/udf.rst @@ -125,11 +125,43 @@ works in ``ewise_mult``, ``ewise_add``, and the other elementwise paths: indexbinary.register_new("discounted_sum", discounted_sum) - bound = indexbinary.discounted_sum[float](0.5) # theta = 0.5 - A = Matrix.from_coo([0, 1], [0, 1], [1.0, 2.0]) - B = Matrix.from_coo([0, 1], [0, 1], [3.0, 4.0]) + 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[0,0] = (1+3)*0.5 = 2; C[1,1] = (2+4)*0.5 = 3 + # 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 @@ -142,12 +174,12 @@ bound IBO directly: sr = semiring.register_anonymous(monoid.plus, bound) C = A.mxm(B, sr).new() -The resulting Semiring is monomorphic in the bound IBO's input/output types -(SuiteSparse builds exactly one ``GrB_Semiring`` for that type pair, rather -than the type-polymorphic table the standard semirings carry). To reuse the -same IBO at a different type, bind theta again under that type and build a -new Semiring. Per SuiteSparse, monoids themselves cannot be built from an -IndexBinaryOp; only the multiplier slot accepts one. +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. diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index ae2cef9e2..2a0b86555 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -74,6 +74,9 @@ Field type rules: 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`` or ``timedelta64``, so + with Numba installed, registering a UDT that holds one raises ``ValueError``. + Use an array of ``uint8`` for ``bool``, and of ``int64`` for the others. - Strings and Python objects are not supported. - Dataclass annotation strings (e.g., ``"int"``) resolve through ``lookup_dtype``. @@ -134,12 +137,152 @@ For example: plus_times = semiring.plus_times[edge_dtype] # field-wise multiply-add semiring total = v.reduce(agg.sum[edge_dtype]).new() # field-wise reduce -The lift is field-by-field for record UDTs and element-by-element for array -UDTs. Field types that don't support the operation raise ``KeyError`` on the -first lookup. ``binary.min``, ``binary.max``, and ``binary.floordiv`` reject -UDTs with any complex leaf (no ordering, no integer modulus); use ``plus``, -``minus``, ``times``, or ``truediv`` for complex arithmetic, or register a -custom binary op. +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; diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 15f66bc2f..e8940ad9b 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -7,8 +7,10 @@ from . import NULL from .descriptor import lookup as descriptor_lookup from .expr import AmbiguousAssignOrExtract, Updater +from .literals import _literal_store_dtype from .mask import 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") @@ -341,6 +343,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( @@ -362,6 +365,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 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) @@ -386,6 +393,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: @@ -399,16 +407,20 @@ 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 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() @@ -418,7 +430,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: @@ -446,6 +458,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) diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 93b3d4d8a..484dcf354 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -3,7 +3,7 @@ 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 @@ -117,7 +117,7 @@ def _get_udt_c_type_info(datatype): word or isn't an identifier. Supports record UDTs (including nested records) and array UDTs (e.g., ``FP64[3]``). """ - from .operator.udt_utils import _udt_c_typedef + from .operator.udt_codegen import _udt_c_typedef if datatype.np_type is None: return None @@ -143,7 +143,12 @@ def _set_udt_jit_c_definition(datatype): Supports both record UDTs (``{'x': float, 'y': float}``) and array UDTs (``np.dtype((np.float64, (3,)))``). """ - from .operator.udt_utils import _has_jit_set + 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 @@ -151,6 +156,11 @@ def _set_udt_jit_c_definition(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: + info = _udt_c_typedef(f"_gbudt_{next(_synthetic_udt_counter)}", datatype.np_type) + 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 @@ -163,6 +173,45 @@ def _set_udt_jit_c_definition(datatype): 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. @@ -283,6 +332,12 @@ def register_anonymous(dtype, name=None): 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 @@ -320,7 +375,6 @@ 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 @@ -566,14 +620,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)) @@ -649,6 +704,43 @@ def _view_if_same_layout(values, 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. diff --git a/graphblas/core/infixmethods.py b/graphblas/core/infixmethods.py index 67248223a..53f9d76f0 100644 --- a/graphblas/core/infixmethods.py +++ b/graphblas/core/infixmethods.py @@ -3,6 +3,7 @@ from .infix import MatrixInfixExpr, ScalarInfixExpr, VectorInfixExpr from .matrix import Matrix, MatrixExpression, MatrixIndexExpr, TransposedMatrix from .scalar import Scalar, ScalarExpression, ScalarIndexExpr +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) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py new file mode 100644 index 000000000..810bde22c --- /dev/null +++ b/graphblas/core/literals.py @@ -0,0 +1,835 @@ +"""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 ..dtypes import FP64, lookup_dtype +from . import utils +from .dtypes import _nats_in_leaf_units, _view_if_same_layout +from .operator import IndexUnaryOp, 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: + 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 = _dict_to_record(np_type, value) if isinstance(value, dict) else 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_type == dtype.np_type: + rv = dtype + else: + try: + rv = lookup_dtype(np_type) + except ValueError: # A complex literal where the backend has no complex types + return None + 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/matrix.py b/graphblas/core/matrix.py index 4625e4258..166f263a0 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -10,6 +10,7 @@ from .base import BaseExpression, BaseType, _check_mask, call from .descriptor import lookup as descriptor_lookup from .expr import _ALL_INDICES, AmbiguousAssignOrExtract, IndexerResolver, InfixExprBase, Updater +from .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype from .mask import Mask, StructuralMask, ValueMask from .operator import ( UNKNOWN_OPCLASS, @@ -19,6 +20,7 @@ get_typed_op, op_from_string, ) +from .operator.select import _value_select_in_own_type from .scalar import ( _COMPLETE, _MATERIALIZE, @@ -27,6 +29,15 @@ 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, @@ -100,6 +111,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 +127,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 +167,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 +186,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): @@ -384,6 +412,11 @@ def isequal(self, other, *, check_dtype=False, **opts): ------- 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 +427,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: @@ -707,7 +742,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 @@ -1499,7 +1538,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, @@ -1921,14 +1961,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 +2147,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") @@ -2419,6 +2433,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 +2459,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 +2495,7 @@ 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, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -2596,8 +2615,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 +2631,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 +2971,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="") @@ -3026,6 +3048,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 +3117,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 +3128,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 +3148,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 @@ -3311,7 +3346,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): diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index e4c41b95e..aea4a6f51 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -1,20 +1,15 @@ -import re from functools import lru_cache from operator import getitem from types import BuiltinFunctionType, ModuleType -import numpy as np - from ... import _STANDARD_OPERATOR_NAMES, backend, op -from ...dtypes import BOOL, INT8, UINT64, _supports_complex, lookup_dtype -from ...exceptions import UdfParseError, check_status_carg +from ...dtypes import BOOL, UINT64, _supports_complex, lookup_dtype +from ...exceptions import UdfParseError from .. import _has_numba, _supports_udfs, ffi, lib -from ..dtypes import _sample_values from ..expr import InfixExprBase from ..utils import output_type if _has_numba: - import numba from numba import NumbaError else: NumbaError = TypeError @@ -98,20 +93,6 @@ def _hasop(module, name): ) -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-05-24: still hits on Numba 0.65. Re-test periodically and - drop the INT8 routing when upstream is fixed. - """ - return INT8 if dtype == BOOL else dtype - - class OpPath: def __init__(self, parent, name): self._parent = parent @@ -182,632 +163,13 @@ def _call_op(op, left, right=None, thunk=None, **kwargs): ) -if _has_numba: - - 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``. - """ - wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) - 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 - if isinstance(numba_ret_type, numba.core.types.BaseTuple): - from .udt_utils import _iter_record_leaves - - n = len(numba_ret_type.types) - - def _leaves(d): - """Yield ``(python_path, c_path, leaf_dtype)`` for 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): - continue - 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-05-24: still hits on Numba 0.65; 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. - from .udt_utils import _iter_record_leaves - - 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" - f"{body_setup}" - f" _result = {call_expr}\n" - f"{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: - 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": - from .udt_utils import _iter_record_leaves - - # 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): - return # a wrong field count, not a wrong shape - 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) - from .udt_utils import _compile_codegen - - 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, *, numba_ret_type=None - ): - """Build a Numba cfunc wrapper for IndexBinaryOp UDFs on UDTs. - - Signature: ``f(x, ix, jx, y, iy, jy, theta) -> z``. ``dtype2`` is the - shared type of ``y`` and ``theta``. - """ - 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(dtype2, "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, dtype2] - _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) - from .udt_utils import _compile_codegen - - wrapper = _compile_codegen( - text, - func_name="wrapper", - source_label=f"", - extra_ns={"numba_func": numba_func}, - ) - return 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``. +# 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", @@ -1241,7 +603,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 2cdab6eff..e17889d2a 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -26,6 +26,10 @@ 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, @@ -35,14 +39,10 @@ _hasop, ) -# Imported unconditionally (plain dict, no numba): ``_compile_udt`` consults it -# even when numba is absent. -from .udt_utils import BUILTIN_UDT_BINARY_OPS as _BUILTIN_UDT_BINARY_OPS - if _has_numba: import numba - from .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _finalize_udt_op, @@ -61,85 +61,100 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): 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``. Returns - ``(wrapper_func, wrapper_sig)``. + 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_utils import _check_udt_pair, _get_udt_info, _iter_record_leaves + 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 - nt = numba.types - cmp_op = "==" if is_eq else "!=" - join = " and " if is_eq else " or " - if udt_info[0] == "array": - # Array UDT: unroll element-by-element so the scalar broadcast - # case can drop ``numba.carray`` on that side entirely. - # ``_get_udt_info`` already flattens the multi-dim shape. - # The wrapper sees the UDT 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. - _base, N = udt_info[1] - base_numba = numba.from_dtype(_base) - x_ptr_type = ( - nt.CPointer(numba.from_dtype(dtype.np_type)) - if x_is_scalar - else nt.CPointer(base_numba) - ) - y_ptr_type = ( - nt.CPointer(numba.from_dtype(dtype2.np_type)) - if y_is_scalar - else nt.CPointer(base_numba) - ) + # 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) - if x_is_scalar: - terms = [f"(x_ptr[0] {cmp_op} y[{i}])" for i in range(N)] - arrays = f" y = numba.carray(y_ptr, {N})\n" - elif y_is_scalar: - terms = [f"(x[{i}] {cmp_op} y_ptr[0])" for i in range(N)] - arrays = f" x = numba.carray(x_ptr, {N})\n" - else: - terms = [f"(x[{i}] {cmp_op} y[{i}])" for i in range(N)] - arrays = f" x = numba.carray(x_ptr, {N})\n y = numba.carray(y_ptr, {N})\n" - body = join.join(terms) - src = f"def wrapper(z_ptr, x_ptr, y_ptr):\n{arrays} z_ptr[0] = {body}\n" + 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. For an - # array-valued sub-field, unroll its elements so the scalar side - # never sees a numpy fancy-broadcast (Numba's record access can't - # express that as a single expression). + # 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: - _base, shape = leaf_dtype.subdtype - sub_N = int(reduce(mul, shape)) - if x_is_scalar: - subterms = [f"(x_ptr[0] {cmp_op} y[0]{py_path}[{i}])" for i in range(sub_N)] - elif y_is_scalar: - subterms = [f"(x[0]{py_path}[{i}] {cmp_op} y_ptr[0])" for i in range(sub_N)] - else: - reducer = ".all()" if is_eq else ".any()" - terms.append(f"(x[0]{py_path} {cmp_op} y[0]{py_path}){reducer}") - continue - terms.append("(" + join.join(subterms) + ")") + reducer = ".all()" if is_eq else ".any()" + terms.append(f"({x_expr} {cmp_op} {y_expr}){reducer}") else: - 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}" 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):"] @@ -389,7 +404,7 @@ def _pair_dtype(op, dtype, dtype2): if _has_numba: - from .udt_utils import ( + from .udt_codegen import ( _compile_codegen, _has_jit_set, _maybe_warn_jit_skipped, @@ -617,21 +632,36 @@ 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) + + # 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") + ) # 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 self.name in _BUILTIN_UDT_BINARY_OPS and not dtype._is_udt and not dtype2._is_udt: + 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 self.name in ("eq", "ne") and not self._anonymous and _has_numba: + 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") ) ret_type = BOOL - elif _has_numba and self.name in _BUILTIN_UDT_BINARY_OPS: + 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 @@ -661,9 +691,15 @@ def _compile_udt(self, dtype, dtype2): ) # Set JIT C definition for auto-generated built-in UDT ops. Only # same-type pairs apply: mixed UDT+scalar ops don't have a - # meaningful single-type C definition. ``_has_jit_set`` gates access - # to ``lib.GrB_BinaryOp_set_String``, which is absent on SS < 9. - is_jittable = self.name in _BUILTIN_UDT_BINARY_OPS or self.name in ("eq", "ne") + # meaningful single-type C definition. The arithmetic kernel also + # declares ``z`` with the operands' type, so it fits only an op whose + # result is that type too, which ``truediv`` on integers is not. + # ``_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( @@ -677,7 +713,7 @@ def _compile_udt(self, dtype, dtype2): else: # eq/ne return BOOL, not the UDT; the kernel is a chain of # per-leaf scalar comparisons. - from .udt_utils import set_jit_c_comparison_on_op + from .udt_codegen import set_jit_c_comparison_on_op op._jit_c_info = set_jit_c_comparison_on_op( op.gb_obj, diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index de2ca5f71..4539387a9 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -1,19 +1,23 @@ 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 UdfParseError, check_status_carg +from ...exceptions import EmptyObject, UdfParseError, check_status_carg from .. import _has_numba, ffi, lib -from ..dtypes import _sample_values +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 .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _get_udt_wrapper_indexbinary, @@ -24,6 +28,7 @@ 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 @@ -42,26 +47,44 @@ 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 infer from the value. A raw user-defined-type - value (array/tuple) can't be inferred, so raise a clear error that names the - value and the working alternatives. + ``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}; pass an explicit dtype " - "(e.g., op(value, dtype=...)) or a typed Scalar. Required for user-defined types." + 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 an IndexBinaryOp. + """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, @@ -69,7 +92,7 @@ class _BoundIndexBinaryOp(TypedOpBase): Semiring that then works in ``mxm`` / ``mxv`` / ``vxm``. """ - __slots__ = ("_theta",) + __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 @@ -90,6 +113,9 @@ class _BoundIndexBinaryOp(TypedOpBase): _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 @@ -99,7 +125,86 @@ class _BoundIndexBinaryOp(TypedOpBase): _custom_dtype = None def __reduce__(self): - return (_rebind_indexbinaryop, (self.parent, self.type, self._theta)) + 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): @@ -116,7 +221,8 @@ def __call__(self, theta=None): Parameters ---------- theta : scalar, optional - The theta parameter to bind. Defaults to 0 (False). + 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 ------- @@ -140,6 +246,7 @@ def __call__(self, theta=None): 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), @@ -156,20 +263,20 @@ def __call__(self, theta=None): dtype2=self._type2, ) rv._theta = theta_value + rv._theta_type = self.thunk_type return rv - @property - def thunk_type(self): - return self.type if self._type2 is None else self._type2 - -class TypedUserIndexBinaryOp(TypedOpBase): - __slots__ = () +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): + 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): @@ -179,8 +286,6 @@ def orig_func(self): def _numba_func(self): return self.parent._numba_func - thunk_type = TypedBuiltinIndexBinaryOp.thunk_type - def __call__(self, theta=None): return TypedBuiltinIndexBinaryOp.__call__(self, theta) @@ -211,11 +316,14 @@ class IndexBinaryOp(OpBase): 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[dtype](theta)``) produces a BinaryOp usable - directly in ``ewise_mult`` and ``ewise_add``. 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. + 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. @@ -225,6 +333,11 @@ class IndexBinaryOp(OpBase): _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 @@ -361,7 +474,7 @@ def indexbinary_wrapper( return new_type_obj raise UdfParseError("Unable to parse function using Numba") - def _compile_udt(self, dtype, dtype2): + 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. @@ -371,9 +484,15 @@ def _compile_udt(self, dtype, dtype2): ) if dtype2 is None: dtype2 = dtype - dtypes = (dtype, dtype2) + 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") @@ -385,15 +504,15 @@ def _compile_udt(self, dtype, dtype2): dtype2.numba_type, UINT64.numba_type, UINT64.numba_type, - dtype2.numba_type, + theta_type.numba_type, ) _compile_udf_for_udt( - numba_func, sig, op_kind="indexbinary", op_name=self.name, dtypes=(dtype, dtype2) + 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, dtype, dtype2) + ret_type = _resolve_udt_return_type(numba_ret_type, *dtypes) indexbinary_wrapper, wrapper_sig = _get_udt_wrapper_indexbinary( - numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type + 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")( @@ -407,7 +526,7 @@ def _compile_udt(self, dtype, dtype2): ret_type._carg, dtype._carg, dtype2._carg, - dtype2._carg, + theta_type._carg, ffi_new("char[]", self.name.encode()), ffi.NULL, ), @@ -421,11 +540,38 @@ def _compile_udt(self, dtype, dtype2): 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. @@ -517,8 +663,8 @@ def __init__(self, name, func=None, *, anonymous=False, is_udt=False, numba_func self._numba_func = numba_func self._is_udt = is_udt if is_udt: - self._udt_types = {} # {(dtype, dtype2): DataType} - self._udt_ops = {} # {(dtype, dtype2): TypedUserIndexBinaryOp} + 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. @@ -526,13 +672,18 @@ def __call__(self, theta=None, dtype=None): Parameters ---------- theta : scalar, optional - The theta parameter to bind. Defaults to 0 (False). + 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 dtype to use. Inferred from ``theta`` when not provided. + 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 ------- - TypedOpBase + BinaryOp A BinaryOp built from this IndexBinaryOp with the given theta. """ @@ -541,22 +692,21 @@ def __call__(self, theta=None, dtype=None): if theta is None: theta = False - if not isinstance(theta, Scalar): - # An explicit dtype supplies a UDT; with no dtype the value is - # inferred (and a raw UDT value raises a clear error). - theta = _theta_to_scalar(theta, _lookup_dtype(dtype) if dtype is not None 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 if dtype is None: - dtype = theta.dtype - else: - dtype = _lookup_dtype(dtype) - typed_op = self[dtype] - return typed_op(theta) + 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 5820e66c5..a8eca2460 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -12,7 +12,7 @@ if _has_numba: import numba - from .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _finalize_udt_op, @@ -207,6 +207,9 @@ 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") @@ -402,8 +405,8 @@ 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} + self._udt_types = {} # {(dtype, thunk_type): DataType} + self._udt_ops = {} # {(dtype, thunk_type): TypedUserIndexUnaryOp} __call__ = TypedBuiltinIndexUnaryOp.__call__ diff --git a/graphblas/core/operator/monoid.py b/graphblas/core/operator/monoid.py index 0b02d4226..72765e2c2 100644 --- a/graphblas/core/operator/monoid.py +++ b/graphblas/core/operator/monoid.py @@ -69,7 +69,8 @@ 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) - if val is None: + # 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" @@ -90,7 +91,10 @@ def _udt_identity_np(monoid_name, np_type, top_dtype): ) vals.append(val) return tuple(vals) - raise KeyError(f"monoid.{monoid_name} does not work with {top_dtype}") + # 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 diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 3cc506d84..1c1138ae5 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -7,7 +7,138 @@ from .indexunary import IndexUnaryOp, TypedBuiltinIndexUnaryOp if _has_numba: - from .base import _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper + from .udf import _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper + + +_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): @@ -129,6 +260,9 @@ 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") @@ -298,8 +432,8 @@ 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} + self._udt_types = {} # {(dtype, thunk_type): DataType} + self._udt_ops = {} # {(dtype, thunk_type): TypedUserSelectOp} __call__ = TypedBuiltinSelectOp.__call__ diff --git a/graphblas/core/operator/semiring.py b/graphblas/core/operator/semiring.py index b333259e3..f9d8603db 100644 --- a/graphblas/core/operator/semiring.py +++ b/graphblas/core/operator/semiring.py @@ -224,17 +224,18 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): raise TypeError(f"monoid must be a Monoid, not {type(monoid)}") # 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`` 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. + # (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 type(binaryop) is not BinaryOp: + if not isinstance(binaryop, BinaryOp): raise TypeError( f"binaryop must be a BinaryOp or a bound IndexBinaryOp " f"(``ibo[dtype](theta)``), not {type(binaryop)}" @@ -351,9 +352,10 @@ def register_anonymous(cls, monoid, binaryop, name=None): The monoid of the semiring (like "plus" 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). A bound IndexBinaryOp (``ibo[dtype](theta)``) - is accepted as well; the resulting semiring is monomorphic in - the bound op's input/output types. + "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}``. @@ -379,9 +381,10 @@ def register_new(cls, name, monoid, binaryop, *, lazy=False): monoid : Monoid or ParameterizedMonoid The monoid of the semiring (like "plus" in the default "plus_times" semiring). binaryop : BinaryOp, ParameterizedBinaryOp, or bound IndexBinaryOp - The multiplier of the semiring. A bound IndexBinaryOp - (``ibo[dtype](theta)``) is also accepted; the resulting - semiring is monomorphic in the bound op's input/output types. + 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..ab8144fc2 --- /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-05-24: still hits on Numba 0.65. 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-05-24: still hits on Numba 0.65; 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..81d78f73f --- /dev/null +++ b/graphblas/core/operator/udt_codegen.py @@ -0,0 +1,1536 @@ +"""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 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 UDT type names or record field names in the +# JIT C source. Covers two failure modes from the JIT include chain +# (``GraphBLAS.h`` transitively pulls in ````, ````, +# ````, ````, ````, ````): +# +# 1. **Macro names**: the preprocessor expands them inside any declarator +# or expression they appear in, so a field declared ``double M_PI ;`` +# becomes ``double 3.14159... ;`` and won't compile. The C-standard +# separate namespace for struct members doesn't help here. +# 2. **Typedef names** at the *outer* type-name position: emitting +# ``typedef struct { ... } FILE ;`` collides with the existing +# ``FILE`` typedef in scope (most compilers reject as a redefinition). +# Struct member names of the same spelling are fine (members live in +# their own namespace), but our gate runs the same check at both +# positions for simplicity. +# +# Either way the JIT compile fails, SuiteSparse swallows the failure +# silently, and the op runs through the slower Numba cfunc path. Block +# these eagerly so the user sees the "without JIT" warning instead. +_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", + } +) + + +def _is_valid_c_identifier(name): + """True if ``name`` is a valid C identifier and not a reserved word.""" + return isinstance(name, str) and name.isidentifier() and name not in _C_RESERVED + + +# Process-global counter for synthesizing C names when the user-supplied +# Python name isn't a valid C identifier (or is a C reserved word, 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 = {} + + +def _pick_c_type_name(python_name): + """Return a valid C identifier for use as ``GxB_JIT_C_NAME`` on a UDT. + + When ``python_name`` is already a valid C identifier (and not reserved), + return it unchanged so introspection and JIT cache filenames remain + readable. Otherwise mint a fresh ``_gbudt_NNN`` name. This lets UDTs + that were anonymous in the Python sense (no ``name=`` passed, name with + dots / special chars, name that collides with a C keyword) still take + the JIT path; users never see the synthetic name unless they read + ``DataType.jit_c_name``. + """ + if _is_valid_c_identifier(python_name): + return python_name + return f"_gbudt_{next(_synthetic_udt_counter)}" + + +def _udt_c_typedef(python_name, np_type): + """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 (may differ from + ``python_name`` if that wasn't a valid C identifier; see + ``_pick_c_type_name``). 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 = _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) + 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 + # synthetic ``_gbudt_NNN`` name (for UDTs whose Python name isn't a valid + # C identifier) is recorded; using it keeps op codegen consistent. + pinned_name = dtype.jit_c_name + typedef_info = _udt_c_typedef(pinned_name or dtype.name, dtype.np_type) + 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 _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 + pinned_name = dtype.jit_c_name + # ``_udt_c_typedef`` returns None when the numpy layout doesn't match a + # C compiler's layout for the same struct. + typedef_info = _udt_c_typedef(pinned_name or dtype.name, np_type) + 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() + set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) + set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) + + +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 + _set_jit_c_strings(gb_obj, *result, set_string_func) + 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 + _set_jit_c_strings(gb_obj, *result, set_string_func) + return result diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py deleted file mode 100644 index ede25b467..000000000 --- a/graphblas/core/operator/udt_utils.py +++ /dev/null @@ -1,1003 +0,0 @@ -"""Shared utilities for auto-generating UDT operator implementations. - -Provides Numba wrapper generators and JIT C code generators for element-wise -operations on record UDTs (struct types) and array UDTs (e.g., FP64[3]). - -Compilation is lazy: functions here are only called from ``_compile_udt`` -methods when an operator is first used with a particular UDT. -""" - -import ast -import itertools -import linecache -from functools import reduce -from operator import mul - -import numpy as np - -from ... import backend -from .. import _has_numba, ffi, lib - -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") - namespace = {"min": min, "max": max, "abs": abs} - if _has_numba: - namespace["numba"] = numba - if extra_ns: - namespace.update(extra_ns) - exec(code, namespace) - return namespace[func_name] - - -BUILTIN_UDT_BINARY_OPS = { - "plus": "+", - "minus": "-", - "times": "*", - "truediv": "/", - "floordiv": "//", - "min": "min", - "max": "max", -} - -BUILTIN_UDT_UNARY_OPS = { - "ainv": "-", - "abs": "abs", -} - -# 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 and are skipped via _op_supports_field_dtype. - np.dtype(np.complex64): "GxB_FC32_t", - np.dtype(np.complex128): "GxB_FC64_t", -} - -# Ops whose JIT C codegen would produce uncompilable kernels on complex fields -# (no ordering for min/max, no integer-mod for floordiv). -_OPS_NOT_FOR_COMPLEX = frozenset({"min", "max", "floordiv"}) - -# C operator equivalents for JIT code generation -_C_INFIX_OPS = {"+": "+", "-": "-", "*": "*", "/": "/", "//": "/"} - -# 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 UDT type names or record field names in the -# JIT C source. Covers two failure modes from the JIT include chain -# (``GraphBLAS.h`` transitively pulls in ````, ````, -# ````, ````, ````, ````): -# -# 1. **Macro names**: the preprocessor expands them inside any declarator -# or expression they appear in, so a field declared ``double M_PI ;`` -# becomes ``double 3.14159... ;`` and won't compile. The C-standard -# separate namespace for struct members doesn't help here. -# 2. **Typedef names** at the *outer* type-name position: emitting -# ``typedef struct { ... } FILE ;`` collides with the existing -# ``FILE`` typedef in scope (most compilers reject as a redefinition). -# Struct member names of the same spelling are fine (members live in -# their own namespace), but our gate runs the same check at both -# positions for simplicity. -# -# Either way the JIT compile fails, SuiteSparse swallows the failure -# silently, and the op runs through the slower Numba cfunc path. Block -# these eagerly so the user sees the "without JIT" warning instead. -_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", - } -) - - -def _is_valid_c_identifier(name): - """True if ``name`` is a valid C identifier and not a reserved word.""" - return isinstance(name, str) and name.isidentifier() and name not in _C_RESERVED - - -# Process-global counter for synthesizing C names when the user-supplied -# Python name isn't a valid C identifier (or is a C reserved word, 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) - - -def _pick_c_type_name(python_name): - """Return a valid C identifier for use as ``GxB_JIT_C_NAME`` on a UDT. - - When ``python_name`` is already a valid C identifier (and not reserved), - return it unchanged so introspection and JIT cache filenames remain - readable. Otherwise mint a fresh ``_gbudt_NNN`` name. This lets UDTs - that were anonymous in the Python sense (no ``name=`` passed, name with - dots / special chars, name that collides with a C keyword) still take - the JIT path; users never see the synthetic name unless they read - ``DataType.jit_c_name``. - """ - if _is_valid_c_identifier(python_name): - return python_name - return f"_gbudt_{next(_synthetic_udt_counter)}" - - -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 _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`` (one side is a scalar broadcast), no check fires. Without this - gate the generated wrapper code raises a cryptic Numba ``TypingError`` - on the first field access. - """ - 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" and detail_x != detail_y: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must share field names; got {list(detail_x)} vs " - f"{list(detail_y)}." - ) - if kind_x == "array" and detail_x != detail_y: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"array UDTs must share base dtype and flat size; " - f"got {detail_x} vs {detail_y}." - ) - - -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: the generated - expression operates on the whole sub-array and the wrapper slice-assigns - it back. The JIT path never sees such a record, because a subarray dtype - has no entry in ``NP_TO_C_TYPES`` and ``_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 _udt_c_typedef(python_name, np_type): - """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 (may differ from - ``python_name`` if that wasn't a valid C identifier; see - ``_pick_c_type_name``). 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 = _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) - 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. - """ - if np_type.names is None: - return np_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 - - -def _op_supports_field_dtypes(op_name, np_type): - """Return True if the JIT codegen for ``op_name`` can handle every field - type in ``np_type``. - - Currently the only restriction is complex fields: ``min`` / ``max`` / - ``floordiv`` have no defined semantics on C99 ``_Complex`` (no ordering, - no integer mod), so we skip JIT and let SuiteSparse use the Numba cfunc - instead. The cfunc itself errors on these combinations, so this gate - only moves the failure earlier and gives a clearer message. - - Recurses into nested record fields so a complex field nested inside an - outer record is still detected. - """ - if op_name not in _OPS_NOT_FOR_COMPLEX: - return True - if np_type.names is not None: - return not any(leaf.kind == "c" for _, _, leaf in _iter_record_leaves(np_type)) - if np_type.subdtype is not None: - return np_type.subdtype[0].kind != "c" - 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): - """Python-source builder; sibling of :func:`_c_expr_binary` for JIT C.""" - if py_op in _FUNC_BINARY_OPS: - return f"{py_op}({x_expr}, {y_expr})" - return f"{x_expr} {py_op} {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: - return f"{py_op}({operand})" - return f"{py_op}{operand}" - - def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scalar=False): - """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 (in base.py) 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. - """ - if arity == 2: - parts = [] - for path in leaf_paths: - x_expr = "x" if x_is_scalar else f"x{path}" - y_expr = "y" if y_is_scalar else f"y{path}" - parts.append(_expr_binary(py_op, x_expr, y_expr)) - sig = "x, y" - else: - parts = [_expr_unary(py_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})" - src = f"def _op({sig}):\n return {ret}\n" - op_func = _compile_codegen( - src, - func_name="_op", - source_label=f"", - ) - return numba.njit(op_func, error_model="numpy") - - def _make_array_wrapper( - size, - base_numba_type, - arity, - py_op, - *, - x_scalar_type=None, - y_scalar_type=None, - ): - """Build a cfunc-ready wrapper for an array UDT (element-by-element). - - When ``x_scalar_type`` or ``y_scalar_type`` is set, that side is a plain - scalar pointer (broadcast to all elements). - - Returns (wrapper_func, wrapper_sig). - """ - nt = numba.types - if arity == 2: - x_ref = "x_ptr[0]" if x_scalar_type else "x[{i}]" - y_ref = "y_ptr[0]" if y_scalar_type else "y[{i}]" - assigns = "\n".join( - f" z[{i}] = {_expr_binary(py_op, x_ref.format(i=i), y_ref.format(i=i))}" - for i in range(size) - ) - params = "z_ptr, x_ptr, y_ptr" - arrays = f" z = numba.carray(z_ptr, {size})\n" - if not x_scalar_type: - arrays += f" x = numba.carray(x_ptr, {size})\n" - if not y_scalar_type: - arrays += f" y = numba.carray(y_ptr, {size})\n" - x_numba = nt.CPointer(x_scalar_type) if x_scalar_type else nt.CPointer(base_numba_type) - y_numba = nt.CPointer(y_scalar_type) if y_scalar_type else nt.CPointer(base_numba_type) - sig = nt.void(nt.CPointer(base_numba_type), x_numba, y_numba) - else: - assigns = "\n".join( - f" z[{i}] = {_expr_unary(py_op, f'x[{i}]')}" for i in range(size) - ) - params = "z_ptr, x_ptr" - arrays = f" z = numba.carray(z_ptr, {size})\n x = numba.carray(x_ptr, {size})\n" - sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(base_numba_type)) - src = f"def _op({params}):\n{arrays}{assigns}\n" - op_func = _compile_codegen( - src, - func_name="_op", - source_label=f"", - ) - 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): - """Compile a built-in element-wise binary op for a UDT. - - Handles two cases: - - - Both sides are the same UDT: 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). - - 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 .base import _get_udt_wrapper, _resolve_udt_return_type - - info_x = _get_udt_info(dtype) - info_y = _get_udt_info(dtype2) - _check_udt_pair(op_name, dtype, dtype2, info_x, info_y) - - # Pick the UDT side. Both sides may be UDTs; the pre-check above - # ensures they share a shape in that case. - if info_x is not None: - udt_dtype = dtype - udt_info = info_x - elif info_y is not None: - udt_dtype = dtype2 - udt_info = info_y - else: - 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])." - ) - - # ``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 not _op_supports_field_dtypes(op_name, udt_dtype.np_type): - raise 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." - ) - - 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 - kind, detail = udt_info - - if kind == "record": - from .base 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']"``. - leaf_paths = [py for py, _c, _d in _iter_record_leaves(udt_dtype.np_type)] - func = _make_record_func( - leaf_paths, - 2, - py_op, - x_is_scalar=x_is_scalar, - y_is_scalar=y_is_scalar, - ) - sig = (dtype.numba_type, dtype2.numba_type) - _compile_udf_for_udt( - func, sig, op_kind="binary", op_name=op_name, dtypes=(dtype, dtype2) - ) - numba_ret_type = func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, udt_dtype) - wrapper, wrapper_sig = _get_udt_wrapper( - func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type - ) - else: - base_dtype, size = detail - ret_type = udt_dtype - wrapper, wrapper_sig = _make_array_wrapper( - size, - numba.from_dtype(base_dtype), - 2, - py_op, - x_scalar_type=numba.from_dtype(dtype.np_type) if x_is_scalar else None, - y_scalar_type=numba.from_dtype(dtype2.np_type) if y_is_scalar else None, - ) - return wrapper, wrapper_sig, ret_type - - def compile_udt_unary_wrapper(op_name, py_op, dtype): - """Compile a built-in element-wise unary op for a UDT. - - Returns (wrapper_func, wrapper_sig, ret_type). - Raises KeyError if the dtype is not supported. - """ - from .base 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 .base 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) - wrapper, wrapper_sig = _get_udt_wrapper( - func, ret_type, dtype, numba_ret_type=numba_ret_type - ) - else: - base_dtype, size = detail - ret_type = dtype - wrapper, wrapper_sig = _make_array_wrapper(size, numba.from_dtype(base_dtype), 1, py_op) - return wrapper, wrapper_sig, ret_type - - -# JIT C code generators below. - - -def _c_expr_binary(py_op, lhs, rhs, field_dtype=None): - """Return a C expression for a binary op: e.g., ``(x->a) + (y->a)``. - - ``field_dtype`` is the numpy dtype of the *result* element. It is only - consulted for ``floordiv`` (``//``), which needs Python ``//`` semantics - rather than C ``/`` (trunc toward zero for ints, true division for - floats). Other ops are type-agnostic at the C level. - """ - if py_op == "min": - # Match Python ``min(a, b) = b if b < a else a`` so NaN propagates - # from the first operand (cfunc / numba follows the same rule). - # The naive ``(a < b ? a : b)`` would silently swallow NaN to the - # right-hand side and disagree with the cfunc path. - return f"(({rhs}) < ({lhs}) ? ({rhs}) : ({lhs}))" - if py_op == "max": - return f"(({rhs}) > ({lhs}) ? ({rhs}) : ({lhs}))" - if py_op == "//": - return _c_floordiv_expr(lhs, rhs, field_dtype) - c_op = _C_INFIX_OPS.get(py_op, py_op) - return f"({lhs}) {c_op} ({rhs})" - - -def _c_floordiv_expr(lhs, rhs, field_dtype): - """Return a C expression for Python-semantics floor division. - - Python ``//`` is floor (rounds toward negative infinity); C ``/`` is - trunc toward zero for ints and true division for floats. The two only - agree for non-negative integer operands; for everything else the JIT - path silently disagreed with the Numba cfunc path before this helper. - - Float fields use ``floor()`` / ``floorf()`` from ````, which is - available in the JIT kernel via SuiteSparse's include chain - (``GraphBLAS.h`` -> ````). Signed integer fields use the - standard trunc-to-floor adjustment. Unsigned integers don't need - adjusting because both operands are non-negative. - """ - if field_dtype is None: - # Caller didn't pass dtype info. The C ``/`` semantics match Python - # ``//`` for non-negative integer operands only. - return f"({lhs}) / ({rhs})" - kind = field_dtype.kind - if kind == "f": - if field_dtype.itemsize == 4: - return f"floorf((float)({lhs}) / (float)({rhs}))" - return f"floor((double)({lhs}) / (double)({rhs}))" - if kind in ("u", "b"): - return f"({lhs}) / ({rhs})" - # Signed integer: trunc-toward-zero is one greater than floor when the - # signs of ``a`` and ``b`` differ and the division has a non-zero - # remainder; subtract 1 in that case. - return f"(({lhs}) / ({rhs}) - ((({lhs}) % ({rhs}) != 0) && ((({lhs}) < 0) != (({rhs}) < 0))))" - - -def _c_expr_unary(py_op, operand, field_dtype=None): - """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 is not None: - 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}") - - -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), - or if the op isn't defined on the dtype's field types (e.g. ``min`` on a - complex field). - """ - np_type = dtype.np_type - if not _op_supports_field_dtypes(op_name, np_type): - return None - # 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 - # synthetic ``_gbudt_NNN`` name (for UDTs whose Python name isn't a valid - # C identifier) is recorded; using it keeps op codegen consistent. - pinned_name = dtype.jit_c_name - typedef_info = _udt_c_typedef(pinned_name or dtype.name, dtype.np_type) - 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 expression builder so - # type-sensitive ops (currently floordiv) can emit correct C. - assigns = " ".join( - f"z->{c} = {_c_expr_binary(py_op, 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) - if arity == 2: - assigns = " ".join( - f"z->v[{i}] = {_c_expr_binary(py_op, f'x->v[{i}]', f'y->v[{i}]', base_dtype)} ;" - for i in range(size) - ) - else: - assigns = " ".join( - f"z->v[{i}] = {_c_expr_unary(py_op, f'x->v[{i}]', base_dtype)} ;" - for i in range(size) - ) - return c_name, f"void {c_name} ({params}) {{ {assigns} }}" - - -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 unroll - their elements at codegen time. - - 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 - pinned_name = dtype.jit_c_name - # ``_udt_c_typedef`` returns None when the numpy layout doesn't match a - # C compiler's layout for the same struct. - typedef_info = _udt_c_typedef(pinned_name or dtype.name, np_type) - 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: unroll all elements. - _base, shape = np_type.subdtype - size = reduce(mul, shape) - terms = [f"((x->v[{i}]) {op} (y->v[{i}]))" for i in range(size)] - elif np_type.names is not None: - terms = [f"((x->{c}) {op} (y->{c}))" for _py, c, _d in _iter_record_leaves(np_type)] - else: - return None - body = join.join(terms) if terms else ("1" if is_eq else "0") - params = f"_Bool *z, const {type_name} *x, const {type_name} *y" - return c_name, f"void {c_name} ({params}) {{ *z = {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() - set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) - set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) - - -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 JIT isn't available or - the dtype can't be expressed in C; callers cache the returned strings - on the op for introspection (see ``TypedUserBinaryOp.jit_c_source``). - """ - if not _has_jit_set: - return None - result = _make_jit_c_comparison_definition(op_name, dtype, is_eq=is_eq) - if result is None: - return None - _set_jit_c_strings(gb_obj, *result, set_string_func) - 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 JIT setters are available and 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``). - """ - if not _has_jit_set: - return None - result = _make_jit_c_definition(op_name, py_op, dtype, arity) - if result is None: - return None - _set_jit_c_strings(gb_obj, *result, set_string_func) - return result diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index fc9c92661..c0bdf0434 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -35,7 +35,7 @@ if _has_numba: import numba - from .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _finalize_udt_op, @@ -120,10 +120,11 @@ def _one(x): if _has_numba: - from .udt_utils import BUILTIN_UDT_UNARY_OPS as _BUILTIN_UDT_UNARY_OPS - from .udt_utils import ( + 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, ) @@ -255,6 +256,9 @@ 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: @@ -494,3 +498,23 @@ def __init__( 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 6f2df5535..37da5cc5f 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -68,6 +68,23 @@ 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 kind == "binary" and 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. Indexing (``monoid.plus[X, Y]``) + # still raises. + return get_typed_op(op.binaryop, dtype, dtype2) + return op._compile_udt(dtype, dtype2) # Generic case: try to unify the two dtypes try: return op[ @@ -219,9 +236,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) @@ -274,6 +292,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] diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index ce191f374..bb7339ce8 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -1,5 +1,4 @@ import itertools -from functools import lru_cache import numpy as np @@ -8,9 +7,28 @@ from ..exceptions import EmptyObject, check_status from . import _has_numba, _supports_udfs, automethods, ffi, lib, utils from .base import BaseExpression, BaseType, call -from .dtypes import _view_if_same_layout 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 .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: @@ -133,15 +151,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) __hash__ = None def __ne__(self, other): - return not self.isequal(other) + return not _scalar_eq(self, other) def __bool__(self): """Truthiness check. @@ -198,6 +220,11 @@ def isequal(self, other, *, check_dtype=False): ------- 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 @@ -206,8 +233,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( @@ -221,6 +255,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: @@ -353,12 +389,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] - base, shape = np_type.subdtype - return rv.view(base).reshape(shape) + return _element_value(self.dtype, scalar.gb_obj) return scalar.gb_obj[0] @value.setter @@ -369,19 +400,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 - # Zeros, not empty: when numpy writes only the fields, a - # record's alignment padding is whatever the allocation - # happened to contain. ``_udt_bytes`` zeros the padding in the - # other direction, when numpy copies whole items instead. - if np_type.subdtype is None: - arr = np.zeros(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.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) - arr[:] = _view_if_same_layout(val, arr.dtype) - self.gb_obj[0 : np_type.itemsize] = _udt_bytes(np_type, arr) + self.gb_obj[0 : self.dtype.np_type.itemsize] = _element_bytes(self.dtype, val) else: self.gb_obj[0] = val self._empty = False @@ -390,7 +409,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]) @@ -468,7 +487,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: @@ -668,7 +691,7 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): other = op(other).new() if type(other) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + other, dtype = _literal_operand(self.dtype, other, op) try: other = Scalar.from_value(other, dtype, is_cscalar=False, name="") except TypeError: @@ -682,6 +705,8 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): ) 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}", @@ -739,7 +764,7 @@ def _ewise_mult(self, other, op=binary.times, is_infix=False): other = op(other).new() if type(other) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + other, dtype = _literal_operand(self.dtype, other, op) try: other = Scalar.from_value(other, dtype, is_cscalar=False, name="") except TypeError: @@ -816,8 +841,9 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): right_dtype = self.dtype dtype = right_dtype if right_dtype._is_udt else None if type(other) is not Scalar: + other, literal_dtype = _literal_operand(right_dtype, other, op) try: - other = Scalar.from_value(other, dtype, is_cscalar=False, name="") + other = Scalar.from_value(other, literal_dtype, is_cscalar=False, name="") except TypeError: other = self._expect_type( other, @@ -828,46 +854,17 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): op=op, ) 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") @@ -1140,6 +1137,47 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): # End auto-generated code: Scalar +def _scalar_eq(scalar, other): + """Return ``scalar == other`` as a bool, by ``binary.eq`` (see ``Scalar.__eq__``).""" + 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) + 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="") @@ -1148,76 +1186,41 @@ 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) - - -@lru_cache -def _padding_ranges(np_type): - """The ``(start, stop)`` byte ranges of ``np_type`` that are alignment padding. +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``. - Walks records, nested records, and array dtypes, and returns the gaps - between the field data. + 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. """ - 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) + 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: - 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) + 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/_utils.py b/graphblas/core/ss/_utils.py index 6ab0c68f7..85017804f 100644 --- a/graphblas/core/ss/_utils.py +++ b/graphblas/core/ss/_utils.py @@ -1,4 +1,6 @@ +import hashlib import itertools +import re from ...exceptions import _error_code_lookup from .. import ffi, lib @@ -45,3 +47,33 @@ def _resolve_serialized_dtype(data_obj, data_nbytes, obj_kind): 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 87c213d32..547d06ccf 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 ---------- @@ -108,6 +114,7 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): 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*") with _compiling_c_source(f"gb.binary.{name}"): check_status_carg( @@ -117,8 +124,8 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): ret_type._carg, left_type._carg, right_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), ), "BinaryOp", gb_obj[0], @@ -126,6 +133,7 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): 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/dtypes.py b/graphblas/core/ss/dtypes.py index 12f74fbce..0ccdd8264 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -29,7 +29,11 @@ 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") # 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 @@ -84,7 +88,7 @@ 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 @@ -96,4 +100,8 @@ def register_new(name, jit_c_definition, *, np_type=None): if save_np_type or np_type not in core.dtypes._registry: core.dtypes._registry[np_type] = 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 return rv diff --git a/graphblas/core/ss/indexbinary.py b/graphblas/core/ss/indexbinary.py index 85b3253e0..5aeed9f11 100644 --- a/graphblas/core/ss/indexbinary.py +++ b/graphblas/core/ss/indexbinary.py @@ -6,27 +6,29 @@ from ..operator.indexbinary import ( IndexBinaryOp, TypedBuiltinIndexBinaryOp, - TypedUserIndexBinaryOp, _has_idxbinop, + _TypedIndexBinaryOpMixin, ) +from ._utils import _jit_c_source ffi_new = ffi.new -class TypedJitIndexBinaryOp(TypedOpBase): - __slots__ = "_jit_c_definition" +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=None): + 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 - thunk_type = TypedUserIndexBinaryOp.thunk_type - def __call__(self, theta=None): return TypedBuiltinIndexBinaryOp.__call__(self, theta) @@ -47,7 +49,12 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_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 ---------- @@ -101,18 +108,22 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): if not isinstance(rv, IndexBinaryOp): IndexBinaryOp._remove_nesting(name) if ( - (x_type, theta_type) in rv.types + (x_type, y_type, theta_type) in rv.types or rv._udt_types is not None - and (x_type, theta_type) in rv._udt_types + and (x_type, y_type, theta_type) in rv._udt_types ): raise TypeError( f"IndexBinaryOp gb.indexbinary.{name} already defined for " - f"({x_type}, {theta_type}) input types" + 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( @@ -123,15 +134,16 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): x_type._carg, y_type._carg, theta_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), + 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, dtype2=theta_type + 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 fa394eb4f..a3a462f59 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. @@ -107,6 +113,7 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): 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*") with _compiling_c_source(f"gb.indexunary.{name}"): check_status_carg( @@ -116,8 +123,8 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): ret_type._carg, input_type._carg, thunk_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), ), "IndexUnaryOp", gb_obj[0], @@ -125,6 +132,7 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): 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 @@ -150,6 +158,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 index 8f3aa4a9a..677b979d8 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -387,7 +387,7 @@ def _compiling_c_source(name): 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_utils._maybe_warn_jit_skipped``, which fires + 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 diff --git a/graphblas/core/ss/select.py b/graphblas/core/ss/select.py index 3ba135eee..48d16b740 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`` diff --git a/graphblas/core/ss/unary.py b/graphblas/core/ss/unary.py index e0be45db9..cee7a16a4 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 ---------- @@ -92,6 +98,7 @@ def register_new(name, jit_c_definition, input_type, ret_type): 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*") with _compiling_c_source(f"gb.unary.{name}"): check_status_carg( @@ -100,13 +107,14 @@ def register_new(name, jit_c_definition, input_type, ret_type): NULL, ret_type._carg, input_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), + 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/udt.py b/graphblas/core/udt.py new file mode 100644 index 000000000..94e7b3136 --- /dev/null +++ b/graphblas/core/udt.py @@ -0,0 +1,911 @@ +"""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: the generated + expression operates on the whole sub-array and the wrapper slice-assigns + it 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 _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): + """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. + """ + 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" + 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. + """ + if (cast_op := _store_cast_op(value.dtype, dtype)) is None: + return value + from .base import BaseType + from .scalar import Scalar + + 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 72f7219bd..9f15254ab 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -4,7 +4,7 @@ from ..dtypes import _INDEX, lookup_dtype from . import ffi, lib -from .dtypes import _view_if_same_layout +from .dtypes import _nats_in_leaf_units, _view_if_same_layout _NP2 = np.__version__.startswith("2.") @@ -49,6 +49,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)] @@ -97,7 +121,7 @@ def values_to_numpy_buffer( np_type = _get_subdtype(dtype.np_type) # https://numpy.org/doc/stable/release/2.0.0-notes.html#new-copy-keyword-meaning-for-array-and-asarray-constructors array = np.array( - _view_if_same_layout(array, np_type), + _nats_in_leaf_units(_view_if_same_layout(array, np_type), np_type), np_type, copy=copy or _NP2 and None, order=order, @@ -132,10 +156,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 diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 5d8c8fef9..f4b72fca4 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -9,6 +9,7 @@ from .base import BaseExpression, BaseType, _check_mask, call from .descriptor import lookup as descriptor_lookup from .expr import _ALL_INDICES, AmbiguousAssignOrExtract, IndexerResolver, InfixExprBase, Updater +from .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype from .mask import Mask, StructuralMask, ValueMask from .operator import ( UNKNOWN_OPCLASS, @@ -18,6 +19,7 @@ get_typed_op, op_from_string, ) +from .operator.select import _value_select_in_own_type from .scalar import ( _COMPLETE, _MATERIALIZE, @@ -26,6 +28,13 @@ 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, @@ -351,6 +360,11 @@ def isequal(self, other, *, check_dtype=False, **opts): ------- 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 +373,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: @@ -593,7 +609,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 @@ -927,14 +947,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, @@ -1026,15 +1047,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, @@ -1204,42 +1233,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") @@ -1420,6 +1415,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: @@ -1441,7 +1441,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: @@ -1477,7 +1477,7 @@ 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, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -1593,8 +1593,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: @@ -1607,6 +1609,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: @@ -1843,6 +1846,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) @@ -1874,7 +1886,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: @@ -1885,6 +1897,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]) @@ -1901,6 +1917,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 @@ -1935,7 +1953,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): diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index a9550ff51..5fb3aad63 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -2,6 +2,7 @@ import pickle import string import sys +import warnings import numpy as np import pytest @@ -10,7 +11,7 @@ from graphblas import core, dtypes from graphblas.core import _supports_udfs as supports_udfs from graphblas.core import lib -from graphblas.core.operator.udt_utils import _has_jit_set +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 @@ -368,6 +369,35 @@ class _Empty: 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 or in a record. 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(("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 @@ -440,6 +470,10 @@ def test_udt_with_macro_field_name_falls_back_to_cfunc(): 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") @@ -539,6 +573,46 @@ def test_udt_jit_c_info_pinned_at_first_register(): 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 == "_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("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") + # The rest compiles kernels for both layouts, which needs a C compiler the JIT can use. + from graphblas.core.ss.jit_config import _enable_jit_for_udt + + if not _enable_jit_for_udt(): + pytest.skip("the JIT cannot compile here (no usable C compiler)") + jit_c_control = gb.ss.config["jit_c_control"] + gb.ss.config["jit_c_control"] = "on" + try: + 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: diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 19b65d720..333c35fcc 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -7,7 +7,7 @@ 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 UdfParseError +from graphblas.exceptions import EmptyObject, UdfParseError pytestmark = [ pytest.mark.skipif(not supports_udfs, reason="requires numba"), @@ -83,6 +83,68 @@ def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) 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.""" @@ -326,19 +388,119 @@ def _combine(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert good[arr_udt].return_type is arr_udt -def test_bind_raw_udt_theta_without_dtype_errors(): - """A raw UDT theta with no dtype can't be inferred; the error is clear and actionable.""" - indexbinary.register_new("raw_no_dtype_op", _ibo_return_x, is_udt=True) +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([1, 2, 3], dtype=np.int64)) - with pytest.raises(TypeError, match="Cannot infer a dtype for theta"): - op((1, 2, 3)) + 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. @@ -450,6 +612,27 @@ def test_bad_type(): 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 @@ -479,6 +662,23 @@ def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) 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.""" @@ -508,6 +708,7 @@ def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) bound = typed(0) assert bound.opclass == "BinaryOp" + assert find_opclass(op(0))[1] == "BinaryOp" def test_udt_tuple_return(): @@ -715,6 +916,84 @@ def weighted(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) 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) diff --git a/graphblas/tests/test_infix.py b/graphblas/tests/test_infix.py index a1a59fd69..94d2555ae 100644 --- a/graphblas/tests/test_infix.py +++ b/graphblas/tests/test_infix.py @@ -497,6 +497,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 diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index ea8a29f26..d1405c49f 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -3696,11 +3696,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) @@ -4340,6 +4341,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) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 36152dfe0..d51993c2e 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1,4 +1,5 @@ import itertools +import sys import numpy as np import pytest @@ -42,7 +43,7 @@ UINT32, UINT64, ) -from graphblas.exceptions import DomainMismatch, UdfParseError +from graphblas.exceptions import DimensionMismatch, DomainMismatch, InvalidValue, UdfParseError from .conftest import shouldhave @@ -1263,9 +1264,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) @@ -1274,7 +1278,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 @@ -1351,10 +1359,11 @@ 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) @@ -1426,6 +1435,470 @@ def _ne_thunk(x, i, j, thunk): # pragma: no cover (numba) 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 _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") +@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): @@ -1579,6 +2052,17 @@ def _three(x, y): # pragma: no cover (numba; raises before execution) 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 @@ -1596,6 +2080,24 @@ def _bad_tuple(x, y): # pragma: no cover (numba; raises before execution) 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(): @@ -1636,7 +2138,7 @@ def _record_expected(udt, 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, 0.1), (0, 0.1), (0, 0.1)]), + ("truediv", [(0.1, 0.1), (0.1, 0.1), (0.1, 0.1)]), ], ) @pytest.mark.skipif("not supports_udfs") @@ -1645,7 +2147,16 @@ 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() - assert result.isequal(_record_expected(record_udt, expected_rows)) + 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: + expected_type = record_udt + assert result.dtype == expected_type + assert result.isequal(_record_expected(expected_type, expected_rows)) @pytest.mark.skipif("not supports_udfs") @@ -1760,7 +2271,7 @@ def test_udt_unary_ainv_abs_array(array_udt): 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_utils import _has_jit_set + from graphblas.core.operator.udt_codegen import _has_jit_set if not _has_jit_set: pytest.skip("JIT not available") @@ -1776,13 +2287,14 @@ def test_udt_jit_typedef(): assert "double jy" in defn assert "JitTypeTest" in defn - # Array UDT - arr_dtype = np.dtype((np.float64, (7,))) - arr_udt = dtypes.register_anonymous(arr_dtype, "Vec7") + # 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 [7]" in defn - assert "Vec7" in defn + assert "double v [31]" in defn + assert "Vec31" in defn # 2D array UDT mat_dtype = np.dtype((np.int32, (5, 5))) @@ -1798,7 +2310,7 @@ def test_udt_jit_typedef(): 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_utils import _has_jit_set + from graphblas.core.operator.udt_codegen import _has_jit_set if not _has_jit_set: pytest.skip("JIT not available") @@ -1824,20 +2336,20 @@ def test_udt_jit_op_definitions(): assert "ainv_JitOpTest" in defn assert "jp" in defn - # Array UDT JIT definitions - arr_dtype = np.dtype((np.float64, (5,))) - arr_udt = dtypes.register_anonymous(arr_dtype, "Vec5Jit") + # 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_Vec5Jit" in defn - assert "v[0]" in defn - assert "v[4]" in defn + assert "plus_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_Vec5Jit" in defn + assert "ainv_Vec37Jit" in defn @pytest.mark.skipif("not supports_udfs") @@ -1953,7 +2465,7 @@ def test_udt_array_wrapper_stays_within_element(shape): import numba - from graphblas.core.operator.base import _get_udt_wrapper + from graphblas.core.operator.udf import _get_udt_wrapper udt = dtypes.register_anonymous(np.dtype((np.float64, shape)), f"_ArrWrapPin{len(shape)}D") @@ -2332,7 +2844,7 @@ def test_udt_broadcast_matches_numba_slice_assign(): """ import numba - from graphblas.core.operator.base import _fits_by_broadcast + from graphblas.core.operator.udf import _fits_by_broadcast @numba.njit def _assign(z, src): # pragma: no cover (numba) @@ -2381,16 +2893,16 @@ def test_udt_udf_shape_check_runs_udf_only_for_built_arrays(monkeypatch): 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 base as _base + from graphblas.core.operator import udf as _udf probed = [] - real_probe = _base._run_udf_probe + 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(_base, "_run_udf_probe", spy) + monkeypatch.setattr(_udf, "_run_udf_probe", spy) udt = dtypes.register_anonymous(np.dtype((np.float32, (18,))), "_ProbeWhen18") rec = dtypes.register_anonymous( np.dtype( @@ -2576,7 +3088,7 @@ def test_udt_array_udf_known_extents_never_match_by_rank(monkeypatch): """ import numba - from graphblas.core.operator import base as _base + 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") @@ -2585,16 +3097,16 @@ def test_udt_array_udf_known_extents_never_match_by_rank(monkeypatch): assert ret.dtype == udt21.numba_type.dtype assert ret.ndim == udt21.numba_type.ndim - real_lookup = _base.lookup_dtype + 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(_base, "lookup_dtype", refuse_nested) + monkeypatch.setattr(_udf, "lookup_dtype", refuse_nested) with pytest.raises(UdfParseError, match="matches no input array UDT"): - _base._resolve_udt_return_type(ret, udt21) + _udf._resolve_udt_return_type(ret, udt21) @pytest.mark.skipif("not supports_udfs") @@ -3073,6 +3585,77 @@ def test_udt_eq_ne_scalar_broadcast_nan_propagates(): 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(): @@ -3167,6 +3750,2612 @@ def test_udt_eq_ne_rejects_incompatible_pairs(): 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") + 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) + 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"), + (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), + ]: + assert (s == literal) is equal + if supports_udfs: + assert (v == literal).new()[0].new().value == equal + assert (v != literal).new()[0].new().value != equal + 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="a UDT casts only to another UDT"): + gb.Scalar(FP64) << v.reduce(agg.sum) + 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. + with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): + added(gb.Scalar.from_value(0.5)) + + +@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. + """ + 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() + with pytest.raises(TypeError, match="Monoid inputs must be the same dtype"): + monoid.plus[X, Y] + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_aggregators(): @@ -3536,7 +6725,6 @@ def test_ops_have_ss(): mod.ss -@pytest.mark.skipif("not supports_udfs") def test_compile_codegen_helper(): """The ``_compile_codegen`` helper validates source and surfaces typos clearly. @@ -3549,7 +6737,7 @@ def test_compile_codegen_helper(): """ import linecache - from graphblas.core.operator.udt_utils import _compile_codegen + from graphblas.core.operator.udt_codegen import _compile_codegen fn = _compile_codegen( "def _op(x, y):\n return x + y\n", diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index b8c00ffec..686e4861a 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -279,6 +279,56 @@ def test_udt_scalar_nested_array_roundtrip(): 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 = None if gb.core._has_numba else dtypes.register_anonymous(np.dtype(("m8[s]", 2))) + 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") + 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() @@ -367,6 +417,53 @@ 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 + + @autocompute def test_neg(): for dtype in sorted( @@ -628,6 +725,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): @@ -660,6 +760,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 diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index eeeeac206..9e083a2a1 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -25,9 +25,10 @@ ) from graphblas.core import _supports_udfs as supports_udfs from graphblas.core.operator.indexbinary import _has_idxbinop -from graphblas.core.operator.udt_utils import _has_jit_set +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 @@ -591,6 +592,26 @@ 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. @@ -786,10 +807,15 @@ def test_jit_indexbinary(v): 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} + 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) @@ -799,6 +825,17 @@ def test_jit_indexbinary(v): # 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(): @@ -814,7 +851,7 @@ def test_jit_indexbinary(v): assert "FP32" in add_theta with pytest.raises( TypeError, - match="IndexBinaryOp gb.indexbinary.ss.add_theta already defined for .FP64, FP64. input", + 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 @@ -825,14 +862,94 @@ def test_jit_indexbinary(v): 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) + (double)(*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", "FP32") in add_theta - assert ("FP32", "FP64") not in 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", "v[0] = (x->v[0]) * (y->v[0])" in src_arr - assert "z->v[10] = (x->v[10]) * (y->v[10])" in src_arr + 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 @@ -1105,41 +1221,181 @@ def test_floordiv_udt_jit_matches_python_semantics(): @pytest.mark.skipif("not supports_udfs") -def test_min_max_udt_jit_propagates_nan(): - """``binary.min``/``max`` on a float UDT must propagate NaN like Python/numba. - - Regression: the JIT codegen used to emit ``(a < b ? a : b)``, which - silently swallows NaN to the right-hand side. Python ``min(a, b)`` (and - numba's ``min``) returns ``a`` when neither comparison is true (NaN - involved), so ``min(NaN, 1.0) == NaN`` and ``min(1.0, NaN) == 1.0``. - The fix swaps the ternary to ``(b < a ? b : a)``. +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 _IS_SSGB7: - pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + 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.float64)]), "_NanJitMM" + 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) - u[i] = (1.0 + i, nan if i % 2 else 3.0 + i) + 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 np.isnan(w[0].new().value[0]) # min(NaN, 1.0) -> NaN - assert w[1].new().value[1] == 3.0 # min(3.0, NaN) -> 3.0 (NaN swallowed) + 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 np.isnan(w[0].new().value[0]) # max(NaN, 1.0) -> NaN - assert w[1].new().value[1] == 3.0 # max(3.0, NaN) -> 3.0 (NaN swallowed) + 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") @@ -1475,6 +1731,44 @@ def test_packed_nested_record_layout_skips_jit(): 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(): diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 52c382742..7d5cab524 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -3,6 +3,7 @@ import pickle import sys import types +import warnings import weakref import numpy as np @@ -1493,6 +1494,292 @@ 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) + # 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 @@ -2097,11 +2384,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) @@ -2641,6 +2929,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) diff --git a/pyproject.toml b/pyproject.toml index 9bd988a10..9c9701dde 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -388,7 +388,7 @@ 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_utils.py" = ["S102"] # exec is used for UDT op codegen +"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 From 9b18bba3c7bab0189a5b2b5d92d8e83f9ecf49c5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 9 Nov 2026 09:45:49 -0600 Subject: [PATCH 25/52] Add Python 3.15; routine maintenance; zizmor in pre-commit and CI (#636) * Pin the CI interpreter and its ABI in the update-env step conda-forge ships two ABIs for CPython 3.13+, and the "Update env" step did not restate ``python``, so the solver was free to replace the interpreter setup-miniconda installed. Three of the last sixteen python 3.14 draws on ubuntu-24.04-arm did exactly that: the log reads ``python 3.14.7-h6bfacdd_100_cp314 --> 3.14.7-hefd156b_0_cp314t`` and pytest then reports "Python: 3.14.7 free-threading build". Two of the three passed only because they also drew pandas 3.0. The third drew pandas 2.3, whose extension raises "The global interpreter lock (GIL) has been enabled to load module 'pandas._libs.pandas_parser'", and ``filterwarnings = error`` turns that into a collection error. It took test_formatting.py and test_recorder.py down with it, the second only because the first had left pandas half-imported. The window opened on 2026-08-19, when conda-forge uploaded python 3.14.7 for linux-aarch64 in both ABIs. Re-solving the failing job's own spec set per platform: linux-aarch64 and osx-arm64 both pick ``cp314t`` today, so macos-latest is on the same fuse and has simply not drawn 3.14 with an unlucky pandas yet. linux-64, osx-64 and win-64 pick the GIL build either way, which is why this has been an arm-only symptom. ``python-gil`` pins ``python_abi``; restating ``python=`` pins the version, since ``python-gil`` alone constrains only the ABI and its own version is left open. setup-miniconda's ``python-version`` is equally ABI-free, so this step is where both get settled. Re-solved with the pair added: every platform lands on a ``cp314`` build, and no draw becomes unsatisfiable or downgrades numpy, numba, pandas or psg, including python 3.11 with numpy 1.24 and scipy 1.9. python-gil has builds for 3.9 through 3.14, so one spec covers every version in the pool. The pypy guard matches the neighbouring lines: ``pypy python-gil`` is unsatisfiable, so an unguarded spec would hard-fail a pypy draw if one is ever added back. 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. An entry like that would be worth having here too, but it is a coverage decision rather than a fix, so it is not in this commit. * Bump the GitHub Actions and pin them to commit SHAs Every action moves to its newest release: checkout 7.0.1, setup-python 7.0.0, setup-miniconda 4.1.0, codecov-action 7.1.1, upload-artifact 7.0.2, gh-action-pypi-publish 1.14.2, and action-tmate 3.24 (pre-commit/action stays at 3.0.1, its newest). This supersedes the open Dependabot PRs #572, #573, #575, #577 and #578. Each is pinned to the commit of its release, with the version in a comment, since a tag can be moved to other code. Dependabot updates the pin and the comment together. The SHAs are the same ones python-suitesparse-graphblas and afar pin. Dependabot now waits seven days after a release before proposing it and groups all action bumps into one PR a week, as those two repositories do. None of the major bumps changes an input this repository uses. Most are the move to Node.js 24; setup-python 7 drops the ``pip-install`` input, checkout 7 refuses fork checkouts under pull_request_target and workflow_run, and upload-artifact 7 adds unzipped single-file uploads. One input changes: * setup-miniconda deprecated ``auto-activate-base`` in favor of ``auto-activate`` (already in 3.3), and every run logged a warning for it. * Update the pre-commit hooks pre-commit autoupdate: isort 9.0.2, pyupgrade 3.22.0, black 26.10.1, ruff 0.17.1, flake8 7.4.1, prettier 3.9.10 and check-jsonschema 0.38.2, with flake8's hand-pinned plugins moved to flake8-bugbear 26.9.30 and flake8-simplify 0.31.1 (and scripts/check_versions.sh with them). This supersedes pre-commit.ci's #579, which predates black 26.10.1, ruff 0.17 and prettier 3.9.10. pyroma stays on 5.0.1: autoupdate offers 5.1b1, and 5.1 is still in beta on PyPI (5.1b2), which is what the comment above it waits for. None of the new versions changes or flags any file, here or at the top of the current stack. * Add numba 0.68, awkward 2.14, networkx 3.7 and psg 10.5.1.1 to the CI pools Each is the newest release on conda-forge, and the audits named all four as newer than any pin. * numba 0.68 caps numpy at <2.6, the same as 0.67 (conda-forge's ``depends``). * awkward 2.14 needs awkward-cpp 57 and Python >=3.11, which every Python in the pool has. * networkx 3.7 requires Python >=3.12, so py3.11's pool stops at 3.6. It constrains numpy, scipy and pandas exactly as 3.5 and 3.6 do, so the existing ``>= (3, 5)`` rule covers it; audit_conda_versions.py now checks 3.7's constraints too. * psg 10.5.1.1 is the newest build of SuiteSparse:GraphBLAS 10.5.1, with numpy 2 builds for every Python we test. numpy (2.5), scipy (1.18), pandas (3.0), sparse (0.19) and pyyaml (6.0) have no newer series. scripts/check_versions.sh moves to the new minimums. The rest of the files with dependency versions need no change: pyproject.toml holds floors, environment.yml and dev-requirements.txt are unpinned, and binder/environment.yml and docs/env.yml pin Pythons that are still supported. `ci_pick_versions.py --validate`, audit_pypi_versions.py --samples 200 and audit_conda_versions.py --samples 100 all pass. * Support Python 3.15 conda-forge has Python 3.15.0 final (cp315 and cp315t), and every dependency CI installs has py315 builds: numba 0.68, numpy 2.5.2+, scipy 1.18.1, pandas 3.0.6, awkward 2.13+, psg 10.5.1.1, matplotlib 3.11.2 and coverage 7.16 (compiled, so the tracer works). PyPI has cp315 wheels of all of them except pyyaml, which pip builds from its sdist. Numba supports 3.15 from the start, so this is a single step rather than the usual two: the numba guards in pyproject.toml move to `python_version<'3.16'`. * pyproject.toml: the 3.15 classifier, the numba guards (all five, sparse's included), and py315 in black's target-version (black 26.10 knows it). * test_and_build.yml and imports.yml draw 3.15. The interpreter pin from the update-env step makes it the GIL build (cp315). * ci_pick_versions.py gets a 3.15 row in every per-Python table, and three rules for what has no py315 build: scipy <1.18, pandas <3.0, and awkward <2.13 (awkward is noarch, but each release pins one awkward-cpp, and 56 is the first with py315 builds). numba starts at 0.68 and psg at 10.5.1.1 for every source: its sdist is the first from after psg added 3.15 support. The validator learned the same three rules and now names the floor each failure needs; disabling any one of the new rules makes `--validate` fail. Checked: `--validate` passes (200000 combos), as do audit_pypi_versions.py and audit_conda_versions.py. A dry-run solve of the update-env install, built the way the workflow builds it, succeeds for 60 draws: 3.15 on all five CI platforms with every source and slowtask (numba skipped too), and 3.11-3.14 on four platforms. Every one lands on a GIL interpreter. The full test suite with --runslow passes on 3.15 with the newest pins (numpy 2.5.4, numba 0.68.0, psg 10.5.1.1, scipy 1.18.1, pandas 3.0.6, networkx 3.7, awkward 2.14.0, sparse 0.19.2). * Run zizmor in pre-commit and CI, and fix what it reports zizmor audits GitHub Actions workflows for security problems. It runs as a pre-commit hook (and so on pre-commit.ci, offline), and security.yml runs it again with a token for its online audits: advisories against the actions we use, pins to commits outside an action's repository, and version comments that no longer match their pin. security.yml also runs pysentry, which pre-commit.ci skips for lack of network. Both run on PRs that touch pyproject.toml, .pre-commit-config.yaml or .github/, and weekly for advisories published since. The setup follows afar's. pysentry now runs when pyproject.toml or .pre-commit-config.yaml changes rather than at every commit, where it fails offline, and with ``--direct-only`` it audits the floors in pyproject.toml instead of the newest versions uv resolves, which would hide a floor that allows a vulnerable release. It finds nothing either way today. What zizmor reported, and what changed: * excessive-permissions: test_and_build.yml, imports.yml, debug.yml and publish_pypi.yml had no ``permissions:``, so their jobs got this repository's default token, which is read-write on everything. They now ask for ``contents: read``, which is all checkout needs; uploads to Codecov and of artifacts do not use the token. * template-injection: the "Update env" and "Build extension module" steps pasted step outputs into their scripts with ``${{ }}``. The values come from fixed lists, so nothing could be injected, but they now reach the scripts as environment variables. That includes the ``format('python={0} python-gil', ...)`` spec, which zizmor does not flag but which is the same pattern. * use-trusted-publishing: suppressed inline on the token, with a TODO. A ``pypi`` environment with required reviewers already exists, but moving to trusted publishing needs a publisher registered on PyPI, which cannot be checked from here. * unpinned-uses and dependabot-cooldown were fixed by the previous commits. .github/zizmor.yml ignored all three of the first findings for whole files; nothing is left for it to ignore, so it is gone. Offline and online, zizmor now reports no findings (one ignored, the inline suppression). The default persona is the one afar uses; the pedantic one would also flag the matrix values in run steps, which are not attacker-controlled. * Re-date the numba#5395 notes: bool pointers still fail on Numba 0.68 The two MAINT notes on the BOOL-to-INT8 routing asked for a re-test. A cfunc that reads ``CPointer(boolean)`` still fails with "cond is not i1: i8", and one that writes it with "cannot store i1 to i8*", on Numba 0.66 and on 0.68, the release that adds Python 3.15. The routing stays; only the notes' dates and versions change. * Publish to PyPI with Trusted Publishing The publish step used a PyPI API token stored as a secret. A trusted publisher is now registered on PyPI for this repository, so the upload moves to its own job that exchanges a short-lived OIDC token instead: only that job gets `id-token: write`, and it runs in the `pypi` environment, which the PyPI registration names and which asks a maintainer to approve the upload. The build job is unchanged and hands the files over as the `releases` artifact it already uploads. This removes zizmor's last inline suppression. * Leave py315 out of black's target versions for now CI's bizarro job installs black from conda-forge, and on Python 3.15 the newest black it can install there is 26.1, because the pytokens >=0.4 that black 26.5 needs has no 3.15 build yet. Black 26.1 does not know "py315" and exits with an error, so the auto-generated code check failed on every 3.15 draw. Any other black older than 26.5 refused to run on this repository for the same reason. Nothing here uses syntax new in 3.15, so the formatting is the same; a MAINT note says when to add it back. * Install the drawn psg version on wheel and source draws The version picker draws a python-suitesparse-graphblas version for every source type, but ``psgver`` was a shell variable of the "Update env" step. "Build extension module", which pip-installs psg for wheel and source draws, runs in a new shell, where ``suitesparse-graphblas${psgver}`` expanded to plain ``suitesparse-graphblas``. So every wheel and source draw installed the newest psg on PyPI, whatever the "versions:" line in its log said; only conda-forge draws, which install psg in "Update env" itself, got the version they drew. It has been this way since wheel and source draws began in #385. It turned up in the CI logs read for why test_ssjit.py's JIT tests skip in CI. In all four wheel and source jobs there that drew a pin (psg 8.0.2.1, 9.4.5.0, 10.1.1.0 and 10.3.1.0), pip installed 10.5.1.1. "Update env" now passes psgver on as a step output, and "Build extension module" reads it through ``env:`` like the other step outputs (zizmor's template-injection audit). An output rather than ``$GITHUB_ENV``, so no other step sees it. Installing the drawn version brings two kinds of draw to life for the first time, and neither could have worked: * source: psg's sdist compiles against conda-forge's GraphBLAS, which was always the newest one (``graphblas>=7.4``). psg 9.4.5.0's sdist does not compile against GraphBLAS 10.5.1 (undeclared GxB_*_GPU_ID constants, changed GxB_Iterator and *_get signatures); against 9.4.5 it builds and imports. A source draw now asks for the GraphBLAS that conda-forge's build of the same psg depends on: python-suitesparse-graphblas X.Y.Z.W requires graphblas >=X.Y.Z,=10.5.2,<10.6, which conda-forge has on all five, so a "source" draw's graphblas pin solves too. The builds need numpy >=1.24 (py3.11) or >=1.25 (py3.12 on), which the numpy 2.x pool this list feeds always meets. No Python floor moves: 10.5.2.0 builds for every Python 3.11 to 3.15 on every source, so PSG_MIN_PYTHON stays as it is. * Give a pinned psg the GraphBLAS it was built for on conda-forge and source draws python-suitesparse-graphblas X.Y.Z.W wraps SuiteSparse:GraphBLAS X.Y.Z, and psg's own test_version_tracks_libgraphblas asserts that the two match. conda-forge's psg builds accept any graphblas >=X.Y.Z,=X.Y.Z,=X.Y.Z on conda-forge, and graphblas X.Y.Z is there for linux-64, linux-aarch64, osx-64, osx-arm64 and win-64. Unpinned psg, upstream and wheel draws are unchanged. Checked with dry-run solves of CI's spec list, 17 of 17 landing on graphblas X.Y.Z: conda-forge draws of psg 10.5.1.1 on py3.15 on all five platforms, of an old psg on py3.11 on each (7.4.0 on linux-64 and osx-arm64, 7.4.3.2 on linux-aarch64, 8.0.2.1 on osx-64, 7.4.1 on win-64), of 8.2.1.0, 9.3.1.0, 9.4.5.0 and 10.5.2.0, and source draws of 9.4.5.0, 10.3.1.0 and 10.5.1.1. Without the pin, a py3.13 osx-arm64 conda-forge draw of 10.5.1.1 solves to graphblas 10.5.2. * Bump pypa/gh-action-pypi-publish to v1.14.3 for GHSA-5xjx-3w67-w9x9 v1.14.2 and earlier trust a certificate authority that Sigstore has revoked, which Trusted Publishing and digital attestations rely on. zizmor's known-vulnerable-actions audit fails on the old pin. * Update the pre-commit hooks ruff 0.17.3, flake8 7.5.0 (with flake8-bugbear 26.10.27), zizmor 1.32.0, validate-pyproject 0.27 and check-jsonschema 0.39.0. The isort and pyroma updates pre-commit autoupdate offers are betas (10.0.0b2, 5.1b1), so they stay where they are. pre-commit run --all-files passes. --- .github/dependabot.yml | 9 ++ .github/workflows/audit_versions.yml | 8 +- .github/workflows/debug.yml | 7 +- .github/workflows/imports.yml | 16 +++- .github/workflows/lint.yml | 6 +- .github/workflows/publish_pypi.yml | 31 +++++-- .github/workflows/security.yml | 58 ++++++++++++ .github/workflows/test_and_build.yml | 131 +++++++++++++++++++++++---- .github/zizmor.yml | 16 ---- .pre-commit-config.yaml | 49 ++++++---- graphblas/core/operator/udf.py | 4 +- pyproject.toml | 15 ++- scripts/audit_conda_versions.py | 3 + scripts/check_versions.sh | 12 +-- scripts/ci_pick_versions.py | 82 ++++++++++++----- 15 files changed, 336 insertions(+), 111 deletions(-) create mode 100644 .github/workflows/security.yml delete mode 100644 .github/zizmor.yml 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 index 646e96f59..162713824 100644 --- a/.github/workflows/audit_versions.yml +++ b/.github/workflows/audit_versions.yml @@ -43,10 +43,10 @@ jobs: name: Audit against PyPI runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false - - uses: actions/setup-python@v6 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: "3.12" - run: python -m pip install --upgrade pip packaging @@ -64,10 +64,10 @@ jobs: run: shell: bash -el {0} steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false - - uses: conda-incubator/setup-miniconda@v3 + - uses: conda-incubator/setup-miniconda@be893c923ea9cf1cf7cd510fbdde27c7e18cbdcb # v4.1.0 with: auto-update-conda: true python-version: "3.12" diff --git a/.github/workflows/debug.yml b/.github/workflows/debug.yml index 45401c704..bb8503570 100644 --- a/.github/workflows/debug.yml +++ b/.github/workflows/debug.yml @@ -9,6 +9,9 @@ on: required: false default: false +permissions: + contents: read + jobs: build: runs-on: ubuntu-latest @@ -26,7 +29,7 @@ jobs: # - "conda-forge" steps: - name: Checkout - uses: actions/checkout@v6 + 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 7f1596c0a..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 @@ -25,7 +33,7 @@ jobs: - name: RNG for Python version id: pyver run: | - choices=(3.11 3.12 3.13 3.14) + 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" @@ -35,14 +43,14 @@ jobs: # runs-on: ${{ matrix.os }} # strategy: # matrix: - # python-version: ["3.11", "3.12", "3.13", "3.14"] + # python-version: ["3.11", "3.12", "3.13", "3.14", "3.15"] # os: ["ubuntu-latest", "macos-latest", "windows-latest"] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false - - uses: actions/setup-python@v6 + - 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 8bceeeadd..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@v6 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false - - uses: actions/setup-python@v6 + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: "3.11" - - uses: pre-commit/action@v3.0.1 + - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 diff --git a/.github/workflows/publish_pypi.yml b/.github/workflows/publish_pypi.yml index ad9284f66..fd2f3af10 100644 --- a/.github/workflows/publish_pypi.yml +++ b/.github/workflows/publish_pypi.yml @@ -5,6 +5,9 @@ on: tags: - "20*" +permissions: + contents: read + jobs: build_and_deploy: runs-on: ubuntu-latest @@ -13,12 +16,12 @@ jobs: shell: bash -l {0} steps: - name: Checkout - uses: actions/checkout@v6 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false - name: Set up Python - uses: actions/setup-python@v6 + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: "3.11" - name: Install build dependencies @@ -27,16 +30,28 @@ jobs: 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.13.0 - if: startsWith(github.ref, 'refs/tags/') && github.repository == 'python-graphblas/python-graphblas' + + 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 8dec76f13..dadb23e36 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: @@ -100,12 +108,14 @@ jobs: # - 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@v6 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 persist-credentials: false @@ -117,7 +127,7 @@ jobs: # 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) + 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" @@ -130,7 +140,7 @@ jobs: 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 @@ -139,17 +149,25 @@ 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: | # 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 "v='${{ steps.pyver.outputs.selected }}'.split('.');print(f'{v[0]}.{v[1]}')") + 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 ${{ steps.sourcetype.outputs.selected }}) ; then + 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 @@ -158,9 +176,44 @@ jobs: # 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}} == "conda-forge" ]] ; then + if [[ $SOURCETYPE == "conda-forge" ]] ; then psg=python-suitesparse-graphblas${psgver} 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" + + # 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,> "$GITHUB_ENV" + fi # Decide whether to install numba and sparse. # Random skip: pypy always, 20% on Windows, 40% on other OSes (unless notebooks). @@ -203,6 +256,22 @@ jobs: echo "versions: np${npver} sp${spver} pd${pdver} ak${akver} nx${nxver} numba${numbaver} yaml${yamlver} sparse${sparsever} psg${psgver}" set -x # echo on + # 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 $(command -v mamba || command -v conda) install -c nodefaults \ packaging pip pytest coverage pytest-randomly cffi donfig tomli c-compiler make \ pyyaml${yamlver} ${sparse} pandas${pdver} scipy${spver} numpy${npver} ${awkward} \ @@ -210,29 +279,32 @@ jobs: ${{ matrix.slowtask == 'pytest_bizarro' && 'black' || '' }} \ ${{ matrix.slowtask == 'notebooks' && '"matplotlib${mplver}" nbconvert jupyter "ipython>=7" drawsvg' || '' }} \ ${{ steps.sourcetype.outputs.selected == 'upstream' && 'cython' || '' }} \ - ${{ steps.sourcetype.outputs.selected != 'wheel' && '"graphblas>=7.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: | # ``python -m pip`` ensures the conda env's interpreter drives the # install, not a system ``pip`` on PATH. - if [[ ${{ steps.sourcetype.outputs.selected }} == "wheel" ]]; then + if [[ $SOURCETYPE == "wheel" ]]; then # Add --pre if installing a pre-release - python -m 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' || '' }} \ # python -m 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 + elif [[ $SOURCETYPE == "source" ]]; then # Add --pre if installing a pre-release - python -m 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' || '' }} \ # 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 [[ ${{ steps.sourcetype.outputs.selected }} == "upstream" ]]; then + elif [[ $SOURCETYPE == "upstream" ]]; then python -m pip install --no-deps git+https://github.com/GraphBLAS/python-suitesparse-graphblas.git@main#egg=suitesparse-graphblas fi python -m pip install --no-deps -e . @@ -245,6 +317,9 @@ jobs: - 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' || '' }} @@ -277,6 +352,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 \ @@ -317,6 +393,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 @@ -340,11 +417,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: | @@ -357,7 +447,7 @@ 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 @@ -365,6 +455,7 @@ jobs: 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 def06f9d5..8515f7f41 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -11,7 +11,9 @@ ci: autoupdate_commit_msg: "chore: update pre-commit hooks" autofix_commit_msg: "style: pre-commit fixes" # pysentry downloads advisory databases at run time, and pre-commit.ci - # runs hooks without network access, so it can only run locally. + # 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: @@ -41,7 +43,7 @@ repos: - id: name-tests-test args: ["--pytest-test-first"] - repo: https://github.com/abravalheri/validate-pyproject - rev: "0.26" + rev: v0.27 hooks: - id: validate-pyproject name: Validate pyproject.toml @@ -54,12 +56,12 @@ repos: # 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: 9.0.1 + 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.21.2 + rev: v3.22.0 hooks: - id: pyupgrade args: [--py311-plus] @@ -69,27 +71,27 @@ repos: - id: auto-walrus args: [--line-length, "100"] - repo: https://github.com/psf/black-pre-commit-mirror - rev: 26.5.1 + rev: 26.10.1 hooks: - id: black - id: black-jupyter - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.16.8 + rev: v0.17.3 hooks: - 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.3.0 + rev: 7.5.0 hooks: - id: flake8 args: ["--config=.flake8"] additional_dependencies: &flake8_dependencies # These versions need updated manually - - flake8==7.3.0 - - flake8-bugbear==26.9.9 - - flake8-simplify==0.30.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.3 hooks: @@ -99,7 +101,7 @@ repos: - tomli; python_version<'3.11' files: ^(graphblas|docs)/ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.16.8 + rev: v0.17.3 hooks: - id: ruff-check # - id: ruff-format # we prefer black; uncomment to compare with ruff's formatter @@ -122,7 +124,7 @@ repos: hooks: - id: shellcheck - repo: https://github.com/rbubley/mirrors-prettier - rev: v3.9.7 + rev: v3.9.10 hooks: - id: prettier args: [--prose-wrap=preserve] @@ -146,8 +148,12 @@ repos: 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.38.0 + rev: 0.39.0 hooks: - id: check-dependabot - id: check-github-workflows @@ -156,15 +162,22 @@ repos: rev: v1.38.0 hooks: - id: yamllint - # - repo: https://github.com/woodruffw/zizmor-pre-commit - # rev: v1.14.2 - # hooks: - # - id: zizmor - repo: https://github.com/pysentry/pysentry-pre-commit rev: v0.5.0 hooks: - id: pysentry - args: ["--format=human", "--fail-on=low", "--sources=pypa,osv,pypi"] + # 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 diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index ab8144fc2..30b3e7ff4 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -39,7 +39,7 @@ def _bool_to_int8(dtype): numba/numba#5395. Routing BOOL through INT8 sidesteps that; GraphBLAS coerces back at the cfunc boundary. - MAINT 2026-05-24: still hits on Numba 0.65. Re-test periodically and + 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 @@ -306,7 +306,7 @@ def _input_operand(dtype, var): if dtype == BOOL: # Numba can't compile bool ptrs (numba/numba#5395); expose them # as int8 and cast on deref. - # MAINT 2026-05-24: still hits on Numba 0.65; re-test periodically. + # 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) diff --git a/pyproject.toml b/pyproject.toml index 9c9701dde..cf0251968 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ classifiers = [ "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", @@ -62,7 +63,7 @@ dependencies = [ # 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", - "numba >=0.57; python_version<'3.15'", # make optional where numba is not supported + "numba >=0.57; python_version<'3.16'", # make optional where numba is not supported ] [project.urls] @@ -83,11 +84,11 @@ suitesparse-udf = [ # udf requires numba repr = ["python-graphblas[pandas]"] io = [ "python-graphblas[networkx,scipy]", - "python-graphblas[numba]; python_version<'3.15'", # make optional where numba is not supported + "python-graphblas[numba]; python_version<'3.16'", # make optional where numba is not supported "awkward >=2.0", # 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.15'", + "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"] @@ -99,11 +100,11 @@ datashade = [ # datashade requires numba 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.15'", # 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.15'", # make optional, b/c datashade needs numba + "python-graphblas[datashade]; python_version<'3.16'", # make optional, b/c datashade needs numba ] [tool.setuptools] @@ -138,6 +139,10 @@ dirty_template = "{tag}+{ccount}.g{sha}.dirty" [tool.black] line-length = 100 +# 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] diff --git a/scripts/audit_conda_versions.py b/scripts/audit_conda_versions.py index 0d2fe4c16..5b837622d 100755 --- a/scripts/audit_conda_versions.py +++ b/scripts/audit_conda_versions.py @@ -55,6 +55,9 @@ # 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"), diff --git a/scripts/check_versions.sh b/scripts/check_versions.sh index f702f2f41..e99cb7ca6 100755 --- a/scripts/check_versions.sh +++ b/scripts/check_versions.sh @@ -10,20 +10,20 @@ # 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.9.9' -conda search 'flake8-simplify[channel=conda-forge]>=0.30.0' +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.6' -conda search 'awkward[channel=conda-forge]>=2.13' +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.67' +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.0' +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 index 8c7d4786c..34f76e9ee 100755 --- a/scripts/ci_pick_versions.py +++ b/scripts/ci_pick_versions.py @@ -24,7 +24,7 @@ # # When updating versions here, also update scripts/check_versions.sh # -# To add a Python version (3.15, say): add it to `choices` in the "RNG for Python +# 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 @@ -32,12 +32,14 @@ # 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. +# 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). @@ -99,16 +101,22 @@ "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.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", ""], + "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)} +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.13": ["3.4", "3.5", "3.6", ""], - "3.14": ["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 = { @@ -145,6 +156,7 @@ "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 @@ -161,6 +173,7 @@ "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 @@ -182,8 +195,10 @@ ], "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": ["7.4.3.2", "8.0.2.1", "8.2.0.1", "8.2.1.0"], + "3.11": ["8.0.2.1", "8.2.0.1", "8.2.1.0"], "3.12": ["8.2.0.1", "8.2.1.0"], }, } @@ -200,6 +215,8 @@ "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. @@ -211,9 +228,14 @@ # 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. -# MAINT: 2026-09-11 checked conda-forge build strings and PyPI wheel tags for psg 9-10 +# +# 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"}, } @@ -320,8 +342,10 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): # --- scipy / Python version availability --- - # scipy <1.14 has no py3.13 builds; scipy <1.16 has no py3.14 builds - if _ver(pyver) >= (3, 14) and v["scipy"] not in ("", "NA") and _ver(v["scipy"]) < (1, 16): + # 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)) @@ -332,8 +356,10 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): # --- pandas constraints --- - # pandas <2.3 has no py3.14 builds - if _ver(pyver) >= (3, 14) and v["pandas"] == "2.2": + # 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 @@ -367,8 +393,12 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): # --- awkward / Python version availability --- - # awkward <2.7 has no py3.13 builds; awkward <2.8 has no py3.14 builds - if _ver(pyver) >= (3, 14) and v["awkward"] not in ("", "NA") and _ver(v["awkward"]) < (2, 8): + # 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", ""]) @@ -388,14 +418,14 @@ def apply_constraints(v, pyver, scipy_pool, numba_pool): # --- networkx / numpy, scipy, pandas floors --- - # conda-forge networkx 3.5 and 3.6 declare `constrains: numpy >=1.25, scipy >=1.11.2, + # 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-08-04 conda-forge networkx 3.6.1 constrains numpy >=1.25, scipy >=1.11.2 + # 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) ): @@ -685,8 +715,10 @@ def validate(v, pyver, source_type=None): errors.append(f"scipy {v['scipy']} requires numpy <{cap}, got {v['numpy'] or 'latest'}") # scipy Python availability - if _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") + 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,))): @@ -694,7 +726,9 @@ def validate(v, pyver, source_type=None): errors.append(f"scipy {v['scipy']} requires Python >={need}, has no py{pyver} build") # pandas Python availability - if _ver(pyver) >= (3, 14) and v["pandas"] == "2.2": + 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`) @@ -711,8 +745,10 @@ def validate(v, pyver, source_type=None): errors.append(f"pandas {v['pandas']} constrains scipy >=1.10.0, got {v['scipy']}") # awkward Python availability - if _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") + 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") From 0f5dc9b33c95de11314bd27b648b4911957eac0e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 9 Nov 2026 09:51:04 -0600 Subject: [PATCH 26/52] Set the SuiteSparse JIT control from config, and JIT-compile lifted UDT ops on mixed operand types (#652) * ss.fix_jit_config: find a conda compiler on Windows (Library/bin, .exe names) The conda repair looked for {baked-in name, cc, clang, gcc} in $CONDA_PREFIX/bin only. Conda on Windows installs compilers in %CONDA_PREFIX%\Library\bin, and their names end in .exe, so the repair found nothing there and left SuiteSparse's build-time compiler, which a user's machine does not have. It now looks in both places, with and without .exe, and takes a Windows path's basename with either slash. Tested with a fake prefix (no Windows machine at hand). * JIT C definitions for lifted ops on a scalar and a UDT, or on two UDT types A lifted binary op (plus, minus, times, truediv, floordiv, min, max) whose operands are a scalar and an array or record UDT, in either order, or two UDTs of one layout with other element types, ran only through its Numba cfunc: the C definition was set for same-type pairs alone, with one C type for x, y and z. So a semiring such as plus_times[FP64, FP64[R]] ran SuiteSparse's generic kernel with a function call per multiply. Such an op now gets a C definition with a type for each of z, x and y (udt_codegen._make_jit_c_mixed_definition), named op_N_x_y with N the length of x's C name, so that no two ops share a name (SuiteSparse finds a compiled kernel in memory by its op's name): a scalar operand is read as (*x) for every element or leaf, and each operand is converted in C to the type the op computes in, as the cfunc converts it (udt._udt_op_types). Ops with no such form keep the cfunc: an array operand that broadcasts, an element not computed in the result's type (integer truediv), a record with an array leaf, and eq/ne on a mixed pair. The whole mxv with plus_times[FP64, FP64[R]] now compiles as one JIT kernel (burble: dot2, jit). test_udt_scalar_broadcast_skips_jit becomes test_udt_scalar_broadcast_jit_routing: mixed arithmetic now JITs, mixed eq/ne still does not. * Test the mixed semiring's JIT kernel by burble's generic, absent and present The burble names a JIT kernel only when it compiles or loads it, so a kernel loaded earlier in the process prints no '(jit: ...)'. A UDT product has no kernel but the JIT one and the generic one, which burble names, so the test checks that 'generic' is absent with the JIT on and present with it off. * Vary the mixed JIT C name test's inputs, so its kernel runs on every SuiteSparse test_udt_mixed_jit_c_name_is_its_own built its vectors from constant values. SuiteSparse before 10.5 makes a vector built from equal values iso, and an ewise_mult of iso vectors gives an iso result computed without the JIT kernel the test is about, so on those versions it checked nothing. On them its to_dense also raised "TypeError: unsupported operand type(s) for +: 'int' and 'tuple'" (wave 2 CI, 17 jobs, e.g. 122633175811 and 122633175893), the iso array-UDT export bug fixed in #591. The values now differ by index. * Set the SuiteSparse JIT control at init, from a new config key Behavior change. Before: python-graphblas left jit_c_control where the SuiteSparse:GraphBLAS build started it ("run" in the conda-forge and PyPI builds, "on" in upstream's own), and UDT work raised "run" to "on" when it needed C compiled, whoever had set it. Turning the JIT on for everything meant setting gb.ss.config["jit_c_control"] after init. After: when the "suitesparse" backend initializes, it sets jit_c_control from the new config key graphblas.config["jit_c_control"]. The key is unset by default, which sets "run": nothing changes for the conda-forge and PyPI builds, and a build that starts at "on" now starts at "run" too, so every build has one documented default. GRAPHBLAS_JIT_C_CONTROL, gb.config.set before init, or a YAML file in ~/.config/graphblas/ set it; the values are "on", "load", "run", "pause" and "off" (YAML's unquoted on/off arrive as bools and are taken as such), and any other value raises ValueError before anything is initialized. UDT work still raises the default "run" to "on" when it needs C compiled, but now keeps a "run" the user set, as it already kept "off", "pause" and "load". A "run" set before init counts: GRAPHBLAS_JIT_C_CONTROL, gb.config.set or a YAML file, one key that donfig merges, so they could not be told apart and need not be; the environment variable is also the form that reaches every worker of a pool or cluster, which is what a "run" is for (GraphBLAS #449, below). A later gb.ss.config["jit_c_control"] = "run" counts when it changes the configured control (a "run" over a configured "on", say), and keeps a configured "run"; with nothing configured it is the default "run" again. The comparison is with the configured control, not the live one, which UDT work may have raised to "on": code that saves the default "run", sets something else and restores it would otherwise turn the default into a chosen "run" by accident. So, at runtime, "load" or "off" keep UDT work from switching to "on", as does GRAPHBLAS_JIT_C_CONTROL=run set before starting. gb.ss.config["jit_c_control"] = None is the default whatever is configured. Under a "run" that was set, UDT ops run their Numba functions (no compiled kernel, typically 2-3x slower for elementwise ops), and types and ops registered from C source (dtypes.ss.register_new and the op modules' ss.register_new) fail with a message that says the control was set to "run" and how to opt in: gb.ss.config["jit_c_control"] = "on", or None for the default. "on" makes SuiteSparse compile kernels it has no precompiled one for: Matrix.reduce_rowwise with lor on BOOL or plus on UINT64 measured 2.7-8x faster than the generic kernel. It stays opt-in because from SuiteSparse 9.4 on a failed compile returns GxB_JIT_ERROR, and processes that share one cache directory race (GraphBLAS #449; 10.1.1 dropped the file locks): on 10.5.2, 127 of 280 concurrent processes compiling the same new kernels into one cache got errors, invalidated or cleared outputs, wrong values or hangs, where 8.2.1, which locked the cache, had none. When "on" is asked for, it is set only where it works (SuiteSparse 7 has no JIT, and nothing is set there); elsewhere init sets "run", with no error or warning: a library built without the JIT (psg's wheels clamp writes to "run"), no compiler on disk after the conda repair, a cache directory it cannot write to (a probe cached there still loads, and the first new kernel would then fail to compile), or a compiler that fails a probe. Without the probe a broken compiler would fail the first user operation that needs a new kernel (a slice, say) with a bare JitError. The probe compiles a trivial type through GxB_Type_new; the kernel is cached, so after the first time per cache directory a probe is a load, about 2 ms against 0.3 s. The type is named after a hash of the compiler name, flags and libraries, so another toolchain sharing the cache gets a probe of its own. A failed probe sets "run", not the "load" SuiteSparse leaves, since "load" still reads kernels from the cache. UDT work makes the same writability check before it raises the default "run" to "on": without it (as before this change), a read-only cache that already held the probe let the raise through, and the first new UDT kernel failed with a bare JitError; now the UDT op runs its Numba function and the control stays "run". _probe_jit now goes through the C API for fix_jit_config and UDT work too, so it registers no dtypes.ss._jit_probe and needs no re-entrancy guard. Init can run in the middle of importing a core module, so it reads and writes these settings through the C API (_CConfig) instead of gb.ss.config, and jit_config imports ffi and lib inside the functions that use them: init imports it, so importing it first would otherwise start init halfway through its own import. Importing gb.ss still changes nothing (#630). * Docs: the SuiteSparse JIT control, how to opt in, and when "on" is unsafe A section in the initialization page of the user guide, about a screen long: what "on" buys (the 2.7-8x reduce numbers), that the default is "run" and that UDT work raises it but leaves a control the user set, how to opt in (GRAPHBLAS_JIT_C_CONTROL, gb.config, gb.ss.config) and the values, that unset is the default and a "run" set later in gb.ss.config is then that default again (at runtime "load" or "off", or GRAPHBLAS_JIT_C_CONTROL=run before starting, keep UDT work from switching to "on"), where "on" falls back to "run" with no error, and the costs of "on". The first is plain: processes running at the same time must not share a cache directory on SuiteSparse 10.1.1 and later, which no longer lock the cache (GraphBLAS #449: JitError, invalid or cleared outputs, wrong values returned with success, hangs). The simple remedy for multiprocessing, dask or pytest-xdist is GRAPHBLAS_JIT_C_CONTROL=run in the workers' environment, which holds for UDT work too; it costs UDT operators their compiled kernels (they run their Numba functions, typically 2-3x slower elementwise) and makes C-source registration fail. A cache per process keeps the kernels. The default will change once SuiteSparse publishes kernels atomically. Then the first-use compile latency, the cache keyed only on the version, and what to do on trouble: set GRAPHBLAS_JIT_C_CONTROL=run, delete the cache, report it. The UDT page says the control starts at the config value, "run" by default, that UDT work and C-source registration raise that default but keep a control the user set (a "run" when set before init), and what a "run" that was set costs there, and links to the section; graphblas.yaml, jit_config's module docstring and gb.ss.config's jit_c_control entry point there, and the entry says what None does. * Exclude the mixed op codegen's array leaf check from coverage A record with an array leaf has no C typedef, so the mixed op's C definition returns before it looks at the leaves. * test_global_config: None sets jit_c_control to the default "run" The test expected every key outside GlobalConfig._defaults to refuse None, but "Set the SuiteSparse JIT control at init, from a new config key" gave jit_c_control a setter in which None is the default "run" whatever graphblas.config has, and not a "run" the user chose (UDT work may still raise it to "on"). Check both, then put the control back as the test found it, and restore the user-set flag, which the round trip earlier in the test can change too. * Skip the unusable-compiler probe check where SuiteSparse compiles with cmake test_set_jit_c_control_at_init points jit_c_compiler_name at a directory and expects the probe to fail. SuiteSparse built with MSVC always compiles through cmake, and cmake's Visual Studio generator chooses the compiler itself, so on Windows the probe succeeds and the control stays 'on'. Every Windows job failed on it. * Say that the JIT stays at "run" on Windows, where SuiteSparse compiles with cmake SuiteSparse built with MSVC, as conda-forge builds it, compiles JIT kernels only through cmake. python-graphblas checks the configured C compiler path before it sets "on", and on conda-forge Windows that is the build machine's cl.exe, which no user has, so a configured "on" becomes "run" there. The user guide's JIT section and the jit_use_cmake entry of gb.ss.config say so. * Keep a UDT op's JIT C name within SuiteSparse's 127 characters SuiteSparse refuses a JIT C name of GxB_MAX_NAME_LEN (128) characters or more, but the definition was set anyway and taken under the empty name. A mixed op on two UDTs whose names add up to about 118 characters then compiled a kernel that calls no function: every output value was left uninitialized, with no error. A single-type op on a UDT named with 122 or more characters (on main too) failed with an empty JitError. Shorten a long mixed-op name to a digest (a length of 0 cannot come from the readable form, so the name stays the pair's own), and set the definition only when SuiteSparse took the name; otherwise the op keeps its cfunc. Found by review (r652f). --- docs/user_guide/init.rst | 71 +++ docs/user_guide/udt.rst | 27 +- graphblas/__init__.py | 32 ++ graphblas/core/operator/binary.py | 29 +- graphblas/core/operator/udt_codegen.py | 127 ++++- graphblas/core/ss/jit_config.py | 309 +++++++++--- graphblas/graphblas.yaml | 6 + graphblas/ss/__init__.py | 10 +- graphblas/ss/_core.py | 29 +- graphblas/tests/test_core.py | 21 + graphblas/tests/test_dtype.py | 13 +- graphblas/tests/test_ss_utils.py | 14 +- graphblas/tests/test_ssjit.py | 641 +++++++++++++++++++++++-- scripts/jit_diagnostics.py | 2 +- 14 files changed, 1174 insertions(+), 157 deletions(-) diff --git a/docs/user_guide/init.rst b/docs/user_guide/init.rst index ffb6a3463..141e1f250 100644 --- a/docs/user_guide/init.rst +++ b/docs/user_guide/init.rst @@ -57,3 +57,74 @@ raised. However, if the arguments do not match, an error will be raised. There is no mechanism to change the initialization arguments after graphblas has been initialized the first time. + +.. _suitesparse_jit: + +The SuiteSparse:GraphBLAS JIT +----------------------------- + +SuiteSparse:GraphBLAS ships precompiled kernels for common operations. For the rest it +can compile a kernel with a C compiler (its JIT), or else run a generic kernel that is +often several times slower (``Matrix.reduce_rowwise`` with ``lor`` on ``BOOL`` or +``plus`` on ``UINT64`` ran 2.7 to 8 times slower on 10 million entries). The JIT control, +``gb.ss.config["jit_c_control"]``, is ``"run"`` by default: run only kernels already +loaded in the process, and compile nothing. UDT operators, and types and operators +registered from C source, raise this default to ``"on"`` when they need a kernel (see +:doc:`udt`), but leave a control you set as it is. To opt in for everything, set the +control to ``"on"`` (compile, load and run): + +.. code-block:: bash + + export GRAPHBLAS_JIT_C_CONTROL=on + +.. code-block:: python + + import graphblas as gb + + gb.config.set(jit_c_control="on") # before graphblas initializes + gb.ss.config["jit_c_control"] = "on" # or later, for the rest of the process + +``graphblas.config["jit_c_control"]`` is read once, when graphblas initializes. It can be +``"on"``, ``"load"`` (load and run cached kernels, never compile), ``"run"``, ``"pause"`` +(run no JIT kernels) or ``"off"`` (run none, and unload them). Set it with the +``GRAPHBLAS_JIT_C_CONTROL`` environment variable, with ``gb.config.set`` before graphblas +initializes, or in a YAML file in ``~/.config/graphblas/``. Left unset, it is the +default ``"run"``, which UDT work raises. With nothing configured, a ``"run"`` written +later to ``gb.ss.config`` is that same default, so to keep UDT work from compiling, set +``"run"`` in one of those three ways, or write ``"load"`` or ``"off"`` later. With +``"on"``, python-graphblas checks the C compiler by compiling a tiny type (once per cache +directory and compiler), and falls back to ``"run"``, with no error, where the JIT cannot +work: SuiteSparse built without the JIT (as in the PyPI wheels), SuiteSparse 7, no working +C compiler (in an activated conda environment, ``conda install c-compiler``), or a cache +directory it cannot write to. On Windows, SuiteSparse built with MSVC (as conda-forge +builds it) compiles JIT kernels only through cmake, and python-graphblas does not turn the +JIT on there yet, so a configured ``"on"`` becomes ``"run"``. + +With the JIT on: + +- **Do not share one cache directory between processes running at the same time** on + SuiteSparse 10.1.1 or later, which no longer locks the cache. They race when they + compile the same new kernel, and calls fail with ``JitError``, leave their outputs + invalid or cleared, return wrong values with no error, or hang + (`GraphBLAS issue #449 `_). + 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. That + 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/udt.rst b/docs/user_guide/udt.rst index 2a0b86555..44af6a1f0 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -442,25 +442,32 @@ 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. 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 +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`` is raised from the SS default ``'run'`` (run cached kernels -only; no compile, no load from disk) to ``'on'`` (compile, load, and run) by -each of these, since each asks SuiteSparse to compile C source: +``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: - a UDT operator that gets C source, such as the first ``binary.plus[udt]``; - ``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 setting other than ``'run'`` is kept, including ``'off'``, ``'pause'``, and -``'load'`` (load and run cached kernels, never compile). SuiteSparse itself -sets ``'load'`` after a compile fails, so keeping it means a failing compile -is tried once, not again at every later request. +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:: diff --git a/graphblas/__init__.py b/graphblas/__init__.py index 86759f570..2de45dda9 100644 --- a/graphblas/__init__.py +++ b/graphblas/__init__.py @@ -137,6 +137,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: @@ -197,6 +200,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/core/operator/binary.py b/graphblas/core/operator/binary.py index e17889d2a..491a3e0ac 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -409,6 +409,7 @@ def _pair_dtype(op, dtype, dtype2): _has_jit_set, _maybe_warn_jit_skipped, compile_udt_binary_wrapper, + set_jit_c_mixed_on_op, set_jit_c_on_op, ) @@ -689,13 +690,13 @@ def _compile_udt(self, dtype, dtype2): op = _finalize_udt_op( self, dtype, dtype2, ret_type, binary_wrapper, wrapper_sig, TypedUserBinaryOp ) - # Set JIT C definition for auto-generated built-in UDT ops. Only - # same-type pairs apply: mixed UDT+scalar ops don't have a - # meaningful single-type C definition. The arithmetic kernel also - # declares ``z`` with the operands' type, so it fits only an op whose - # result is that type too, which ``truediv`` on integers is not. - # ``_has_jit_set`` gates access to ``lib.GrB_BinaryOp_set_String``, - # which is absent on SS < 9. + # 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: @@ -723,6 +724,20 @@ def _compile_udt(self, dtype, dtype2): 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 diff --git a/graphblas/core/operator/udt_codegen.py b/graphblas/core/operator/udt_codegen.py index 81d78f73f..41fba757a 100644 --- a/graphblas/core/operator/udt_codegen.py +++ b/graphblas/core/operator/udt_codegen.py @@ -17,6 +17,7 @@ """ import ast +import hashlib import itertools import linecache from functools import lru_cache, reduce @@ -1419,6 +1420,102 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): 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 = _udt_c_typedef(dtype.jit_c_name or dtype.name, dtype.np_type) + return None if typedef_info is None else typedef_info[0] + + +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]``. @@ -1499,8 +1596,15 @@ def _set_jit_c_strings(gb_obj, c_name, c_defn, set_string_func): from ..ss.jit_config import _enable_jit_for_udt _enable_jit_for_udt() - set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) + # 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): @@ -1516,7 +1620,23 @@ def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq result = _make_jit_c_comparison_definition(op_name, dtype, is_eq=is_eq) if result is None: return None - _set_jit_c_strings(gb_obj, *result, set_string_func) + 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 @@ -1532,5 +1652,6 @@ def set_jit_c_on_op(gb_obj, op_name, py_op, dtype, set_string_func, arity=2): result = _make_jit_c_definition(op_name, py_op, dtype, arity) if result is None: return None - _set_jit_c_strings(gb_obj, *result, set_string_func) + if not _set_jit_c_strings(gb_obj, *result, set_string_func): + return None return result diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py index 677b979d8..62d78a44a 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -18,16 +18,27 @@ ``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 -from .. import lib # noqa: F401 (sets up cffi) +# ``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. @@ -53,7 +64,10 @@ def jit_compiler_is_usable(): Cheap (no compile attempt); use before suggesting ``fix_jit_config()``. """ - cfg = _ss_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() @@ -73,8 +87,8 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): 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, try to JIT-register a trivial UDT to verify - the compiler actually works. A failed compile makes SuiteSparse lower + 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. @@ -94,8 +108,8 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): rv = _repair_jit_compiler(cfg, use_sysconfig=use_sysconfig) if rv is None: return None - # Enabling compilation is what the caller asked for. The import-time - # repair deliberately stops short of this; see ``_auto_fix_jit_at_import``. + # 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 @@ -127,14 +141,19 @@ def _repair_jit_compiler(cfg, *, use_sysconfig=True): def _fix_from_conda(cfg, conda_prefix): - """Conda-aware fix; swap the compiler path with one from $CONDA_PREFIX/bin/.""" + """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(): - cc_basename = pathlib.Path(jit_cc).name - bin_dir = pathlib.Path(conda_prefix) / "bin" - for candidate in [cc_basename, "cc", "clang", "gcc"]: - local_cc = bin_dir / candidate - if local_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: @@ -195,39 +214,158 @@ def _strip_mismatched_arch(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): - """Probe a trivial JIT compile to verify the config works. - - ``jit_c_control`` is left wherever SuiteSparse puts it, and that 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. + """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 dtypes as _dtypes - - # The probe dtype is installed at ``dtypes.ss._jit_probe`` on first - # success; reuse it on a second probe so a repeat call doesn't appear to - # fail. Without the hasattr short-circuit, ``register_new`` would raise - # ``ValueError("name unavailable")`` on the second call. - probe_name = "_jit_probe" - if hasattr(_dtypes.ss, probe_name): - return True - global _probing_jit - _probing_jit = True - try: - _dtypes.ss.register_new(probe_name, "typedef struct { int _probe ; } _jit_probe ;") - except Exception: - # ``register_new`` can raise ``JitError`` (bad path / flags / arch / - # SDK), ``RuntimeError`` (SS<8 has no JIT), or one of its input - # validation ``ValueError``s. The probe's contract is "did this - # work?", so absorb every failure mode here. + 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 - finally: - _probing_jit = False - return True + 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(): @@ -237,12 +375,14 @@ def _auto_fix_jit_at_import(): breaks ``import graphblas.ss``. The body sticks to dict ops and the string rewriting in :func:`_repair_jit_compiler`. - Deliberately leaves ``jit_c_control`` at whatever SuiteSparse set. 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. - 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. + 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: @@ -251,57 +391,61 @@ def _auto_fix_jit_at_import(): _repair_jit_compiler(cfg) -# Tri-state: ``None`` until 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. +# 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 -# True only while ``_probe_jit`` is registering its own UDT. That registration -# goes through ``dtypes.ss.register_new``, which asks to enable the JIT, which -# would probe again; the second probe would register the same name a second -# time and the first one would then die on cffi's "multiple declarations". -_probing_jit = False +# ``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. - SuiteSparse defaults ``jit_c_control`` to ``'run'``, 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. + ``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. An explicit ``'off'``, ``'pause'``, - or ``'load'`` is honored, 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. + 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 - if _probing_jit: - # Re-entered from the probe's own registration. The probe is the thing - # deciding this answer, so say yes and let it finish rather than - # starting a second one inside it. - return True cfg = _ss_config() control = cfg.get("jit_c_control") - if control in (None, "off", "pause", "load"): - # ``None`` is a library with no JIT at all. The others are honored as - # set and not cached, so a saved ``'run'`` restored afterwards is - # still raised. + 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) - if _jit_enabled_for_udt := jit_compiler_is_usable(): + # 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 @@ -366,6 +510,15 @@ def _compiling_c_source(name): "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: " 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/ss/__init__.py b/graphblas/ss/__init__.py index 33a342373..87bc3e2cd 100644 --- a/graphblas/ss/__init__.py +++ b/graphblas/ss/__init__.py @@ -14,11 +14,11 @@ # 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 SuiteSparse set it, 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. + # 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 f036b3bb7..40c7695f0 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -143,9 +143,17 @@ class GlobalConfig(BaseConfig): "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. @@ -254,6 +262,25 @@ 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 = { diff --git a/graphblas/tests/test_core.py b/graphblas/tests/test_core.py index a15da6f54..bfad8279d 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") diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 5fb3aad63..8ca033625 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -498,8 +498,8 @@ def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): 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 SuiteSparse set, and - # blaming it would send the reader to a setting that cannot fix anything. + # 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 @@ -521,8 +521,8 @@ def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): 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, SuiteSparse leaves ``jit_c_control`` - at ``'run'`` and ``_enable_jit_for_udt`` caches its ``False``. Codegen + 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``. @@ -592,13 +592,14 @@ def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): with pytest.raises(ValueError, match="unavailable"): dtypes.ss.register_new("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") # 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 - if not _enable_jit_for_udt(): - pytest.skip("the JIT cannot compile here (no usable C compiler)") 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) diff --git a/graphblas/tests/test_ss_utils.py b/graphblas/tests/test_ss_utils.py index b1469b7a0..4e553f51b 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -246,9 +246,14 @@ def test_openmp_enabled(): ) -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 @@ -271,6 +276,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 diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 9e083a2a1..3721eee1f 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -47,12 +47,22 @@ # 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) -# ``_auto_fix_jit_at_import`` repairs the compiler path but leaves -# ``jit_c_control`` alone, so compilation is enabled on demand by -# ``_enable_jit_for_udt``. Ask it once here, before the autouse ``_setup_jit`` -# fixture mutates the live config: it answers whether this environment can -# JIT-compile at all, and caches that answer for the rest of the process. -_JIT_WORKS_AT_IMPORT = not _IS_SSGB7 and backend == "suitesparse" and _enable_jit_for_udt() + +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) @@ -114,19 +124,28 @@ def _jit_mode(mode): gb.ss.config["jit_c_control"] = prev -def _run_in_fresh_process(script): +_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. + 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=dict(os.environ, PYTHONPATH=str(repo_root)), + env=env, ) report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" assert result.returncode == 0, report @@ -159,17 +178,115 @@ def control(): @pytest.mark.skipif("_IS_SSGB7") -def test_ss_import_leaves_jit_c_control_alone(): +@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. + 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) + 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 @@ -187,11 +304,11 @@ def test_ss_import_leaves_jit_c_control_alone(): def test_fix_jit_config_first_call_agrees_with_the_second(): """``fix_jit_config()`` must report on the environment, not on itself. - Its probe registers a UDT, and registering a UDT asks to enable the JIT, - which would start a second probe inside the first. The inner probe would - register the probe type and the outer one would then die redeclaring it, - so a working environment would answer ``False`` the first time and - ``True`` every time after. + 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 @@ -199,8 +316,8 @@ def test_fix_jit_config_first_call_agrees_with_the_second(): # 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 ``'run'``, so one registrar raising the control -# cannot cover for another that does not. +# 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 @@ -238,7 +355,7 @@ def test_fix_jit_config_first_call_agrees_with_the_second(): "FP64", "FP64", "FP64", "FP64", ) for kind, register in registrars.items(): - gb.ss.config["jit_c_control"] = "run" + gb.ss.config["jit_c_control"] = None try: register() except Exception as exc: @@ -250,7 +367,7 @@ def test_fix_jit_config_first_call_agrees_with_the_second(): @pytest.mark.skipif("_IS_SSGB7") def test_ss_register_new_enables_the_jit(): - """Registering an op from C source raises SuiteSparse's default ``'run'``. + """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 @@ -298,9 +415,9 @@ def test_ss_register_new_says_why_the_jit_cannot_build_it(monkeypatch): 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 SuiteSparse's default in place. + # 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"] = "run" + 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. @@ -314,13 +431,141 @@ def failed_probe(cfg): 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"] = "run" + 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'``.""" @@ -377,12 +622,18 @@ def test_enable_jit_for_udt_re_raises_a_lowered_control(monkeypatch): 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") @@ -390,9 +641,10 @@ 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 a saved ``'run'``, restored afterwards, unraised for + 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'`` and ``'load'``. + say so. The same goes for ``'pause'``, ``'load'``, and a ``'run'`` the user + chose. """ from graphblas.core.ss import jit_config @@ -401,19 +653,36 @@ def test_enable_jit_for_udt_does_not_cache_an_opt_out(monkeypatch): 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) - for start in ("off", "pause", "load"): + 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(): - """The same raise, through SuiteSparse's own config rather than a stand-in.""" +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'") @@ -423,9 +692,23 @@ def test_enable_jit_for_udt_re_raises_the_live_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 @@ -465,7 +748,7 @@ def test_failed_compile_is_not_retried_at_the_next_arming(): 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"] = "run" + 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. @@ -507,8 +790,8 @@ def test_public_fix_jit_config_repairs_a_broken_compiler(): assert result is True assert gb.ss.jit_compiler_is_usable() # Calling fix_jit_config a second time must not disable JIT. The probe - # tries to register a one-shot ``_jit_probe`` UDT and previously - # interpreted the duplicate-name error as a probe failure. + # 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 @@ -2046,15 +2329,11 @@ def test_eq_ne_udt_cfunc_works_for_unrepresentable_dtypes(): @pytest.mark.skipif("not supports_udfs") -def test_udt_scalar_broadcast_skips_jit(): - """Mixed UDT+scalar arithmetic and eq/ne stay on the cfunc path. - - The JIT codegen helpers (``set_jit_c_on_op`` and - ``set_jit_c_comparison_on_op``) only run when ``dtype == dtype2``. A - mixed pair would need a per-side type in the C signature plus one C - name per (op, left_type, right_type) triple; for now we keep the - cfunc fallback for broadcast and skip JIT. This test pins that - routing so a future change doesn't accidentally regress it (either +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: @@ -2068,14 +2347,195 @@ def test_udt_scalar_broadcast_skips_jit(): # 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 pairs do not JIT (cfunc handles broadcast correctly; this is - # the documented limit, not an oversight). - assert binary.plus[record, dtypes.INT64].jit_c_source is None - assert binary.plus[dtypes.INT64, record].jit_c_source is None + # Mixed arithmetic does, with the scalar read once for every leaf + assert binary.plus[record, dtypes.INT64].jit_c_source == ( + "void plus_11__BcastJitUV_int64_t (_BcastJitUV *z, const _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"): @@ -2085,3 +2545,94 @@ 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) + t = dtypes.register_new("TooLongName" + "t" * 120, rec) + 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() diff --git a/scripts/jit_diagnostics.py b/scripts/jit_diagnostics.py index 3c94f19bd..6959bfb2d 100755 --- a/scripts/jit_diagnostics.py +++ b/scripts/jit_diagnostics.py @@ -68,7 +68,7 @@ def main(): print("\nSuiteSparse:GraphBLAS 7.x: JIT not available") return 0 - print("\n--- GraphBLAS JIT defaults (from compiled C library) ---") + print("\n--- GraphBLAS JIT settings after graphblas init ---") for key in [ "jit_c_control", "jit_c_compiler_name", From 8f4519ee813210df5ac7f422a8654e9d764cf2e9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 10:43:22 -0600 Subject: [PATCH 27/52] Follow-ups to #591: Monoids on UDT pairs, out-of-range ints, select.row bounds, wording (#661) * Cover date-field text, a Vector accumulated into a UDT Scalar and more Five tests reach lines of #591 that no test ran (local coverage of the full suite with --runslow: suitesparse and vanilla, with and without Numba): - Text in a record's duration field, or text numpy cannot read as a date or a duration, takes the field's type, and converting it raises numpy's error (literals.py: _as_numpy_date and leaf_dtype). - Accumulating a Vector into a UDT Scalar raises the TypeError a built-in Scalar raises (udt.py: _scalar_accum_operand). - A NaT of no unit does not fit a one-byte field: b"N" and NaT cast to each other neither way (literals.py: _converts_back). - The error for an array of records holding an array of bool names the bool array (dtypes.py: _numba_untyped_array_element). - Without Numba, a NaT duration in an array field of a record stays NaT in a coarser unit (dtypes.py: _nats_in_leaf_units). The try/except around lookup_dtype in _weak_builtin_literal_dtype is removed: where the backend has no complex types (vanilla), lookup_dtype registers a complex literal's dtype as a UDT rather than raise, so the except that returned None never ran on any backend. * Index a Monoid with two UDTs as with two built-in types monoid.plus["INT8", "INT16"] is monoid.plus["INT16"]: the Monoid takes the type the pair promotes to. Before #591 the same held for two UDTs (they were unified with numpy), but since #591 stopped unifying UDT pairs, every UDT pair went to Monoid._compile_udt, which raises TypeError "Monoid inputs must be the same dtype". The comment beside that branch said indexing "still raises". Now an indexed Monoid on a pair with a UDT takes the type its BinaryOp gives the pair, the per-leaf promotion that v.ewise_mult(w, monoid.plus) already uses. Behavior change (before -> after), with R8 and R16 records of one int8 and one int16 field of the same name: - monoid.plus[R8, R16] and monoid.plus[R16, R8]: TypeError -> monoid.plus[R16], as before #591. - monoid.plus[R8, "INT16"]: TypeError -> monoid.plus[R16], as binary.plus[R8, "INT16"] gives R16 (TypeError before #591 too). - Records whose fields do not pair: TypeError -> the KeyError the BinaryOp gives for that pair. * Take an out-of-range int exactly beside any op whose result has another type An int beyond the operand's range is typed exactly, not refused, when the op keeps no result in the operand's type: a comparison, or an op that gives another type anyway (int8_vec / 300 is FP64). _gives_another_type decided the second only for a named op it could index, so it said no for a parameterized op (binary.isclose has no __getitem__), for the ops a parameterized op makes (binary.isclose(0.1) is anonymous), and for a user's anonymous op. Those raised OverflowError. Before #591 the int was typed as INT64, giving False, as numpy's isclose does. Now any op counts by the type it gives on the operand's dtype, and a parameterized op is typed with its default parameters, as get_typed_op types it. An op that gives the operand's type (plus, or a user's x + y) still raises. Behavior change (before -> after), with v an INT8 Vector of [100, -100]: - v.apply(binary.isclose, right=300), left=-300, and with binary.isclose(0.1): OverflowError -> [False, False], as before #591. - v.ewise_union(w, binary.isclose, 0, 300): OverflowError -> False where the default is used. - v.apply(gt, right=300) with gt a user's anonymous x > y: OverflowError -> [False, False]. - int64_vec.apply(binary.isclose, right=2**63), which raised OverflowError before #591: now compared in FP64, as numpy's isclose compares it. * Say in udt.rst what a UDT of text or dates supports The field rules ended with "Strings and Python objects are not supported", though a bare text, datetime64 or timedelta64 UDT registers, stores, and compares with ==, != and isequal, and text fields compare in records. The bullet on arrays Numba cannot type left out arrays of str and bytes, which register_anonymous refuses with the same ValueError. * Say how lifted ops compute a record's array field, in _iter_record_leaves Its docstring said the generated expression operates on the whole sub-array. Binary ops now compute an array field element by element, as they compute an array UDT; only unary ops still apply the expression to the whole sub-array. * Name Scalar.__eq__ in the error for a value that is not a literal Scalar == [1] raised "Bad type for keyword argument `other=` in Scalar.ewise_mult(op=binary.eq)", the method == calls inside; before #591 it named Scalar.isequal, also not what the user wrote. Now == converts a literal as ewise_mult does, and the error names Scalar.__eq__ (or Scalar.__ne__ for !=), for built-in and UDT Scalars. * Call a stored or accumulated object a value in the store-cast error Storing or accumulating an object whose type does not cast, such as s(accum=binary.plus) << Scalar.from_value(0.5) into a UDT Scalar, raised "cannot store a result of type FP64 ... Use .new() for an object of the result's type, or compute the result in the type you want to store". The FP64 Scalar is a value, not a result, and .new() does not help it. Now _store_cast_op takes is_value, which the stores of an object as it is pass (w << v, C << A.T, assigning or accumulating a Scalar, and _cast_for_store's other callers): "cannot store a value of type FP64 in an object of type INT8[3]: a UDT casts only to another UDT." A computed result keeps the old message and its advice. * Index monoid.any with a UDT pair by the type the pair promotes to "Index a Monoid with two UDTs as with two built-in types" typed an indexed Monoid by its BinaryOp's result on the pair. For plus, times, min and max that is the promoted type, as before #591, but binary.any on two UDTs keeps the second operand's type, so monoid.any[R16, R8] became monoid.any[R8], and v16.reduce(monoid.any[R16, R8]) raised DomainMismatch where before #591 it gave an R16 Scalar: the pair was unified in either order. monoid.any now takes the type lifted arithmetic (plus) gives the pair. With R8 and R16 records of one int8 and one int16 field: - monoid.any[R16, R8]: monoid.any[R8] -> monoid.any[R16], as before #591 (the other order already gave it). - monoid.any[R8, "INT16"]: the built-in monoid.any[INT16] -> monoid.any[R16], as monoid.plus gives (TypeError before #591). - Records that do not pair: monoid.any[R8] -> plus's KeyError, as for monoid.plus, min and max (TypeError before #591). * Compare with infinity only for a built-in comparison, not an op named like one An int beyond an integer operand's range is compared as infinity of its sign by the built-in comparisons. The check was the op's name alone, and register_anonymous names an op after its function, so a user's op made from a function named ge or eq got the int replaced by infinity. Until "Take an out-of-range int exactly beside any op whose result has another type", such an op raised OverflowError first; since then, int8_vec.apply(register_anonymous(ge), right=300), with ge returning x * 1.0 - y, gave [-inf, -inf], and an eq returning y - x > 250 gave [True, True]. Before #591 they gave [-200.0, -400.0] and [False, True], which they give again: the check now also asks that the op is not anonymous, as the exact typing beside it does. * Raise for a user's op that would meet an int beyond 64 bits in a float type The rule from "Take an out-of-range int exactly beside any op whose result has another type" covers a user's op too. Beside INT64 or UINT64, the int's exact type (UINT64 for 2**63, FP64 for 2**64) unifies with the operand to FP64, which rounds the int: a register_anonymous x < y gave [False, True] for int64_vec [2**63 - 1, 0] with 2**63, where numpy and binary.lt give [True, True]. #591 already did so for a register_new op; before #591 it raised OverflowError. Built-in comparisons are given infinity instead, and the other built-ins that give another type (truediv, isclose, the float ufuncs) compute in float anyway, as numpy does. Now a user's op raises OverflowError, as #591 did, where the exact type and the operand would unify to a float type. Where they share an integer type (300 beside INT8 is INT16) it still takes the int exactly. Built-in ops are told apart by _STANDARD_OPERATOR_NAMES, and the ops a built-in parameterized op makes (binary.isclose(0.1), which is anonymous) by their parent. types.rst states the rule. Behavior change (before -> after), with an anonymous or register_new x < y or x != y: - int64_vec.apply(op, right=2**63) and right=-2**63 - 1, uint64_vec with 2**64, the Scalar, Matrix and ewise_union paths: a wrong answer -> OverflowError, as before #591. - int64_vec.apply(binary.isclose, right=2**63): unchanged, [True, False] as numpy's isclose. * Read a dict as a record only beside a record UDT in == An array UDT Scalar or Vector compared with a dict (a == {}) raised AttributeError: 'NoneType' object has no attribute 'items', since the comparison read the dict as a record of the array dtype's fields, which it has none of. Before #591 it raised the TypeError for a value that is not a literal. Now the dict is read as a record only for a record dtype, as _literal_type does, and a == {} raises "Bad type for argument `other` in Scalar.__eq__(...)" (Vector: the usual Vector.apply TypeError). * Say that operators on a text or date UDT need Numba udt.rst's bullet on text, datetime64 and timedelta64 UDTs listed ==, !=, isequal, first, second, pair and any as working. Without Numba, a Vector's ==, isequal, first, second and any raise KeyError; only a Scalar's ==, != and isequal (compared in numpy) and pair work. * Compare indices with the bound given in select.row, column and index select.row(A < k) reads its thunk from the comparison A < k, which since #591 types a literal beside the matrix's values (NEP 50): beside INT8 or UINT8 values, 300 is infinity, and beside FP32, 2**24 + 1 rounds. That is right for a value comparison, but the helpers used the same thunk for an index. On a 400-row INT8 matrix, select.row(B < 300) kept all 400 rows and select.row(B > 299) none; on FP32 with 2**24 + 3 rows, select.row(F < 2**24 + 1) dropped row 2**24. Before #591 the helpers were right. Now an expression remembers the right operand as the caller gave it (_right_literal, set by Vector.apply, Matrix.apply, Scalar.apply and Scalar's ewise_mult, which Scalar comparisons use), and the row, column and index helpers compare indices with that number, or with the thunk Scalar's value when a Scalar was given. As indices are whole numbers, the bound is made one: i < t is i <= ceil(t) - 1, i >= t is i > ceil(t) - 1, and i <= t and i > t take floor(t). GraphBLAS truncated a float thunk toward zero, so select.row(A > -0.5) dropped row 0 and select.row(A < 2.5) row 2 (before #591 too). A whole bound beyond INT64 is clamped to it, which holds every index. The Scalar unwrap now comes first, so select.row(s < 1) works (it raised AttributeError). Covered: Vector, Matrix (also transposed, through select.column) and Scalar operands, INT8, UINT8 and FP32, fractional and infinite bounds, and a Scalar given as the bound. * Count every op python-graphblas defines at init as standard, truediv included _is_builtin_op, from "Raise for a user's op that would meet an int beyond 64 bits in a float type", asks _STANDARD_OPERATOR_NAMES whether an op is python-graphblas's own, and that set held only the names register_new recorded during initialization. BinaryOp._initialize makes cdiv, truediv and rtruediv directly, with the semirings on cdiv, and binom is registered on first use, so none of them was there. A user's op beside a 64-bit int raises OverflowError where the exact type would compute in float, and truediv was taken for one: uint64_vec / -1, int64_vec / 2**63 and int8_vec / 2**70 raised, where #591 (and, for the first, the code before it) and numpy compute them; binary.numpy.true_divide and divide are truediv too. Now, once the operators are initialized, every op the namespaces hold, and every name queued for first use, is recorded as standard. That covers whatever _initialize makes, now or later, with no list to keep in step; it adds exactly cdiv, truediv, rtruediv, binom and the five semirings on cdiv. They now also pickle by name, as the other built-ins do; their old pickles still load. * Take the thunk's value as an index bound that is not a number A bound given as a Scalar expression (w[0] with autocompute), a 0-d array or np.True_ was recorded as the caller's literal, then found not to be a number, so the helpers shifted lt to le without shifting the bound: select.row(A < np.array(3)) kept rows 0 to 3, and select.row(A >= w[0]) dropped row 3. Use the caller's literal only when it is a real number; otherwise the thunk, made from that operand with its own type, holds the bound. Found in review. * Give a select helper's left-operand comparison the op's TypeError again "Take the thunk's value as an index bound that is not a number" read the thunk's value whatever it was, and for A.apply(binary.le, left=3) the second operand is the Matrix: select.row(A.apply(binary.le, left=3)) before: AttributeError: 'Matrix' object has no attribute 'value' after: TypeError: Bad type for keyword argument `thunk=` in Vector.select(op=select.rowle), as before #591 _index_bound reads a value only from a Scalar thunk; any other goes to the op as it is, which says what is wrong. Found in review. * Select no index for a NaN bound in the row, column and index helpers Every comparison with NaN is False, but NaN went to rowle or rowgt as it was, and its cast to the INT64 thunk made it 0: A of 4 rows: select.row(A < nan), (A <= nan) before: row 0 kept (also before #591) select.row(A > nan), (A >= nan) before: rows 1, 2 and 3 kept (also before #591) after: nothing kept, also for Scalar(nan) and np.float32(nan) _index_bound now gives a NaN bound as le -1 or gt INT64_MAX, which no index meets. Found in review. * Give the NaT in test_udt_literal_converts_only_when_exact a unit numpy 2.5 deprecates the generic unit, so np.datetime64("NaT") warns, and CI (which turns warnings into errors) failed this test on every job that drew numpy 2.5. A NaT in seconds reaches the same check: the literal still does not fit the S1 field, with the same error. * Ignore _test_weak_lt in test_operator_types like other test-made ops test_python_number_is_weak_beside_builtin_dtypes registers it in the binary namespace, so test_binarytypes failed whenever pytest-randomly ran it after that test. * Give text read as NaT the date field's unit A date given as text is parsed by np.datetime64, which reads "NaT" (in any case) and empty text as NaT with no unit. numpy 2.5 deprecates that unit, so comparing a date UDT, or a record's date field, with "NaT" warned (an error under the test suite's filters), where the code before #591 compared without a warning: s == "NaT" # s a Scalar of an M8[s] UDT before (numpy 2.5): DeprecationWarning: The 'generic' unit for NumPy timedelta is deprecated after: False, as numpy's NaT equals nothing Such text is now NaT in the field's unit. Other text keeps the unit it is written to, so a finer date still does not fit a coarser field. Found in review of the numpy 2.5 CI fixes. * Type a Monoid indexed with a UDT pair without compiling plus monoid.any[X, Y] found the type the pair promotes to by compiling binary.plus on the pair (a Numba cfunc and a GrB_BinaryOp that nothing used) and reading its return type. _monoid_pair_type, in udt.py beside the other UDT result-type rules, runs the same check and promotion (_check_udt_pair, _udt_op_types) and compiles nothing. It gives the type the compiled ops give on every pair tried. Without Numba, monoid.any[X, Y] now raises the KeyError monoid.any raises on the promoted type ("any does not work with ..."), not one that names plus. Also drop a test comment that said monoid.any[R16, R8] was monoid.any[R8] before #591: it was monoid.any[R16], as it is now. * Keep the result message for a computed result assigned to part of an object "Call a stored or accumulated object a value in the store-cast error" had _cast_for_store pass is_value=True for everything, but indexed assignment hands it expressions too: v[:] << w.apply(unary.ainv) # w is FP64, v of an INT8[3] UDT before: cannot store a value of type FP64 in an object of type ... after: cannot store a result of type FP64 in an object of type ... Use .new() for an object of the result's type, or ... A computed result keeps the result message, as that commit meant; only a Scalar, Vector or Matrix stored as it is is a value. --- docs/user_guide/types.rst | 9 +++- docs/user_guide/udt.rst | 17 +++++-- graphblas/core/base.py | 10 +++- graphblas/core/literals.py | 63 +++++++++++++++++-------- graphblas/core/matrix.py | 6 ++- graphblas/core/operator/utils.py | 24 ++++++++-- graphblas/core/scalar.py | 34 +++++++++++--- graphblas/core/udt.py | 41 +++++++++++++---- graphblas/core/vector.py | 6 ++- graphblas/select/__init__.py | 64 ++++++++++++++++++++------ graphblas/tests/test_dtype.py | 6 ++- graphblas/tests/test_matrix.py | 51 ++++++++++++++++++++ graphblas/tests/test_op.py | 62 ++++++++++++++++++++++--- graphblas/tests/test_operator_types.py | 1 + graphblas/tests/test_scalar.py | 32 ++++++++++++- graphblas/tests/test_vector.py | 58 +++++++++++++++++++++++ 16 files changed, 414 insertions(+), 70 deletions(-) diff --git a/docs/user_guide/types.rst b/docs/user_guide/types.rst index 36d93e8a4..65b941cf9 100644 --- a/docs/user_guide/types.rst +++ b/docs/user_guide/types.rst @@ -48,7 +48,14 @@ 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 ``v / 300`` raise, since its result is FP64 whatever the integer. An +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 diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 44af6a1f0..96d8e1f91 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -74,10 +74,19 @@ Field type rules: 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`` or ``timedelta64``, so - with Numba installed, registering a UDT that holds one raises ``ValueError``. - Use an array of ``uint8`` for ``bool``, and of ``int64`` for the others. -- Strings and Python objects are not supported. +- 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``. diff --git a/graphblas/core/base.py b/graphblas/core/base.py index e8940ad9b..a3562f922 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -365,7 +365,7 @@ 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 not None: + 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 @@ -407,7 +407,9 @@ 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: - if (cast_op := _store_cast_op(expr.dtype, self.dtype)) is not None: + 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: @@ -570,6 +572,7 @@ class BaseExpression: "expr_repr", "dtype", "_value", + "_right_literal", "__weakref__", ) output_type = None @@ -605,6 +608,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): diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 810bde22c..d43c97aa8 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -22,10 +22,11 @@ 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, SelectOp, TypedOpBase, binary_from_string +from .operator import IndexUnaryOp, ParameterizedUdf, SelectOp, TypedOpBase, binary_from_string from .operator.binary import BinaryOp from .operator.monoid import Monoid from .udt import ( @@ -97,7 +98,7 @@ def _literal_operand(dtype, value, op, *, left=False): # 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: + 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: @@ -114,7 +115,9 @@ def _literal_operand(dtype, value, op, *, left=False): exact = is_comparison or op.name == "truediv" if is_comparison and not isinstance(value, (np.generic, np.ndarray)): np_type = dtype.np_type - parts = _dict_to_record(np_type, value) if isinstance(value, dict) else value + 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: @@ -179,13 +182,15 @@ def _weak_builtin_literal_dtype(dtype, value, op=None): # 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_type == dtype.np_type: - rv = dtype - else: - try: - rv = lookup_dtype(np_type) - except ValueError: # A complex literal where the backend has no complex types - return None + 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: @@ -204,15 +209,29 @@ def _weak_builtin_literal_dtype(dtype, value, op=None): return rv +def _is_builtin_op(op): + """Whether python-graphblas defines ``op``, or a built-in parameterized op made it.""" + if op._anonymous: + info = getattr(getattr(op, "orig_func", None), "_parameterized_info", None) + return info is not None and _is_builtin_op(info[0]) + module = op._op_class._modname if isinstance(op, ParameterizedUdf) else op._modname + return f"{module}.{op.name}" in _STANDARD_OPERATOR_NAMES + + def _gives_another_type(op, dtype): - """Whether built-in ``op`` on two operands of built-in ``dtype`` gives another type. + """Whether ``op`` on two operands of built-in ``dtype`` gives another type. - ``truediv`` on integers gives FP64, and the numpy ufuncs that compute in - float give a float type, as numpy computes them in a float loop. + ``truediv`` on integers gives FP64, ``isclose`` gives BOOL, and the numpy + ufuncs that compute in float give a float type, as numpy computes them in + a float loop. A user's op counts by the type it gives too. A + parameterized op is typed with its default parameters, as + ``get_typed_op`` types it. """ - if op is None or getattr(op, "_anonymous", True) or not hasattr(op, "__getitem__"): + if op is None: return False try: + if isinstance(op, ParameterizedUdf): + op = op() return op[dtype].return_type != dtype except Exception: # An op that does not take this dtype reports it itself, later. return False @@ -257,7 +276,7 @@ def _literal_type(dtype, value, *, exact=False): if _keeps_value(value, np_type): return dtype return lookup_dtype(np.dtype(f"{np_type.kind}{max(len(value), 1)}")) - if np_type.kind in "mM" and (date := _as_numpy_date(value, np_type.kind)) is not None: + if np_type.kind in "mM" and (date := _as_numpy_date(value, np_type)) is not None: if date.dtype.kind == np_type.kind and _keeps_value(date, np_type): return dtype return lookup_dtype(date.dtype) @@ -611,7 +630,7 @@ def leaf_dtype(leaf, part, *, in_record=False): # takes the field's type, and converting it reports what is wrong. if isinstance(part, (str, bytes)) and base.kind in "SU": own = np.dtype(f"{base.kind}{max(len(part), 1)}") - elif (date := _as_numpy_date(part, base.kind)) is not None: + elif (date := _as_numpy_date(part, base)) is not None: part = date own = date.dtype if own.kind != base.kind: @@ -674,21 +693,27 @@ def record_leaves(np_type, part): return dtype if literal_type == np_type else lookup_dtype(literal_type) -def _as_numpy_date(value, kind): +def _as_numpy_date(value, leaf): """Return ``value`` as numpy's ``datetime64`` or ``timedelta64``, or ``None`` if it is neither. A numpy date or duration is itself; a Python ``datetime.date`` or ``datetime.datetime`` is a date in days or microseconds and a ``datetime.timedelta`` a duration in microseconds, as numpy makes them; text is a date in the unit it is written to (``"2020-01-01T12:00"`` is in - minutes) when ``kind``, the leaf's, is ``"M"`` and numpy can read it. + minutes) when ``leaf``, the numpy dtype it meets, is a date and numpy can + read it. Text numpy reads as NaT (``"NaT"`` in any case, or empty) is NaT + in the leaf's unit: numpy would give it no unit, which numpy 2.5 deprecates. """ if isinstance(value, (np.datetime64, np.timedelta64)): return value try: if isinstance(value, datetime.timedelta): return np.timedelta64(value) - if isinstance(value, datetime.date) or (kind == "M" and isinstance(value, (str, bytes))): + if isinstance(value, datetime.date): + return np.datetime64(value) + if leaf.kind == "M" and isinstance(value, (str, bytes)): + if value.lower() in ("nat", "", b"nat", b""): + return np.datetime64("NaT", np.datetime_data(leaf)[0]) return np.datetime64(value) except ValueError: # text that is not a date; converting it reports that pass diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 166f263a0..049b39772 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -2426,6 +2426,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." @@ -2495,6 +2496,7 @@ 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: + right_literal = right right, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") @@ -2538,7 +2540,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, @@ -2549,6 +2551,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 diff --git a/graphblas/core/operator/utils.py b/graphblas/core/operator/utils.py index 37da5cc5f..c3a8f5b8f 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -1,6 +1,7 @@ from types import BuiltinFunctionType, FunctionType, ModuleType from ... import ( + _STANDARD_OPERATOR_NAMES, backend, binary, config, @@ -14,6 +15,7 @@ ) from ...dtypes import UINT64, lookup_dtype, unify from ..expr import InfixExprBase +from ..udt import _monoid_pair_type from .base import ( _SS_OPERATORS, OpBase, @@ -49,6 +51,18 @@ # 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 @@ -78,12 +92,14 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala return op[UINT64] if op._udt_types is None: raise KeyError(f"{op.name} does not work with ({dtype}, {dtype2})") - if kind == "binary" and isinstance(op, Monoid): + 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. Indexing (``monoid.plus[X, Y]``) - # still raises. - return get_typed_op(op.binaryop, dtype, dtype2) + # 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: diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index bb7339ce8..a7cf2189b 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -158,12 +158,12 @@ def __eq__(self, other): Without Numba, which ``binary.eq`` needs for UDTs, a UDT's values are compared in numpy, which broadcasts them the same way. """ - return _scalar_eq(self, other) + return _scalar_eq(self, other, "__eq__") __hash__ = None def __ne__(self, other): - return not _scalar_eq(self, other) + return not _scalar_eq(self, other, "__ne__") def __bool__(self): """Truthiness check. @@ -750,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 @@ -778,7 +779,7 @@ def _ewise_mult(self, other, op=binary.times, is_infix=False): ) 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()], @@ -786,6 +787,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, @@ -940,7 +944,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, @@ -950,6 +954,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) @@ -1137,8 +1143,11 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): # End auto-generated code: Scalar -def _scalar_eq(scalar, other): - """Return ``scalar == other`` as a bool, by ``binary.eq`` (see ``Scalar.__eq__``).""" +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: @@ -1175,6 +1184,19 @@ def _scalar_eq(scalar, other): # 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) diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index 94e7b3136..27ebc7b90 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -80,11 +80,12 @@ def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): ``("['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: the generated - expression operates on the whole sub-array and the wrapper slice-assigns - it 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. + 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] @@ -300,6 +301,20 @@ def _udt_op_types(op_name, dtype, dtype2): 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. @@ -430,7 +445,7 @@ def _casts_safely_into(literal_type, dtype): 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): +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 @@ -442,6 +457,8 @@ def _store_cast_op(result_dtype, output_dtype): 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 @@ -453,6 +470,11 @@ def _store_cast_op(result_dtype, output_dtype): 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 " @@ -530,11 +552,14 @@ def _cast_for_store(value, dtype): 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. """ - if (cast_op := _store_cast_op(value.dtype, dtype)) is None: - return value 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}: " diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index f4b72fca4..8e07f2bc2 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -1408,6 +1408,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." @@ -1477,6 +1478,7 @@ 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: + right_literal = right right, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") @@ -1520,7 +1522,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, @@ -1528,6 +1530,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 diff --git a/graphblas/select/__init__.py b/graphblas/select/__init__.py index b55766ff8..876b3ec6a 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 = {} @@ -42,21 +45,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 +53,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. diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 8ca033625..e2df81017 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -378,10 +378,12 @@ def test_register_anonymous_zero_size(): def test_register_anonymous_array_numba_cannot_type(): - # Numba cannot type an array of these, alone or in a record. Its own error was - # "AttributeError: 'Boolean' object has no attribute 'bitwidth'" and the like. + # 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"), diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index d1405c49f..306f9837b 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -1273,6 +1273,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) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index d51993c2e..3c72e36c2 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4576,6 +4576,7 @@ def test_udt_literal_converts_only_when_exact(): 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( @@ -4593,6 +4594,8 @@ def test_udt_literal_converts_only_when_exact(): 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) @@ -4656,6 +4659,8 @@ def test_udt_literal_converts_only_when_exact(): (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}")): @@ -5004,11 +5009,32 @@ def test_udt_record_date_literals_compare_as_numpy(): ((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) @@ -5839,8 +5865,11 @@ def test_udt_store_casts_reach_every_store_path(): 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="a UDT casts only to another UDT"): + 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] @@ -5902,9 +5931,14 @@ def added_to_int8(value): 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. - with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): + # 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") @@ -6333,7 +6367,9 @@ def test_udt_pair_is_never_unified(): 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. + 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") @@ -6352,8 +6388,22 @@ def test_udt_pair_is_never_unified(): 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() - with pytest.raises(TypeError, match="Monoid inputs must be the same dtype"): - monoid.plus[X, Y] + # 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") diff --git a/graphblas/tests/test_operator_types.py b/graphblas/tests/test_operator_types.py index 027f02fcc..8c250bc25 100644 --- a/graphblas/tests/test_operator_types.py +++ b/graphblas/tests/test_operator_types.py @@ -204,6 +204,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_scalar.py b/graphblas/tests/test_scalar.py index 686e4861a..3ade2fb25 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -293,7 +293,10 @@ def test_udt_nat_duration_in_another_unit_stays_nat(): 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 = None if gb.core._has_numba else dtypes.register_anonymous(np.dtype(("m8[s]", 2))) + 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) @@ -321,6 +324,9 @@ def test_udt_nat_duration_in_another_unit_stays_nat(): 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"]) @@ -462,6 +468,17 @@ def test_eq_types_a_literal_as_binary_eq_does(): 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 @@ -717,6 +734,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 diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 7d5cab524..3fa48e661 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -751,6 +751,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 @@ -1608,6 +1613,59 @@ def test_python_number_is_weak_beside_builtin_dtypes(): 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 From 0bf7cf9891dd84197f1de8a05c29e613c422dcae Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 10:47:32 -0600 Subject: [PATCH 28/52] Fix gh-559 memory growth: build .ss lazily instead of storing it (#598) * Fix gh-559 memory growth: build .ss lazily instead of storing it Every Matrix and Vector was born inside two reference cycles: the stored self.ss held _parent back to the object, and so did ss.config. Instances therefore died by the cyclic garbage collector rather than by refcount, and their C-side GrB buffers accumulated between gc sweeps. With gc disabled, a scaled version of the gh-559 batched-mxm loop grew without bound. The .ss namespace is now built per access (a property inside the existing class_property), nothing is stored on the parent, no cycle forms, and objects free as soon as their refcount drops. Two user-visible behavior changes: A.ss is A.ss is now False. It was True, because .ss was a stored attribute; each access now returns a fresh namespace object. Code that compares .ss by identity, or that caches attributes on it, will see the difference. Assigning A.ss now raises AttributeError. It previously succeeded and silently replaced the namespace, because "ss" was in __slots__. Class-level Matrix.ss and Vector.ss still resolve to the ss class, so the import_* classmethods are unchanged. Building the namespace per access costs roughly 215ns against roughly 78ns for the stored attribute (timeit), so about 1% of a single small mxm plus to_dense iteration, which runs about 15us. The .ss namespace is typically touched once per user operation. Six of the eight new tests fail without the fix. They assert on the reference graph rather than on process memory, which a functional suite cannot see and which an RSS delta would measure only statistically: a weakref must be dead the instant the last strong reference drops with the cyclic collector switched off, and a batched mxm loop must not raise the number of live Matrix objects reported by gc.get_objects(). The remaining two cover invariants the fix has to preserve, class-level access and views whose _parent is set, so they pass either way. * Exclude the gc-was-off arm of a test's cleanup from branch coverage test_batched_mxm_loop_does_not_accumulate_matrices restores the cyclic collector only if it was on when the test started. It always is under pytest, so the other arm of that if never runs and Codecov reported a partial branch. The same check in _dies_by_reference_count is not reported: its try block returns, and coverage.py gives a finally block reached that way only one arc. --- graphblas/core/matrix.py | 20 ++-- graphblas/core/vector.py | 17 ++-- graphblas/tests/test_ss_refcount.py | 144 ++++++++++++++++++++++++++++ 3 files changed, 167 insertions(+), 14 deletions(-) create mode 100644 graphblas/tests/test_ss_refcount.py diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 049b39772..55920d4eb 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -208,7 +208,7 @@ class Matrix(BaseType): """ - __slots__ = "_nrows", "_ncols", "_parent", "ss" + __slots__ = "_nrows", "_ncols", "_parent" ndim = 2 _is_transposed = False _name_counter = itertools.count() @@ -226,8 +226,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 @@ -239,8 +237,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): @@ -3571,10 +3567,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 ) diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 8e07f2bc2..9ae745063 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -161,7 +161,7 @@ class Vector(BaseType): """ - __slots__ = "_size", "_parent", "ss" + __slots__ = "_size", "_parent" ndim = 1 _name_counter = itertools.count() @@ -174,8 +174,6 @@ def __new__(cls, dtype=FP64, size=0, *, name=None): 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 @@ -186,8 +184,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): @@ -2127,10 +2123,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/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() From 4bdc169990e55811cada696e148665a5f063ccc9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 10:49:04 -0600 Subject: [PATCH 29/52] Hold a Matrix by column before casting it to a Vector (#637) Matrix.ss.flatten reshapes a copy to n x 1 and casts the GrB_Matrix to a GrB_Vector in place (_as_vector). SuiteSparse holds a GrB_Vector as a one-column matrix by column, but GxB_Matrix_reshapeDup keeps the input's orientation, so flattening a Matrix held by row (python-graphblas's default) gave a Vector that SuiteSparse sees as size 1: v = Matrix.from_coo([0, 1], [1, 2], [1.0, 2.0]).ss.flatten() v.shape # (6,) v.size # 1 v.get(5) # InvalidIndex isequal and to_coo read the Python-side size and looked right, which is why test_ss_flatten passed. Reproduced on 03cf0bcc with SuiteSparse 10.5.1 as well as on this stack with 10.5.0. _as_vector now sets the matrix's format to by_col when it is held by row, so the cast is valid for every caller; flatten is the only one today. The new test flattens a 2 x 3 Matrix imported in each of the eight formats, in both orders, and checks size, every element through get, and dup; the four by-row formats failed with "assert 1 == 6". --- graphblas/core/matrix.py | 8 +++++++- graphblas/tests/test_matrix.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 55920d4eb..f67414438 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -252,13 +252,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, diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 306f9837b..8d8abf82b 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -3116,6 +3116,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 = [ From 1649d995ae2dc27708e4f0a3415e2adf5a2c2c71 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 11:02:00 -0600 Subject: [PATCH 30/52] Return an empty Scalar from ss.iso_value of an empty iso object (#638) * Hold a Matrix by column before casting it to a Vector (#637) Matrix.ss.flatten reshapes a copy to n x 1 and casts the GrB_Matrix to a GrB_Vector in place (_as_vector). SuiteSparse holds a GrB_Vector as a one-column matrix by column, but GxB_Matrix_reshapeDup keeps the input's orientation, so flattening a Matrix held by row (python-graphblas's default) gave a Vector that SuiteSparse sees as size 1: v = Matrix.from_coo([0, 1], [1, 2], [1.0, 2.0]).ss.flatten() v.shape # (6,) v.size # 1 v.get(5) # InvalidIndex isequal and to_coo read the Python-side size and looked right, which is why test_ss_flatten passed. Reproduced on 03cf0bcc with SuiteSparse 10.5.1 as well as on this stack with 10.5.0. _as_vector now sets the matrix's format to by_col when it is held by row, so the cast is valid for every caller; flatten is the only one today. The new test flattens a 2 x 3 Matrix imported in each of the eight formats, in both orders, and checks size, every element through get, and dup; the four by-row formats failed with "assert 1 == 6". * Return an empty Scalar from ss.iso_value of an empty iso object SuiteSparse marks a Matrix or Vector with no entries iso in ordinary use: Matrix.from_coo([], [], 7, nrows=3, ncols=2) is one, and so can be the result of a select that keeps no entries. ss.iso_value read the value with next(self.itervalues()), so on such an object it leaked StopIteration, which turns into "RuntimeError: generator raised StopIteration" inside a generator. graphblas-algorithms reached it through every edgeless graph its operators build (difference(G, G), for one): pagerank, katz_centrality and the Bellman-Ford names raised StopIteration. An object with no entries holds no value, so return an empty Scalar of its dtype. That is how python-graphblas reports a missing value elsewhere: an empty reduction (reduce with the default allow_empty=True) and the extraction of a missing element both give an empty Scalar. ValueError stays for an object that is not iso, which is a different answer from "iso, with nothing in it". Both properties gain a docstring saying so. * Check ss.iso_value on an empty non-iso object in the same loop The empty-iso tests loop over objects and check iso_value only when SuiteSparse marks one iso. Whether an empty select result is iso depends on the SuiteSparse version, so the not-iso arm ran on some CI draws only and asserted nothing when it did. A new empty Matrix or Vector, which is never iso, now goes through the loop too, and the not-iso arm checks that iso_value raises. The separate check that did this is folded in. --- graphblas/core/ss/matrix.py | 11 ++++++++++- graphblas/core/ss/vector.py | 11 ++++++++++- graphblas/tests/test_matrix.py | 22 ++++++++++++++++++++++ graphblas/tests/test_vector.py | 19 +++++++++++++++++++ 4 files changed, 61 insertions(+), 2 deletions(-) diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index 2c9a4d99e..e7887311c 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -201,9 +201,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 diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index 80932da73..f4cfa0cdc 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -120,9 +120,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 diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 8d8abf82b..1144e4109 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -177,6 +177,28 @@ 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 + + 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]) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 3fa48e661..bd4b9660e 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -168,6 +168,25 @@ 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 + + def test_from_coo_default_values(): u = Vector.from_coo([1, 3]) v = Vector.from_coo([1, 3], [1.0, 1.0]) From f5d1763e7943409abc7f3a25398c04fbd0dd9371 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 11:15:16 -0600 Subject: [PATCH 31/52] Return {} from Matrix.to_dicts of a Matrix with no entries (#639) to_dicts paired consecutive indptr entries with sliding_window_view(indptr, 2). A Matrix with no entries compresses to indptr == [0], one element, and sliding_window_view refuses a window longer than its input: Matrix(int, 3, 2).to_dicts() ValueError: window shape cannot be larger than input array shape so to_dicts failed on every empty Matrix (and its transpose), in both orders, while from_dicts({}) built one. Reproduced on 03cf0bcc as well. Pair indptr[:-1] with indptr[1:] instead, as lists: an empty Matrix gives {}, and the rest is unchanged (on 100000 x 100000 with 10**6 entries, 90 ms against 96 ms before, best of 5 at load average 4.6). --- graphblas/core/matrix.py | 13 ++++++------- graphblas/tests/test_matrix.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index f67414438..889878364 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -1868,18 +1868,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( diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 1144e4109..1ffb1919c 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -4200,6 +4200,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) From 8e3c34366c188fd03dec9781e118416f09502c25 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 11:23:56 -0600 Subject: [PATCH 32/52] Say why a full export or unpack of a non-full object fails (#640) Vector.ss.export("full"), Matrix.ss.export("fullr") and "fullc", and the same formats through ss.unpack, on an object with missing entries raised InvalidValue with an empty message: SuiteSparse refuses with GrB_INVALID_VALUE and sets no error string, so there was nothing to show. The object itself was left intact. Check the entry count first and raise InvalidValue (the same type as before) with the count and the format: Matrix with 1 of 4 entries is not full, so it cannot export to format 'fullr' --- graphblas/core/ss/matrix.py | 8 +++++++- graphblas/core/ss/vector.py | 8 +++++++- graphblas/tests/test_matrix.py | 12 ++++++++++++ graphblas/tests/test_vector.py | 11 +++++++++++ 4 files changed, 37 insertions(+), 2 deletions(-) diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index e7887311c..9f9bfccae 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -7,7 +7,7 @@ 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 InvalidValue, _error_code_lookup, check_status, check_status_carg from .. import NULL, _has_numba, ffi, lib from ..base import call from ..operator import get_typed_op @@ -1126,6 +1126,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: diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index f4cfa0cdc..a9515c448 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -7,7 +7,7 @@ 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 ..operator import get_typed_op @@ -604,6 +604,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, diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 1ffb1919c..16e48cef0 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -199,6 +199,18 @@ def test_ss_iso_value_of_an_empty_iso_matrix(): 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]) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index bd4b9660e..5edebd235 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -187,6 +187,17 @@ def test_ss_iso_value_of_an_empty_iso_vector(): 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]) From c6348d539cd668b0034e9fe162b02eea15f03818 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 19:57:55 -0600 Subject: [PATCH 33/52] Hand GraphBLAS only arrays it can free, and free each once (#641) Imports that hand arrays to GraphBLAS (take_ownership=True, and every secure import) could free memory with the wrong allocator, or twice: - owndata says an array owns its data, not which allocator made it. An array from a custom NumPy memory handler (NEP 49) owns its data too, yet GraphBLAS frees what it takes with its own allocator. With suitesparse-graphblas 10.5.1.1, NumPy 2.5 and PYTHONMALLOC=debug, unclaim_buffer raised ValueError after GraphBLAS had taken such an array, and the process aborted when the object was freed. With memory_manager="c", every ss import raised that ValueError, take_ownership=False included. - An array given as two arguments of one import (the same uint64 array as indptr and col_indices of import_csr, say) went to GraphBLAS twice and was freed twice. Only an index array given as the values was copied. - GraphBLAS moves the arrays into the new object, clearing the caller's pointers, before a secure import's check, and frees them with the object when the check fails. The arrays were unclaimed only after a successful status, so NumPy freed them again and the process aborted. Arrays under 1 KiB hid this: NumPy caches those blocks. - A secure pack whose check fails frees the packed GraphBLAS object itself (a SuiteSparse:GraphBLAS bug in 10.5.2 and earlier), so the Python object kept a freed handle: A.nvals read freed memory, and the process aborted when A was freed again. Now an import asks suitesparse-graphblas's can_unclaim_buffer (10.5.1.1 and later; older versions keep the owndata and writeable check) whether GraphBLAS can take an array, and copies one it can't, once, into memory GraphBLAS can own: NumPy's when GraphBLAS can own that (the default), else memory from suitesparse_graphblas.utils.empty. An array given twice is copied for its second use. Each of the eleven imports unclaims every array whose pointer GraphBLAS cleared, right after the call and before checking the status. pack_*(secure_import=True) imports the data securely into a new object and moves its content into the target with an unpack and a fast pack, so a failed check leaves the target unchanged. A secure import now always copies, even with take_ownership=True: a failed check would otherwise leave the caller holding read-only views of memory GraphBLAS had freed, and a secure import reads every entry anyway. Values given without a dtype, whose record dtype nests subarrays, are viewed with the flattened dtype before they are copied, not after, which copied them twice. The thirteen import docstrings say which arrays can be taken as they are: those allocated by GraphBLAS (from export, unpack and suitesparse_graphblas.utils.empty), and those from NumPy's default allocator unless memory_manager="c". Tests import arrays from a GraphBLAS arena given libc's allocator (a real foreign memory handler), refuse can_unclaim_buffer, and check the rule for older suitesparse-graphblas. Subprocess tests, where a double free aborts only the child, give an array twice and import and pack data that only the secure check rejects, in every format with such a check. Another runs imports under memory_manager="c" with PYTHONMALLOC=debug and checks with tracemalloc that they make no NumPy copy. A secure import and pack of valid data runs in all eleven formats; no test used secure_import before. With suitesparse-graphblas 10.5.1.1 and later, the test session turns on NumPy's warning for freeing owned data that has no memory handler, which filterwarnings = error makes a failure. suitesparse-graphblas's give_buffer context manager could replace the unclaim logic once 10.5.1.1 is the minimum version. --- graphblas/core/ss/_utils.py | 37 +++ graphblas/core/ss/matrix.py | 257 ++++++++++++++----- graphblas/core/ss/vector.py | 87 +++++-- graphblas/core/utils.py | 92 +++++-- graphblas/tests/conftest.py | 29 +++ graphblas/tests/test_ss_utils.py | 409 +++++++++++++++++++++++++++++++ 6 files changed, 815 insertions(+), 96 deletions(-) diff --git a/graphblas/core/ss/_utils.py b/graphblas/core/ss/_utils.py index 85017804f..ca4fc87d3 100644 --- a/graphblas/core/ss/_utils.py +++ b/graphblas/core/ss/_utils.py @@ -2,11 +2,48 @@ 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. diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index 9f9bfccae..9fc535924 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -1,7 +1,7 @@ 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 @@ -14,6 +14,7 @@ from ..scalar import Scalar, _as_scalar, _scalar_index from ..utils import ( _CArray, + _distinct_buffers, _MatrixArray, _Pointer, get_order, @@ -24,7 +25,7 @@ values_to_numpy_buffer, wrapdoc, ) -from ._utils import _resolve_serialized_dtype +from ._utils import _pack_securely, _resolve_serialized_dtype, _unclaim_taken from .config import BaseConfig from .descriptor import get_descriptor @@ -1214,12 +1215,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``. @@ -1312,7 +1318,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" ) @@ -1324,8 +1343,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)) @@ -1349,6 +1367,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, @@ -1358,9 +1377,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 @@ -1403,12 +1419,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``. @@ -1501,7 +1522,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" ) @@ -1513,8 +1547,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)) @@ -1538,6 +1571,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, @@ -1547,9 +1581,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 @@ -1598,12 +1629,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``. @@ -1704,7 +1740,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" ) @@ -1719,8 +1770,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)) @@ -1760,6 +1810,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, @@ -1769,10 +1820,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 @@ -1821,12 +1868,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``. @@ -1927,7 +1979,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" ) @@ -1940,8 +2007,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)) @@ -1983,6 +2049,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, @@ -1992,10 +2059,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 @@ -2043,12 +2106,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``. @@ -2137,7 +2205,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" ) @@ -2146,8 +2226,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: @@ -2177,6 +2256,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, @@ -2186,8 +2266,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 @@ -2235,12 +2313,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``. @@ -2329,7 +2412,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" ) @@ -2338,8 +2433,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: @@ -2369,6 +2463,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, @@ -2378,8 +2473,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 @@ -2420,12 +2513,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``. @@ -2506,7 +2604,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( @@ -2528,6 +2636,7 @@ def _import_fullr( 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, @@ -2537,7 +2646,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 @@ -2578,12 +2686,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``. @@ -2664,7 +2777,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( @@ -2685,6 +2808,7 @@ def _import_fullc( 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, @@ -2694,7 +2818,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 @@ -2926,12 +3049,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. @@ -3094,12 +3222,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. diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index a9515c448..9cd4e22d0 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -1,7 +1,7 @@ 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 @@ -14,13 +14,14 @@ from ..scalar import Scalar, _as_scalar from ..utils import ( _CArray, + _distinct_buffers, _MatrixArray, ints_to_numpy_buffer, normalize_chunks, values_to_numpy_buffer, wrapdoc, ) -from ._utils import _resolve_serialized_dtype +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 @@ -879,12 +880,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``. @@ -974,15 +980,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: @@ -1011,6 +1029,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, @@ -1020,8 +1039,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 @@ -1063,12 +1080,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``. @@ -1154,7 +1176,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 @@ -1162,8 +1196,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)) @@ -1197,6 +1230,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, @@ -1206,8 +1240,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 @@ -1243,12 +1275,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``. @@ -1326,7 +1363,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 @@ -1355,6 +1402,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, @@ -1364,7 +1412,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) diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index 9f15254ab..f076da0c9 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -6,6 +6,12 @@ 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.") @@ -87,11 +93,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): @@ -119,17 +179,21 @@ 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( - _nats_in_leaf_units(_view_if_same_layout(array, np_type), np_type), - 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 @@ -142,8 +206,8 @@ def values_to_numpy_buffer( 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 diff --git a/graphblas/tests/conftest.py b/graphblas/tests/conftest.py index 4270fec17..f68a597fd 100644 --- a/graphblas/tests/conftest.py +++ b/graphblas/tests/conftest.py @@ -13,6 +13,14 @@ 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 @@ -162,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_ss_utils.py b/graphblas/tests/test_ss_utils.py index 4e553f51b..6e7a86b8e 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -1,14 +1,28 @@ +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 @@ -348,3 +362,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 From d9edd77dbe0550a0a8e8f2d2e1b73d645d895145 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 21:23:03 -0600 Subject: [PATCH 34/52] Check that values hold whole elements before building from them (#642) * Check that values hold whole elements before building from them Values for a sub-array dtype are converted to its scalar type, and build counted only the leading dimension, so `Vector.from_coo([0, 1], np.ones(2), "FP64[2]")` handed GraphBLAS 2 numbers where it reads 4 and built from whatever followed the array. Raise unless the trailing shape is the element's, in build, from_csr/from_csc, and from_dense. Without SuiteSparse, a scalar value now broadcasts to whole elements too (it read past the end the same way), and from_csr/from_csc with a scalar no longer fail for a sub-array dtype on SuiteSparse. * Skip the values shape tests where a UDT cannot be built Without SuiteSparse, build gives a UDT the dup_op binary.any, and that needs Numba, so test_subarray_dtype_values_shape failed with a KeyError on the vanilla backend without Numba. Skip it there only; without Numba on SuiteSparse it still runs. * Give the number of elements in the build lengths error build and from_coo compare the number of indices with the number of elements (values.shape[0]), but the message gave values.size, the number of numbers: Vector.from_coo([0, 1], np.ones((3, 2)), "FP64[2]") said "lengths must match: 2 != 6". It now says 2 != 3, and the Matrix message gives 3 as well. --- graphblas/core/matrix.py | 12 +++++++++--- graphblas/core/utils.py | 14 ++++++++++++++ graphblas/core/vector.py | 10 ++++++---- graphblas/tests/test_matrix.py | 30 ++++++++++++++++++++++++++++++ graphblas/tests/test_vector.py | 34 ++++++++++++++++++++++++++++++++++ 5 files changed, 93 insertions(+), 7 deletions(-) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 889878364..335c1f3e8 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -42,6 +42,7 @@ from .utils import ( _CArray, _Pointer, + check_element_shape, class_property, get_order, ints_to_numpy_buffer, @@ -678,8 +679,9 @@ def build(self, rows, columns, values, *, dup_op=None, clear=False, nrows=None, 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: @@ -927,7 +929,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 @@ -1060,6 +1063,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: @@ -1083,7 +1087,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: @@ -1476,6 +1481,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) diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index f076da0c9..7da0e0645 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -211,6 +211,20 @@ def values_to_numpy_buffer( 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 normalize_values(self, values, dtype, shape=None, is_iso=False): """Expand and/or update dtype of values array.""" if dtype is not None: diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 9ae745063..61f820cef 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -39,6 +39,7 @@ from .utils import ( _CArray, _Pointer, + check_element_shape, class_property, ints_to_numpy_buffer, normalize_values, @@ -543,9 +544,8 @@ def build(self, indices, values, *, dup_op=None, clear=False, size=None): 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: @@ -744,7 +744,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) @@ -897,6 +898,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: diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 16e48cef0..30d2a5fa0 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -4508,6 +4508,36 @@ 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 + # 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]") + + def test_power(A): expected = A.dup() for i in range(1, 50): diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 5edebd235..de1d2a53d 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -3075,3 +3075,37 @@ 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) + + +# 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]") From d6b6bbfaa8023715ed52a638c9593ba7734c575a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 21:25:47 -0600 Subject: [PATCH 35/52] Check that imported values have an element for every entry (#643) * Check that imported values have an element for every entry SuiteSparse compares the byte size of an imported values buffer with the number of entries, not with the bytes they need (GB_import.c sets Ax_memsize_for_non_iso to nvals or vlen * vdim), so any buffer of at least nvals bytes got through and was read past its end: Vector.ss.import_full(values=np.ones(3), size=5) built a Vector whose last two values were heap garbage, and so did values too short for a sub-array dtype, as in Vector.ss.import_sparse(size=3, indices=[0, 1], values=np.ones(2), dtype="FP64[2]") Every ss import and pack (sparse, bitmap and full Vectors; csr, csc, hypercsr, hypercsc, bitmapr, bitmapc, fullr and fullc Matrices, and so from_dcsr and from_dcsc) now raises ValueError: values must have 5 or more elements of dtype FP64, not 3 before handing GraphBLAS the buffer. An iso import needs one element. * Note that SuiteSparse:GraphBLAS 10.5.2 checks the size of imported values SuiteSparse:GraphBLAS 10.5.2 fixed the import size check that check_values_size works around: GB_import.c compared the byte size of the values with the number of entries. With check_values_size bypassed, the five short imports tried (import_full, import_sparse of FP64[2], import_bitmapr, and two import_csr) now raise InvalidValue with an empty message on 10.5.2, where 10.5.1 built them from bytes past the end. check_values_size still runs on every version: older ones need it, and its ValueError says what is wrong. Its docstring and the comments of the two test_ss_import_values_size tests now give the versions. Comments only. * Fill one element with a scalar in from_dcsr and from_dcsc on SuiteSparse Matrix.from_dcsr([0, 1], [0, 1, 2], [0, 1], 1.5, "FP64[2]") gave [[1.5, 1.5], [1.5, 1.5]] without SuiteSparse, and on SuiteSparse raised "values must have 1 or more elements of dtype FP64[2], not 0": its SuiteSparse branch imported the scalar as it was, 8 bytes where one element is 16. It now broadcasts a scalar to one whole element and imports it iso, as from_csr and from_csc do; the same for from_dcsc. check_values_size also says what is wrong in two cases it misreported: - A scalar that does not fill one element of a sub-array dtype, as in Vector.ss.import_full(values=1.5, size=3, dtype="FP64[2]", is_iso=True): "a scalar does not fill one element of dtype FP64[2], of shape (2,)" (it said "not 0"). - Index pointers that give more entries than there are indices, as in import_csr(nrows=1, ncols=3, indptr=[0, 2**62], col_indices=[0], values=[1.0]): "index pointers give 4611686018427387904 entries, more than the number of indices, 1" (it asked for that many values). GraphBLAS rejects such an indptr itself, with an InvalidValue and no message, so the check comes first. * Say a scalar does not fill an element only when it is shorter than one import_fullr(values=1.5, nrows=2, ncols=2) said "a scalar does not fill one element of dtype FP64", though 1.5 is a whole FP64 element and the problem is that four are needed. Keep that message for a scalar shorter than one element of a sub-array dtype. Found by review (r47abc). --- graphblas/core/matrix.py | 10 +++++ graphblas/core/ss/matrix.py | 11 +++++ graphblas/core/ss/vector.py | 4 ++ graphblas/core/utils.py | 28 +++++++++++++ graphblas/tests/test_matrix.py | 73 ++++++++++++++++++++++++++++++++++ graphblas/tests/test_vector.py | 27 +++++++++++++ 6 files changed, 153 insertions(+) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 335c1f3e8..e1876c58f 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -1269,11 +1269,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, @@ -1354,11 +1359,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, diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index 9fc535924..443019557 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -17,6 +17,7 @@ _distinct_buffers, _MatrixArray, _Pointer, + check_values_size, get_order, get_shape, ints_to_numpy_buffer, @@ -1353,6 +1354,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, @@ -1557,6 +1560,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, @@ -1793,6 +1798,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, @@ -2032,6 +2038,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, @@ -2244,6 +2251,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, @@ -2451,6 +2459,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, @@ -2632,6 +2641,7 @@ 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 @@ -2804,6 +2814,7 @@ 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 diff --git a/graphblas/core/ss/vector.py b/graphblas/core/ss/vector.py index 9cd4e22d0..8ee259510 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -16,6 +16,7 @@ _CArray, _distinct_buffers, _MatrixArray, + check_values_size, ints_to_numpy_buffer, normalize_chunks, values_to_numpy_buffer, @@ -1016,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, @@ -1218,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, @@ -1393,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, diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index 7da0e0645..0b6b46a28 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -225,6 +225,34 @@ def check_element_shape(values, dtype, ndim): ) +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: diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 30d2a5fa0..9d5d43f52 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -2178,6 +2178,7 @@ def test_ss_import_export(A, do_iso, methods): 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) @@ -4533,11 +4534,83 @@ def test_subarray_dtype_values_shape(): 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_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): diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index de1d2a53d..3e4c9cf40 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -3109,3 +3109,30 @@ def test_subarray_dtype_values_shape(): # 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 From 8aef9c400b4d34ee46d859c90e75a15476fd7a18 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 23:09:20 -0600 Subject: [PATCH 36/52] Fix Scalar infix expressions with a literal operand (#654) * Keep a literal on its own side of a Scalar's `&` and `|` Behavior fix. Since 2023.1.0, `x & s` and `x | s`, with a literal x and a Scalar s, built the expression with s on the left and x on the right, so the op applied to it ran with its operands swapped, and every op that does not commute gave wrong values, silently: s = Scalar.from_value(5) binary.minus(1 & s) # was 4, now -4 binary.truediv(1 | s) # was 5.0, now 0.2 binary.first(1 & s) # was 5, now 1 binary.lt(1 & s) # was False, now True So did rminus, cdiv, rdiv, pow, the other comparisons and is-comparisons, and the rest of the ops that do not commute. So did the defaults of an ewise_union, each of which went with the other operand: with an empty Scalar e, `binary.minus(1 | e, left_default=0, right_default=10)` was -1 and is now -9. And so did chains: `binary.minus(1 & (2 & s))` was 2 and is now 4. The expression now keeps the literal where it was written, as it is, and the op applied to it types it: Scalar's ewise_mult, ewise_add and ewise_union take a literal on the left as they take one on the right, typed beside the Scalar by literals._literal_operand. A Python number is weak, as beside `-`: `binary.minus(2 & int8_s)` is INT8 and `binary.minus(0.5 & int8_s)` is FP64, as `2 - int8_s` and `0.5 - int8_s` are. Beside a UDT, a lifted op such as minus types the literal as numpy would and `first` makes it an element of the UDT. How a literal is typed depends on the op (300 beside INT8 is an OverflowError for `-` and is compared as it is by `<`), so the literal is not converted when `&` or `|` builds the expression. Converting it then, with its own dtype, would get the typing wrong: the 2 of `int8_s & 2` would be INT64 under every op, where `int8_s - 2` keeps it INT8, and `0.5 & int8_udt_s` would cut 0.5 to 0 in the UDT. `op(x & y)` now asks the expression to call its left operand's method (InfixExprBase._call_left), so a Scalar expression can pass a literal there. Two more errors with a literal operand are fixed: - Computing a BOOL expression with a literal by itself, as `(s & True).new()` and `(False | s).new()` do, raised AttributeError: 'bool' object has no attribute 'dtype'. It now types the literal for LAND or LOR and computes, or raises the usual TypeError for a dtype other than BOOL. - `s & object()` gave an expression that failed only when an op was applied; it now raises TypeError at once, as `s + object()` does. The check types the literal as `==` would and drops it. Beside an infix expression, as in `(s & t) & object()`, it raises when the op is applied. * Apply `lit - s` and the like as minus(lit, s), not as rminus(s, lit) Behavior fix. Since 2023.1.0, a Scalar's reflected operators (`lit - s`, `lit / s`, `lit ** s` and so on, with a literal on the left) computed the op's commutes_to with the Scalar first, rminus(s, lit). That fails where no such op exists: 1 - udt_s, 1 / udt_s, 1 // udt_s, (1, 2) - udt_s # KeyError: 'rminus does not work with ...'; no reflected op # takes a UDT, though `udt_s - 1` and `1 - udt_vector` work 7 % s # TypeError for every dtype: numpy.mod has no reflected op 2 ** s # TypeError without Numba: rpow is a Numba function Now the literal becomes a Scalar on the left, typed beside s as in `binary.minus(1 & s)`, and the op itself is applied: `1 - udt_s` is [-4, -6] for udt_s = [5, 7], and `7 % int8_s` is 2 for 5. `lit ** s` now gives what `binary.pow(lit, s)` and `lit ** vector` give. That changes four kinds of result, where rpow, a Numba function, differed: - an integer power out of range saturates, as GraphBLAS's pow does, where rpow wrapped: `3 ** int8_s` for 5 was -13 and is now 127, and `0 ** int8_s` for -3 was 0 and is now 127; - `False ** bool_s` is BOOL, not INT8; - `0 ** complex_s` is 0, not nan (for 1+2j), and other complex powers may differ in the last bit; - a Python int beyond INT64 beside a BOOL Scalar, as in `(2**64 - 1) ** bool_s`, raises OverflowError, as `bool_s ** (2**64 - 1)` and `(2**64 - 1) ** bool_vector` do, where it gave a UINT64. Nothing else changes: for `+ - * / // % ** ^`, nine Scalars (INT8, INT64, UINT8, FP32, FP64, BOOL, empty) and eight literals, every result the previous code gave is the same but those of `**`; and for the seven that are not empty, `lit OP s` now matches `binary.op(lit, s)` for each of `+ - * / // % **` (it did not in 66 of 392 cases). * Show a Scalar's `&` and `|` with a literal or a non-BOOL dtype Behavior fix. Since 2023.1.0, showing a Scalar `&` or `|` expression raised in two cases: repr(s & 1), repr(1 | s), repr(t & True), and their _repr_html_ # AttributeError: 'int' object has no attribute 'name' repr(s & s) # s not BOOL, with autocompute on (the default) # TypeError: Bad dtypes for `x & y`! Automatic computation ... A literal operand now shows its value, as an unnamed Scalar does in an expression, and its own type where a dtype goes, since it has no dtype until an op types it: gb.ScalarEwiseMultExpr left_dtype right_dtype s_1 & 1 INT8 int The text repr now computes the expression to show its result only where the HTML repr, and the reprs of Vector and Matrix infix expressions, do: for `x @ y`, and for `x & y` and `x | y` of BOOL operands. A literal counts as BOOL when it is a bool, so `t & True` shows its result. A chain such as `(s & t) & u` still raises in a repr, as it does for Vector and Matrix. * Raise DimensionMismatch, not TypeError, for a size mismatch beside an infix expression A behavior fix. When an operand of `&`, `|` or `@` was itself an infix expression and the sizes did not match, building the error crashed in _mismatched, which called _dummy(obj) while _dummy took a second parameter, obj_type, that it never used (since #517). Before: >>> (v & v) & w # v of size 3, w of size 4 TypeError: _dummy() missing 1 required positional argument: 'obj_type' After: >>> (v & v) & w DimensionMismatch: ... The same held with the expression on the right (w & (v & v)), for `|`, for Matrix operands ((A & A) & B), and for `@` ((A @ A) @ w). _dummy now takes only obj. Tests cover Vector and Matrix with the expression on either side, for `|`, `&` and `@`. * Let numpy's operators defer to graphblas objects Behavior change. A numpy scalar or 0-d array on the left of an infix operator was computed by numpy, which converted a Scalar with __array__ and could not convert a Vector or Matrix. graphblas objects now set __array_priority__, so numpy's operators return NotImplemented and Python calls the graphblas reflected method, which takes the numpy value as a literal operand, typed as it is on the right: s = Scalar.from_value(5, dtypes.INT8) np.int16(2) - s # was np.int16(-3), now an INT16 ScalarExpression np.int16(2) < s # was np.True_, now a BOOL ScalarExpression np.int16(5) == s # was np.True_, now True (Scalar's ==) np.int16(2) - v # was a TypeError, now an INT16 VectorExpression np.int16(2) == A # was a TypeError, now a BOOL MatrixExpression The same holds for expressions, A.T and index expressions. `&`, `|` and `^` now go to graphblas too: `np.int16(2) & s` was numpy's bitwise and, np.int16(0), and is now the TypeError of `2 & s` (only BOOL computes). With an empty Scalar, numpy computed with nan; the result is now what a Python float on the left gives: e = Scalar(float) np.float64(5) + e # was nan, now 5.0, as 5.0 + e is (and `-`) np.float64(5) * e # was nan, now an empty Scalar, as 5.0 * e is np.float64(5) < e # was np.False_, now an empty Scalar np.float64(5) == e # was np.False_, now False What numpy computed and graphblas does not now raises TypeError: np.array([1, 2]) - s # was array([-4, -3]) arr -= s # was array([-4, -3]), in place np.int16(2) << s # was np.int16(64); graphblas has no `<<` On numpy < 2, where a size-1 array beside a Scalar is a literal (on the right too), `np.array([1]) - s` is now an INT64 ScalarExpression. Ufuncs still use __array__: np.sin(s), np.subtract(np.int16(2), s) and np.asarray(s) give numpy values as before. * Type a literal on the left of a Scalar op as the op's left operand _literal_scalar types a literal on either side of a Scalar op, as in op(1 & s) or 1 - s, through _literal_operand, which now takes left=True for a left operand: a typed op takes each operand in its own input type ("Type a literal beside a typed op in the op's own input type", in #591). _literal_scalar did not pass it, so with an op typed for two different types, a literal on the left was typed as the right operand. Every left-side caller passes argname="left", so _literal_scalar passes left=argname == "left". An untyped op, as every infix operator is, types the literal the same either way. Behavior change (before -> after), with A a record of two float64 fields, s = Scalar.from_value((3.0, 4.0), A) and add_a a UDT BinaryOp of x + y.ll_a: - add_a[INT64, A](2 & s): GraphBLAS's DomainMismatch (2 became an A element where the op takes an INT64) -> 5.0. * Tidy the helpers for a literal operand of a Scalar infix expression No behavior change. - ScalarInfixExpr._call_left is a method that falls back to InfixExprBase._call_left, in place of a module function assigned on ScalarEwiseAddExpr and ScalarEwiseMultExpr. - _typed_operands asks _is_literal of each operand, without a check for a Scalar expression first: a Vector or Matrix infix expression never holds a literal. - formatting._is_bool, which had one caller, is folded into _computes_itself. - call_op names the method for an error only where a literal is on the left. --- docs/user_guide/types.rst | 4 +- graphblas/core/base.py | 7 + graphblas/core/expr.py | 6 + graphblas/core/formatting.py | 41 ++++-- graphblas/core/infix.py | 96 ++++++++++++-- graphblas/core/infixmethods.py | 20 ++- graphblas/core/matrix.py | 1 + graphblas/core/operator/base.py | 4 +- graphblas/core/operator/binary.py | 4 +- graphblas/core/scalar.py | 67 +++++----- graphblas/tests/test_formatting.py | 58 ++++++++ graphblas/tests/test_infix.py | 205 ++++++++++++++++++++++++++++- graphblas/tests/test_scalar.py | 41 ++++++ 13 files changed, 473 insertions(+), 81 deletions(-) diff --git a/docs/user_guide/types.rst b/docs/user_guide/types.rst index 65b941cf9..2a0649e38 100644 --- a/docs/user_guide/types.rst +++ b/docs/user_guide/types.rst @@ -41,7 +41,9 @@ 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. python-graphblas applies this rule itself, +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`). diff --git a/graphblas/core/base.py b/graphblas/core/base.py index a3562f922..76a2ae3cc 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -190,6 +190,12 @@ 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 __call__( self, @@ -577,6 +583,7 @@ class BaseExpression: ) output_type = None _is_scalar = False + __array_priority__ = BaseType.__array_priority__ # numpy's operators defer to ours def __init__( self, diff --git a/graphblas/core/expr.py b/graphblas/core/expr.py index 47ff18a0b..236e29534 100644 --- a/graphblas/core/expr.py +++ b/graphblas/core/expr.py @@ -290,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 @@ -484,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 @@ -511,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 5fe9b6972..471b264d3 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 @@ -717,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) @@ -740,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}" @@ -764,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 73b3b91b5..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,21 @@ 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 @@ -485,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="") @@ -498,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) @@ -530,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 53f9d76f0..dfe539468 100644 --- a/graphblas/core/infixmethods.py +++ b/graphblas/core/infixmethods.py @@ -2,7 +2,7 @@ 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 @@ -32,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): diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index e1876c58f..23d65f4fa 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -3893,6 +3893,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" diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index aea4a6f51..b9cd34e52 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -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" diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 491a3e0ac..b2860b3d7 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -203,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 diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index a7cf2189b..5f1ac7c01 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -689,20 +689,11 @@ 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: - other, dtype = _literal_operand(self.dtype, other, op) - 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): @@ -763,20 +754,11 @@ 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: - other, dtype = _literal_operand(self.dtype, other, op) - 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") expr = ScalarExpression( @@ -841,22 +823,13 @@ 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: - other, literal_dtype = _literal_operand(right_dtype, other, op) - try: - other = Scalar.from_value(other, literal_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. @@ -1208,6 +1181,28 @@ def _as_scalar(scalar, dtype=None, *, is_cscalar): return scalar +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``. diff --git a/graphblas/tests/test_formatting.py b/graphblas/tests/test_formatting.py index 2b412ac1d..75393afa0 100644 --- a/graphblas/tests/test_formatting.py +++ b/graphblas/tests/test_formatting.py @@ -4010,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)") diff --git a/graphblas/tests/test_infix.py b/graphblas/tests/test_infix.py index 94d2555ae..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): @@ -763,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_scalar.py b/graphblas/tests/test_scalar.py index 3ade2fb25..4bd5be4ec 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -905,3 +905,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 From 526a64262b4f91da1ba9234190e882380a4553fc Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 23:10:02 -0600 Subject: [PATCH 37/52] Raise clearer errors for constructor, mask, and attribute misuse (#606) * Raise a helpful TypeError when data is passed to Vector()/Matrix() Vector([1, 2, 3]) previously failed with "Unknown dtype: [1, 2, 3] of type ". The constructors take a dtype first, so passing data is a common newcomer mistake; now list/tuple/ndarray first arguments that fail dtype lookup raise a TypeError pointing to from_coo/from_dense. Valid list/tuple dtype specs (structured and subarray dtypes) still work: the hint only engages after lookup_dtype has rejected the argument. Other bad dtypes keep the original ValueError. * Raise a clear TypeError for a wrong-kind mask on a Matrix operation A Vector mask on a full-Matrix operation (ewise, extract, .new) leaked a raw cffi error, "initializer for ctype 'struct GB_Matrix_opaque'", which says nothing about masks. The Vector-output direction was already guarded; _check_mask now guards the Matrix-output direction the same way and reports "Mask object must be type Matrix; got ...". The check is gated behind a strict_kind flag rather than applied everywhere, because assignment is a genuine exception: a Vector mask on a Matrix row/column assign is valid, and Matrix.__setitem__ validates it separately. So strict_kind is False for __setitem__ updates and True for .new() and other full-tensor operations. * Hint the right API when .new or .transpose is accessed as an attribute Two common misuses failed with a bare "no attribute" message that gave no direction: A.new() is the expression resolver, not an instance method A.transpose() transpose is the .T property, not a method A __getattr__ on BaseType supplies the hint on the attribute miss. It raises AttributeError, exactly as before, so hasattr() and getattr() with a default keep answering the way they did and duck-typing probes are unaffected; only the message changes. Names with no hint keep the plain message. .new and .transpose ride in one commit because they are one mechanism: a single __getattr__ consulting a single hint table. Splitting them would mean adding the same function twice. __getattr__ fires only on a genuine attribute miss, since slots and methods resolve first, so the attribute hot path is untouched. There is deliberately no class-level hint (Matrix.new): that needs a metaclass, and both candidates break something real. A plain type metaclass makes mixing Matrix into any abc-based class die with "metaclass conflict", and an ABCMeta-derived one leaks register, __abstractmethods__ and _abc_impl into dir(Matrix), which the expression-surface guard rightly reports as drift. Class-level misuse keeps Python's default AttributeError, and test_abc_mixin_subclass pins the abc composition. * Instantiate the Matrix subclass with an abc mixin in its test test_abc_mixin_subclass checked only issubclass, so the subclass's __len__ never ran. Make an instance and check that it is a Sized whose len is its nvals, 0 and then 2 after two stores (a __len__ that returned a constant would fail): the subclass must work as a Matrix and as the abc, not just be declared. * Show the repr of a Matrix subclass Behavior fix; main has it. The repr's info row picks the storage format with get_format(matrix) for a Matrix and from matrix._matrix for a TransposedMatrix, and it chose with type(matrix) is Matrix. Any subclass of Matrix took the TransposedMatrix branch, so on the suitesparse backend repr() and _repr_html_() raised AttributeError: 'Sub' object has no attribute '_matrix', for a plain subclass as for one with an abc mixin, and so did the reprs of its masks (A.S, ~A.V). Choose with _is_transposed, which Matrix and TransposedMatrix both define. The test checks that the repr and HTML repr of a subclass and of its structural and complemented value masks are a Matrix's but for the type's name, and that its .T is a TransposedMatrix whose reprs match a Matrix's .T; it fails before this with the AttributeError. * Find the error text for a failing call on a Matrix or Vector subclass Behavior fix; main has it. check_status takes an error's text from GrB__error, with the type's own name, and a subclass of Matrix or Vector has no such function. Any failing call on one, for example on a subclass instance whose data was exported with give_ownership=True, raised the right exception with check_status's fallback text, "Unable to get the error string for type SubMatrix. This is most likely a bug in graphblas. Please report this as an issue". A class without an error function of its own now uses that of the nearest base class in its MRO that has one. Only classes in graphblas.core count, which is where all eleven such types are: a user's mixin may share a GraphBLAS type's name, and listed first (a class Context mixed into a Vector) it would hand GxB_Context_error a GrB_Vector, which cffi rejects with a raw TypeError. A class with an error function of its own keeps it, and TransposedMatrix, which does not derive from Matrix, keeps the fallback. Two of graphblas's own classes change besides subclasses, for the better: BuiltinDescriptor and GlobalContext now get their C types' error text (Descriptor, Context) instead of the fallback. The test makes a failing call on a given-away subclass of Matrix, of Vector, and of Vector with such a Context mixin listed first, and checks the text is the base class's; it fails before this with the fallback text, and the mixin case with the TypeError without the graphblas.core filter. * Report an error on a TransposedMatrix against its Matrix Behavior fix; main has it. A TransposedMatrix shares Matrix.wait and Matrix._nvals, so when one of those calls failed, check_status was handed the TransposedMatrix and looked for GrB_TransposedMatrix_error, which GraphBLAS does not have. On a Matrix whose data was exported with give_ownership=True, A.T.wait() and repr(A.T) raised NullPointer with check_status's fallback text, "Unable to get the error string for type TransposedMatrix. This is most likely a bug in graphblas. Please report this as an issue", the same kind of message the previous commit fixed for subclasses. TransposedMatrix now gives its Matrix as _exc_arg, which check_status already uses in place of the object it is handed (as for _MatrixArray), so the error and its text are the Matrix's. The test checks that both calls raise NullPointer with A.wait()'s text; it fails before this with the fallback text. * Exclude the error-type walk's unmatched arms from coverage With a TransposedMatrix now reporting against its Matrix, every object a failing call can hand check_status is a Matrix, Vector, Scalar, Descriptor or Context, or a subclass of one, so _error_type_name always finds a base with an error function. A graphblas.core class without one is passed over only for BuiltinDescriptor and GlobalContext, whose calls cannot fail: the global context engages and disengages any number of times, and builtin descriptors are never freed. Mark the skip and the no-match return "defensive", and say why in a comment. * Keep a raising property's own message, and the dtype error's cause A property that raises AttributeError sends Python to __getattr__, whose generic "no attribute" message replaced the property's own: v.ss on the vanilla backend lost "ss attribute is only available with "suitesparse" backend", and a float Scalar's __index__ lost "only has __index__ for integral dtypes". Look the name up again after the hints, so the property raises its own message. The constructor hint for data passed as the dtype also fired on a malformed dtype spec such as Vector([("a", "zzz")], 3) and dropped the lookup's reason with "from None". Keep it as the cause, and word the hint as "If you meant to pass values". Found by review (r606s). --- graphblas/core/base.py | 51 +++++++++- graphblas/core/dtypes.py | 22 +++++ graphblas/core/formatting.py | 2 +- graphblas/core/matrix.py | 11 ++- graphblas/core/vector.py | 6 +- graphblas/exceptions.py | 25 ++++- graphblas/tests/test_matrix.py | 168 +++++++++++++++++++++++++++++++++ graphblas/tests/test_scalar.py | 2 + graphblas/tests/test_vector.py | 37 ++++++++ 9 files changed, 315 insertions(+), 9 deletions(-) diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 76a2ae3cc..b1347d355 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -168,7 +168,7 @@ def _expect_op(self, op, values, *, within, **kwargs): AmbiguousAssignOrExtract._expect_type = _expect_type -def _check_mask(mask, output=None): +def _check_mask(mask, output=None, strict_kind=False): if not isinstance(mask, Mask): # Convert bool objects to value masks if output_type(mask).__name__ in {"Vector", "Matrix"}: @@ -180,11 +180,34 @@ def _check_mask(mask, output=None): 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)}") + 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 +# 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: # pylint: disable=assigning-non-slot __slots__ = "gb_obj", "dtype", "name", "__weakref__" @@ -197,6 +220,18 @@ class BaseType: # 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, *optional_mask_accum_replace, @@ -495,7 +530,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 @@ -651,7 +690,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 484dcf354..2a85f46a1 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -612,6 +612,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. diff --git a/graphblas/core/formatting.py b/graphblas/core/formatting.py index 471b264d3..d4ac8cfe4 100644 --- a/graphblas/core/formatting.py +++ b/graphblas/core/formatting.py @@ -379,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)) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 23d65f4fa..bd4f1d8d7 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -9,6 +9,7 @@ from . import _supports_udfs, automethods, ffi, lib, utils from .base import BaseExpression, BaseType, _check_mask, 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 .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype from .mask import Mask, StructuralMask, ValueMask @@ -218,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 @@ -3954,6 +3958,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" diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 61f820cef..bb9bea32a 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -8,6 +8,7 @@ from . import _supports_udfs, automethods, ffi, lib, utils from .base import BaseExpression, BaseType, _check_mask, 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 .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype from .mask import Mask, StructuralMask, ValueMask @@ -168,7 +169,10 @@ class Vector(BaseType): 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*") diff --git a/graphblas/exceptions.py b/graphblas/exceptions.py index df674b55e..2c034c212 100644 --- a/graphblas/exceptions.py +++ b/graphblas/exceptions.py @@ -162,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/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 9d5d43f52..7ae48ae72 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -3018,6 +3018,7 @@ def test_expr_is_like_matrix(A): "__call__", "__del__", "__delitem__", + "__getattr__", "__lshift__", "__setitem__", "_assign_element", @@ -3084,6 +3085,7 @@ def test_index_expr_is_like_matrix(A): expected = { "__del__", "__delitem__", + "__getattr__", "__setitem__", "_assign_element", "_delete_element", @@ -4780,3 +4782,169 @@ 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_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_scalar.py b/graphblas/tests/test_scalar.py index 4bd5be4ec..2d5974032 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -542,6 +542,7 @@ def test_expr_is_like_scalar(s): expected = { "__call__", "__del__", + "__getattr__", "__imatmul__", "__lshift__", "_carg", @@ -586,6 +587,7 @@ def test_index_expr_is_like_scalar(s): # Should we make any of these raise informative errors? expected = { "__del__", + "__getattr__", "__imatmul__", "_carg", "_deserialize", diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 3e4c9cf40..541a892f9 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -2033,6 +2033,7 @@ def test_expr_is_like_vector(v): "__call__", "__del__", "__delitem__", + "__getattr__", "__lshift__", "__setitem__", "_assign_element", @@ -2083,6 +2084,7 @@ def test_index_expr_is_like_vector(v): expected = { "__del__", "__delitem__", + "__getattr__", "__setitem__", "_assign_element", "_delete_element", @@ -3077,6 +3079,41 @@ def test_subarray_dtypes(): 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(): From 306176a5fafd61374eedfcbe4fef520ec049ec96 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 23:10:43 -0600 Subject: [PATCH 38/52] Make config attribute assignment raise instead of silently no-op (#608) Measured before this change: `graphblas.config.autocompute = False` left `config["autocompute"]` at True and added a dead `autocompute` entry to `vars(config)`. donfig.Config defines no __setattr__ for options, so the write landed on the instance and the real option never moved. The mistake is easy to make and leaves no trace. graphblas.config is now an instance of a donfig.Config subclass whose __setattr__ rejects writes that are not donfig's own. For a known option the AttributeError names the canonical idioms (config.set(name=value), optionally as a `with` block, and config[name] to read). For an unknown name it says so and lists the known options, deliberately WITHOUT advising config.set(): donfig's set() accepts arbitrary keys, so pointing there would trade a loud error for a silently-created bogus option, the same class of trap this change exists to close. The allowlist of writable attributes is derived rather than hardcoded: the subclass snapshots vars(self) once donfig's __init__ returns, so a donfig release that adds or renames an internal field cannot break attribute access here. Today that snapshot is name, env_prefix, env, main_path, paths, defaults, deprecations, config, and config_lock. An instance built without running __init__ has no allowlist yet and stays permissive, so the guard cannot raise from a half-built object. What this does not change: reading options, config.set() as a call or as a context manager, config[name], get(), update(), refresh(), to_dict(), and writes to donfig's own attributes all behave as before. Item assignment (config[name] = value) was never supported by donfig and still raises TypeError. Under the pinned config (--backend suitesparse --blocking --no-mapnumpy), the suite goes from 1073 to 1075 passed by the two tests added here; 141 skipped is unchanged. --- graphblas/__init__.py | 40 +++++++++++++++++++++++++++++++- graphblas/tests/test_core.py | 45 ++++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 1 deletion(-) diff --git a/graphblas/__init__.py b/graphblas/__init__.py index 2de45dda9..e5038563e 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) diff --git a/graphblas/tests/test_core.py b/graphblas/tests/test_core.py index bfad8279d..c324c7d7d 100644 --- a/graphblas/tests/test_core.py +++ b/graphblas/tests/test_core.py @@ -115,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 From a6cc5a2ed96dbd333e1b1dc5b01d0b75789f3838 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 23:12:34 -0600 Subject: [PATCH 39/52] Reject building a user monoid from a built-in BinaryOp (#617) --- graphblas/core/operator/monoid.py | 28 ++++++++++++++++++++++++++++ graphblas/tests/test_op.py | 18 ++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/graphblas/core/operator/monoid.py b/graphblas/core/operator/monoid.py index 72765e2c2..bcf06d323 100644 --- a/graphblas/core/operator/monoid.py +++ b/graphblas/core/operator/monoid.py @@ -29,6 +29,16 @@ _BUILTIN_UDT_MONOIDS = {"plus", "times", "min", "max"} +def _is_builtin_binaryop(binaryop): + """True if ``binaryop`` wraps built-in GraphBLAS operators rather than a UDF. + + Built-ins are the only binary ops 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. + """ + return binaryop.orig_func is None and binaryop._numba_func is None + + def _scalar_identity(monoid_name, scalar_dtype): """Return the identity for a single numeric (or bool) numpy dtype, or ``None``. @@ -262,6 +272,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( diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 3c72e36c2..c76c895df 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -422,6 +422,24 @@ 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") + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_semiring_parameterized(): From ed1bacb88bfb699952d9855cbd0b25bd37a33b38 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 23:13:11 -0600 Subject: [PATCH 40/52] Extend the monoid guard to built-ins that carry a UDT fallback func (#619) * Extend the monoid guard to built-ins that carry a UDT fallback func The guard added with the monoid-from-built-in rejection keyed on "constructed with neither a Python func nor a numba func", which is true for most built-ins but not all of them: binary.any, first, pair, and second carry a Python orig_func as their UDT fallback while their typed ops are still TypedBuiltinBinaryOp. All four sailed past the guard and reproduced exactly what it exists to prevent: an opaque AttributeError partway through registration, with binary.._monoid left pointing at the half-built monoid (binary.any.monoid is monoid.any flipped from True to False after the rejected call). binary.any is even named in the guard's own rationale as an op whose identity SuiteSparse ignores. Key the probe on the typed ops as well: any TypedBuiltinBinaryOp among them marks the op built-in. The two signals agree everywhere else, including under mapnumpy=True, where the mapped numpy binary ops hold the built-in's typed ops and no funcs, and are flagged by both tests alike; the numpy monoids themselves are mapped without passing through Monoid._build, so their registration is unaffected (verified under both mapnumpy settings). A UDF's typed ops are never TypedBuiltinBinaryOp, so nothing legitimate is newly rejected. * Give a timedelta64 field the same KeyError as a datetime64 one in min/max Behavior fix. monoid.min and monoid.max generate a UDT's identity field by field, and _scalar_identity picked the integer arm with np.issubdtype(dtype, np.integer). numpy files timedelta64 under np.integer, so a record with a timedelta64 field reached np.iinfo, which rejects it: rec = dtypes.register_anonymous(np.dtype([("n", np.int64), ("t", "m8[s]")])) monoid.min[rec] # before: ValueError: Invalid integer data type 'm'. rec in agg.min # before: the same ValueError, not False A datetime64 field already got "KeyError: monoid.min does not work with ...: field 't' has unsupported dtype datetime64[s]", and `in` answered False. Test the dtype's kind instead, so a timedelta64 field takes that path too: monoid.min[rec] # after: KeyError ... field 't' has unsupported dtype timedelta64[s] rec in agg.min # after: False Real identities for timedelta64 fields (its own max and min, with NaT in mind) are a separate enhancement. * Say why a UDT monoid fails when its op returns another type Behavior fix. Monoid._compile_udt handed a binary op whose return type differs from its input type straight to GrB_Monoid_new_UDT, which rejects it with a DomainMismatch that has no message, also from `in`: op = BinaryOp.register_anonymous(lambda x, y: x["a"] + y["a"], is_udt=True) m = Monoid.register_anonymous(op, 0) m[rec] # before: DomainMismatch (empty message) rec in m # before: the same DomainMismatch Check the compiled op's return type first and raise a TypeError, the error the neighbouring check (two different input types) raises: m[rec] # after: TypeError: monoid. does not work with rec: a # monoid's binary operator must return its input type, rec, not INT64 rec in m # after: False * Keep is_idempotent through Monoid.register_new(lazy=True) Behavior fix. The lazy branch stored register_new's arguments for the delayed call but left is_idempotent out, so the monoid built on first access was never idempotent: Monoid.register_new("m", op, 0, is_idempotent=True, lazy=True) monoid.m.is_idempotent # before: False # after: True The eager branch already passed it. The test registers lazy_idem in the binary and monoid namespaces, so test_operator_types ignores it like the other names test_lazy_op makes. * Generate a UDT identity for the built-in monoids only, not by name Behavior fix. Monoid._compile_udt generated a field-wise identity whenever the monoid had no identity and its name was plus, times, min or max. A user monoid gets its name from its BinaryOp, which gets it from the function, so a user op called `min` (or a monoid given name="plus") silently took the built-in's identity and gave wrong results. With an op that takes the field-wise max and identity None: op = BinaryOp.register_anonymous(fieldwise_max, "min", is_udt=True) v = [(1, 2), (5, -3)]; w = [(-1, -2), (-5, -3)] Monoid.register_anonymous(op, None) # name "min" before: identity (INT64_MAX, INT64_MAX); v.reduce -> (INT64_MAX, INT64_MAX) Monoid.register_anonymous(op, None, "plus") before: identity (0, 0); v.reduce -> (5, 2), w.reduce -> (0, 0) both, after: KeyError: monoid. does not work with ...: no identity value Generate the identity only for a built-in monoid, the ones with no _binaryop of their own. A user monoid with an identity is unaffected. * Match the UDT monoid error by the record UDT's current name test_udt_monoid_op_must_return_its_input_type matched the error against "_BuiltinOpsRec", the name the record_udt fixture registers. But register_anonymous returns the one DataType for an np.dtype and gives it the newest name, and test_udt_op_pickle registers the same record dtype as "_PickleUdt". In random order, with a record_udt test, then test_udt_op_pickle, then this test, the message said "_PickleUdt" and the match failed (wave 2 CI: 13 jobs, e.g. 122634118659 and 122634118406, on both backends and with mapnumpy on or off). Match the fixture's name as it is when the error is raised. --- graphblas/core/operator/monoid.py | 43 +++++++++--- graphblas/tests/test_op.py | 92 ++++++++++++++++++++++++++ graphblas/tests/test_operator_types.py | 1 + 3 files changed, 128 insertions(+), 8 deletions(-) diff --git a/graphblas/core/operator/monoid.py b/graphblas/core/operator/monoid.py index bcf06d323..e6d2ce07e 100644 --- a/graphblas/core/operator/monoid.py +++ b/graphblas/core/operator/monoid.py @@ -32,11 +32,22 @@ def _is_builtin_binaryop(binaryop): """True if ``binaryop`` wraps built-in GraphBLAS operators rather than a UDF. - Built-ins are the only binary ops 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. + 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 + 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): @@ -53,7 +64,9 @@ def _scalar_identity(monoid_name, scalar_dtype): return scalar_dtype.type(1) if monoid_name not in {"min", "max"}: return None - if np.issubdtype(scalar_dtype, np.integer): + # 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): @@ -340,9 +353,18 @@ def _compile_udt(self, dtype, dtype2): from ..scalar import Scalar ret_type = binaryop.return_type + 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.name in _BUILTIN_UDT_MONOIDS: - # Auto-generate the identity for built-in monoids on UDTs. + 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( @@ -439,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) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index c76c895df..594b196a6 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -440,6 +440,24 @@ def test_monoid_rejects_builtin_binaryop(): 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(): @@ -1188,6 +1206,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) @@ -1981,6 +2002,27 @@ def _add_udt(x, y): 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): @@ -2456,6 +2498,56 @@ def test_udt_auto_monoid(): 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): diff --git a/graphblas/tests/test_operator_types.py b/graphblas/tests/test_operator_types.py index 8c250bc25..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", From bf5eec34ec83521bb8d849e2bc46c5c5e999c7a8 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 23:14:09 -0600 Subject: [PATCH 41/52] Stop building floordiv at import time when seeding UDT auto-lift (#610) * Stop building floordiv at import time when seeding UDT auto-lift The first access to any operator namespace paid a numba compile it did not need. BinaryOp._initialize ended with a loop over _BUILTIN_UDT_BINARY_OPS that seeded _udt_types/_udt_ops/_custom_dtype via getattr(binary, op_name). floordiv is the one name in that set registered with lazy=True, so the getattr materialized it, compiling every dtype signature on every process start. Seed those three attributes in __init__ instead. That covers every path that creates a BinaryOp (builtin enumeration, the specials in _initialize, delayed UDF materialization, anonymous), so the seeding is not lost, and laziness survives: floordiv still compiles on first use and UDT lift behavior is unchanged. Measured here, fresh process, first `binary.plus` access: ~648 ms before, ~152 ms after (medians of 3 and 5 runs). Treat these as indicative, not benchmark-grade; load average was ~4 on a shared machine throughout, so the absolute numbers move but the ~4x gap does not. This does not make imports free. The ~152 ms that remains is the rest of _initialize plus namespace setup and is untouched here. It also does not change what floordiv costs once you use it; the compile is deferred, not removed. The other four lazily-registered UDFs (rfloordiv, absfirst, abssecond, rpow) were never in _BUILTIN_UDT_BINARY_OPS and so were never force-built by this loop. test_initialize_does_not_build_lazy_udfs covers the invariant in a subprocess, since _initialize has already run by the time any in-process test executes. * Lift only the built-in ops to UDTs, not an anonymous op of the same name Seeding the UDT lift in BinaryOp.__init__ keyed on the name alone, so BinaryOp.register_anonymous(f, "plus") was lifted field by field like binary.plus, and a record operand gave a UdfParseError from inside the lift instead of the KeyError a user's op gives. The old loop in _initialize seeded only the real binary. objects; skip anonymous ops to match it. --- graphblas/core/operator/binary.py | 23 ++++++++----- graphblas/tests/test_op.py | 56 +++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 9 deletions(-) diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index b2860b3d7..ace225d46 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -1053,15 +1053,12 @@ def _initialize(cls): binary.eq._udt_ops = {} binary.ne._udt_types = {} binary.ne._udt_ops = {} - # Element-wise arithmetic ops are auto-generated from per-field / - # per-element scalar ops. - if _has_numba: - for op_name in _BUILTIN_UDT_BINARY_OPS: - binop = getattr(binary, op_name, None) - if binop is not None: - binop._udt_types = {} - binop._udt_ops = {} - binop._custom_dtype = _udt_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__( @@ -1086,6 +1083,14 @@ def __init__( if is_udt: self._udt_types = {} # {(dtype, dtype): DataType} self._udt_ops = {} # {(dtype, dtype): TypedUserBinaryOp} + 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 diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 594b196a6..0de6ca3e1 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1,5 +1,8 @@ import itertools +import os +import subprocess import sys +from pathlib import Path import numpy as np import pytest @@ -6926,3 +6929,56 @@ def test_compile_codegen_helper(): assert "Source:" in msg assert bad_src in msg assert isinstance(exc_info.value.__cause__, SyntaxError) + + +# 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 From 00ec75bf330e3a3fb4b12784a8a5162837c70fb0 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Nov 2026 23:41:50 -0600 Subject: [PATCH 42/52] Persist the numba compile of the built-in UDF binops to disk (#611) * Persist the numba compile of the built-in UDF binops to disk floordiv, rfloordiv, absfirst, abssecond, and rpow are UDFs, so every process that touched one paid its numba compile again. They are plain module-level functions, which is exactly the case numba can cache, so that cost belongs once per machine rather than once per process. Thread the flag through as a private keyword-only `_cache` on register_new into _build, where it reaches numba.njit alongside the existing error_model. Keeping it private avoids widening the public API for something no caller outside this module should set. Ops registered by users stay uncached: numba keys a cache entry on a stable on-disk source, which lambdas and interactively defined functions do not have. Measured here, fresh process, medians of 3, load average ~4 on a shared machine so read these as indicative. First-touch floordiv alone: 434 ms uncached, 347 ms warm (-20%). First-touch of all five: 1337 ms uncached, 733 ms warm (-45%). Writing a cold cache costs about 107 ms once for all five (1444 ms on the run that populates it). What stays: the remaining ~733 ms is numba/LLVM startup and per-process object linking, which no on-disk cache removes. In-process dispatch once an op is built is untouched. The first process on a machine is slightly slower, not faster. Read-only installs are safe, verified rather than assumed. With graphblas/core/operator and its __pycache__ chmod'd a-w, the run wrote zero files in-tree and 66 .nbi/.nbc files under ~/Library/Caches/numba, with empty stderr under `python -W always`. The NumbaWarning in numba/core/caching.py is a source-content check, not a filesystem one. docs/user_guide/udf.rst gains a "Compilation caching" section covering the cache location, the read-only fallback, and why user ops are excluded. * Exclude a test UDF's body from coverage test_builtin_udfs_are_disk_cached registers _uncached_probe only to check that Numba gives it no disk cache, so its body never runs in Python. Mark it as the file marks other UDF bodies. * Compile the built-in UDF binops uncached when Numba cannot cache With a read-only install and no writable per-user cache directory either (a read-only container or HPC image run by a user without a writable HOME), numba.njit(cache=True) raises RuntimeError: "cannot cache function '_floordiv': no locator available". floordiv, rfloordiv, absfirst, abssecond and rpow then failed outright, and agg.L1norm reported a misleading "Bad semiring string: 'plus_absfirst'". Fall back to an uncached compile, as for an op a user registers, and say so in udf.rst. * Retry the built-in UDF compile uncached only when it asked for a cache numba.njit(cache=False) never enables caching, so a RuntimeError from it is not the missing-locator case: retrying the same call could only fail again, with a chained traceback that blames the cache. Wrap only the cache=True call. Every RuntimeError Numba raises there comes from setting up the cache ("no locator available", or a bad NUMBA_CACHE_LOCATOR_CLASSES), so catching RuntimeError needs no match on the message. --- docs/user_guide/udf.rst | 25 ++++++++++++++++++ graphblas/core/operator/binary.py | 42 ++++++++++++++++++++++++------- graphblas/tests/test_op.py | 35 ++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 9 deletions(-) diff --git a/docs/user_guide/udf.rst b/docs/user_guide/udf.rst index 6255c1c56..b9ad278b8 100644 --- a/docs/user_guide/udf.rst +++ b/docs/user_guide/udf.rst @@ -220,3 +220,28 @@ time: # 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/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index ace225d46..0db22748c 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -533,7 +533,7 @@ 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: @@ -546,7 +546,21 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): # 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. - binary_udf = numba.njit(func, error_model="numpy") + # + # ``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 @@ -779,7 +793,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 @@ -846,13 +862,14 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal "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) @@ -905,13 +922,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 diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 0de6ca3e1..eb2675f46 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -6982,3 +6982,38 @@ def test_initialize_does_not_build_lazy_udfs(): 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])) From db36c4ced8151c5614f4b825396f53266bda67cc Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 11:09:53 -0600 Subject: [PATCH 43/52] Test the fmm engine guards and config key completions; combine repeated indices in csr, csc and GCXS imports (#609) * Test the fmm engine guards and config key completions Two behaviors had no test at all in the fast tier. Asking for the deprecated "fmm"/"fast_matrix_market" Matrix Market engine when fast_matrix_market is not installed must warn about the deprecation before it fails, and then fail with an ImportError naming the engine; neither read nor write had coverage for that, since the existing engine tests skip themselves when the package is missing. And the config mapping's _ipython_key_completions_ hook, which drives `config[]` in IPython, was never called. Each test was checked by perturbing the line it covers and confirming it goes red. * Combine repeated indices in an imported csr, csc, or GCXS array Behavior fix. A scipy csr or csc array outside canonical format, and a pydata GCXS array, may repeat an index within a row (column), and from_scipy_sparse imported the arrays unchecked. The Matrix was invalid (GxB_Matrix_fprint: GrB_INVALID_OBJECT), with nvals counting both entries. Before: from_pydata_sparse of a 2x3 GCXS with indices [1, 1, 0] in row 0 (values 1, 2, 3) and [2] in row 1 gave nvals 4 with two (0, 1) entries and to_dense 2.0 at (0, 1), where pydata's todense gives 3.0. from_scipy_sparse of the same csr array did the same. After: from_scipy_sparse builds such an array through coo, so repeats need dup_op as they do in coo input (ValueError "Duplicate indices found" without it). from_pydata_sparse sums them by default (with scipy's sum_duplicates, on a copy, so any dtype works), as pydata's todense does and its COO does when built; dup_op combines them otherwise. Unsorted but distinct indices give the same Matrix as before, now built from coo. A 1-d GCXS already summed its repeats (through asformat("coo")). --- graphblas/core/ss/config.py | 2 +- graphblas/io/_scipy.py | 8 +++- graphblas/io/_sparse.py | 13 +++++- graphblas/tests/test_io.py | 80 ++++++++++++++++++++++++++++++++ graphblas/tests/test_ss_utils.py | 16 +++++++ 5 files changed, 115 insertions(+), 4 deletions(-) diff --git a/graphblas/core/ss/config.py b/graphblas/core/ss/config.py index 70a7dd196..e1ccd15a0 100644 --- a/graphblas/core/ss/config.py +++ b/graphblas/core/ss/config.py @@ -210,5 +210,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/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/tests/test_io.py b/graphblas/tests/test_io.py index 962595aaa..a655bb5e0 100644 --- a/graphblas/tests/test_io.py +++ b/graphblas/tests/test_io.py @@ -355,6 +355,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): @@ -388,6 +413,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) @@ -531,3 +574,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_ss_utils.py b/graphblas/tests/test_ss_utils.py index 6e7a86b8e..86d9abd51 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -305,6 +305,22 @@ def test_global_config(monkeypatch): 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] + + @pytest.mark.skipif("gb.core.ss._IS_SSGB7") def test_context(): context = gb.ss.Context() From 42847adfb2caa72605f914d9e5794f0aaa3577a4 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:29:31 -0600 Subject: [PATCH 44/52] Install CI's packages from the configured channels, without -c nodefaults (#671) In an environment file, nodefaults is a keyword that drops the defaults channel. On the command line it is an ordinary channel name, and anaconda.org has a channel called nodefaults (empty for now), so `conda install -c nodefaults` searched it first, ahead of conda-forge, with strict channel priority: a package uploaded there would be what CI installed. setup-miniconda already makes conda-forge the only channel (conda-forge and defaults for PyPy), and the install now uses those. --- .github/workflows/test_and_build.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index dadb23e36..ecd82e2ab 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -272,7 +272,10 @@ jobs: else interpreter="python=$PYVER python-gil" fi - $(command -v mamba || command -v conda) install -c nodefaults \ + # 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} ${psg} \ From c7fa29a3f236f0b8c63994d0c2b69641fcb2a590 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:29:54 -0600 Subject: [PATCH 45/52] Keep nodes without edges in to_networkx (#646) * Keep nodes without edges in to_networkx to_networkx built its DiGraph from the matrix's entries alone, so an index whose row and column are both empty (a node with no edges) vanished, and the nodes came out in the order the sorted entries first named them. A 5x5 matrix with entries (0, 3), (2, 0) and (3, 3) became a graph with nodes [0, 3, 2]; from_networkx then gave back a 3x3 matrix with rows and columns permuted. viz.draw, which draws through to_networkx, left such nodes out of the picture. Add range(nrows) as nodes before the edges, so every row is a node and the nodes iterate in index order, and from_networkx(to_networkx(M)) gives back M. The user guide's line on to_networkx said it always labels the values "weight"; it now says what the nodes are and that edge_attribute names (or drops) the values. * Add create_isolates to to_networkx and require a square Matrix Making every row a node cannot work for a Matrix with huge dimensions: for a 2**60 x 2**60 hypersparse Matrix, to_networkx would run out of memory adding nodes (a DiGraph spends about 350 bytes on each). create_isolates=False, from #476, makes nodes only of the indices that have entries. They are still added in index order, so from_networkx gives back the Matrix without its empty rows and columns. With the default create_isolates=True, more than 2**40 rows raise ValueError naming the keyword instead of exhausting memory; a graph that large could never be built. Behavior change: a non-square Matrix raises ValueError. Row i and column i are both node i, which needs a square Matrix; a non-square one gave a graph that mixed row and column indices. The docstring already asked for a square Matrix. Co-authored-by: Sultan Orazbayev <20208402+SultanOrazbayev@users.noreply.github.com> * Give from_networkx the dtype in the to_networkx round-trip tests With numpy 1 on Windows, from_networkx infers INT32 from the Python int weights networkx stores, so the round trips compared with check_dtype=True failed against the INT64 Matrix. These tests are about the nodes and edges, so they now give the dtype. --------- Co-authored-by: Sultan Orazbayev <20208402+SultanOrazbayev@users.noreply.github.com> --- docs/user_guide/io.rst | 6 ++++- graphblas/io/_networkx.py | 34 +++++++++++++++++++++++-- graphblas/tests/test_io.py | 51 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 88 insertions(+), 3 deletions(-) diff --git a/docs/user_guide/io.rst b/docs/user_guide/io.rst index f144661d2..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 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/tests/test_io.py b/graphblas/tests/test_io.py index a655bb5e0..1536fa4c7 100644 --- a/graphblas/tests/test_io.py +++ b/graphblas/tests/test_io.py @@ -160,6 +160,57 @@ 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): From 3e37f836b564a668247ce354c2c2323ffbf25830 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:30:33 -0600 Subject: [PATCH 46/52] Make ss config attribute assignment raise instead of silently no-op (#615) Assigning an attribute on a SuiteSparse config object used to write a plain instance attribute and leave the real config unchanged: gb.ss.config.nthreads = 1 config["nthreads"] stayed 18 (its real value), while attribute reads of gb.ss.config.nthreads then returned the dead 1 v.ss.config.sparsity_control = "bitmap" vanished entirely, since v.ss.config is built fresh per access BaseConfig now defines __setattr__ that raises AttributeError: - known writable option: points at item assignment, the supported write idiom (config["nthreads"] = value) - known read-only option: "is read-only", matching the ValueError text that item assignment gives - unknown name: "Unknown config option ...; known options are [...]". Item assignment of an unknown key already raises KeyError, so there is no advice that could silently create junk. The allowlist of internal attributes is a snapshot of the instance __dict__ taken when __init__ finishes, not a hardcoded list, so a subclass that adds instance attributes cannot silently fall out of sync. BaseConfig does not use __slots__, so the snapshot is the natural source of truth. Context.__init__ now sets gb_obj and _prev_context before calling super().__init__() so they exist by the time the snapshot is taken (Context._from_obj already set them before init). Properties such as Context._context pass through to their setters, so assigning the context itself keeps working and assigning a different context keeps raising from the property. About (gb.ss.about) had the same trap: about.mode = "junk" left about["mode"] alone but made attribute reads return "junk". It never sets instance attributes at all, so its new __setattr__ raises unconditionally: "is read-only" for known keys, "Unknown About option" otherwise. Item assignment, reads, iteration, repr, and IPython key completions are unchanged, as is the donfig-based gb.config guard. Pinned suite (suitesparse, blocking, no-mapnumpy): the suite gains exactly the 4 new tests, 145 skipped unchanged (each test, each verified to fail with the guard reverted). test_ss_utils.py still module-skips on the suitesparse-vanilla backend. --- graphblas/core/ss/config.py | 30 +++++++++++++++++ graphblas/core/ss/context.py | 6 ++-- graphblas/ss/_core.py | 7 ++++ graphblas/tests/test_ss_utils.py | 58 ++++++++++++++++++++++++++++++++ 4 files changed, 99 insertions(+), 2 deletions(-) diff --git a/graphblas/core/ss/config.py b/graphblas/core/ss/config.py index e1ccd15a0..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.") diff --git a/graphblas/core/ss/context.py b/graphblas/core/ss/context.py index 67c7a6c20..99026e16f 100644 --- a/graphblas/core/ss/context.py +++ b/graphblas/core/ss/context.py @@ -33,9 +33,12 @@ 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 @@ -49,7 +52,6 @@ def __init__(self, engage=True, *, stack=True, nthreads=None, chunk=None, gpu_id self["chunk"] = chunk 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() diff --git a/graphblas/ss/_core.py b/graphblas/ss/_core.py index 40c7695f0..9fd6eecb6 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -347,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/test_ss_utils.py b/graphblas/tests/test_ss_utils.py index 86d9abd51..e76f6e45d 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -321,6 +321,64 @@ def test_global_config_key_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() From 4e8bc3de4e7eeac5f81da2bb489a8ee6c1c9e1e7 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:31:21 -0600 Subject: [PATCH 47/52] Extend the wrong-kind mask guard to dup and whole-object update (#614) A.dup(mask=v.S) and C(mask=v.S) << A assign with C(mask)[...] = A, which builds a __setitem__ expression, so the strict mask-kind check in BaseType._update did not apply and a Vector mask on a Matrix output reached GrB_Matrix_assign, leaking a raw cffi error: TypeError: Error calling GrB_Matrix_assign: ... initializer for ctype 'struct GB_Matrix_opaque *' must be a pointer to same type, not cdata 'struct GB_Vector_opaque *' Validate the mask kind in the Matrix-valued branch of Matrix._prep_for_assign, mirroring the guard the Vector-valued branch already has. This also covers C(v.S)[:, :] << A and the submask form C[:, :](v.S) << A. Row and column assignment with a Vector mask (C(v.S)[0, :] << w and friends) is handled in earlier branches and remains valid. The guard keys on mask.parent.ndim rather than exact type, the same idiom _check_mask uses, so a Vector subclass's mask is caught too. Two pre-existing guards in the Vector-valued branch and one more in the scalar path still use exact-type checks and carry the same subclass hole; they predate this commit and are left for their own change. --- graphblas/core/matrix.py | 7 +++++ graphblas/tests/test_matrix.py | 55 ++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index bd4f1d8d7..3abdcd124 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -3342,6 +3342,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 = ( diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 7ae48ae72..2079897ec 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -4828,6 +4828,61 @@ def test_wrong_kind_mask_on_matrix_raises(): 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. From e4b09cc5a7138c76142c698432e9149c5c39f5c4 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:31:50 -0600 Subject: [PATCH 48/52] Docs: fix examples, fill API reference gaps, document gotchas (#602) * Fix the README UDF example and add a Types API reference page The README UDF example imported unary but then used Vector without importing it, so the snippet as published raised NameError: name 'Vector' is not defined. Add Vector to the import; the copy in docs/user_guide/udf.rst was already correct. Also a new docs/api_reference/types.rst, wired into the reference toctree, surfacing DataType's JIT introspection properties (jit_c_name, jit_c_definition) and the dtypes register_new / register_anonymous helpers. None of that had an API reference entry. * Document the from_coo shape-inference gotcha and lazy backend init Matrix.from_coo and Vector.from_coo infer the shape from the largest index when nrows/ncols/size are omitted, so trailing all-empty rows, columns, or positions are silently dropped. Add a warning to both docstrings pointing at the explicit nrows/ncols/size fix. Also a note on the module-level backend global that reading gb.backend does not trigger initialization: it is a plain global, not a special attribute, so it stays None on a fresh import until a special-attribute access or init() runs. * Document the plus_first and plus_second semirings (GH #497) Add a user-guide subsection for plus_first and plus_second. They sum one operand's values while ignoring the other's, so when one operand is a boolean or iso adjacency they skip a plus_times multiply that would only multiply by one or rescale every term by the same constant. On a 0/1 matrix plus_first counts length-two paths. Includes a runnable example and lists both among the common semirings. * Add missing entries to the API reference (GH #410, #541) Document six public symbols absent from the API reference: IndexBinaryOp and Aggregator (operators), viz.spy and viz.datashade (visualization), and graphblas.init and graphblas.Recorder in a new utilities page wired into the reference toctree. The io converters and the collection classes' methods were already covered by existing autoclass :members: entries. * Add a class docstring to Aggregator The Aggregator class had no docstring, so its API reference entry rendered with no description. Add a numpydoc class docstring in the style of the sibling operator classes: what an aggregator is (a reduction operator for reduce, reduce_rowwise, reduce_columnwise, reduce_scalar), the graphblas.agg namespace, the summary-versus-position split, and the note that an aggregator is composed from a monoid or semiring plus an optional finalize step per dtype rather than being a single GraphBLAS object. The position aggregators are named as agg.ss.argmin and friends. The bare agg.argmin spellings still resolve but raise DeprecationWarning pointing at agg.ss, so a new docstring should not teach them. Docstring-only, no code change. * Fix a README constructor typo and a run-on apply error message The README's "Creating new Vectors / Matrices" block showed Matrix.new(dtype, num_rows, num_cols). There is no Matrix.new; the constructor is Matrix(dtype, num_rows, num_cols). Vector.apply built its error message from two adjacent string literals with no space at the join, so an invalid op reported "... `right` scalaror IndexUnaryOp with `right` thunk." * Say what Matrix.wait does in blocking mode The docstring said wait has no effect in blocking mode. With SuiteSparse:GraphBLAS, GrB_Matrix_wait with GrB_MATERIALIZE also builds the hyper_hash of a hypersparse matrix that has none yet, in blocking mode too: results of operations already have one, but a matrix from from_coo or an import may not. SuiteSparse 10.5 builds it only when the number of non-empty vectors exceeds GxB_HYPER_HASH (1024 by default; 1025 builds it, 1024 does not), and 8.2 for any hypersparse matrix. It took 0.7 to 1 s on one thread for a hypersparse matrix the size of the FROSTT NELL-2 tensor's matricization. Matrix.ss.unpack_hyperhash(compute=True) relies on it. wait("complete") never builds it. Vector.wait and Scalar.wait keep the sentence, which holds for them. * Say that plus_first on two boolean matrices needs an integer type plus has no BOOL form, so A.mxm(A, semiring.plus_first) on a boolean adjacency matrix raises KeyError; semiring.plus_first[int] counts the paths. Found by review (rmisc1). --- README.md | 4 ++-- docs/api_reference/index.rst | 2 ++ docs/api_reference/io.rst | 4 ++++ docs/api_reference/operators.rst | 12 ++++++++++ docs/api_reference/types.rst | 22 ++++++++++++++++++ docs/api_reference/utilities.rst | 13 +++++++++++ docs/user_guide/operators.rst | 39 ++++++++++++++++++++++++++++++++ graphblas/__init__.py | 4 ++++ graphblas/core/matrix.py | 12 +++++++++- graphblas/core/operator/agg.py | 21 +++++++++++++++++ graphblas/core/vector.py | 7 +++++- 11 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 docs/api_reference/types.rst create mode 100644 docs/api_reference/utilities.rst diff --git a/README.md b/README.md index 81df4fe64..3b84afef2 100644 --- a/README.md +++ b/README.md @@ -178,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 ``` @@ -225,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: 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/user_guide/operators.rst b/docs/user_guide/operators.rst index fb827afe3..3a3b2845a 100644 --- a/docs/user_guide/operators.rst +++ b/docs/user_guide/operators.rst @@ -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 -------------------- diff --git a/graphblas/__init__.py b/graphblas/__init__.py index e5038563e..e15604058 100644 --- a/graphblas/__init__.py +++ b/graphblas/__init__.py @@ -74,6 +74,10 @@ def __setattr__(self, key, value): 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 = { diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 3abdcd124..365f36527 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -822,7 +822,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() @@ -876,6 +880,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 diff --git a/graphblas/core/operator/agg.py b/graphblas/core/operator/agg.py index 7afeb9e46..574f4ec1b 100644 --- a/graphblas/core/operator/agg.py +++ b/graphblas/core/operator/agg.py @@ -29,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__( diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index bb9bea32a..5ea058e7c 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -700,6 +700,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 @@ -1412,7 +1417,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" + "apply only accepts UnaryOp with no scalars or BinaryOp with `left` or `right` scalar " "or IndexUnaryOp with `right` thunk." ) if isinstance(op, str): From 2d1f2c9893a59e565c556512f2e6d17d7e36fa51 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:32:26 -0600 Subject: [PATCH 49/52] Suggest close operator names in namespace AttributeErrors (#603) binary.pluss now raises "module 'graphblas.binary' has no attribute 'pluss'. Did you mean 'plus'?" via difflib.get_close_matches over the module's __dir__(), so lazily-registered operators are suggested without forcing them to build (binary._delayed is untouched by a typo lookup, which the test pins). Applies to all eight operator namespaces; difflib is imported only on the error path. --- graphblas/agg/__init__.py | 4 +++- graphblas/binary/__init__.py | 4 +++- graphblas/core/utils.py | 17 ++++++++++++++++ graphblas/indexunary/__init__.py | 4 +++- graphblas/monoid/__init__.py | 4 +++- graphblas/op/__init__.py | 4 +++- graphblas/select/__init__.py | 4 +++- graphblas/semiring/__init__.py | 4 +++- graphblas/tests/test_op.py | 34 ++++++++++++++++++++++++++++++++ graphblas/unary/__init__.py | 4 +++- 10 files changed, 75 insertions(+), 8 deletions(-) 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/utils.py b/graphblas/core/utils.py index 0b6b46a28..2d3de41cb 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -40,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, 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/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/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 876b3ec6a..df598667e 100644 --- a/graphblas/select/__init__.py +++ b/graphblas/select/__init__.py @@ -32,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): 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/tests/test_op.py b/graphblas/tests/test_op.py index eb2675f46..abb2bfd48 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -6931,6 +6931,40 @@ def test_compile_codegen_helper(): 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. 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 From 52dc9b498ed8c0e47ac04598e92fcea027b7a10e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:33:00 -0600 Subject: [PATCH 50/52] Move _check_mask into core/mask.py and drop dead numpy 1.21/1.22 guards (#607) Two cleanups that share matrix.py and vector.py in disjoint regions. _check_mask lived in base.py, so mask.py had to defer-import it inside three methods (Mask.new, Mask.__and__, Mask.__or__) to dodge the import cycle. It is mask logic, so it now lives in mask.py and those three calls are direct. base.py, matrix.py, and vector.py import it from mask.py; no new cycle appears, because mask.py has no module-level import of base. expr.py still needs a function-level import (the import cycle closes through mask.py's own module-level import of graphblas.binary names), now pointed at mask.py. Matrix.from_dicts and Vector.from_dict each carried a np.__version__ fallback for subarray dtypes on numpy 1.21 and 1.22, where np.fromiter could not build them. pyproject.toml requires numpy >=1.24, so neither branch was reachable. The np.fromiter path stays. What this does not change: same function, same signature, same error messages. The wrong-kind-mask TypeError and its assignment carve-out move across intact. The suite count under the pinned config (--backend suitesparse --blocking --no-mapnumpy) is identical before and after (1073 passed, 141 skipped). --- graphblas/core/base.py | 28 +--------------------------- graphblas/core/expr.py | 2 +- graphblas/core/mask.py | 31 +++++++++++++++++++++++++------ graphblas/core/matrix.py | 9 +++------ graphblas/core/vector.py | 9 +++------ 5 files changed, 33 insertions(+), 46 deletions(-) diff --git a/graphblas/core/base.py b/graphblas/core/base.py index b1347d355..95d89c670 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -2,13 +2,12 @@ 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 .literals import _literal_store_dtype -from .mask import Mask +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 @@ -168,31 +167,6 @@ def _expect_op(self, op, values, *, within, **kwargs): AmbiguousAssignOrExtract._expect_type = _expect_type -def _check_mask(mask, output=None, strict_kind=False): - 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: - 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 - - # 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. diff --git a/graphblas/core/expr.py b/graphblas/core/expr.py index 236e29534..1e16c44b4 100644 --- a/graphblas/core/expr.py +++ b/graphblas/core/expr.py @@ -322,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) diff --git a/graphblas/core/mask.py b/graphblas/core/mask.py index 4a8412a60..f812e2119 100644 --- a/graphblas/core/mask.py +++ b/graphblas/core/mask.py @@ -75,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)] @@ -95,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 @@ -121,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) @@ -202,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 365f36527..f92a4b028 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -7,12 +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 .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype -from .mask import Mask, StructuralMask, ValueMask +from .mask import Mask, StructuralMask, ValueMask, _check_mask from .operator import ( UNKNOWN_OPCLASS, _get_typed_op_from_exprs, @@ -1662,10 +1662,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 ) diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 5ea058e7c..ddb69fcdf 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -6,12 +6,12 @@ from ..dtypes import _INDEX, FP64, INT64, 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 .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype -from .mask import Mask, StructuralMask, ValueMask +from .mask import Mask, StructuralMask, ValueMask, _check_mask from .operator import ( UNKNOWN_OPCLASS, _get_typed_op_from_exprs, @@ -2107,10 +2107,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) From 63830fdd1892012fab8f2913e6cd624e519baa9a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:36:52 -0600 Subject: [PATCH 51/52] Fix wrong answers from select by A.T, coo import and isclose (#668) * Select by a transposed comparison as by its mask, not by the untransposed one select._match_expr rewrote A.select(A.T == x) as select.valueeq(A, x): its clause "args[0] is parent or args[0].T is parent" matched the transpose of the parent, and the rewrite then compared A itself, dropping the transpose and, for a matrix that is not square, the dimension check: S = Matrix.from_coo([0, 1], [1, 0], [1, 5], nrows=2, ncols=2) S.select(S.T == 1) before: (0, 1) = 1 after: (1, 0) = 5, as S.select((S.T == 1).new()) gives B (2 x 3): B.select(B.T == 1) before: a 2 x 3 result after: DimensionMismatch The clause is now "args[0] is parent"; anything else takes the mask path. On main since #252. * Clear an exported object's handle as soon as GraphBLAS frees it GraphBLAS frees the Matrix or Vector as it hands the arrays over in an export. _export gave GraphBLAS a copy of the handle and cleared the object's own handle only at the end, so an error in between, such as psg's claim_buffer_2d raising for an array UDT, left that handle pointing at the freed object: A.ss.export(format, give_ownership=True) raising after the export, then A.nvals before: UninitializedObject (a read of freed memory) after: NullPointer, as after an export that succeeds A plain export, of a copy, also freed that copy a second time when it was collected, which can abort the process (it did on SuiteSparse 7.4). GraphBLAS now gets the object's own handle, which it sets to NULL. * Say an infinity is close only to itself in isclose, as math.isclose does isclose used math.isclose's formula without its infinity check. With an infinite x or y, abs(x - y) is inf and so is rel_tol * max(abs(x), abs(y)), so infinity was close to every number, in binary.isclose (Vector and Matrix isclose with Numba, Scalar.isclose) and in _isclose_recipe (without Numba): Scalar.from_value(inf).isclose(1.0), .isclose(-inf) Vector.from_coo([0], [inf]).isclose(Vector.from_coo([0], [1.0])) before: True (all three) after: False, as math.isclose and numpy.isclose say Two infinities of the same sign are equal, and so still close; NaN is still close to nothing. * Compute isclose on integers without wrapping in x - y or abs(x) isclose computed abs(x - y) and abs(x) in the integer type, where they wrap: in UINT64, and with Numba in uint64 for any unsigned type (Numba widens to 64 bits); in INT64 at its extremes, as is abs(-2**63); and without Numba in every integer type. So, for Vector, Matrix and Scalar: UINT64 [0] and [2**64 - 1]: before True, after False UINT8 or UINT32 5 and 6, rel_tol=0.5: before False, after True INT64 -2**63 and 2**63 - 1: before True, after False INT8 -128 and 127, rel_tol=0.5: before True without Numba, after False and Scalar.isclose raised OverflowError for a UINT64 value above INT64's maximum, because Numba typed the Python int as int64. The difference is now exact (in uint64, for an integer type), the magnitudes are in FP64, and Scalar.isclose passes its values in the op's input types, as ewise_mult does. Mixed types are compared in the type they unify to, as before, so INT64 with UINT64 is in FP64. * Keep the imaginary part of x - y in isclose without Numba _isclose_recipe, which Vector and Matrix isclose use without Numba, stored x - y in an FP64 Vector before taking abs, and so kept only the real part of a complex difference: Vector.from_coo([0], [0j]).isclose(Vector.from_coo([0], [1j])) before: True without Numba (False with it) after: False, as cmath.isclose says abs is now taken in the complex type first. abs(x) and abs(y) are taken in the type x - y is in, as the Numba op does, so an FC32 operand beside an FP64 one is not rounded to FP32 first. * Count each row's entries when importing coo, so empty rows keep indptr whole import_coor and import_cooc build CSR (or CSC) index pointers with indices_to_indptr, which set indptr only at the rows that appear. Every empty row after the first one left a 0 in indptr, so an empty row in the middle moved entries to another row, and an empty last row lost every entry. Exporting as "coo" with sort=True, as "coor" or as "cooc" and importing back hit it, on main since 2021: A = Matrix.from_coo([0, 2], [1, 0], [1.0, 2.0], nrows=3, ncols=2) B = Matrix.ss.import_any(**A.ss.export("coo", sort=True)) before: B holds (2, 1) and (2, 0), and A.isequal(B) said True after: B holds (0, 1) and (2, 0), as A does A = Matrix.from_coo([0, 0], [0, 1], [1.0, 2.0], nrows=3, ncols=2) before: the round trip gives an empty Matrix after: both entries It now counts the entries of each row, then sums them. The indices are converted to uint64 first, so lists work, and one at or beyond nrows (or ncols) raises IndexOutOfBound, where it was written past the end of indptr with no bounds check under Numba. * Read 2-D coo indices as a flat buffer, and count them with numpy without Numba The new indptr kernel for import_coor and import_cooc indexed a 2-D rows (or cols) array by its rows, past the end of its first axis: rows = np.array([[0, 1], [1, 2]], np.uint64) Matrix.ss.import_coor(rows=rows, cols=[0, 0, 1, 1], values=[1, 2, 3, 4], nrows=3, ncols=2) before: Segmentation fault with Numba; IndexError without it after: the Matrix from_coo(rows.ravel(), ...) gives, as other imports read 2-D indices Without Numba the kernel ran in Python, once per entry and once per row (1.1 s for 10**7 rows and 10**5 entries, where the old kernel took 12 ms). It now counts with np.bincount and sums with np.cumsum (20 ms). * Compare in the type the dtypes unify to in Scalar.isclose without Numba Without Numba, Scalar.isclose passed the Python values straight to the isclose function. Every other path compares in the type the two dtypes unify to: Vector and Matrix isclose, and Scalar.isclose with Numba, which passes its values in the op's input types. INT64 and UINT64 unify to FP64, where 2**62 + 1 rounds to 2**62: Scalar.from_value(2**62 + 1).isclose( Scalar.from_value(2**62, dtype="UINT64"), rel_tol=0) with Numba, and Vector.isclose of the same: True Scalar.isclose without Numba: False, now True The values are converted to that type and back to Python numbers, so x - y still cannot wrap. * Describe the coo indptr helpers by what they return, not by the old bug The indices_to_indptr docstring told how the old kernel lost entries, which the commit that fixed it already says. It now states the invariant: every empty row gets an empty range. The select test's docstring no longer names main as a reference point. --- graphblas/core/matrix.py | 3 +- graphblas/core/operator/binary.py | 51 ++++++++++++- graphblas/core/operator/unary.py | 6 +- graphblas/core/scalar.py | 21 ++++-- graphblas/core/ss/matrix.py | 52 +++++++++---- graphblas/core/ss/vector.py | 6 +- graphblas/core/vector.py | 43 ++++++++--- graphblas/select/__init__.py | 2 +- graphblas/tests/test_matrix.py | 118 ++++++++++++++++++++++++++++++ graphblas/tests/test_scalar.py | 40 ++++++++++ graphblas/tests/test_vector.py | 73 ++++++++++++++++++ 11 files changed, 375 insertions(+), 40 deletions(-) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index f92a4b028..299d98d8c 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -459,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 ---------- diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 0db22748c..6abccb39e 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -339,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 @@ -843,11 +883,16 @@ def register_new( >>> 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) diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index c0bdf0434..e8f0870e0 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -319,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 @@ -357,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 diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 5f1ac7c01..0b711c258 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -17,6 +17,7 @@ _literal_type, ) from .operator import get_typed_op +from .operator.binary import _isclose from .udt import ( _element_bytes, _element_shapes_differ, @@ -34,7 +35,7 @@ if _supports_udfs: from ..binary import isclose else: - from .operator.binary import _isclose as isclose + isclose = _isclose ffi_new = ffi.new @@ -271,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 ---------- @@ -318,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, @@ -327,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, @@ -336,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. diff --git a/graphblas/core/ss/matrix.py b/graphblas/core/ss/matrix.py index 443019557..d39eb9145 100644 --- a/graphblas/core/ss/matrix.py +++ b/graphblas/core/ss/matrix.py @@ -7,7 +7,13 @@ from ... import binary, monoid from ...dtypes import _INDEX, BOOL, INT64, UINT64, lookup_dtype -from ...exceptions import InvalidValue, _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 ..operator import get_typed_op @@ -863,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*") @@ -1170,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 @@ -3171,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, @@ -3344,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, @@ -4558,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/vector.py b/graphblas/core/ss/vector.py index 8ee259510..4ad1c3eb9 100644 --- a/graphblas/core/ss/vector.py +++ b/graphblas/core/ss/vector.py @@ -515,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_) @@ -642,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 diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index ddb69fcdf..82fb5270f 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -3,7 +3,7 @@ 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, call @@ -110,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 @@ -124,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) @@ -398,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 ---------- diff --git a/graphblas/select/__init__.py b/graphblas/select/__init__.py index df598667e..a2d94a0a0 100644 --- a/graphblas/select/__init__.py +++ b/graphblas/select/__init__.py @@ -116,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/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 2079897ec..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, ) @@ -251,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() @@ -1269,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() @@ -1935,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 = [ @@ -4546,6 +4626,44 @@ def test_subarray_dtype_values_shape(): 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 diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 2d5974032..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 @@ -199,6 +201,44 @@ 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") diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 541a892f9..509b74173 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1,5 +1,7 @@ +import cmath import inspect import itertools +import math import pickle import sys import types @@ -1137,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 From 4cd73f16cc9ec066124c0cc59bf954560eb9a1ea Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Nov 2026 22:39:53 -0600 Subject: [PATCH 52/52] Give UDTs C names that no JIT kernel identifier collides with (#669) * Skip the JIT for a UDT field named like a standard C macro A record field named like an object-like macro of the C headers the JIT kernels include, such as INT_MAX, SIZE_MAX, BUFSIZ, EXIT_FAILURE or PRId64, is expanded by the preprocessor inside the struct, so the kernel did not compile. SuiteSparse 9.4 and later raised JitError (with no message) for that op and set jit_c_control to 'load', so the process compiled no more kernels. Main did the same. The comment beside _C_RESERVED said SuiteSparse swallowed the failure silently. _C_RESERVED now also has the object-like macros the C standard defines in , , , , , , and , so such a field skips the JIT: the op runs through the Numba cfunc with a NoJITWarning, as udt.rst says. Macros a platform adds (POSIX errno codes and the like) are not listed and still fail to compile. * Give each UDT the C name gbudt_ and its name, which no kernel identifier collides with SuiteSparse's JIT puts a UDT's C name in each kernel beside its own identifiers, and that name was the UDT's Python name. A UDT named like one of them did not compile: SuiteSparse raised JitError (with no message) for that op and set jit_c_control to 'load', so the process compiled no more kernels, although the code comment and udt.rst say such a UDT falls back to the cfunc. Main did the same. Names that failed, each in a fresh process: z and x (parameters of our op definitions), A, B, C, M, p, i, j, k, Ax, Cx, cnz, tid and result (locals of SuiteSparse's templates; which kernels fail depends on the name), div, abs, exp, round, free, index and printf (declared by the C headers), and GB_void and GrB_Info (SuiteSparse's). A list of reserved names cannot cover the template locals. Now the C name of a named UDT is gbudt_ and its name (_pick_c_type_name), so every one of those compiles and computes; a C keyword as a name (class) now takes the JIT too, where it got a synthetic name. A name that is not a C identifier, or would make a C name of 128 characters or more, gets a synthetic _gbudt_NNN name: SuiteSparse refuses such a long name and keeps the type with an empty one, and its first kernel (A.T.new() say) raised JitError, on main too for a UDT name of 128 characters or more. jit_c_name is unreleased (it came with #576, after 2025.2.0), so changing it breaks no release. udt.rst's example and the tests that pin C names say gbudt_. The test of an op whose C name SuiteSparse refuses now uses a 119-character UDT name: its C name (125 characters) fits, and plus_ and that does not. New tests: UDTs named z, A, p, i, div and result take the JIT for ewise_mult, reduce and mxv; UDT names of 121, 122 and 200 characters get C names that fit, and transpose under the JIT. * Say in the JIT docs how mixed UDT ops compile and how long the JIT stays on udt.rst said SuiteSparse compiles a kernel for each op[udt]. A lifted op on a UDT and a built-in type, or on two UDT types that pair, compiles too, and these mixed cases keep the Numba cfunc without a warning: an array UDT that broadcasts to another shape, a record with an array field, eq and ne on two types, and an op on one UDT that gives another type (v / v on INT8[3]). The JIT section now says so, with the mixed C name. UDT work raises the default 'run' to 'on' the first time it needs a kernel, for the rest of the process and for every operation, built-in ones included; a mixed op does too now that it has C source. init.rst and udt.rst read as if the raise were per op. The concurrency bullet now says such processes race too. init.rst says python-graphblas checks the C compiler under "on" and falls back to "run". Only a configured "on" is checked: an "on" written later to gb.ss.config is taken as it is, and with a compiler that does not work the first new kernel raises JitError and SuiteSparse lowers the control to "load" (checked with /usr/bin/false as the compiler). It says so, and that gb.ss.fix_jit_config() checks. * Keep a name starting with _gb as a UDT's C name, unprefixed python-graphblas names the types it makes itself with _gb (the synthetic _gbudt_NNN and _gbnest_NNN, and the tensor engine's _gbtensor_fp64xR), and no kernel identifier starts with it. Only _gbudt_ and _gbnest_ were kept as they are, so the tensor engine's type became gbudt__gbtensor_fp64x16, which its test of the name it claims caught. Keep every _gb name. * Name a C-source type by its own C name in the ops lifted to it The prefix that keeps a numpy UDT's C name clear of kernel identifiers was put on the name of every type an op definition names, and a type registered from C source (dtypes.ss.register_new) has no jit_c_name, so its Python name stood in, and got the prefix too: binary.plus[t] on a type cpt_t defined plus_gbudt_cpt_t on a type gbudt_cpt_t that SuiteSparse does not know. The kernel failed to compile, JitError, and the JIT was 'load' for the rest of the process. Before the prefix it was plus_cpt_t, and right. Now a C-source type's jit_c_name is the name it was registered under (and jit_c_definition its definition), and the op code generators name a type by its jit_c_name as it is (_dtype_c_typedef, is_c_name=True), so only a numpy UDT's Python name goes through _pick_c_type_name. That function no longer needs to map a pinned gbudt_ name to itself, so a user's name gbudt_Foo is now gbudt_gbudt_Foo and cannot collide with one made here (rBj N1). A name starting with _gb, which python-graphblas keeps for the types it makes itself (_gbudt_NNN, _gbnest_NNN, the tensor engine's _gbtensor_...), is still kept as it is. * Give a UDT a synthetic C name when its ops' names would not fit A numpy UDT kept the C name gbudt_ and its name while that was under SuiteSparse's 128 characters. But its one-type ops are named {op}_{type}, so a name of 113 to 121 characters left a C name its ops did not fit beside: SuiteSparse refused floordiv_ or plus_ and it, the op kept its cfunc, and the NoJITWarning said the UDT was not expressible as a C struct, blaming its fields. Now the C name may be at most _max_type_c_name_len() characters, which leaves room for the longest prefix of an op built from it (GxB_MAX_NAME_LEN - 1 - len("floordiv_") = 118, from the lifted ops, eq and ne), and a longer name gets a synthetic _gbudt_NNN, so every op on the UDT keeps its JIT. Mixed ops shorten their own names already. The test of a refused op name now registers a type from C source, whose 125-character name is the user's, so plus_ and eq_ and it still do not fit. The long-name test checks names of 112 (C name 118, kept), 113, 119 and 200 characters: floordiv, plus, eq and ainv all have C source, and transpose and plus run under the JIT. * Say that C source names a numpy UDT by its jit_c_name A UDT from numpy is known to SuiteSparse as gbudt_ and its name, so C source passed to gb.binary.ss.register_new and its siblings must name it so: written with its Python name, the operator failed to compile, raised JitError and left the JIT at 'load' for the rest of the process (rBj L1). The five register_new docstrings and udt.rst now say to use jit_c_name, and that a type registered from C source keeps its name. udt.rst also gives the new threshold for a synthetic C name (over 118 characters) and the exception for a name starting with _gb, which is kept as it is. * Skip the JIT for a field named like a signal macro or with GraphBLAS's prefix A record field named like one of 's nine C-standard macros (SIGINT, SIG_DFL, ...), which the C library headers bring into the kernels on macOS, or like one of GraphBLAS's own macros (GrB_INDEX_MAX, GxB_IMPLEMENTATION_MAJOR) was expanded by the preprocessor inside the struct, so the kernel failed to compile: JitError, then 'load' for the rest of the process (rBj N2). udt.rst said only macros beyond the C standard's still failed. Now _C_RESERVED has the nine, and a field that starts with GB_, GrB_ or GxB_ skips the JIT too: GraphBLAS keeps those prefixes for its own names, and each version adds more of them. The op runs through the Numba cfunc with a NoJITWarning, as for the other macros. * Say how long a C-source type's name may be, and that gbudt_ is usual The five ss.register_new docstrings said a numpy UDT's C name is gbudt_ and its name; a synthetic _gbudt_NNN is the exception, so they now say usually. A type registered from C source keeps the name given, and with 119 to 127 characters the built-in ops lifted to it get C names SuiteSparse refuses, so they lose the JIT, with a NoJITWarning that blames the fields. udt.rst now says to keep such a name to 118 characters (rBj round 2). --- docs/user_guide/init.rst | 17 ++- docs/user_guide/udt.rst | 58 +++++--- graphblas/core/dtypes.py | 25 ++-- graphblas/core/operator/udt_codegen.py | 192 +++++++++++++++++++------ graphblas/core/ss/binary.py | 2 + graphblas/core/ss/dtypes.py | 2 + graphblas/core/ss/indexbinary.py | 2 + graphblas/core/ss/indexunary.py | 2 + graphblas/core/ss/select.py | 2 + graphblas/core/ss/unary.py | 2 + graphblas/tests/test_dtype.py | 10 +- graphblas/tests/test_op.py | 8 +- graphblas/tests/test_ssjit.py | 151 +++++++++++++++++-- 13 files changed, 375 insertions(+), 98 deletions(-) diff --git a/docs/user_guide/init.rst b/docs/user_guide/init.rst index 141e1f250..306b596a4 100644 --- a/docs/user_guide/init.rst +++ b/docs/user_guide/init.rst @@ -69,8 +69,9 @@ often several times slower (``Matrix.reduce_rowwise`` with ``lor`` on ``BOOL`` o ``plus`` on ``UINT64`` ran 2.7 to 8 times slower on 10 million entries). The JIT control, ``gb.ss.config["jit_c_control"]``, is ``"run"`` by default: run only kernels already loaded in the process, and compile nothing. UDT operators, and types and operators -registered from C source, raise this default to ``"on"`` when they need a kernel (see -:doc:`udt`), but leave a control you set as it is. To opt in for everything, set the +registered from C source, raise this default to ``"on"`` the first time they need a +kernel, for the rest of the process and for every operation (see :doc:`udt`), but leave +a control you set as it is. To opt in for everything, set the control to ``"on"`` (compile, load and run): .. code-block:: bash @@ -96,7 +97,11 @@ later to ``gb.ss.config`` is that same default, so to keep UDT work from compili directory and compiler), and falls back to ``"run"``, with no error, where the JIT cannot work: SuiteSparse built without the JIT (as in the PyPI wheels), SuiteSparse 7, no working C compiler (in an activated conda environment, ``conda install c-compiler``), or a cache -directory it cannot write to. On Windows, SuiteSparse built with MSVC (as conda-forge +directory it cannot write to. That check is for a configured ``"on"`` only: an ``"on"`` +written later to ``gb.ss.config`` is taken as it is, and with a C compiler that does not +work, the first operation that needs a new kernel raises ``JitError``, and SuiteSparse +lowers the control to ``"load"``. ``gb.ss.fix_jit_config()`` repairs the compiler setting +where it can and checks it. On Windows, SuiteSparse built with MSVC (as conda-forge builds it) compiles JIT kernels only through cmake, and python-graphblas does not turn the JIT on there yet, so a configured ``"on"`` becomes ``"run"``. @@ -109,8 +114,10 @@ With the JIT on: (`GraphBLAS issue #449 `_). 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. That - costs UDTs their compiled kernels: UDT operators call their Numba functions instead + 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 diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 96d8e1f91..88b77086a 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -376,9 +376,11 @@ from Numba's typing pass instead of a 200-line traceback. JIT and introspection --------------------- -When the UDT name and 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. The kernel is cached on disk +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 @@ -389,11 +391,11 @@ Inspect what SuiteSparse is JIT-ing from Python: .. code-block:: python - udt.jit_c_name # 'Edge', or None if not JIT-able - udt.jit_c_definition # 'typedef struct { ... } Edge ;' + 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_Edge' + 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 @@ -404,10 +406,11 @@ Inspect what SuiteSparse is JIT-ing from Python: JIT is skipped when: -- A field name (or the UDT name, if you supplied one) is a C reserved word - (``class``, ``return``, ...) or a stdlib macro/typedef pulled in by - ``GraphBLAS.h`` (``NULL``, ``FILE``, ``M_PI``, ``complex``, ...). The op - falls through to the Numba cfunc path. +- 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 @@ -418,11 +421,25 @@ JIT is skipped when: 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). - -Anonymous UDTs (no ``name=`` argument) still take the JIT path: a synthetic -``_gbudt_NNN`` C name is minted so SuiteSparse always has a registerable -identifier. Pass ``name=`` only for readable JIT cache filenames and -introspection. +- 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 @@ -439,7 +456,9 @@ category or by message:: 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. +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 ~~~~~~~~~~~~~~~~~~~~~ @@ -461,9 +480,12 @@ 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: +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]``; +- 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. diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 2a85f46a1..97cbf1f62 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -108,14 +108,14 @@ def _deserialize(name, dtype, is_anonymous): 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`` may differ from ``datatype.name``: if the user - didn't supply a name, or supplied one that isn't a valid C identifier, - ``_udt_c_typedef`` synthesizes a fresh ``_gbudt_NNN`` 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]``). + 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 @@ -159,7 +159,8 @@ def _set_udt_jit_c_definition(datatype): # 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: - info = _udt_c_typedef(f"_gbudt_{next(_synthetic_udt_counter)}", datatype.np_type) + 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( @@ -276,9 +277,9 @@ def register_anonymous(dtype, name=None): string, and nested-UDT fields are not. name : str, optional A human-readable name for this UDT. Used in error messages and - ``repr``. Must be a valid C identifier (and not a C reserved word) - to be usable on the SuiteSparse JIT path; otherwise JIT is skipped - silently and ops fall back to the Numba cfunc path. + ``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 ------- diff --git a/graphblas/core/operator/udt_codegen.py b/graphblas/core/operator/udt_codegen.py index 41fba757a..0ac54404a 100644 --- a/graphblas/core/operator/udt_codegen.py +++ b/graphblas/core/operator/udt_codegen.py @@ -147,25 +147,20 @@ def _is_signed_int(np_dtype): # ``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 UDT type names or record field names in the -# JIT C source. Covers two failure modes from the JIT include chain -# (``GraphBLAS.h`` transitively pulls in ````, ````, -# ````, ````, ````, ````): +# 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.) # -# 1. **Macro names**: the preprocessor expands them inside any declarator -# or expression they appear in, so a field declared ``double M_PI ;`` -# becomes ``double 3.14159... ;`` and won't compile. The C-standard -# separate namespace for struct members doesn't help here. -# 2. **Typedef names** at the *outer* type-name position: emitting -# ``typedef struct { ... } FILE ;`` collides with the existing -# ``FILE`` typedef in scope (most compilers reject as a redefinition). -# Struct member names of the same spelling are fine (members live in -# their own namespace), but our gate runs the same check at both -# positions for simplicity. -# -# Either way the JIT compile fails, SuiteSparse swallows the failure -# silently, and the op runs through the slower Numba cfunc path. Block -# these eagerly so the user sees the "without JIT" warning instead. +# 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 @@ -280,17 +275,80 @@ def _is_signed_int(np_dtype): # "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.""" - return isinstance(name, str) and name.isidentifier() and name not in _C_RESERVED + """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 isn't a valid C identifier (or is a C reserved word, or the -# UDT was registered without ``name=`` at all). The Python-side ``DataType.name`` +# 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 @@ -304,29 +362,61 @@ def _is_valid_c_identifier(name): _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. - - When ``python_name`` is already a valid C identifier (and not reserved), - return it unchanged so introspection and JIT cache filenames remain - readable. Otherwise mint a fresh ``_gbudt_NNN`` name. This lets UDTs - that were anonymous in the Python sense (no ``name=`` passed, name with - dots / special chars, name that collides with a C keyword) still take - the JIT path; users never see the synthetic name unless they read - ``DataType.jit_c_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 _is_valid_c_identifier(python_name): - return python_name + 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): +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 (may differ from - ``python_name`` if that wasn't a valid C identifier; see - ``_pick_c_type_name``). Returns ``None`` when the UDT still can't be + 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). @@ -345,7 +435,7 @@ def _udt_c_typedef(python_name, np_type): # 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 = _pick_c_type_name(python_name) + 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 @@ -366,7 +456,7 @@ def _udt_c_typedef(python_name, np_type): # 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) + 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 @@ -1374,10 +1464,9 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): # 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 - # synthetic ``_gbudt_NNN`` name (for UDTs whose Python name isn't a valid - # C identifier) is recorded; using it keeps op codegen consistent. - pinned_name = dtype.jit_c_name - typedef_info = _udt_c_typedef(pinned_name or dtype.name, dtype.np_type) + # ``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 @@ -1426,10 +1515,22 @@ def _c_type_name(dtype): """ 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 = _udt_c_typedef(dtype.jit_c_name or dtype.name, 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 @@ -1542,10 +1643,9 @@ def _make_jit_c_comparison_definition(op_name, dtype, *, is_eq): semantic comparison. """ np_type = dtype.np_type - pinned_name = dtype.jit_c_name # ``_udt_c_typedef`` returns None when the numpy layout doesn't match a # C compiler's layout for the same struct. - typedef_info = _udt_c_typedef(pinned_name or dtype.name, np_type) + typedef_info = _dtype_c_typedef(dtype) if typedef_info is None: return None type_name, _typedef = typedef_info diff --git a/graphblas/core/ss/binary.py b/graphblas/core/ss/binary.py index 547d06ccf..d940ac342 100644 --- a/graphblas/core/ss/binary.py +++ b/graphblas/core/ss/binary.py @@ -61,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 diff --git a/graphblas/core/ss/dtypes.py b/graphblas/core/ss/dtypes.py index 0ccdd8264..54d3fb812 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -104,4 +104,6 @@ def register_new(name, jit_c_definition, *, np_type=None): # 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 index 5aeed9f11..8fd53a27e 100644 --- a/graphblas/core/ss/indexbinary.py +++ b/graphblas/core/ss/indexbinary.py @@ -67,6 +67,8 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): ``"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 diff --git a/graphblas/core/ss/indexunary.py b/graphblas/core/ss/indexunary.py index a3a462f59..77eafe067 100644 --- a/graphblas/core/ss/indexunary.py +++ b/graphblas/core/ss/indexunary.py @@ -60,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 diff --git a/graphblas/core/ss/select.py b/graphblas/core/ss/select.py index 48d16b740..323af9f40 100644 --- a/graphblas/core/ss/select.py +++ b/graphblas/core/ss/select.py @@ -57,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 cee7a16a4..695dec5c1 100644 --- a/graphblas/core/ss/unary.py +++ b/graphblas/core/ss/unary.py @@ -55,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 diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index e2df81017..a7edc6d7b 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -565,13 +565,13 @@ def test_udt_jit_c_info_pinned_at_first_register(): # 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 == "_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 == "_PinUDT_first" + assert udt.jit_c_name == "gbudt__PinUDT_first" assert "_PinUDT_first" in udt.jit_c_definition @@ -587,12 +587,14 @@ def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): """ 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 == "_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("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") + 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 diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index abb2bfd48..64bc1880c 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2387,7 +2387,7 @@ def test_udt_jit_op_definitions(): 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}_JitOpTest" in defn + assert f"{op_name}_gbudt_JitOpTest" in defn assert "jp" in defn assert "jq" in defn assert expected_c_op in defn @@ -2396,7 +2396,7 @@ def test_udt_jit_op_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_JitOpTest" in defn + 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) @@ -2405,14 +2405,14 @@ def test_udt_jit_op_definitions(): 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_Vec37Jit" in defn + 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_Vec37Jit" in defn + assert "ainv_gbudt_Vec37Jit" in defn @pytest.mark.skipif("not supports_udfs") diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 3721eee1f..573d0f236 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1327,7 +1327,7 @@ def test_udt_jit_c_source_introspection(): udt = dtypes.register_anonymous(record_dtype, "_IntrospectUDT") # dtype-level introspection - assert udt.jit_c_name == "_IntrospectUDT" + assert udt.jit_c_name == "gbudt__IntrospectUDT" assert "typedef struct" in udt.jit_c_definition assert "int64_t introsp_a" in udt.jit_c_definition assert "double introsp_b" in udt.jit_c_definition @@ -1338,7 +1338,7 @@ def test_udt_jit_c_source_introspection(): # Auto-lifted record UDT binary op plus_udt = binary.plus[udt] - assert plus_udt.jit_c_name == "plus__IntrospectUDT" + assert plus_udt.jit_c_name == "plus_gbudt__IntrospectUDT" src = plus_udt.jit_c_source assert src is not None assert "z->introsp_a = (x->introsp_a) + (y->introsp_a)" in src @@ -1346,7 +1346,7 @@ def test_udt_jit_c_source_introspection(): # Auto-lifted record UDT unary op ainv_udt = unary.ainv[udt] - assert ainv_udt.jit_c_name == "ainv__IntrospectUDT" + 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`` @@ -1392,12 +1392,12 @@ def test_op_jit_signature_uses_pinned_type_name(): udt2 = dtypes.register_anonymous(record_dtype, "_PinnedRenamed") assert udt2 is udt assert udt.name == "_PinnedRenamed" - assert udt.jit_c_name == "_PinnedOrig" + 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__PinnedOrig" - assert "_PinnedOrig *" in op.jit_c_source + assert op.jit_c_name == "plus_gbudt__PinnedOrig" + assert "gbudt__PinnedOrig *" in op.jit_c_source assert "_PinnedRenamed" not in op.jit_c_source @@ -2349,8 +2349,9 @@ def test_udt_scalar_broadcast_jit_routing(): 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_11__BcastJitUV_int64_t (_BcastJitUV *z, const _BcastJitUV *x, " - "const int64_t *y) { z->bcast_u = (x->bcast_u) + (((double) ((*y)))) ; " + "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 @@ -2619,7 +2620,13 @@ def test_udt_jit_c_name_too_long_keeps_the_cfunc(): pytest.skip("requires SuiteSparse:GraphBLAS >= 9") _require_jit_on() rec = np.dtype([("tl_x", np.float64), ("tl_y", np.int64)], align=True) - t = dtypes.register_new("TooLongName" + "t" * 120, rec) + # 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 @@ -2636,3 +2643,129 @@ def test_udt_jit_c_name_too_long_keeps_the_cfunc(): 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"