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})")