From 8088a6490e99e2e98ede071ef97e69b6e6209bb0 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:21:33 -0600 Subject: [PATCH 01/77] Add graphblas.array_api: an Array API namespace over Tensor (skeleton) An Array is a Tensor of stored entries plus a fill value for every other position, as in pydata/sparse. This first step has the Array class, the dtypes, creation functions, the element-wise functions, astype and the data type functions, __array_namespace_info__, basic indexing (ints, slices, ..., None) and reshape, permute_dims, expand_dims and squeeze. Element-wise results follow numpy on every position, the fill included; numpy maps the stored values where no GraphBLAS operator gives its answer exactly. --- graphblas/array_api/__init__.py | 85 ++++ graphblas/array_api/_array.py | 355 ++++++++++++++ graphblas/array_api/_constants.py | 11 + graphblas/array_api/_creation.py | 387 +++++++++++++++ graphblas/array_api/_dtypes.py | 172 +++++++ graphblas/array_api/_elementwise.py | 693 +++++++++++++++++++++++++++ graphblas/array_api/_indexing.py | 99 ++++ graphblas/array_api/_info.py | 53 ++ graphblas/array_api/_linalg.py | 3 + graphblas/array_api/_manipulation.py | 116 +++++ graphblas/array_api/_searching.py | 3 + graphblas/array_api/_set.py | 3 + graphblas/array_api/_sorting.py | 3 + graphblas/array_api/_statistical.py | 3 + graphblas/array_api/_utils.py | 169 +++++++ graphblas/graphblas.yaml | 2 + 16 files changed, 2157 insertions(+) create mode 100644 graphblas/array_api/__init__.py create mode 100644 graphblas/array_api/_array.py create mode 100644 graphblas/array_api/_constants.py create mode 100644 graphblas/array_api/_creation.py create mode 100644 graphblas/array_api/_dtypes.py create mode 100644 graphblas/array_api/_elementwise.py create mode 100644 graphblas/array_api/_indexing.py create mode 100644 graphblas/array_api/_info.py create mode 100644 graphblas/array_api/_linalg.py create mode 100644 graphblas/array_api/_manipulation.py create mode 100644 graphblas/array_api/_searching.py create mode 100644 graphblas/array_api/_set.py create mode 100644 graphblas/array_api/_sorting.py create mode 100644 graphblas/array_api/_statistical.py create mode 100644 graphblas/array_api/_utils.py diff --git a/graphblas/array_api/__init__.py b/graphblas/array_api/__init__.py new file mode 100644 index 000000000..e9db10cf7 --- /dev/null +++ b/graphblas/array_api/__init__.py @@ -0,0 +1,85 @@ +"""An Array API (2025.12) namespace over ``graphblas.Tensor``: ``import graphblas.array_api as xp``. + +**THIS API IS EXPERIMENTAL AND MAY CHANGE.** + +An Array is logically dense: its value at a position is its stored entry there, or its +``fill_value`` where it stores none (pydata/sparse's model). Results equal numpy's on the same +dense data, and keep the stored entries sparse where the fill values allow. +""" + +from . import ( + _constants, + _creation, + _dtypes, + _elementwise, + _indexing, + _info, + _linalg, + _manipulation, + _searching, + _set, + _sorting, + _statistical, +) +from ._array import Array +from ._constants import * # noqa: F403 +from ._creation import * # noqa: F403 +from ._dtypes import * # noqa: F403 +from ._dtypes import DTYPES as _DTYPES +from ._elementwise import * # noqa: F403 +from ._indexing import * # noqa: F403 +from ._info import * # noqa: F403 +from ._linalg import * # noqa: F403 +from ._manipulation import * # noqa: F403 +from ._searching import * # noqa: F403 +from ._set import * # noqa: F403 +from ._sorting import * # noqa: F403 +from ._statistical import * # noqa: F403 +from ._utils import DensifyError + +__array_api_version__ = "2025.12" + +bool = _DTYPES["bool"] # noqa: A001 (the Array API's name) +int8 = _DTYPES["int8"] +int16 = _DTYPES["int16"] +int32 = _DTYPES["int32"] +int64 = _DTYPES["int64"] +uint8 = _DTYPES["uint8"] +uint16 = _DTYPES["uint16"] +uint32 = _DTYPES["uint32"] +uint64 = _DTYPES["uint64"] +float32 = _DTYPES["float32"] +float64 = _DTYPES["float64"] +complex64 = _DTYPES["complex64"] +complex128 = _DTYPES["complex128"] + +__all__ = [ + "Array", + "DensifyError", + "__array_api_version__", + "bool", + "complex64", + "complex128", + "float32", + "float64", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", +] +__all__ += _constants.__all__ +__all__ += _creation.__all__ +__all__ += _dtypes.__all__ +__all__ += _elementwise.__all__ +__all__ += _indexing.__all__ +__all__ += _info.__all__ +__all__ += _linalg.__all__ +__all__ += _manipulation.__all__ +__all__ += _searching.__all__ +__all__ += _set.__all__ +__all__ += _sorting.__all__ +__all__ += _statistical.__all__ diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py new file mode 100644 index 000000000..6b1b55ed7 --- /dev/null +++ b/graphblas/array_api/_array.py @@ -0,0 +1,355 @@ +"""The Array: a Tensor of stored entries plus a fill value for every other position.""" + +import math + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._utils import np_dtype + +CPU_DEVICE = "cpu" +API_VERSIONS = ("2021.12", "2022.12", "2023.12", "2024.12", "2025.12") + + +def check_device(device): + if device is not None and device != CPU_DEVICE: + raise ValueError(f"graphblas.array_api has only the device {CPU_DEVICE!r}, not {device!r}") + + +class Array: + """An n-dimensional array whose value at a position is its stored entry there, or + ``fill_value`` where it stores none (pydata/sparse's model). + + Create Arrays with the namespace's functions (``asarray``, ``zeros``, ``from_coo``, ...). + The backing ``graphblas.Tensor`` holds the stored entries; stored entries may equal the + fill value. + """ + + __slots__ = "_tensor", "_fill", "_facts", "__weakref__" + # numpy hands its operators over to ours (``ndarray + Array``), as it does for other + # array types, rather than converting the Array to a dense ndarray + __array_ufunc__ = None + + def __init__(self, *args, **kwargs): + raise TypeError( + "Array is not created directly; use graphblas.array_api.asarray (or zeros, full, " + "from_coo, ...)" + ) + + @classmethod + def _new(cls, tensor, fill): + """An Array of ``tensor``'s entries and ``fill`` (cast to its dtype) elsewhere. + + A 0-d Array stores nothing: its value is its fill. + """ + if tensor.ndim == 0 and tensor.nvals: + fill = tensor.to_scalar().value + tensor = Tensor(tensor.dtype, ()) + self = object.__new__(cls) + self._tensor = tensor + self._fill = _utils.scalar(fill, np_dtype(tensor.dtype)) + self._facts = None + return self + + def _T(self): + """The Tensor of stored entries. Every read of it goes through here, so that an Array + may later hold an unevaluated expression and compute it on first use. + """ + return self._tensor + + def _replace(self, other): + """Make this Array hold ``other``'s values (for in-place operators and assignment).""" + self._tensor = other._T() + self._fill = other._fill + self._facts = None + + def _fact(self, name, compute): + """A cached property of the stored values; assignment clears the cache.""" + facts = self._facts + if facts is None: + facts = self._facts = {} + if name not in facts: + facts[name] = compute(self) + return facts[name] + + def _copy(self): + return Array._new(self._T().dup(), self._fill) + + # Attributes + @property + def dtype(self): + return np_dtype(self._T().dtype) + + @property + def device(self): + return CPU_DEVICE + + @property + def shape(self): + return self._T().shape + + @property + def ndim(self): + return self._T().ndim + + @property + def size(self): + return math.prod(self._T().shape) + + @property + def T(self): + """The axes reversed (numpy's ``.T``; the Array API defines it for 2-d arrays).""" + from ._manipulation import permute_dims + + return permute_dims(self, tuple(range(self.ndim - 1, -1, -1))) + + @property + def mT(self): + from ._manipulation import permute_dims + + if self.ndim < 2: + raise ValueError("mT needs at least 2 dimensions") + axes = list(range(self.ndim)) + axes[-2], axes[-1] = axes[-1], axes[-2] + return permute_dims(self, tuple(axes)) + + # Extensions + @property + def fill_value(self): + """The value at every position the Array does not store (a numpy scalar).""" + return self._fill + + @property + def nnz(self): + """The number of stored entries (some may equal the fill value).""" + return self._T().nvals + + @property + def tensor(self): + """The stored entries as a ``graphblas.Tensor`` (positions holding the fill value are + missing); a copy that shares the data until either side changes. + """ + return self._T().dup() + + def todense(self): + """A numpy array of every position (raises past ``array_api_densify_limit``).""" + _utils.check_dense(self.size, f"densifying an array of shape {self.shape}") + dt = self.dtype + T = self._T() + if T.ndim == 0 or T.nvals == 0: + return np.full(self.shape, self._fill, dtype=dt) + if T.nvals == self.size: + return T.to_dense(dtype=T.dtype).astype(dt, copy=False) + coords, values = T.to_coo() + out = np.full(self.shape, self._fill, dtype=dt) + out[coords] = values + return out + + def to_coo(self): + """``(coords, values)`` of the stored entries: a tuple of index arrays (one per axis) + and the values, in C order. + """ + return self._T().to_coo() + + def __repr__(self): + fill = self._fill.item() + if self.ndim == 0: + return f"Array({fill!r}, dtype={self.dtype})" + return f"Array(shape={self.shape}, dtype={self.dtype}, nnz={self.nnz}, fill_value={fill!r})" + + # Protocols + def __array_namespace__(self, *, api_version=None): + if api_version is not None and api_version not in API_VERSIONS: + raise ValueError(f"unsupported Array API version {api_version!r}") + from .. import array_api + + return array_api + + def __array__(self, dtype=None, copy=None): + if copy is False: + raise ValueError("converting an Array to a numpy array always copies") + out = self.todense() + return out if dtype is None else out.astype(dtype, copy=False) + + def __dlpack__(self, *, stream=None, max_version=None, dl_device=None, copy=None): + if copy is False: + raise BufferError("exporting an Array through DLPack always copies") + kwargs = {} if max_version is None else {"max_version": max_version} + if dl_device is not None: + kwargs["dl_device"] = dl_device + return self.todense().__dlpack__(stream=stream, **kwargs) + + def __dlpack_device__(self): + return np.empty(0).__dlpack_device__() + + def to_device(self, device, /, *, stream=None): + check_device(device) + return self + + def _value(self, what): + if self.ndim != 0: + raise TypeError(f"{what} needs a 0-d array; this one has shape {self.shape}") + return self._fill + + def __bool__(self): + return bool(self._value("bool()")) + + def __int__(self): + value = self._value("int()") + if self.dtype.kind == "c": + raise TypeError("int() of a complex array") + return int(value) + + def __float__(self): + value = self._value("float()") + if self.dtype.kind == "c": + raise TypeError("float() of a complex array") + return float(value) + + def __complex__(self): + return complex(self._value("complex()")) + + def __index__(self): + value = self._value("an index") + if self.dtype.kind not in "iu": + raise TypeError(f"only integer arrays are indices, not {self.dtype}") + return int(value) + + def __len__(self): + if self.ndim == 0: + raise TypeError("len() of a 0-d array") + return self.shape[0] + + def __iter__(self): + if self.ndim == 0: + raise TypeError("iteration over a 0-d array") + return (self[i] for i in range(self.shape[0])) + + __hash__ = None + + # Indexing (graphblas.array_api._indexing) + def __getitem__(self, key): + from ._indexing import getitem + + return getitem(self, key) + + def __setitem__(self, key, value): + from ._indexing import setitem + + setitem(self, key, value) + + # Operators + def __matmul__(self, other): + if not _is_operand(other): + return NotImplemented + from ._linalg import matmul + + return matmul(self, other) + + def __rmatmul__(self, other): + if not _is_operand(other): + return NotImplemented + from ._linalg import matmul + + return matmul(other, self) + + def __imatmul__(self, other): + from ._linalg import matmul + + return self._inplace(matmul(self, other)) + + def _inplace(self, result): + if result.shape != self.shape: + raise ValueError( + f"an in-place operation would change the shape {self.shape} to {result.shape}" + ) + if result.dtype != self.dtype: + raise TypeError( + f"an in-place operation would change the data type {self.dtype} to " + f"{result.dtype}" + ) + self._replace(result) + return self + + def __neg__(self): + from ._elementwise import negative + + return negative(self) + + def __pos__(self): + from ._elementwise import positive + + return positive(self) + + def __abs__(self): + from ._elementwise import abs as abs_ + + return abs_(self) + + def __invert__(self): + from ._elementwise import bitwise_invert + + return bitwise_invert(self) + + +def _is_operand(x): + return isinstance(x, (Array, bool, int, float, complex, np.generic, np.ndarray)) + + +def _binary_methods(name, func_name): + def forward(self, other): + if not _is_operand(other): + return NotImplemented + from . import _elementwise + + return getattr(_elementwise, func_name)(self, other) + + def reflected(self, other): + if not _is_operand(other): + return NotImplemented + from . import _elementwise + + return getattr(_elementwise, func_name)(other, self) + + def inplace(self, other): + from . import _elementwise + + return self._inplace(getattr(_elementwise, func_name)(self, other)) + + forward.__name__ = f"__{name}__" + reflected.__name__ = f"__r{name}__" + inplace.__name__ = f"__i{name}__" + return forward, reflected, inplace + + +for _name, _func in [ + ("add", "add"), + ("sub", "subtract"), + ("mul", "multiply"), + ("truediv", "divide"), + ("floordiv", "floor_divide"), + ("mod", "remainder"), + ("pow", "pow"), + ("and", "bitwise_and"), + ("or", "bitwise_or"), + ("xor", "bitwise_xor"), + ("lshift", "bitwise_left_shift"), + ("rshift", "bitwise_right_shift"), +]: + _f, _r, _i = _binary_methods(_name, _func) + setattr(Array, f"__{_name}__", _f) + setattr(Array, f"__r{_name}__", _r) + setattr(Array, f"__i{_name}__", _i) + +for _name, _func in [ + ("eq", "equal"), + ("ne", "not_equal"), + ("lt", "less"), + ("le", "less_equal"), + ("gt", "greater"), + ("ge", "greater_equal"), +]: + _f, _r, _i = _binary_methods(_name, _func) + setattr(Array, f"__{_name}__", _f) +del _name, _func, _f, _r, _i diff --git a/graphblas/array_api/_constants.py b/graphblas/array_api/_constants.py new file mode 100644 index 000000000..18a5473f8 --- /dev/null +++ b/graphblas/array_api/_constants.py @@ -0,0 +1,11 @@ +"""Constants.""" + +import math + +__all__ = ["e", "inf", "nan", "newaxis", "pi"] + +e = math.e +inf = math.inf +nan = math.nan +newaxis = None +pi = math.pi diff --git a/graphblas/array_api/_creation.py b/graphblas/array_api/_creation.py new file mode 100644 index 000000000..dc7ee3732 --- /dev/null +++ b/graphblas/array_api/_creation.py @@ -0,0 +1,387 @@ +"""Creation functions. + +``zeros``, ``ones``, ``full``, ``empty`` and their ``_like`` forms store nothing: their value is +the fill, so they cost O(1). Conversions from dense data store every position whose value is +not identical to the fill (0 of the dtype unless ``fill_value=`` is given): ``-0.0`` is stored +when the fill is ``0.0``, and NaN is stored unless the fill is NaN. +""" + +import builtins +import math + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._array import Array, check_device +from ._dtypes import as_dtype, check_supported +from ._utils import gb_dtype + +__all__ = [ + "arange", + "asarray", + "empty", + "empty_like", + "eye", + "from_coo", + "from_dlpack", + "full", + "full_like", + "linspace", + "meshgrid", + "ones", + "ones_like", + "tril", + "triu", + "zeros", + "zeros_like", +] + +_DEFAULT_FLOAT = np.dtype("float64") + + +def _shape(shape): + if isinstance(shape, (int, np.integer)): + shape = (shape,) + shape = tuple(int(n) for n in shape) + if any(n < 0 for n in shape): + raise ValueError(f"negative dimensions are not allowed: {shape}") + return shape + + +def _default_fill(dtype): + return _utils.scalar(0, dtype) + + +def _from_numpy(a, fill_value=None): + """An Array of the numpy array ``a`` storing what differs from the fill.""" + dt = check_supported(a.dtype, "an array") + fill = _default_fill(dt) if fill_value is None else _utils.scalar(fill_value, dt) + if a.ndim == 0: + return Array._new(Tensor(gb_dtype(dt), ()), a[()]) + if a.size == 0: + return Array._new(Tensor(gb_dtype(dt), a.shape), fill) + keep = _utils.differs(a, fill) + if keep.all(): + T = Tensor.from_dense(a, dtype=gb_dtype(dt)) + else: + coords = np.nonzero(keep) + T = Tensor.from_coo(coords, a[coords], gb_dtype(dt), shape=a.shape) + return Array._new(T, fill) + + +def _from_graphblas(obj, dtype, fill_value): + """A Tensor, Matrix, Vector or Scalar as an Array whose missing entries hold the fill.""" + from .. import Matrix, Scalar, Vector + from ..core.matrix import TransposedMatrix + + if isinstance(obj, Scalar): + dt = _utils.np_dtype(obj.dtype) + value = obj.value if obj.nvals else fill_value if fill_value is not None else 0 + out = Array._new(Tensor(gb_dtype(dt), ()), value) + else: + if isinstance(obj, Tensor): + T = obj.dup() + elif isinstance(obj, Vector): + T = Tensor.from_vector(obj) + elif isinstance(obj, (Matrix, TransposedMatrix)): + T = Tensor.from_matrix(obj) + else: # pragma: no cover (safety) + raise TypeError(type(obj)) + dt = check_supported(_utils.np_dtype(T.dtype), "an array") + fill = _default_fill(dt) if fill_value is None else fill_value + out = Array._new(T, fill) + if dtype is not None and as_dtype(dtype) != out.dtype: + from ._elementwise import cast + + out = cast(out, as_dtype(dtype)) + return out + + +def _from_sparse(obj, dtype, fill_value): + """A scipy.sparse array or matrix, or a pydata/sparse array.""" + if hasattr(obj, "tocoo") and hasattr(obj, "format") and not hasattr(obj, "fill_value"): + coo = obj.tocoo() + coords = tuple(np.asarray(c) for c in coo.coords) + data, shape, fill = np.asarray(coo.data), coo.shape, None + if fill_value is not None: + # scipy.sparse's missing entries are zeros + raise ValueError("a scipy.sparse array has fill value 0; convert it first") + if not getattr(coo, "has_canonical_format", True): + coo.sum_duplicates() + coords, data = tuple(np.asarray(c) for c in coo.coords), np.asarray(coo.data) + else: + if obj.format != "coo": + obj = obj.asformat("coo") + coords = tuple(np.asarray(c) for c in obj.coords) + data, shape, fill = np.asarray(obj.data), obj.shape, obj.fill_value + if fill_value is not None and not _utils.same_value(fill_value, fill): + raise ValueError(f"fill_value={fill_value!r} differs from the sparse array's {fill!r}") + return from_coo(coords, data, shape, fill_value=fill, dtype=dtype) + + +def _is_sparse(obj): + module = type(obj).__module__ + return module.startswith(("scipy.sparse", "sparse")) + + +def asarray(obj, /, *, dtype=None, device=None, copy=None, fill_value=None): + """Convert ``obj`` to an Array. + + ``obj`` is an Array, a Python scalar or (nested) sequence, an object numpy converts (a + buffer, an ndarray), or (extensions) a ``graphblas`` Tensor, Matrix, Vector or Scalar, a + scipy.sparse array or a pydata/sparse array, whose missing entries hold the fill. + ``fill_value`` (an extension) is the value of unstored positions: 0 of the dtype by + default, or the sparse array's own fill value. + """ + check_device(device) + if isinstance(obj, Array): + if fill_value is not None and not _utils.same_value(fill_value, obj._fill): + raise ValueError("asarray of an Array keeps its fill value; use full_like to refill") + if dtype is not None and as_dtype(dtype) != obj.dtype: + if copy is False: + raise ValueError("changing the data type needs a copy, and copy=False") + from ._elementwise import cast + + return cast(obj, as_dtype(dtype)) + return obj._copy() if copy else obj + if copy is False: + raise ValueError("asarray stores its own copy of anything but an Array; copy=False") + from .. import Matrix, Scalar, Vector + from ..core.matrix import TransposedMatrix + + if isinstance(obj, (Tensor, Matrix, Vector, Scalar, TransposedMatrix)): + return _from_graphblas(obj, dtype, fill_value) + if _is_sparse(obj) and (hasattr(obj, "tocoo") or hasattr(obj, "coords")): + return _from_sparse(obj, dtype, fill_value) + dt = None if dtype is None else as_dtype(dtype) + if isinstance(obj, (list, tuple)): + obj = _unwrap_nested(obj) + if dt is None and isinstance(obj, builtins.bool): + dt = np.dtype(bool) + elif dt is None and isinstance(obj, int) and not isinstance(obj, np.generic): + dt = np.dtype("int64") + a = np.asarray(obj, dtype=dt) + if dt is None and a.dtype.kind == "f" and a.dtype.itemsize < 4: + a = a.astype(np.float32) + return _from_numpy(a, fill_value) + + +def _unwrap_nested(obj): + """A (nested) sequence with its 0-d Arrays replaced by their values (numpy would densify + every Array through ``__array__``, which works too, but slower). + """ + if isinstance(obj, Array): + return obj._fill if obj.ndim == 0 else obj.todense() + if isinstance(obj, (list, tuple)): + return [_unwrap_nested(o) for o in obj] + return obj + + +def from_coo(coords, data, shape=None, *, fill_value=None, dtype=None): + """An Array storing ``data`` at ``coords`` (one index array per axis, or an (ndim, nnz) + array), and ``fill_value`` (0 by default) elsewhere. ``data`` may be a scalar. + """ + data = np.asarray(data) + dt = as_dtype(dtype) if dtype is not None else check_supported(data.dtype, "from_coo data") + if dtype is not None: + data = data.astype(dt, copy=False) + if isinstance(coords, np.ndarray) and coords.ndim == 2: + coords = tuple(coords) + coords = tuple(np.asarray(c) for c in coords) + if shape is None: + shape = tuple(int(c.max()) + 1 if c.size else 0 for c in coords) + shape = _shape(shape) + fill = _default_fill(dt) if fill_value is None else _utils.scalar(fill_value, dt) + if len(shape) == 0: + value = data.reshape(-1)[0] if data.size else fill + return Array._new(Tensor(gb_dtype(dt), ()), value) + if coords and coords[0].size: + T = Tensor.from_coo(coords, data, gb_dtype(dt), shape=shape) + else: + T = Tensor(gb_dtype(dt), shape) + return Array._new(T, fill) + + +def _full(shape, value, dtype): + shape = _shape(shape) + return Array._new(Tensor(gb_dtype(dtype), shape), value) + + +def _fill_dtype(fill_value): + if isinstance(fill_value, Array): + return fill_value.dtype + if isinstance(fill_value, np.generic): + return check_supported(fill_value.dtype, "full") + if isinstance(fill_value, builtins.bool): + return np.dtype(bool) + if isinstance(fill_value, int): + return np.dtype("int64") + if isinstance(fill_value, float): + return _DEFAULT_FLOAT + if isinstance(fill_value, complex): + return np.dtype("complex128") + raise TypeError(f"fill_value must be a Python scalar, not {type(fill_value).__name__}") + + +def full(shape, fill_value, *, dtype=None, device=None): + check_device(device) + dt = _fill_dtype(fill_value) if dtype is None else as_dtype(dtype) + if isinstance(fill_value, Array): + fill_value = fill_value._value("full's fill_value") + if isinstance(fill_value, int) and dt.kind in "iu": + info = np.iinfo(dt) + if not info.min <= fill_value <= info.max: + raise OverflowError(f"fill_value {fill_value} is out of range for {dt}") + return _full(shape, fill_value, dt) + + +def zeros(shape, *, dtype=None, device=None): + check_device(device) + dt = _DEFAULT_FLOAT if dtype is None else as_dtype(dtype) + return _full(shape, 0, dt) + + +def ones(shape, *, dtype=None, device=None): + check_device(device) + dt = _DEFAULT_FLOAT if dtype is None else as_dtype(dtype) + return _full(shape, 1, dt) + + +def empty(shape, *, dtype=None, device=None): + """Zeros: an empty array stores nothing, and its fill is 0 of the dtype.""" + return zeros(shape, dtype=dtype, device=device) + + +def _like(x, dtype, device): + if not isinstance(x, Array): + raise TypeError(f"needs an Array, not {type(x).__name__}") + check_device(device) + return x.shape, x.dtype if dtype is None else as_dtype(dtype) + + +def full_like(x, /, fill_value, *, dtype=None, device=None): + shape, dt = _like(x, dtype, device) + if isinstance(fill_value, Array): + fill_value = fill_value._value("full_like's fill_value") + return _full(shape, fill_value, dt) + + +def zeros_like(x, /, *, dtype=None, device=None): + shape, dt = _like(x, dtype, device) + return _full(shape, 0, dt) + + +def ones_like(x, /, *, dtype=None, device=None): + shape, dt = _like(x, dtype, device) + return _full(shape, 1, dt) + + +def empty_like(x, /, *, dtype=None, device=None): + return zeros_like(x, dtype=dtype, device=device) + + +def eye(n_rows, n_cols=None, /, *, k=0, dtype=None, device=None): + check_device(device) + dt = _DEFAULT_FLOAT if dtype is None else as_dtype(dtype) + n_cols = n_rows if n_cols is None else n_cols + shape = _shape((n_rows, n_cols)) + i = np.arange(builtins.max(0, -k), builtins.max(0, builtins.min(n_rows, n_cols - k))) + if i.size == 0: + return _full(shape, 0, dt) + T = Tensor.from_coo((i, i + k), _utils.scalar(1, dt), gb_dtype(dt), shape=shape) + return Array._new(T, 0) + + +def arange(start, /, stop=None, step=1, *, dtype=None, device=None): + check_device(device) + if stop is None: + start, stop = 0, start + dt = None if dtype is None else as_dtype(dtype) + if step == 0: + raise ValueError("arange's step must not be 0") + n = builtins.max(0, math.ceil((stop - start) / step)) + _utils.check_dense(n, "arange") + if dt is None: + if builtins.all(isinstance(v, int) for v in (start, stop, step)): + dt = np.dtype("int64") + else: + dt = _DEFAULT_FLOAT + a = np.arange(start, stop, step, dtype=dt) + return _from_numpy(a[:n] if a.size > n else a) + + +def linspace(start, stop, /, num, *, dtype=None, device=None, endpoint=True): + check_device(device) + _utils.check_dense(num, "linspace") + if dtype is None: + is_complex = builtins.any(isinstance(v, complex) for v in (start, stop)) + dt = np.dtype("complex128") if is_complex else _DEFAULT_FLOAT + else: + dt = as_dtype(dtype) + a = np.linspace(start, stop, num, endpoint=endpoint, dtype=dt) + return _from_numpy(a) + + +def from_dlpack(x, /, *, device=None, copy=None): + check_device(device) + if copy is False: + raise ValueError("from_dlpack stores its own copy; copy=False") + if isinstance(x, Array): + return x._copy() + return _from_numpy(np.from_dlpack(x)) + + +def meshgrid(*arrays, indexing="xy"): + from ._manipulation import _broadcast_to, reshape + + if indexing not in {"xy", "ij"}: + raise ValueError(f"indexing must be 'xy' or 'ij', not {indexing!r}") + for x in arrays: + if not isinstance(x, Array) or x.ndim != 1: + raise ValueError("meshgrid needs 1-d arrays") + if len({x.dtype for x in arrays}) > 1: + raise ValueError("meshgrid needs arrays of one data type") + n = len(arrays) + shape = [x.shape[0] for x in arrays] + if indexing == "xy" and n >= 2: + shape[0], shape[1] = shape[1], shape[0] + out = [] + for i, x in enumerate(arrays): + axis = i + if indexing == "xy" and n >= 2 and i < 2: + axis = 1 - i + s = [1] * n + s[axis] = x.shape[0] + out.append(_broadcast_to(reshape(x, tuple(s)), tuple(shape))) + return out + + +def _tri(x, k, lower): + if not isinstance(x, Array): + raise TypeError(f"needs an Array, not {type(x).__name__}") + if x.ndim < 2: + raise ValueError("tril and triu need at least 2 dimensions") + zero = _default_fill(x.dtype) + if not _utils.same_value(x._fill, zero): + # The kept triangle holds the fill where nothing is stored: store it there first + _utils.check_dense(x.size, "tril or triu of an array whose fill is not 0") + x = _from_numpy(x.todense()) + T = x._T() + if T.nvals == 0: + return _full(x.shape, 0, x.dtype) + coords, values = T.to_coo() + i, j = coords[-2], coords[-1] + keep = (j - i <= k) if lower else (j - i >= k) + if not keep.any(): + return _full(x.shape, 0, x.dtype) + coords = tuple(c[keep] for c in coords) + return Array._new(Tensor.from_coo(coords, values[keep], T.dtype, shape=T.shape), zero) + + +def tril(x, /, *, k=0): + return _tri(x, k, lower=True) + + +def triu(x, /, *, k=0): + return _tri(x, k, lower=False) diff --git a/graphblas/array_api/_dtypes.py b/graphblas/array_api/_dtypes.py new file mode 100644 index 000000000..f107a3afa --- /dev/null +++ b/graphblas/array_api/_dtypes.py @@ -0,0 +1,172 @@ +"""Data types, type promotion and the data type functions. + +The namespace's dtypes are numpy dtype instances (``xp.float64 is np.dtype("float64")``), and +an Array's ``dtype`` is one of them, so they compare and hash alike. Type promotion is numpy's, +with Python scalars weak (NEP 50). +""" + +import builtins +from typing import NamedTuple + +import numpy as np + +__all__ = ["astype", "can_cast", "finfo", "iinfo", "isdtype", "result_type"] + +_NAMES = ( + "bool", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "float32", + "float64", + "complex64", + "complex128", +) +DTYPES = {name: np.dtype(name) for name in _NAMES} +_SUPPORTED = frozenset(DTYPES.values()) + +_KINDS = { + "bool": frozenset({DTYPES["bool"]}), + "signed integer": frozenset(DTYPES[n] for n in ("int8", "int16", "int32", "int64")), + "unsigned integer": frozenset(DTYPES[n] for n in ("uint8", "uint16", "uint32", "uint64")), + "real floating": frozenset({DTYPES["float32"], DTYPES["float64"]}), + "complex floating": frozenset({DTYPES["complex64"], DTYPES["complex128"]}), +} +_KINDS["integral"] = _KINDS["signed integer"] | _KINDS["unsigned integer"] +_KINDS["numeric"] = _KINDS["integral"] | _KINDS["real floating"] | _KINDS["complex floating"] + + +def as_dtype(dtype): + """``dtype`` (a dtype of this namespace, or anything numpy takes for one) as a supported + numpy dtype; TypeError otherwise. + """ + if dtype is None: + raise TypeError("dtype must not be None here") + try: + dt = np.dtype(dtype) + except TypeError: + raise TypeError(f"{dtype!r} is not a data type") from None + if dt not in _SUPPORTED: + raise TypeError(f"data type {dt} is not supported by graphblas.array_api") + return dt + + +def check_supported(dt, what): + if dt not in _SUPPORTED: + raise TypeError(f"{what} would have data type {dt}, which graphblas.array_api lacks") + return dt + + +def _dtype_of(x): + from ._array import Array + + if isinstance(x, Array): + return x.dtype + return as_dtype(x) + + +def isdtype(dtype, kind): + if not isinstance(dtype, np.dtype) or dtype not in _SUPPORTED: + try: + dtype = as_dtype(dtype) + except TypeError: + raise TypeError(f"isdtype needs a data type, not {dtype!r}") from None + if isinstance(kind, tuple): + return builtins.any(isdtype(dtype, k) for k in kind) + if isinstance(kind, str): + try: + return dtype in _KINDS[kind] + except KeyError: + raise ValueError(f"unknown data type kind {kind!r}") from None + return dtype == as_dtype(kind) + + +def result_type(*arrays_and_dtypes): + from ._array import Array + + if not arrays_and_dtypes: + raise ValueError("result_type needs at least one array or data type") + args = [] + has_typed = False + for x in arrays_and_dtypes: + if isinstance(x, Array): + args.append(x.dtype) + has_typed = True + elif isinstance(x, (builtins.bool, int, float, complex)) and not isinstance(x, np.generic): + args.append(x) # weak, as numpy takes a Python scalar + else: + args.append(as_dtype(x)) + has_typed = True + if not has_typed: + raise ValueError("result_type needs at least one array or data type, not only scalars") + # numpy promotes many arguments pairwise in an order of its own: fold left to right, as the + # Array API's promotion table does, so that the result does not depend on the order + rt = None + for arg in args: + if isinstance(arg, np.dtype): + rt = arg if rt is None else np.result_type(rt, arg) + for arg in args: + if not isinstance(arg, np.dtype): + rt = np.result_type(rt, arg) + return check_supported(rt, "result_type") + + +def can_cast(from_, to, /): + return builtins.bool(np.can_cast(_dtype_of(from_), as_dtype(to), casting="safe")) + + +class finfo_object(NamedTuple): + bits: int + eps: float + max: float + min: float + smallest_normal: float + dtype: np.dtype + + +class iinfo_object(NamedTuple): + bits: int + max: int + min: int + dtype: np.dtype + + +def finfo(type, /): # noqa: A002 (the Array API's name) + dt = _dtype_of(type) + if dt.kind not in "fc": + raise TypeError(f"finfo needs a floating-point data type, not {dt}") + info = np.finfo(dt) + return finfo_object( + builtins.int(info.bits), + float(info.eps), + float(info.max), + float(info.min), + float(info.smallest_normal), + np.dtype(info.dtype), + ) + + +def iinfo(type, /): # noqa: A002 (the Array API's name) + dt = _dtype_of(type) + if dt.kind not in "iu": + raise TypeError(f"iinfo needs an integer data type, not {dt}") + info = np.iinfo(dt) + return iinfo_object(builtins.int(info.bits), builtins.int(info.max), builtins.int(info.min), dt) + + +def astype(x, dtype, /, *, copy=True, device=None): + from ._array import Array, check_device + from ._elementwise import cast + + check_device(device) + if not isinstance(x, Array): + raise TypeError(f"astype needs an Array, not {type(x)}") + dt = as_dtype(dtype) + if dt == x.dtype: + return x._copy() if copy else x + return cast(x, dt) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py new file mode 100644 index 000000000..f0483f340 --- /dev/null +++ b/graphblas/array_api/_elementwise.py @@ -0,0 +1,693 @@ +"""Element-wise functions: numpy's answer, bit for bit, on every position. + +A function of one array maps the stored entries and the fill alike. A function of two arrays +takes the union of their entries, an unstored position of one operand taking its fill there, +and the new fill is the function of the two fills. Python scalars are weak (NEP 50). + +Where a GraphBLAS operator gives numpy's answer exactly for the dtypes at hand, the entries +are computed in GraphBLAS; otherwise numpy maps the stored values (the pattern stays as it is, +see ``_utils.map_values``). numpy also computes every fill, so an entry equal to the fill maps +to the new fill. +""" + +import builtins + +import numpy as np + +from .. import binary, unary +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._dtypes import check_supported +from ._utils import gb_dtype, gb_scalar + +__all__ = [ + "abs", + "acos", + "acosh", + "add", + "asin", + "asinh", + "atan", + "atan2", + "atanh", + "bitwise_and", + "bitwise_invert", + "bitwise_left_shift", + "bitwise_or", + "bitwise_right_shift", + "bitwise_xor", + "ceil", + "clip", + "conj", + "copysign", + "cos", + "cosh", + "divide", + "equal", + "exp", + "expm1", + "floor", + "floor_divide", + "greater", + "greater_equal", + "hypot", + "imag", + "isfinite", + "isinf", + "isnan", + "less", + "less_equal", + "log", + "log1p", + "log2", + "log10", + "logaddexp", + "logical_and", + "logical_not", + "logical_or", + "logical_xor", + "maximum", + "minimum", + "multiply", + "negative", + "nextafter", + "not_equal", + "positive", + "pow", + "real", + "reciprocal", + "remainder", + "round", + "sign", + "signbit", + "sin", + "sinh", + "sqrt", + "square", + "subtract", + "tan", + "tanh", + "trunc", +] + +_FLOAT = "f" +_INEXACT = "fc" +_INTEGRAL = "iu" +_REAL = "biuf" + + +# Operands +def _as_operand(x, name): + """An Array; a Python scalar stays as it is (weak); a numpy scalar or array becomes an + Array (strong, as numpy treats it). + """ + if isinstance(x, Array): + return x + if isinstance(x, (np.generic, np.ndarray)): + from ._creation import asarray + + return asarray(x) + if isinstance(x, (builtins.bool, int, float, complex)): + return x + raise TypeError(f"{name} needs arrays or Python scalars, not {type(x).__name__}") + + +def _probe(x): + """What numpy's type resolution sees: an empty array of an Array's dtype, or the weak + Python scalar itself. + """ + if isinstance(x, Array): + return np.empty(0, x.dtype) + return x + + +def _fill_operand(x): + """The value of ``x`` at an unstored position, typed as numpy would type ``x``.""" + if isinstance(x, Array): + return np.asarray(x._fill) + return x + + +def _result(T, fill): + return Array._new(T, fill) + + +def _zero_d(dtype, fill): + return Array._new(Tensor(gb_dtype(dtype), ()), fill) + + +# Unary functions +def _unary(name, ufunc, x, gb=None, *, const=None): + """``ufunc`` of every position of ``x``. + + ``gb(dtype)`` gives a GraphBLAS UnaryOp equal to ``ufunc`` on that dtype, a callable + taking the Tensor, or None; ``const(dtype)`` gives the value of every position when the + result does not depend on the entries (``isnan`` of an integer array), else None. + """ + x = _as_operand(x, name) + if not isinstance(x, Array): + from ._creation import asarray + + x = asarray(x) + with np.errstate(all="ignore"): + rdt = ufunc(np.empty(0, x.dtype)).dtype + fill = ufunc(np.asarray(x._fill)) + check_supported(rdt, f"{name} of {x.dtype}") + T = x._T() + if T.ndim == 0: + return _zero_d(rdt, fill) + if const is not None and (value := const(x.dtype)) is not None: + return _result(_utils.empty_tensor(rdt, T.shape), value) + if T.nvals == 0: + return _result(_utils.empty_tensor(rdt, T.shape), fill) + op = gb(x.dtype) if gb is not None else None + if op is None: + R = _utils.map_values(ufunc, [T], rdt) + elif op == "copy": + R = T.dup(dtype=gb_dtype(rdt)) + elif callable(op) and not hasattr(op, "opclass"): + R = op(T) + else: + R = T.apply(op).new(dtype=gb_dtype(rdt)) + return _result(R, fill) + + +def _typed(op, kinds): + """``gb`` for _unary: ``op`` typed for dtypes of these kinds, else numpy.""" + + def get(dtype): + return op[gb_dtype(dtype)] if dtype.kind in kinds else None + + return get + + +def _copy_if(kinds): + def get(dtype): + return "copy" if dtype.kind in kinds else None + + return get + + +def _first_of(*getters): + def get(dtype): + for g in getters: + if (op := g(dtype)) is not None: + return op + return None + + return get + + +def abs(x, /): # noqa: A001 (the Array API's name) + # (GraphBLAS's integer abs of the most negative value is undefined in C; numpy wraps) + return _unary("abs", np.abs, x, _first_of(_typed(unary.abs, _FLOAT), _copy_if("ub"))) + + +def negative(x, /): + return _unary("negative", np.negative, x, _typed(unary.ainv, "iufc")) + + +def positive(x, /): + return _unary("positive", np.positive, x, _copy_if("iufc")) + + +def sqrt(x, /): + # IEEE square root is correctly rounded, in C as in numpy + return _unary("sqrt", np.sqrt, x, _typed(unary.sqrt, _FLOAT)) + + +def floor(x, /): + return _unary("floor", np.floor, x, _first_of(_typed(unary.floor, _FLOAT), _copy_if("biu"))) + + +def ceil(x, /): + return _unary("ceil", np.ceil, x, _first_of(_typed(unary.ceil, _FLOAT), _copy_if("biu"))) + + +def trunc(x, /): + return _unary("trunc", np.trunc, x, _first_of(_typed(unary.trunc, _FLOAT), _copy_if("biu"))) + + +def round(x, /): # noqa: A001 (the Array API's name) + # numpy rounds half to even; C's round (GraphBLAS's) rounds half away from zero + return _unary("round", np.round, x, _copy_if("biu")) + + +def _const_unless(kinds, value): + def get(dtype): + return None if dtype.kind in kinds else value + + return get + + +def isnan(x, /): + return _unary( + "isnan", np.isnan, x, _typed(unary.isnan, _INEXACT), const=_const_unless("fc", False) + ) + + +def isinf(x, /): + return _unary( + "isinf", np.isinf, x, _typed(unary.isinf, _INEXACT), const=_const_unless("fc", False) + ) + + +def isfinite(x, /): + return _unary( + "isfinite", + np.isfinite, + x, + _typed(unary.isfinite, _INEXACT), + const=_const_unless("fc", True), + ) + + +def logical_not(x, /): + return _unary("logical_not", np.logical_not, x, _typed(unary.lnot, "b")) + + +def bitwise_invert(x, /): + def gb(dtype): + if dtype.kind == "b": + return unary.lnot[gb_dtype(dtype)] + if dtype.kind in _INTEGRAL: + return unary.bnot[gb_dtype(dtype)] + return None + + return _unary("bitwise_invert", np.invert, x, gb) + + +def square(x, /): + def gb(dtype): + if dtype.kind not in "iuf": + return None + # x * x on the stored entries, as numpy's square computes it + return lambda T: T.ewise_mult(T, binary.times[gb_dtype(dtype)]).new() + + return _unary("square", np.square, x, gb) + + +def reciprocal(x, /): + return _unary("reciprocal", np.reciprocal, x, _typed(unary.minv, _FLOAT)) + + +def conj(x, /): + return _unary("conj", np.conj, x, _first_of(_typed(unary.conj, "c"), _copy_if(_REAL))) + + +def real(x, /): + return _unary("real", np.real, x, _first_of(_typed(unary.creal, "c"), _copy_if(_REAL))) + + +def imag(x, /): + def const(dtype): + return None if dtype.kind == "c" else 0 + + return _unary("imag", np.imag, x, _typed(unary.cimag, "c"), const=const) + + +def sign(x, /): + return _unary("sign", np.sign, x) + + +def signbit(x, /): + return _unary("signbit", np.signbit, x) + + +def _numpy_unary(name, ufunc): + def func(x, /): + return _unary(name, ufunc, x) + + func.__name__ = name + func.__qualname__ = name + return func + + +acos = _numpy_unary("acos", np.arccos) +acosh = _numpy_unary("acosh", np.arccosh) +asin = _numpy_unary("asin", np.arcsin) +asinh = _numpy_unary("asinh", np.arcsinh) +atan = _numpy_unary("atan", np.arctan) +atanh = _numpy_unary("atanh", np.arctanh) +cos = _numpy_unary("cos", np.cos) +cosh = _numpy_unary("cosh", np.cosh) +exp = _numpy_unary("exp", np.exp) +expm1 = _numpy_unary("expm1", np.expm1) +log = _numpy_unary("log", np.log) +log1p = _numpy_unary("log1p", np.log1p) +log2 = _numpy_unary("log2", np.log2) +log10 = _numpy_unary("log10", np.log10) +sin = _numpy_unary("sin", np.sin) +sinh = _numpy_unary("sinh", np.sinh) +tan = _numpy_unary("tan", np.tan) +tanh = _numpy_unary("tanh", np.tanh) + + +# Binary functions +_MANTISSA = {"f": {4: 24, 8: 53}, "c": {8: 24, 16: 53}} + + +def _exact_cast(src, dst): + """Whether every value of dtype ``src`` converts to ``dst`` without change.""" + if src == dst or src.kind == "b": + return True + if src.kind in _INTEGRAL: + if dst.kind in _INTEGRAL: + return np.can_cast(src, dst, casting="safe") + if dst.kind in _INEXACT and dst != np.dtype(bool): + return src.itemsize * 8 - (src.kind == "i") <= _MANTISSA[dst.kind][dst.itemsize] + return False + if src.kind == "f" and dst.kind in _INEXACT: + return _MANTISSA["f"][src.itemsize] <= _MANTISSA[dst.kind][dst.itemsize] + if src.kind == "c" and dst.kind == "c": + return src.itemsize <= dst.itemsize + return False + + +def _python_scalar_fits(value, dtype): + """Whether numpy converts the Python scalar ``value`` to ``dtype`` without changing it + (out-of-range Python ints get special treatment in numpy's comparisons). + """ + if dtype.kind in _INTEGRAL and isinstance(value, int): + info = np.iinfo(dtype) + return info.min <= value <= info.max + return True + + +class _Binary: + """How to compute one binary function. + + ``kind`` is "arith" (the op works in the result dtype), "cmp" (in the promoted dtype of + the operands, with a bool result) or "logical" (in bool). ``gb(dtype)`` gives the + GraphBLAS operator that equals ``ufunc`` on operands of the dtype it works in, or None. + ``annihilator(dtype)`` gives the fill value that, held by both operands, makes the union + equal to the intersection (``ewise_mult``), or None. + """ + + def __init__(self, name, ufunc, kind="arith", gb=None, annihilator=None): + self.name = name + self.ufunc = ufunc + self.kind = kind + self.gb = gb + self.annihilator = annihilator + + +def _binary(spec, x1, x2): + a = _as_operand(x1, spec.name) + b = _as_operand(x2, spec.name) + ufunc = spec.ufunc + if not isinstance(a, Array) and not isinstance(b, Array): + from ._creation import asarray + + a = asarray(a) + with np.errstate(all="ignore"): + rdt = ufunc(_probe(a), _probe(b)).dtype + fill = ufunc(_fill_operand(a), _fill_operand(b)) + check_supported(rdt, f"{spec.name} of {_dtype_name(a)} and {_dtype_name(b)}") + a_arr = isinstance(a, Array) and a.ndim > 0 + b_arr = isinstance(b, Array) and b.ndim > 0 + if not a_arr and not b_arr: + return _zero_d(rdt, fill) + if spec.kind == "arith": + ct = rdt + elif spec.kind == "cmp": + try: + ct = np.result_type(_probe(a), _probe(b)) + except (OverflowError, TypeError): + ct = None # (a Python int out of range: numpy compares it exactly) + else: + ct = np.dtype(bool) + op = None if ct is None else _gb_op(spec, a, b, ct) + if a_arr and b_arr: + return _binary_arrays(spec, a, b, op, ct, rdt, fill) + if a_arr: + T, other, left = a._T(), b, False + else: + T, other, left = b._T(), a, True + if T.nvals == 0: + return _result(_utils.empty_tensor(rdt, T.shape), fill) + if op is not None: + s = gb_scalar(_fill_operand(other), ct) + expr = T.apply(op, left=s) if left else T.apply(op, right=s) + return _result(expr.new(dtype=gb_dtype(rdt)), fill) + s = _fill_operand(other) + if left: + R = _utils.map_values(lambda v: ufunc(s, v), [T], rdt) + else: + R = _utils.map_values(lambda v: ufunc(v, s), [T], rdt) + return _result(R, fill) + + +def _dtype_name(x): + return str(x.dtype) if isinstance(x, Array) else type(x).__name__ + + +def _gb_op(spec, a, b, ct): + """The typed GraphBLAS op for ``spec`` working in ``ct``, when it gives numpy's answer; + else None. + """ + if spec.gb is None: + return None + for x in (a, b): + if isinstance(x, Array): + if spec.kind != "arith" and not _exact_cast(x.dtype, ct): + return None + elif not _python_scalar_fits(x, ct): + return None + return spec.gb(ct) + + +def _binary_arrays(spec, a, b, op, ct, rdt, fill): + Ta, Tb = a._T(), b._T() + try: + shape = np.broadcast_shapes(Ta.shape, Tb.shape) + except ValueError: + raise ValueError( + f"{spec.name}: shapes {Ta.shape} and {Tb.shape} do not broadcast" + ) from None + if Ta.nvals == 0 and Tb.nvals == 0: + return _result(_utils.empty_tensor(rdt, shape), fill) + if op is not None: + fa = _utils.scalar(a._fill, ct) + fb = _utils.scalar(b._fill, ct) + zero = spec.annihilator(ct) if spec.annihilator is not None else None + if zero is not None and _utils.same_value(fa, zero) and _utils.same_value(fb, zero): + # Where either operand holds its fill the result is the fill: the intersection + R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) + return _result(R, fill) + _utils.expansion_check(Ta, Tb, shape) + R = Ta.ewise_union(Tb, op, gb_scalar(fa, ct), gb_scalar(fb, ct), allow_expansion=True).new( + dtype=gb_dtype(rdt) + ) + return _result(R, fill) + # numpy over the values: both operands' values on the union of entries, aligned + _utils.expansion_check(Ta, Tb, shape) + da, db = a.dtype, b.dtype + A = Ta.ewise_union( + Tb, binary.first[gb_dtype(da)], gb_scalar(a._fill, da), gb_scalar(0, da), + allow_expansion=True, + ).new() # fmt: skip + B = Ta.ewise_union( + Tb, binary.second[gb_dtype(db)], gb_scalar(0, db), gb_scalar(b._fill, db), + allow_expansion=True, + ).new() # fmt: skip + return _result(_utils.map_values(spec.ufunc, [A, B], rdt), fill) + + +def _typed_for(op_by_kind): + """``gb`` for _Binary: ``op_by_kind[dtype.kind]`` typed for the dtype, else None.""" + + def get(dtype): + op = op_by_kind.get(dtype.kind) + return None if op is None else op[gb_dtype(dtype)] + + return get + + +def _zero(dtype): + return _utils.scalar(0, dtype) + + +_SPECS = { + "add": _Binary( + "add", np.add, gb=_typed_for({"b": binary.lor, **dict.fromkeys("iufc", binary.plus)}) + ), + "subtract": _Binary( + "subtract", np.subtract, gb=_typed_for(dict.fromkeys("iufc", binary.minus)) + ), + # Complex products and quotients are computed by numpy: C's formulas round otherwise. + # A float zero fill is not an annihilator: 0.0 * -1.0 is -0.0, and 0.0 * inf is NaN. + "multiply": _Binary( + "multiply", + np.multiply, + gb=_typed_for({"b": binary.land, **dict.fromkeys("iuf", binary.times)}), + annihilator=lambda dt: _zero(dt) if dt.kind in "biu" else None, + ), + "divide": _Binary("divide", np.divide, gb=_typed_for({"f": binary.truediv})), + "equal": _Binary("equal", np.equal, "cmp", gb=_typed_for(dict.fromkeys("biufc", binary.eq))), + "not_equal": _Binary( + "not_equal", np.not_equal, "cmp", gb=_typed_for(dict.fromkeys("biufc", binary.ne)) + ), + "less": _Binary("less", np.less, "cmp", gb=_typed_for(dict.fromkeys("biuf", binary.lt))), + "less_equal": _Binary( + "less_equal", np.less_equal, "cmp", gb=_typed_for(dict.fromkeys("biuf", binary.le)) + ), + "greater": _Binary( + "greater", np.greater, "cmp", gb=_typed_for(dict.fromkeys("biuf", binary.gt)) + ), + "greater_equal": _Binary( + "greater_equal", np.greater_equal, "cmp", gb=_typed_for(dict.fromkeys("biuf", binary.ge)) + ), + "logical_and": _Binary( + "logical_and", + np.logical_and, + "logical", + gb=lambda dt: binary.land[gb_dtype(dt)], + annihilator=_zero, + ), + "logical_or": _Binary( + "logical_or", np.logical_or, "logical", gb=lambda dt: binary.lor[gb_dtype(dt)] + ), + "logical_xor": _Binary( + "logical_xor", np.logical_xor, "logical", gb=lambda dt: binary.lxor[gb_dtype(dt)] + ), + "bitwise_and": _Binary( + "bitwise_and", + np.bitwise_and, + gb=_typed_for({"b": binary.land, "i": binary.band, "u": binary.band}), + annihilator=_zero, + ), + "bitwise_or": _Binary( + "bitwise_or", + np.bitwise_or, + gb=_typed_for({"b": binary.lor, "i": binary.bor, "u": binary.bor}), + ), + "bitwise_xor": _Binary( + "bitwise_xor", + np.bitwise_xor, + gb=_typed_for({"b": binary.lxor, "i": binary.bxor, "u": binary.bxor}), + ), + # GraphBLAS's max and min of floats drop NaN (C's fmax); numpy's propagate it + "maximum": _Binary("maximum", np.maximum, gb=_typed_for(dict.fromkeys("biu", binary.max))), + "minimum": _Binary("minimum", np.minimum, gb=_typed_for(dict.fromkeys("biu", binary.min))), + "copysign": _Binary("copysign", np.copysign, gb=_typed_for({"f": binary.copysign})), + # numpy's own definitions: floored division and remainder (with its results for division + # by zero), integer powers, shifts past the width, and libm-free special values + "floor_divide": _Binary("floor_divide", np.floor_divide), + "remainder": _Binary("remainder", np.remainder), + "pow": _Binary("pow", np.power), + "atan2": _Binary("atan2", np.arctan2), + "hypot": _Binary("hypot", np.hypot), + "logaddexp": _Binary("logaddexp", np.logaddexp), + "nextafter": _Binary("nextafter", np.nextafter), + "bitwise_left_shift": _Binary("bitwise_left_shift", np.left_shift), + "bitwise_right_shift": _Binary("bitwise_right_shift", np.right_shift), +} + + +def _make_binary(name): + spec = _SPECS[name] + + def func(x1, x2, /): + return _binary(spec, x1, x2) + + func.__name__ = name + func.__qualname__ = name + return func + + +add = _make_binary("add") +subtract = _make_binary("subtract") +multiply = _make_binary("multiply") +divide = _make_binary("divide") +equal = _make_binary("equal") +not_equal = _make_binary("not_equal") +less = _make_binary("less") +less_equal = _make_binary("less_equal") +greater = _make_binary("greater") +greater_equal = _make_binary("greater_equal") +logical_and = _make_binary("logical_and") +logical_or = _make_binary("logical_or") +logical_xor = _make_binary("logical_xor") +bitwise_and = _make_binary("bitwise_and") +bitwise_or = _make_binary("bitwise_or") +bitwise_xor = _make_binary("bitwise_xor") +maximum = _make_binary("maximum") +minimum = _make_binary("minimum") +copysign = _make_binary("copysign") +floor_divide = _make_binary("floor_divide") +remainder = _make_binary("remainder") +pow = _make_binary("pow") # noqa: A001 (the Array API's name) +atan2 = _make_binary("atan2") +hypot = _make_binary("hypot") +logaddexp = _make_binary("logaddexp") +nextafter = _make_binary("nextafter") +bitwise_left_shift = _make_binary("bitwise_left_shift") +bitwise_right_shift = _make_binary("bitwise_right_shift") + + +def clip(x, /, min=None, max=None): + """``minimum(maximum(x, min), max)`` in the dtype of ``x`` (NaN propagates).""" + if not isinstance(x, Array): + raise TypeError(f"clip needs an Array, not {type(x).__name__}") + out = x + if min is not None: + out = maximum(out, _clip_bound(min, x)) + if max is not None: + out = minimum(out, _clip_bound(max, x)) + if out is x: + return x._copy() + if out.dtype != x.dtype: + out = cast(out, x.dtype) + return out + + +def _clip_bound(bound, x): + """A Python scalar bound in the dtype of ``x`` when it fits (the Array API's clip keeps the + dtype of ``x``). + """ + if isinstance(bound, (builtins.bool, int, float)) and not isinstance(bound, np.generic): + if x.dtype.kind in _INTEGRAL and isinstance(bound, float): + return bound + if x.dtype.kind in _INTEGRAL and not _python_scalar_fits(bound, x.dtype): + info = np.iinfo(x.dtype) + return builtins.min(builtins.max(bound, int(info.min)), int(info.max)) + return bound + + +# Casting (astype) +def _gb_cast_ok(src, dst): + """Whether GraphBLAS's typecast from ``src`` to ``dst`` gives numpy's ``astype``. + + C converts integers and floats as numpy does; float to integer differs where the value + does not fit (NaN, inf, out of range: numpy's result is undefined there, GraphBLAS + saturates), and complex to real is numpy's own rule. + """ + if src.kind == "c": + return dst.kind in "cb" + if src.kind == "f": + return dst.kind in "fcb" + return True + + +def _astype(values, dt): + # numpy's astype from complex to real keeps the real part (and warns that it does) + if values.dtype.kind == "c" and dt.kind not in "cb": + values = values.real + return values.astype(dt) + + +def cast(x, dtype): + """``x.astype(dtype)`` by numpy's rules, entries and fill alike.""" + dt = np.dtype(dtype) + with np.errstate(all="ignore"): + fill = _astype(np.asarray(x._fill), dt) + T = x._T() + if T.ndim == 0: + return _zero_d(dt, fill) + if T.nvals == 0: + return _result(_utils.empty_tensor(dt, T.shape), fill) + if _gb_cast_ok(x.dtype, dt): + return _result(T.dup(dtype=gb_dtype(dt)), fill) + return _result(_utils.map_values(lambda v: _astype(v, dt), [T], dt), fill) diff --git a/graphblas/array_api/_indexing.py b/graphblas/array_api/_indexing.py new file mode 100644 index 000000000..49bfef70b --- /dev/null +++ b/graphblas/array_api/_indexing.py @@ -0,0 +1,99 @@ +"""Indexing: ``x[key]`` (and later ``x[key] = value``, ``take``, ``take_along_axis``). + +Basic keys so far: integers, slices, ``...`` and ``None``, as numpy and the Array API index. +""" + +import operator + +import numpy as np + +from ..core.tensor import Tensor +from ._array import Array + +__all__ = [] + + +def _basic_key(x, key): + """``key`` as a tuple of ints and slices for every axis, and the positions of the new + axes (None) in the result; None when the key is not basic. + """ + if not isinstance(key, tuple): + key = (key,) + items = [] + for k in key: + if k is None or k is Ellipsis or isinstance(k, slice): + items.append(k) + elif isinstance(k, (bool, np.bool_)): + return None # (a boolean index is a mask) + elif isinstance(k, (int, np.integer)): + items.append(int(k)) + elif isinstance(k, Array) and k.ndim == 0 and k.dtype.kind in "iu": + items.append(operator.index(k)) + else: + return None + if sum(k is Ellipsis for k in items) > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + nindexed = sum(k is not None and k is not Ellipsis for k in items) + if nindexed > x.ndim: + raise IndexError(f"too many indices: {nindexed} for an array of {x.ndim} dimensions") + if Ellipsis in items: + i = items.index(Ellipsis) + items[i : i + 1] = [slice(None)] * (x.ndim - nindexed) + else: + items.extend([slice(None)] * (x.ndim - nindexed)) + axes_key = [] + new_axes = [] + out_axis = 0 + axis = 0 + for k in items: + if k is None: + new_axes.append(out_axis) + out_axis += 1 + continue + n = x.shape[axis] + if isinstance(k, int): + if not -n <= k < n: + raise IndexError(f"index {k} is out of bounds for axis {axis} with size {n}") + axes_key.append(k % n) + else: + axes_key.append(k) + out_axis += 1 + axis += 1 + return tuple(axes_key), new_axes + + +def getitem(x, key): + basic = _basic_key(x, key) + if basic is None: + return _advanced_getitem(x, key) + axes_key, new_axes = basic + T = x._T() + if x.ndim == 0: + out = x._copy() + elif all(k == slice(None) for k in axes_key): + out = Array._new(T.dup(), x._fill) + else: + shape = tuple( + len(range(*k.indices(n))) + for k, n in zip(axes_key, T.shape, strict=True) + if isinstance(k, slice) + ) + if T.nvals == 0 or 0 in shape: + out = Array._new(Tensor(T.dtype, shape), x._fill) + else: + out = Array._new(T[axes_key].new(), x._fill) + if new_axes: + from ._manipulation import expand_dims + + out = expand_dims(out, tuple(new_axes)) + return out + + +def _advanced_getitem(x, key): + raise NotImplementedError( + f"indexing with {key!r}: only integers, slices, ... and None are supported so far" + ) + + +def setitem(x, key, value): + raise NotImplementedError("assignment by index is not supported yet") diff --git a/graphblas/array_api/_info.py b/graphblas/array_api/_info.py new file mode 100644 index 000000000..aab48b339 --- /dev/null +++ b/graphblas/array_api/_info.py @@ -0,0 +1,53 @@ +"""``__array_namespace_info__``: what the namespace supports.""" + +from ._array import CPU_DEVICE +from ._dtypes import _KINDS, DTYPES + +__all__ = ["__array_namespace_info__"] + + +class Info: + __module__ = "graphblas.array_api" + + def capabilities(self): + return { + "boolean indexing": True, + "data-dependent shapes": True, + "max dimensions": None, + } + + def default_device(self): + return CPU_DEVICE + + def devices(self): + return [CPU_DEVICE] + + def default_dtypes(self, *, device=None): + _check(device) + return { + "real floating": DTYPES["float64"], + "complex floating": DTYPES["complex128"], + "integral": DTYPES["int64"], + "indexing": DTYPES["int64"], + } + + def dtypes(self, *, device=None, kind=None): + _check(device) + if kind is None: + return dict(DTYPES) + kinds = kind if isinstance(kind, tuple) else (kind,) + selected = set() + for k in kinds: + if k not in _KINDS: + raise ValueError(f"unknown data type kind {k!r}") + selected |= _KINDS[k] + return {name: dt for name, dt in DTYPES.items() if dt in selected} + + +def _check(device): + if device is not None and device != CPU_DEVICE: + raise ValueError(f"unknown device {device!r}") + + +def __array_namespace_info__(): + return Info() diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py new file mode 100644 index 000000000..a22e74f0a --- /dev/null +++ b/graphblas/array_api/_linalg.py @@ -0,0 +1,3 @@ +"""Linear algebra in the main namespace: matmul, matrix_transpose, tensordot, vecdot.""" + +__all__ = [] diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py new file mode 100644 index 000000000..8053f5f97 --- /dev/null +++ b/graphblas/array_api/_manipulation.py @@ -0,0 +1,116 @@ +"""Manipulation functions: shape changes share the stored data (copy-on-write).""" + +import math + +import numpy as np + +from .. import binary +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._utils import gb_dtype + +__all__ = ["expand_dims", "permute_dims", "reshape", "squeeze"] + + +def _check_array(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + + +def _normalize_axis(axis, ndim): + axis = int(axis) + if not -ndim <= axis < ndim: + raise IndexError(f"axis {axis} is out of bounds for {ndim} dimensions") + return axis % ndim + + +def _new_shape(shape, size): + if isinstance(shape, (int, np.integer)): + shape = (shape,) + shape = tuple(int(n) for n in shape) + unknown = [i for i, n in enumerate(shape) if n == -1] + if len(unknown) > 1: + raise ValueError("reshape takes at most one -1") + if any(n < -1 for n in shape): + raise ValueError(f"invalid shape {shape}") + if unknown: + known = math.prod(n for n in shape if n != -1) + if known == 0 or size % known: + raise ValueError(f"cannot reshape an array of size {size} into shape {shape}") + i = unknown[0] + shape = (*shape[:i], size // known, *shape[i + 1 :]) + if math.prod(shape) != size: + raise ValueError(f"cannot reshape an array of size {size} into shape {shape}") + return shape + + +def reshape(x, /, shape, *, copy=None): + _check_array(x, "reshape") + shape = _new_shape(shape, x.size) + T = x._T() + if x.ndim == 0 or T.nvals == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + if shape == (): + # (one position: its entry, as a 0-d Array's value) + return Array._new(Tensor(T.dtype, ()), T.to_coo()[1][0]) + if shape == T.shape: + return x._copy() if copy else Array._new(T.dup(), x._fill) + return Array._new(T.reshape(shape, copy=copy), x._fill) + + +def permute_dims(x, /, axes): + _check_array(x, "permute_dims") + axes = tuple(_normalize_axis(a, x.ndim) for a in axes) + if sorted(axes) != list(range(x.ndim)): + raise ValueError(f"axes {axes} must be a permutation of range({x.ndim})") + if x.ndim == 0: + return x._copy() + return Array._new(x._T().permute_dims(axes), x._fill) + + +def expand_dims(x, /, axis=0): + _check_array(x, "expand_dims") + axes = axis if isinstance(axis, tuple) else (axis,) + ndim = x.ndim + len(axes) + norm = sorted(_normalize_axis(a, ndim) for a in axes) + if len(set(norm)) != len(norm): + raise ValueError(f"repeated axis in {axis}") + T = x._T() + if x.ndim == 0: + return Array._new(Tensor(T.dtype, (1,) * ndim), x._fill) + return Array._new(T.expand_dims(tuple(norm)), x._fill) + + +def squeeze(x, /, axis): + _check_array(x, "squeeze") + axes = axis if isinstance(axis, tuple) else (axis,) + norm = [_normalize_axis(a, x.ndim) for a in axes] + if len(set(norm)) != len(norm): + raise ValueError(f"repeated axis in {axis}") + for a in norm: + if x.shape[a] != 1: + raise ValueError(f"cannot squeeze axis {a} of size {x.shape[a]}") + if not norm: + return x._copy() + T = x._T() + shape = tuple(n for i, n in enumerate(x.shape) if i not in norm) + if not shape: + return reshape(x, ()) + return Array._new(T.squeeze(tuple(norm)), x._fill) + + +def _broadcast_to(x, shape): + """``x`` broadcast to ``shape``: its entries replicated over the broadcast axes.""" + shape = tuple(shape) + if np.broadcast_shapes(x.shape, shape) != shape: + raise ValueError(f"cannot broadcast shape {x.shape} to {shape}") + if x.shape == shape: + return x._copy() + T = x._T() + if x.ndim == 0 or T.nvals == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + E = Tensor(T.dtype, shape) + _utils.expansion_check(T, E, shape) + first = binary.first[gb_dtype(x.dtype)] + return Array._new(T.ewise_add(E, first, allow_expansion=True).new(), x._fill) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py new file mode 100644 index 000000000..cc6f15e5b --- /dev/null +++ b/graphblas/array_api/_searching.py @@ -0,0 +1,3 @@ +"""Searching functions: argmax, argmin, count_nonzero, nonzero, searchsorted, where.""" + +__all__ = [] diff --git a/graphblas/array_api/_set.py b/graphblas/array_api/_set.py new file mode 100644 index 000000000..48c4f080d --- /dev/null +++ b/graphblas/array_api/_set.py @@ -0,0 +1,3 @@ +"""Set functions: isin, unique_all, unique_counts, unique_inverse, unique_values.""" + +__all__ = [] diff --git a/graphblas/array_api/_sorting.py b/graphblas/array_api/_sorting.py new file mode 100644 index 000000000..156e958c6 --- /dev/null +++ b/graphblas/array_api/_sorting.py @@ -0,0 +1,3 @@ +"""Sorting functions: argsort, sort.""" + +__all__ = [] diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py new file mode 100644 index 000000000..c56c2f017 --- /dev/null +++ b/graphblas/array_api/_statistical.py @@ -0,0 +1,3 @@ +"""Statistical functions, and the utility functions all, any and diff.""" + +__all__ = [] diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py new file mode 100644 index 000000000..6736def17 --- /dev/null +++ b/graphblas/array_api/_utils.py @@ -0,0 +1,169 @@ +"""Private helpers shared by the modules of ``graphblas.array_api``.""" + +import contextlib +import math + +import numpy as np + +from .. import config +from ..core import tensor_layout as tl +from ..core.scalar import Scalar +from ..core.tensor import Tensor +from ..core.tensor_ewise import _broadcast_factor +from ..dtypes import lookup_dtype + + +class DensifyError(ValueError): + """An operation would allocate a dense array larger than the namespace allows.""" + + +def densify_limit(): + return config.get("array_api_densify_limit", 10**8) + + +def check_dense(n, what): + """Raise DensifyError when ``what`` would hold ``n`` elements densely past the limit.""" + if n > (limit := densify_limit()): + raise DensifyError( + f"{what} would hold {n} elements densely, more than the limit of {limit}. Raise " + "`graphblas.config` option array_api_densify_limit " + "(`with gb.config.set(array_api_densify_limit=...)`) to allow it." + ) + + +def gb_dtype(dtype): + return lookup_dtype(np.dtype(dtype)) + + +def np_dtype(gbdtype): + return np.dtype(gbdtype.np_type) + + +def scalar(value, dtype): + """``value`` as a numpy scalar of ``dtype`` (numpy's cast, warnings silenced).""" + with np.errstate(all="ignore"): + return np.asarray(value).astype(dtype, copy=False)[()] + + +def gb_scalar(value, dtype): + """A GraphBLAS Scalar of the numpy ``dtype`` holding ``value`` cast as numpy casts it.""" + return Scalar.from_value(scalar(value, dtype), gb_dtype(dtype)) + + +def same_value(a, b): + """Whether two scalars are identical as values: NaN matches NaN, -0.0 differs from 0.0.""" + a = np.asarray(a) + b = np.asarray(b) + if a.dtype.kind == "c" or b.dtype.kind == "c": + return same_value(a.real, b.real) and same_value(a.imag, b.imag) + if a.dtype.kind == "f" or b.dtype.kind == "f": + if np.isnan(a) or np.isnan(b): + return bool(np.isnan(a) and np.isnan(b)) + return bool(a == b and np.signbit(a) == np.signbit(b)) + return bool(np.equal(a, b)) + + +def differs(values, fill): + """Boolean array: where ``values`` is not identical to the scalar ``fill`` (invariant 4).""" + kind = values.dtype.kind + if kind == "c": + return differs(values.real, np.real(fill)) | differs(values.imag, np.imag(fill)) + if kind == "f": + if np.isnan(fill): + return ~np.isnan(values) + return (values != fill) | (np.signbit(values) != np.signbit(fill)) + return values != fill + + +def expansion_check(Ta, Tb, shape): + """Before a union of ``Ta`` and ``Tb`` broadcast to ``shape``: raise past the limit. + + The union replicates each broadcast operand's entries over the axes it broadcasts along. + """ + replicated = 0 + for T in (Ta, Tb): + factor = _broadcast_factor(T.shape, shape) + if factor > 1: + replicated += factor * T.nvals + check_dense(replicated, f"broadcasting shapes {Ta.shape} and {Tb.shape} to {shape}") + + +def empty_tensor(dtype, shape): + return Tensor(gb_dtype(dtype), tuple(shape)) + + +def full_tensor(dtype, shape, value): + """A Tensor holding ``value`` at every position (iso and full: O(1) memory).""" + T = Tensor(gb_dtype(dtype), tuple(shape)) + if math.prod(shape) > 0: + T << gb_scalar(value, dtype) + return T + + +def tensor_values(T): + """The stored values of ``T`` in storage order (a copy).""" + return T._M.to_coo(rows=False, columns=False, sort=False)[2] + + +def _x_array(cont, dtype, stack): + """The value array of an unloaded Container as a numpy view (None when it is empty).""" + return stack.enter_context(tl._VectorArray(cont.x, np.dtype(dtype))).array + + +def map_values(func, tensors, dtype): + """A Tensor with the entries of ``tensors`` (all of one shape, layout and pattern) whose + values are ``func(*value_arrays)``, computed by numpy over the stored values. + + The pattern is copied in GraphBLAS and the values written into the copy's own array + through a GxB_Container, so the entries are never rebuilt from coordinates. + """ + T0 = tensors[0] + shape, layout = T0._shape, T0._layout + gbdt = gb_dtype(dtype) + Ms = [T._M for T in tensors] + if T0._M._nvals == 0: + return Tensor._from_M(tl.new_matrix(gbdt, Ms[0]._nrows, Ms[0]._ncols), shape, layout) + if any(T._layout != layout for T in tensors[1:]): + Ms = [Ms[0]] + [tl.relayout(T._M, shape, T._layout, layout)[0] for T in tensors[1:]] + # An iso Matrix has one value; the template for the result is one that holds them all, + # if any does, so that its value array has the length of func's output. + isos = [M.ss.is_iso for M in Ms] + template = Ms[isos.index(False)] if False in isos else Ms[0] + W = tl.new_matrix(gbdt, template._nrows, template._ncols) + W << template + with contextlib.ExitStack() as stack: + unloaded = {} + for M in Ms: + if id(M) not in unloaded: + unloaded[id(M)] = stack.enter_context(tl._Unloaded(M)) + conts = [unloaded[id(M)] for M in Ms] + wcont = stack.enter_context(tl._Unloaded(W)) + formats = {c.format for c in conts} | {wcont.format} + nvals = {c.nvals for c in conts} | {wcont.nvals} + if len(formats) == 1 and len(nvals) == 1: + arrays = [_x_array(c, np_dtype(M.dtype), stack) for c, M in zip(conts, Ms, strict=True)] + out = _x_array(wcont, dtype, stack) + with np.errstate(all="ignore"): + result = func(*arrays) + result = np.asarray(result) + if out is not None and result.shape == out.shape: + out[...] = result + return Tensor._from_M(W, shape, layout) + # (formats differ: a fallback through coordinates, which aligns the values by position) + return _map_values_coo(func, tensors, dtype) # pragma: no cover (safety) + + +def _map_values_coo(func, tensors, dtype): + coords = None + arrays = [] + for T in tensors: + idx, vals = T.to_coo() + if coords is None: + coords = idx + arrays.append(vals) + with np.errstate(all="ignore"): + result = np.asarray(func(*arrays), dtype=dtype) + if result.ndim == 0: + result = np.broadcast_to(result, coords[0].shape if coords else (1,)) + T0 = tensors[0] + return Tensor.from_coo(coords, result, gb_dtype(dtype), shape=T0._shape) diff --git a/graphblas/graphblas.yaml b/graphblas/graphblas.yaml index f3885d73a..4711f66b2 100644 --- a/graphblas/graphblas.yaml +++ b/graphblas/graphblas.yaml @@ -8,3 +8,5 @@ mapnumpy: True jit_c_control: null # Tensor: largest number of entries an implicit broadcast may replicate (ewise_add / union) tensor_max_expansion: 1000000000 +# graphblas.array_api: most elements an array may hold densely (np.asarray, cumulative_sum, ...) +array_api_densify_limit: 100000000 From 0149fbbd0727ddaf9035c40bad29cc8474c24191 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:24:35 -0600 Subject: [PATCH 02/77] array_api: test the element-wise functions against numpy, and match numpy's clip Every element-wise function is compared with numpy on small arrays of each dtype, with zero and nonzero fills, NaN, infinities, signed zeros, stored entries equal to the fill, broadcasting, 0-d arrays and Python scalars. clip with scalar bounds is numpy's clip, which keeps -0.0 against a bound of 0.0; with array bounds it is minimum(maximum(x, min), max), as numpy's clip with array bounds is. --- graphblas/array_api/_elementwise.py | 35 +- graphblas/tests/test_array_api_elementwise.py | 568 ++++++++++++++++++ 2 files changed, 592 insertions(+), 11 deletions(-) create mode 100644 graphblas/tests/test_array_api_elementwise.py diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index f0483f340..20cfc7ef7 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -628,26 +628,39 @@ def func(x1, x2, /): def clip(x, /, min=None, max=None): - """``minimum(maximum(x, min), max)`` in the dtype of ``x`` (NaN propagates).""" + """``x`` clipped to ``[min, max]`` in the dtype of ``x``; NaN propagates. + + With scalar bounds this is numpy's ``clip`` (which keeps ``-0.0`` against a bound of + ``0.0``); with array bounds it is ``minimum(maximum(x, min), max)``, as numpy's ``clip`` + with array bounds is. + """ if not isinstance(x, Array): raise TypeError(f"clip needs an Array, not {type(x).__name__}") - out = x - if min is not None: - out = maximum(out, _clip_bound(min, x)) - if max is not None: - out = minimum(out, _clip_bound(max, x)) - if out is x: - return x._copy() + lo = _clip_bound(min, x) + hi = _clip_bound(max, x) + if lo is None and hi is None: + out = x._copy() + elif all(b is None or not isinstance(b, Array) or b.ndim == 0 for b in (lo, hi)): + lo_v, hi_v = (b._fill if isinstance(b, Array) else b for b in (lo, hi)) + out = _unary("clip", lambda v: np.clip(v, lo_v, hi_v), x) + else: + out = x + if lo is not None: + out = maximum(out, lo) + if hi is not None: + out = minimum(out, hi) if out.dtype != x.dtype: out = cast(out, x.dtype) return out def _clip_bound(bound, x): - """A Python scalar bound in the dtype of ``x`` when it fits (the Array API's clip keeps the - dtype of ``x``). + """A Python integer bound clamped to the range of an integer ``x`` (the Array API's clip + keeps the dtype of ``x``, and numpy refuses an integer out of its range). """ - if isinstance(bound, (builtins.bool, int, float)) and not isinstance(bound, np.generic): + if isinstance(bound, (np.generic, np.ndarray)): + return _as_operand(bound, "clip") + if isinstance(bound, (builtins.bool, int, float)): if x.dtype.kind in _INTEGRAL and isinstance(bound, float): return bound if x.dtype.kind in _INTEGRAL and not _python_scalar_fits(bound, x.dtype): diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py new file mode 100644 index 000000000..eadaae9b4 --- /dev/null +++ b/graphblas/tests/test_array_api_elementwise.py @@ -0,0 +1,568 @@ +"""graphblas.array_api element-wise functions against numpy, position by position. + +The helpers here (``make``, ``assert_array_equal``, ``inputs``) are shared by the other +``test_array_api_*`` modules. +""" + +import math +import operator +import warnings + +import numpy as np +import pytest + +import graphblas.array_api as xp +from graphblas import backend +from graphblas.core.tensor import _has_container as tensor_has_container + +if backend != "suitesparse": + pytest.skip("graphblas.array_api needs the suitesparse backend", allow_module_level=True) +if not tensor_has_container(): + pytest.skip("graphblas.array_api needs SuiteSparse:GraphBLAS 10", allow_module_level=True) + + +DTYPES = [ + np.dtype(name) + for name in ( + "bool", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "float32", + "float64", + "complex64", + "complex128", + ) +] +FLOAT_SPECIALS = [np.nan, np.inf, -np.inf, -0.0, 0.0, 1.0, -1.0, 0.5, -2.5, 3.0, 1e30] + + +def random_values(rng, dtype, shape): + """Dense values of ``dtype`` with the special values of the kind.""" + if dtype.kind == "b": + return rng.random(shape) < 0.5 + if dtype.kind in "iu": + info = np.iinfo(dtype) + vals = rng.integers(-4 if dtype.kind == "i" else 0, 5, shape).astype(dtype) + edge = rng.random(shape) < 0.1 + extremes = np.where(rng.random(shape) < 0.5, np.array(info.min, dtype), info.max) + return np.where(edge, extremes, vals) + real = rng.choice(np.array(FLOAT_SPECIALS), size=shape) + real = np.where(rng.random(shape) < 0.5, real, np.round(rng.normal(size=shape) * 3, 2)) + if dtype.kind == "c": + imag = rng.choice(np.array(FLOAT_SPECIALS), size=shape) + imag = np.where(rng.random(shape) < 0.6, 0.0, imag) + out = np.empty(shape, dtype) + with np.errstate(all="ignore"): + out.real = real + out.imag = imag + return out + with np.errstate(all="ignore"): + return real.astype(dtype) + + +def fills_for(dtype): + if dtype.kind == "b": + return [False, True] + if dtype.kind in "iu": + return [0, 3] + if dtype.kind == "c": + return [0, complex(1.5, -2), complex(np.nan, 0)] + return [0.0, -0.0, 2.5, np.nan, -np.inf] + + +def same_value(a, b): + a, b = np.asarray(a), np.asarray(b) + if a.dtype.kind == "c": + return same_value(a.real, b.real) and same_value(a.imag, b.imag) + if a.dtype.kind == "f": + if np.isnan(a) or np.isnan(b): + return bool(np.isnan(a) and np.isnan(b)) + return bool(a == b and np.signbit(a) == np.signbit(b)) + return bool(np.equal(a, b)) + + +def make(dense, fill, rng=None, *, explicit=0.3): + """An Array whose value is ``dense`` at every position: the positions that differ from + ``fill`` are stored, and (with ``rng``) some that equal it are stored too. + """ + dense = np.asarray(dense) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", np.exceptions.ComplexWarning) + fill = np.asarray(fill).astype(dense.dtype)[()] + if dense.ndim == 0: + return xp.asarray(dense) + flat = dense.reshape(-1) + keep = np.array([not same_value(v, fill) for v in flat], dtype=bool) + if rng is not None: + keep |= rng.random(flat.size) < explicit + idx = np.nonzero(keep.reshape(dense.shape)) + out = xp.from_coo(idx, dense[idx], dense.shape, fill_value=fill, dtype=dense.dtype) + assert_array_equal(out, dense) + return out + + +def check_invariants(x): + assert isinstance(x, xp.Array) + assert np.asarray(x.fill_value).dtype == x.dtype + T = x._T() + assert T.shape == x.shape + assert np.dtype(T.dtype.np_type) == x.dtype + if x.ndim == 0: + assert T.nvals == 0 + T._check_invariants() + + +def assert_array_equal(x, expected): + """``x`` holds ``expected`` exactly: dtype, shape, NaN, signed zeros.""" + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape, (x.shape, expected.shape) + dense = x.todense() + assert dense.dtype == expected.dtype + _assert_identical(dense, expected) + + +def _assert_identical(a, b): + if a.dtype.kind == "c": + _assert_identical(a.real, b.real) + _assert_identical(a.imag, b.imag) + return + if a.dtype.kind == "f": + nan_a, nan_b = np.isnan(a), np.isnan(b) + bad = (nan_a != nan_b) | (~nan_a & ((a != b) | (np.signbit(a) != np.signbit(b)))) + else: + bad = a != b + if bad.any(): + idx = np.argwhere(bad)[:5] + detail = [(tuple(int(i) for i in ix), a[tuple(ix)], b[tuple(ix)]) for ix in idx] + raise AssertionError(f"{bad.sum()} positions differ (position, got, expected): {detail}") + + +def np_call(func, *args): + with np.errstate(all="ignore"), warnings.catch_warnings(): + warnings.simplefilter("ignore") + return func(*args) + + +def inputs(dtype, shape, seed=0): + """(dense, Array) pairs over the fills of the dtype, some with explicit fills.""" + rng = np.random.default_rng(seed) + for i, fill in enumerate(fills_for(dtype)): + dense = random_values(rng, dtype, shape) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + fillv = np.asarray(fill).astype(dtype)[()] + if dense.ndim: + sparse = rng.random(shape) < 0.5 + dense = np.where(sparse, fillv, dense).astype(dtype) + yield dense, make(dense, fillv, rng if i % 2 else None) + + +# (name, numpy function) +UNARY = [ + ("abs", np.abs), + ("acos", np.arccos), + ("acosh", np.arccosh), + ("asin", np.arcsin), + ("asinh", np.arcsinh), + ("atan", np.arctan), + ("atanh", np.arctanh), + ("bitwise_invert", np.invert), + ("ceil", np.ceil), + ("conj", np.conj), + ("cos", np.cos), + ("cosh", np.cosh), + ("exp", np.exp), + ("expm1", np.expm1), + ("floor", np.floor), + ("imag", np.imag), + ("isfinite", np.isfinite), + ("isinf", np.isinf), + ("isnan", np.isnan), + ("log", np.log), + ("log1p", np.log1p), + ("log2", np.log2), + ("log10", np.log10), + ("logical_not", np.logical_not), + ("negative", np.negative), + ("positive", np.positive), + ("real", np.real), + ("reciprocal", np.reciprocal), + ("round", np.round), + ("sign", np.sign), + ("signbit", np.signbit), + ("sin", np.sin), + ("sinh", np.sinh), + ("sqrt", np.sqrt), + ("square", np.square), + ("tan", np.tan), + ("tanh", np.tanh), + ("trunc", np.trunc), +] + +BINARY = [ + ("add", np.add), + ("atan2", np.arctan2), + ("bitwise_and", np.bitwise_and), + ("bitwise_left_shift", np.left_shift), + ("bitwise_or", np.bitwise_or), + ("bitwise_right_shift", np.right_shift), + ("bitwise_xor", np.bitwise_xor), + ("copysign", np.copysign), + ("divide", np.divide), + ("equal", np.equal), + ("floor_divide", np.floor_divide), + ("greater", np.greater), + ("greater_equal", np.greater_equal), + ("hypot", np.hypot), + ("less", np.less), + ("less_equal", np.less_equal), + ("logaddexp", np.logaddexp), + ("logical_and", np.logical_and), + ("logical_or", np.logical_or), + ("logical_xor", np.logical_xor), + ("maximum", np.maximum), + ("minimum", np.minimum), + ("multiply", np.multiply), + ("nextafter", np.nextafter), + ("not_equal", np.not_equal), + ("pow", np.power), + ("remainder", np.remainder), + ("subtract", np.subtract), +] + +SUPPORTED = set(DTYPES) + + +def expected_or_error(func, *args): + """Numpy's result, or the exception class numpy raises; None when numpy's result dtype is + one graphblas.array_api lacks (float16). + """ + try: + out = np_call(func, *args) + except (TypeError, ValueError, OverflowError) as exc: + return type(exc) + out = np.asarray(out) + if out.dtype not in SUPPORTED: + return None + return out + + +def test_all_elementwise_functions_are_tested(): + names = {name for name, _ in UNARY + BINARY} | {"clip"} + assert names == set(xp._elementwise.__all__) + + +@pytest.mark.parametrize(("name", "func"), UNARY, ids=[n for n, _ in UNARY]) +@pytest.mark.parametrize("shape", [(), (7,), (3, 4), (2, 3, 4)]) +def test_unary(name, func, shape): + for dtype in DTYPES: + for dense, x in inputs(dtype, shape, seed=hash(name) % 1000): + expected = expected_or_error(func, dense) + if expected is None: + with pytest.raises(TypeError): + getattr(xp, name)(x) + continue + if isinstance(expected, type): + with pytest.raises(expected): + getattr(xp, name)(x) + continue + out = getattr(xp, name)(x) + try: + assert_array_equal(out, expected) + except AssertionError as exc: + raise AssertionError(f"{name} of {dtype} fill {x.fill_value!r}: {exc}") from None + + +BINARY_SHAPES = [ + ((5,), (5,)), + ((3, 4), (3, 4)), + ((3, 4), (4,)), + ((3, 1), (1, 4)), + ((2, 3, 4), (3, 1)), + ((), (3, 4)), + ((3, 4), ()), + ((), ()), +] +DTYPE_PAIRS = [(dt, dt) for dt in DTYPES] + [ + (np.dtype("int8"), np.dtype("int64")), + (np.dtype("uint8"), np.dtype("int8")), + (np.dtype("int32"), np.dtype("float32")), + (np.dtype("int64"), np.dtype("float64")), + (np.dtype("float32"), np.dtype("float64")), + (np.dtype("float32"), np.dtype("complex128")), + (np.dtype("bool"), np.dtype("int16")), + (np.dtype("uint64"), np.dtype("int64")), +] + + +def _binary_cases(dt1, dt2, shapes, seed): + for i, (dense1, x1) in enumerate(inputs(dt1, shapes[0], seed=seed)): + for j, (dense2, x2) in enumerate(inputs(dt2, shapes[1], seed=seed + 1)): + if (i + j) % 2 == 0 or i == j: + yield dense1, x1, dense2, x2 + + +@pytest.mark.parametrize(("name", "func"), BINARY, ids=[n for n, _ in BINARY]) +@pytest.mark.parametrize("shapes", BINARY_SHAPES, ids=str) +def test_binary(name, func, shapes): + for dt1, dt2 in DTYPE_PAIRS: + for dense1, x1, dense2, x2 in _binary_cases(dt1, dt2, shapes, hash(name) % 1000): + expected = expected_or_error(func, dense1, dense2) + if expected is None: + with pytest.raises(TypeError): + getattr(xp, name)(x1, x2) + continue + if isinstance(expected, type): + with pytest.raises(expected): + np_call(getattr(xp, name), x1, x2) + continue + out = getattr(xp, name)(x1, x2) + try: + assert_array_equal(out, expected) + except AssertionError as exc: + raise AssertionError( + f"{name} of {dt1} fill {x1.fill_value!r} and {dt2} fill {x2.fill_value!r}: " + f"{exc}" + ) from None + + +PY_SCALARS = [0, 1, 2, -3, 200, 2**40, True, False, 0.0, -0.0, 0.5, -1.5, np.nan, np.inf, 1j] + + +@pytest.mark.parametrize(("name", "func"), BINARY, ids=[n for n, _ in BINARY]) +@pytest.mark.parametrize("shape", [(), (3, 4)]) +def test_binary_python_scalar(name, func, shape): + for dtype in DTYPES: + dense, x = next(iter(inputs(dtype, shape, seed=7))) + for s in PY_SCALARS: + for args, dargs in (((x, s), (dense, s)), ((s, x), (s, dense))): + expected = expected_or_error(func, *dargs) + if expected is None: + with pytest.raises(TypeError): + getattr(xp, name)(*args) + continue + if isinstance(expected, type): + with pytest.raises(expected): + np_call(getattr(xp, name), *args) + continue + out = np_call(getattr(xp, name), *args) + try: + assert_array_equal(out, expected) + except AssertionError as exc: + raise AssertionError(f"{name}{tuple(map(repr, dargs))}: {exc}") from None + + +OPERATORS = [ + (operator.add, "add"), + (operator.sub, "subtract"), + (operator.mul, "multiply"), + (operator.truediv, "divide"), + (operator.floordiv, "floor_divide"), + (operator.mod, "remainder"), + (operator.pow, "pow"), + (operator.and_, "bitwise_and"), + (operator.or_, "bitwise_or"), + (operator.xor, "bitwise_xor"), + (operator.lshift, "bitwise_left_shift"), + (operator.rshift, "bitwise_right_shift"), + (operator.eq, "equal"), + (operator.ne, "not_equal"), + (operator.lt, "less"), + (operator.le, "less_equal"), + (operator.gt, "greater"), + (operator.ge, "greater_equal"), +] + + +@pytest.mark.parametrize(("op", "name"), OPERATORS, ids=[n for _, n in OPERATORS]) +def test_operators(op, name): + rng = np.random.default_rng(1) + a = rng.integers(1, 5, (3, 4)) + b = rng.integers(1, 5, (4,)) + x, y = make(a, 0), make(b, 2) + assert_array_equal(op(x, y), np_call(op, a, b)) + assert_array_equal(op(x, 3), np_call(op, a, 3)) + assert_array_equal(op(3, y), np_call(op, 3, b)) + # numpy arrays defer to the Array + out = op(a, y) + assert isinstance(out, xp.Array) + assert_array_equal(out, np_call(op, a, b)) + + +def test_unary_operators(): + a = np.array([1, -2, 0, 3]) + x = make(a, 0) + assert_array_equal(-x, -a) + assert_array_equal(+x, +a) + assert_array_equal(abs(x), abs(a)) + assert_array_equal(~x, ~a) + + +@pytest.mark.parametrize( + ("op", "name"), + [(operator.iadd, "add"), (operator.imul, "multiply"), (operator.isub, "subtract")], +) +def test_inplace_operators(op, name): + a = np.array([[1.0, 0.0, 2.0], [0.0, -0.0, np.nan]]) + x = make(a, 0.0) + orig = x + x = op(x, make(np.array([1.0, 2.0, 0.0]), 0.0)) + assert x is orig + assert_array_equal(x, np_call(op, a.copy(), np.array([1.0, 2.0, 0.0]))) + x = op(x, 2) + assert x is orig + with pytest.raises(TypeError): + op(make(np.array([1, 2]), 0), 1.5) + with pytest.raises(ValueError, match="shape"): + op(make(np.array([1.0, 2.0]), 0.0), make(np.ones((3, 2)), 0.0)) + + +@pytest.mark.parametrize("dtype", ["int8", "uint8", "float32", "float64", "int64"]) +def test_clip(dtype): + rng = np.random.default_rng(3) + dt = np.dtype(dtype) + for dense, x in inputs(dt, (3, 4), seed=4): + for lo, hi in [(1, None), (None, 2), (0, 3), (-1, 1)]: + if dt.kind == "u" and lo == -1: + lo = 0 + expected = np.clip(dense, lo, hi) if (lo, hi) != (None, None) else dense + out = xp.clip(x, min=lo, max=hi) + assert_array_equal(out, np.asarray(expected).astype(dt)) + lo_dense = random_values(rng, dt, (4,)) + lo_arr = make(lo_dense, 0) + with np.errstate(all="ignore"): + expected = np.maximum(dense, lo_dense).astype(dt) + assert_array_equal(xp.clip(x, min=lo_arr), expected) + x = make(np.array([1, 2, 3], dtype=np.int8), 0) + assert_array_equal(xp.clip(x, min=-1000, max=1000), np.array([1, 2, 3], dtype=np.int8)) + assert_array_equal(xp.clip(x), np.array([1, 2, 3], dtype=np.int8)) + + +def test_signed_zero_union(): + # 0.0 fill: an unstored position is +0.0, a stored -0.0 differs + x = make(np.array([-0.0, 0.0, 1.0, -1.0]), 0.0) + y = make(np.array([0.0, -0.0, 0.0, 0.0]), 0.0) + assert x.nnz == 3 + assert y.nnz == 1 + a, b = x.todense(), y.todense() + for name, func in [("add", np.add), ("multiply", np.multiply), ("subtract", np.subtract)]: + assert_array_equal(getattr(xp, name)(x, y), func(a, b)) + assert_array_equal(getattr(xp, name)(y, x), func(b, a)) + + +def test_integer_multiply_intersects(): + # Both fills 0: integer products are 0 where either side is unstored, so the result + # stores only the intersection of entries. + x = make(np.array([1, 0, 2, 0, 3]), 0) + y = make(np.array([0, 4, 5, 0, 6]), 0) + out = x * y + assert out.nnz == 2 + assert_array_equal(out, np.array([0, 0, 10, 0, 18])) + + +def test_float_multiply_keeps_union(): + x = make(np.array([-1.0, 0.0, 2.0, np.inf]), 0.0) + y = make(np.array([0.0, 4.0, 5.0, 0.0]), 0.0) + assert_array_equal(x * y, np_call(np.multiply, x.todense(), y.todense())) + + +def test_broadcast_expansion_is_guarded(): + import graphblas as gb + + x = make(np.arange(1, 101, dtype=np.float64).reshape(100, 1), 0.0) + y = make(np.arange(1, 101, dtype=np.float64).reshape(1, 100), 0.0) + with gb.config.set(array_api_densify_limit=1000), pytest.raises(xp.DensifyError): + x + y + assert_array_equal(x + y, x.todense() + y.todense()) + + +def test_values_mapped_on_every_storage_format(): + # numpy maps the stored values through the Container in place: check bitmap, full, + # hypersparse and iso storage + import graphblas as gb + + rng = np.random.default_rng(5) + for density in [0.0, 0.05, 0.5, 1.0]: + for shape in [(40,), (30, 20), (6, 5, 4)]: + a = np.round(rng.normal(size=shape), 2) + a[rng.random(shape) >= density] = 0.0 + x = make(a, 0.0) + assert_array_equal(xp.sin(x), np.sin(a)) + assert_array_equal(xp.pow(x, x), np_call(np.power, a, a)) + for fmt in ["bitmap", "sparse", "hypersparse", "full"]: + M = x._T()._M + if fmt == "full" and M.nvals != M.nrows * M.ncols: + continue + M.ss.config["sparsity_control"] = fmt + assert_array_equal(xp.exp(x), np.exp(a)) + assert_array_equal(xp.floor_divide(x, 0.3), np_call(np.floor_divide, a, 0.3)) + iso = xp.from_coo(([0, 3, 5],), 2.5, (8,), fill_value=0.0) + assert iso._T().ss.is_iso + assert_array_equal(xp.log(iso), np_call(np.log, iso.todense())) + other = make(np.arange(8.0), 0.0) + assert_array_equal(xp.atan2(iso, other), np.arctan2(iso.todense(), other.todense())) + assert_array_equal(xp.atan2(other, iso), np.arctan2(other.todense(), iso.todense())) + del gb + + +def test_zero_d_results_store_nothing(): + x = xp.asarray(2.5) + y = xp.add(x, 1) + check_invariants(y) + assert float(y) == 3.5 + z = make(np.array([1.0, 2.0]), 0.0)[1] + check_invariants(z) + assert float(z) == 2.0 + assert z._T().nvals == 0 + + +def test_results_do_not_share_writes(): + x = make(np.array([1.0, 0.0, 2.0]), 0.0) + y = +x + y += 1 + assert_array_equal(x, np.array([1.0, 0.0, 2.0])) + assert_array_equal(y, np.array([2.0, 1.0, 3.0])) + + +def test_math_consistency_of_fill(): + # An entry equal to the fill maps to the new fill, bit for bit + for name, _ in UNARY: + if name in {"logical_not", "bitwise_invert"}: + continue + x = xp.from_coo(([0, 1],), [0.7, 0.7], (4,), fill_value=0.7) + out = getattr(xp, name)(x) + dense = out.todense() + assert all(same_value(v, out.fill_value) for v in dense), name + + +def test_isnan_of_integers_stores_nothing(): + x = make(np.array([1, 0, 2]), 0) + for name, value in [("isnan", False), ("isinf", False), ("isfinite", True)]: + out = getattr(xp, name)(x) + assert out.nnz == 0 + assert out.fill_value == value + + +def test_math_constants(): + assert xp.e == math.e + assert xp.pi == math.pi + assert math.isinf(xp.inf) + assert math.isnan(xp.nan) + assert xp.newaxis is None + + +def test_unsupported_operands(): + x = make(np.array([1.0, 2.0]), 0.0) + with pytest.raises(TypeError): + xp.add(x, "a") + assert (x == "a") is False + with pytest.raises(TypeError): + xp.sin(make(np.array([True, False]), False)) # numpy would give float16 From 33252984025dfa3450f519a549fe42f781f9b491 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:30:27 -0600 Subject: [PATCH 03/77] array_api: statistical functions, all, any and diff sum, prod, min, max, any and all reduce the stored entries in GraphBLAS and then account for each fiber's unstored positions from a count of its entries: k * fill for sum, fill ** k for prod, the fill once for the others. No count is needed when the fill is the monoid's identity or every position is stored. min and max propagate NaN, as numpy's do. mean divides the sum; var and std take two passes (the stored entries' squared deviations from their fiber's mean in GraphBLAS, plus the fill's times its count). cumulative_sum and cumulative_prod go through numpy under the densify limit, since their results are dense. diff subtracts shifted slices, with prepend and append joined along the axis. --- graphblas/array_api/_elementwise.py | 7 + graphblas/array_api/_manipulation.py | 45 ++ graphblas/array_api/_statistical.py | 432 +++++++++++++++++- graphblas/tests/test_array_api_elementwise.py | 4 +- graphblas/tests/test_array_api_statistical.py | 248 ++++++++++ 5 files changed, 732 insertions(+), 4 deletions(-) create mode 100644 graphblas/tests/test_array_api_statistical.py diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 20cfc7ef7..b7b71db26 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -586,6 +586,13 @@ def _zero(dtype): } +def binary_numpy(name, func, x1, x2): + """``func`` (a numpy function of two arrays, which broadcasts) of two operands with the + fill rules of a binary function, computed by numpy over the values (for internal use). + """ + return _binary(_Binary(name, func), x1, x2) + + def _make_binary(name): spec = _SPECS[name] diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py index 8053f5f97..4b6d8d8f2 100644 --- a/graphblas/array_api/_manipulation.py +++ b/graphblas/array_api/_manipulation.py @@ -114,3 +114,48 @@ def _broadcast_to(x, shape): _utils.expansion_check(T, E, shape) first = binary.first[gb_dtype(x.dtype)] return Array._new(T.ewise_add(E, first, allow_expansion=True).new(), x._fill) + + +def _concat(arrays, axis): + """The arrays joined along ``axis``. Parts whose fill differs from the first part's store + their fill explicitly (densely, under the limit). + """ + from ._creation import _from_numpy + from ._dtypes import result_type + from ._elementwise import cast + + arrays = list(arrays) + if not arrays: + raise ValueError("need at least one array to concatenate") + ndim = arrays[0].ndim + if ndim == 0 or any(a.ndim != ndim for a in arrays): + raise ValueError("concatenation needs arrays of one positive number of dimensions") + axis = _normalize_axis(axis, ndim) + shape = list(arrays[0].shape) + for a in arrays[1:]: + if any(a.shape[i] != n for i, n in enumerate(shape) if i != axis): + raise ValueError(f"shapes {arrays[0].shape} and {a.shape} do not concatenate") + shape[axis] = sum(a.shape[axis] for a in arrays) + dtype = result_type(*arrays) + arrays = [a if a.dtype == dtype else cast(a, dtype) for a in arrays] + fill = arrays[0]._fill + coords, values = [], [] + offset = 0 + for a in arrays: + if not _utils.same_value(a._fill, fill): + _utils.check_dense(a.size, "concatenating arrays with different fill values") + a = _from_numpy(a.todense(), fill) + T = a._T() + if T.nvals: + idx, vals = T.to_coo() + idx = list(idx) + idx[axis] = idx[axis] + offset + coords.append(idx) + values.append(vals) + offset += a.shape[axis] + shape = tuple(shape) + if not values: + return Array._new(Tensor(gb_dtype(dtype), shape), fill) + coords = tuple(np.concatenate([c[i] for c in coords]) for i in range(ndim)) + T = Tensor.from_coo(coords, np.concatenate(values), gb_dtype(dtype), shape=shape) + return Array._new(T, fill) diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index c56c2f017..bb897021b 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -1,3 +1,431 @@ -"""Statistical functions, and the utility functions all, any and diff.""" +"""Statistical functions, and the utility functions ``all``, ``any`` and ``diff``. -__all__ = [] +A reduction over a fiber of ``L`` positions, ``k`` of them unstored, is the monoid over the +stored entries combined with the fill ``k`` times: ``plus`` adds ``k * fill``, ``times`` +multiplies by ``fill ** k``, and ``min``, ``max``, ``lor`` and ``land`` take the fill once +when ``k > 0``. The counts cost a second pass, skipped when the fill is the monoid's identity +or every position is stored. A fiber with no entries holds the result's fill. Floating-point +sums and products may round differently from numpy's (the order of the terms differs); +everything else is exact. +""" + +import builtins +import math + +import numpy as np + +from .. import agg, binary, monoid, unary +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._dtypes import as_dtype, check_supported +from ._utils import gb_dtype + +__all__ = [ + "all", + "any", + "cumulative_prod", + "cumulative_sum", + "diff", + "max", + "mean", + "min", + "prod", + "std", + "sum", + "var", +] + + +def _check_array(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + + +def _axes(x, axis): + """The reduced axes, normalized and sorted.""" + if axis is None: + return tuple(range(x.ndim)) + axes = axis if isinstance(axis, tuple) else (axis,) + norm = [] + for a in axes: + a = int(a) + if not -x.ndim <= a < x.ndim: + raise IndexError(f"axis {a} is out of bounds for {x.ndim} dimensions") + norm.append(a % x.ndim) + if len(set(norm)) != len(norm): + raise ValueError(f"repeated axis in {axis}") + return tuple(sorted(norm)) + + +def _out_shape(shape, axes, keepdims): + if keepdims: + return tuple(1 if i in axes else n for i, n in enumerate(shape)) + return tuple(n for i, n in enumerate(shape) if i not in axes) + + +def _keepdims(out, x, axes, keepdims): + if keepdims and axes: + from ._manipulation import expand_dims + + return expand_dims(out, axes) + return out + + +def _reduce_tensor(T, op, axes): + axis = None if len(axes) == T.ndim else axes + return T.reduce(op, axis=axis).new() + + +def _combine(func, tensors, dtype, fill): + """An Array whose entries are ``func`` of the aligned values of ``tensors`` (one + pattern), and ``fill`` elsewhere; 0-d tensors give a 0-d Array. + """ + T0 = tensors[0] + if T0.ndim == 0: + if T0.nvals == 0: + return Array._new(Tensor(gb_dtype(dtype), ()), fill) + values = [T.to_coo()[1] for T in tensors] + with np.errstate(all="ignore"): + return Array._new(Tensor(gb_dtype(dtype), ()), np.asarray(func(*values))[0]) + return Array._new(_utils.map_values(func, tensors, dtype), fill) + + +def _int_power(base, exponent, dtype): + """``base ** exponent`` wrapped to the integer ``dtype``, as repeated products wrap.""" + bits = dtype.itemsize * 8 + value = builtins.pow(int(base), int(exponent), 1 << bits) + if dtype.kind == "i" and value >= 1 << (bits - 1): + value -= 1 << bits + return np.asarray(value).astype(dtype)[()] + + +class _Monoid: + def __init__(self, name, gb, npfunc, identity): + self.name = name + self.gb = gb + self.npfunc = npfunc + self.identity = identity # (dtype -> the identity) + + def repeat(self, fill, L, ct): + """The monoid over ``L`` copies of ``fill``.""" + with np.errstate(all="ignore"): + if self.name == "sum": + if ct.kind in "iu": + return np.multiply(fill, np.asarray(L % (1 << 64)).astype(ct), dtype=ct) + return np.multiply(fill, L, dtype=ct) + if self.name == "prod": + if ct.kind in "iu": + return _int_power(fill, L, ct) + if ct.kind == "b": + return fill + return np.power(fill, ct.type(L), dtype=ct) + return fill + + def combine(self, L, fill, ct): + """``func(r, count)``: a fiber's reduction of its stored entries ``r`` combined with + the fill at its ``L - count`` unstored positions. + """ + name = self.name + if name == "sum": + if ct.kind in "iu": + return lambda r, c: r + (L - c).astype(ct) * fill + # (a fiber with no unstored position must not see 0 * inf) + return lambda r, c: np.where(c < L, r + (L - c) * fill, r) + if name == "prod": + if ct.kind in "iu": + big = np.dtype("uint64" if ct.kind == "u" else "int64") + return lambda r, c: r * np.power(big.type(fill), (L - c).astype(big)).astype(ct) + if ct.kind == "b": + return lambda r, c: np.where(c < L, r & fill, r) + return lambda r, c: np.where(c < L, r * np.power(fill, (L - c).astype(ct)), r) + func = { + "min": np.minimum, + "max": np.maximum, + "any": np.logical_or, + "all": np.logical_and, + }[name] + return lambda r, c: np.where(c < L, func(r, fill), r) + + +def _min_identity(dt): + if dt.kind == "f": + return np.inf + if dt.kind == "b": + return True + return np.iinfo(dt).max + + +def _max_identity(dt): + if dt.kind == "f": + return -np.inf + if dt.kind == "b": + return False + return np.iinfo(dt).min + + +_MONOIDS = { + "sum": _Monoid("sum", monoid.plus, np.sum, lambda dt: 0), + "prod": _Monoid("prod", monoid.times, np.prod, lambda dt: 1), + "min": _Monoid("min", monoid.min, np.min, _min_identity), + "max": _Monoid("max", monoid.max, np.max, _max_identity), + "any": _Monoid("any", monoid.lor, np.any, lambda dt: False), + "all": _Monoid("all", monoid.land, np.all, lambda dt: True), +} + + +def _reduce(name, x, axis, keepdims, dtype=None): + _check_array(x, name) + spec = _MONOIDS[name] + if name in {"min", "max"} and x.dtype.kind == "c": + raise TypeError(f"{name} needs a real-valued array, not {x.dtype}") + if name in {"any", "all"}: + ct = np.dtype(bool) + elif name in {"sum", "prod"}: + if dtype is not None: + ct = as_dtype(dtype) + else: + ct = check_supported(spec.npfunc(np.zeros(1, x.dtype)).dtype, name) + else: + ct = x.dtype + axes = _axes(x, axis) + if x.ndim == 0: + kwargs = {"dtype": ct} if name in {"sum", "prod"} else {} + with np.errstate(all="ignore"): + value = spec.npfunc(np.asarray(x._fill), **kwargs) + return Array._new(Tensor(gb_dtype(ct), ()), value) + if not axes: + return _reduce_no_axes(name, x, ct) + T = x._T() + shape = _out_shape(x.shape, axes, keepdims=False) + L = math.prod(x.shape[a] for a in axes) + with np.errstate(all="ignore"): + fill = _utils.scalar(x._fill, ct) + if L == 0: + if name in {"min", "max"}: + raise ValueError(f"{name} of an empty axis, which has no identity") + out = Array._new(Tensor(gb_dtype(ct), shape), spec.identity(ct)) + return _keepdims(out, x, axes, keepdims) + fill_out = spec.repeat(fill, L, ct) + if T.nvals == 0 or math.prod(shape) == 0: + return _keepdims(Array._new(Tensor(gb_dtype(ct), shape), fill_out), x, axes, keepdims) + R = _reduce_tensor(T, spec.gb[gb_dtype(ct)], axes) + nan_fix = name in {"min", "max"} and ct.kind == "f" and _has_nan(x) + full = T.nvals == x.size + if not nan_fix and (full or _utils.same_value(fill, spec.identity(ct))): + out = Array._new(R, fill if full else fill_out) + else: + C = _reduce_tensor(T, agg.count, axes) + combine = spec.combine(L, fill, ct) + if nan_fix: + # GraphBLAS's min and max skip NaN; numpy's propagate it + N = _reduce_tensor(T.apply(unary.isnan).new(), monoid.lor[gb_dtype(bool)], axes) + + def func(r, c, n): + return np.where(n, np.nan, combine(r, c)) + + out = _combine(func, [R, C, N], ct, fill_out) + else: + out = _combine(combine, [R, C], ct, fill_out) + return _keepdims(out, x, axes, keepdims) + + +def _reduce_no_axes(name, x, ct): + """A reduction over no axes, as numpy computes it: the identity combined with each + element (so ``sum`` turns ``-0.0`` into ``0.0``), or a copy for ``min`` and ``max``. + """ + from ._elementwise import add, cast, multiply + + if name in {"sum", "prod"}: + identity = Array._new(Tensor(gb_dtype(ct), ()), 0 if name == "sum" else 1) + out = (add if name == "sum" else multiply)(identity, x if x.dtype == ct else cast(x, ct)) + return out if out.dtype == ct else cast(out, ct) + return cast(x, ct) if x.dtype != ct else x._copy() + + +def _has_nan(x): + def compute(x): + T = x._T() + if T.nvals == 0: + return False + return bool(T.apply(unary.isnan).new().reduce_scalar(monoid.lor).new().value) + + return x._fact("has_nan", compute) + + +def sum(x, /, *, axis=None, dtype=None, keepdims=False): # noqa: A001 (the Array API's name) + return _reduce("sum", x, axis, keepdims, dtype) + + +def prod(x, /, *, axis=None, dtype=None, keepdims=False): + return _reduce("prod", x, axis, keepdims, dtype) + + +def min(x, /, *, axis=None, keepdims=False): + return _reduce("min", x, axis, keepdims) + + +def max(x, /, *, axis=None, keepdims=False): + return _reduce("max", x, axis, keepdims) + + +def any(x, /, *, axis=None, keepdims=False): # noqa: A001 (the Array API's name) + return _reduce("any", x, axis, keepdims) + + +def all(x, /, *, axis=None, keepdims=False): # noqa: A001 (the Array API's name) + return _reduce("all", x, axis, keepdims) + + +def _mean_dtype(x, name): + return check_supported(np.mean(np.zeros(1, x.dtype)).dtype, name) + + +def mean(x, /, *, axis=None, keepdims=False): + from ._elementwise import divide + + _check_array(x, "mean") + rdt = _mean_dtype(x, "mean") + L = math.prod(x.shape[a] for a in _axes(x, axis)) + total = _reduce("sum", x, axis, keepdims, rdt) + if L == 0: + return Array._new(Tensor(total._T().dtype, total.shape), np.nan) + return divide(total, L) + + +def _squared_magnitude(x): + from ._elementwise import conj, multiply, real + + if x.dtype.kind == "c": + return real(multiply(x, conj(x))) + return multiply(x, x) + + +def var(x, /, *, axis=None, correction=0.0, keepdims=False): + """Two passes: the means, then the squared deviations from them (the stored entries' + in GraphBLAS, and the fill's once per fiber, times its count of unstored positions). + """ + from ._elementwise import add, binary_numpy, cast, divide, subtract + + _check_array(x, "var") + rdt = _mean_dtype(x, "var") + real_dt = np.dtype(rdt.char.lower()) if rdt.kind == "c" else rdt + axes = _axes(x, axis) + L = math.prod(x.shape[a] for a in axes) + if L == 0: + shape = _out_shape(x.shape, axes, keepdims) + return Array._new(Tensor(gb_dtype(real_dt), shape), np.nan) + denominator = builtins.max(L - correction, 0) + if not axes: + xr = cast(x, rdt) + out = divide(_squared_magnitude(subtract(xr, xr)), denominator) + return out if out.dtype == real_dt else cast(out, real_dt) + T = x._T() + m = mean(x, axis=axes) + keep_shape = _out_shape(x.shape, axes, keepdims=True) + if T.nvals: + _utils.check_dense(math.prod(keep_shape), "var or std (the means)") + M = _utils.full_tensor(rdt, keep_shape, m._fill) + Tm = m._T() + if Tm.nvals: + M(accum=binary.second) << Tm.expand_dims(axes) + D = T.ewise_mult(M, binary.minus[gb_dtype(rdt)]).new() + S2 = _utils.map_values(lambda d: np.real(d * np.conj(d)).astype(real_dt), [D], real_dt) + S2 = _reduce_tensor(S2, monoid.plus[gb_dtype(real_dt)], axes) + C = _reduce_tensor(T, agg.count, axes) + else: + S2 = Tensor(gb_dtype(real_dt), m.shape) + C = Tensor(gb_dtype(np.int64), m.shape) + k = subtract(L, Array._new(C, 0)) + dev = subtract(Array._new(Tensor(gb_dtype(rdt), ()), x._fill), m) + # k unstored positions, each (fill - mean) ** 2 away; none for k == 0 (the fill may be inf) + missing = binary_numpy("var", _times_count, k, _squared_magnitude(dev)) + out = divide(add(Array._new(S2, 0), missing), denominator) + if out.dtype != real_dt: + out = cast(out, real_dt) + return _keepdims(out, x, axes, keepdims) + + +def _times_count(k, d): + return np.where(k == 0, d.dtype.type(0), k * d) + + +def std(x, /, *, axis=None, correction=0.0, keepdims=False): + from ._elementwise import sqrt + + return sqrt(var(x, axis=axis, correction=correction, keepdims=keepdims)) + + +def _cumulative(name, npfunc, x, axis, dtype, include_initial): + _check_array(x, name) + if x.ndim == 0: + raise ValueError(f"{name} needs an array of at least one dimension") + if axis is None: + if x.ndim > 1: + raise ValueError(f"{name} of an array of more than one dimension needs an axis") + axis = 0 + axis = _axes(x, axis)[0] + if dtype is None: + reduce_name = "sum" if name == "cumulative_sum" else "prod" + ct = check_supported(_MONOIDS[reduce_name].npfunc(np.zeros(1, x.dtype)).dtype, name) + else: + ct = as_dtype(dtype) + _utils.check_dense(x.size + (x.size // builtins.max(x.shape[axis], 1)), name) + with np.errstate(all="ignore"): + out = npfunc(x.todense(), axis=axis, dtype=ct) + if include_initial: + shape = list(out.shape) + shape[axis] = 1 + initial = np.zeros(shape, ct) if name == "cumulative_sum" else np.ones(shape, ct) + out = np.concatenate([initial, out], axis=axis) + from ._creation import _from_numpy + + return _from_numpy(out) + + +def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False): + return _cumulative("cumulative_sum", np.cumsum, x, axis, dtype, include_initial) + + +def cumulative_prod(x, /, *, axis=None, dtype=None, include_initial=False): + return _cumulative("cumulative_prod", np.cumprod, x, axis, dtype, include_initial) + + +def diff(x, /, *, axis=-1, n=1, prepend=None, append=None): + from ._elementwise import not_equal, subtract + from ._manipulation import _concat + + _check_array(x, "diff") + if x.ndim == 0: + raise ValueError("diff needs an array of at least one dimension") + axis = _axes(x, axis)[0] + n = int(n) + if n < 0: + raise ValueError(f"diff's n must be nonnegative, not {n}") + parts = [_diff_part(p, x, axis) for p in (prepend, x, append) if p is not None] + if len(parts) > 1: + x = _concat(parts, axis) + if n == 0: + return x._copy() + # (numpy's diff of booleans is not_equal: subtracting booleans is an error) + op = not_equal if x.dtype.kind == "b" else subtract + before = (slice(None),) * axis + for _ in range(n): + if x.shape[axis] == 0: + break + x = op(x[(*before, slice(1, None))], x[(*before, slice(None, -1))]) + return x + + +def _diff_part(p, x, axis): + """A ``prepend`` or ``append`` of ``diff``: an Array, or a scalar broadcast to ``x``'s + shape with length 1 along ``axis`` (as numpy takes it). + """ + if isinstance(p, Array) and p.ndim: + return p + from ._creation import asarray + from ._manipulation import _broadcast_to + + shape = list(x.shape) + shape[axis] = 1 + return _broadcast_to(asarray(p), tuple(shape)) diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index eadaae9b4..d3a0d0c5c 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -145,10 +145,10 @@ def _assert_identical(a, b): raise AssertionError(f"{bad.sum()} positions differ (position, got, expected): {detail}") -def np_call(func, *args): +def np_call(func, *args, **kwargs): with np.errstate(all="ignore"), warnings.catch_warnings(): warnings.simplefilter("ignore") - return func(*args) + return func(*args, **kwargs) def inputs(dtype, shape, seed=0): diff --git a/graphblas/tests/test_array_api_statistical.py b/graphblas/tests/test_array_api_statistical.py new file mode 100644 index 000000000..f8c2ed90a --- /dev/null +++ b/graphblas/tests/test_array_api_statistical.py @@ -0,0 +1,248 @@ +"""graphblas.array_api statistical and utility functions against numpy.""" + +import warnings + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import ( + DTYPES, + assert_array_equal, + check_invariants, + inputs, + make, + np_call, +) + +REDUCTIONS = [ + ("sum", np.sum), + ("prod", np.prod), + ("min", np.min), + ("max", np.max), + ("any", np.any), + ("all", np.all), + ("mean", np.mean), + ("var", np.var), + ("std", np.std), +] +SHAPES_AXES = [ + ((6,), [None, 0, -1, (0,), ()]), + ((3, 4), [None, 0, 1, -2, (0, 1), (1, 0), ()]), + ((2, 3, 4), [None, 0, 2, (0, 2), (1, 2), (0, 1, 2)]), + ((), [None, ()]), +] + + +def assert_close(x, expected, scale=1.0): + """``x`` equals ``expected`` up to rounding of floating-point sums (exactly otherwise); + ``scale`` is the size of the largest terms summed, which bounds the rounding. + """ + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape + got = x.todense() + if got.dtype.kind in "fc": + rtol = 1e-4 if got.dtype in (np.float32, np.complex64) else 1e-10 + atol = rtol * scale + np.testing.assert_allclose(got, expected, rtol=rtol, atol=atol, equal_nan=True) + # an infinity's sign, and NaN where numpy has NaN, are not rounding + assert (np.isnan(got) == np.isnan(expected)).all() + assert (np.isinf(got) == np.isinf(expected)).all() + else: + np.testing.assert_array_equal(got, expected) + + +@pytest.mark.parametrize(("name", "func"), REDUCTIONS, ids=[n for n, _ in REDUCTIONS]) +@pytest.mark.parametrize(("shape", "axes"), SHAPES_AXES, ids=str) +@pytest.mark.parametrize("keepdims", [False, True]) +def test_reductions(name, func, shape, axes, keepdims): + for dtype in DTYPES: + for dense, x in inputs(dtype, shape, seed=11): + for axis in axes: + kwargs = {"axis": axis, "keepdims": keepdims} + try: + expected = np_call(func, dense, **kwargs) + except (TypeError, ValueError) as exc: + with pytest.raises(type(exc)): + getattr(xp, name)(x, **kwargs) + continue + if name in {"min", "max"} and dtype.kind == "c": + with pytest.raises(TypeError): + getattr(xp, name)(x, **kwargs) + continue + if np.asarray(expected).dtype not in set(DTYPES): + continue + out = getattr(xp, name)(x, **kwargs) + exact = name in {"min", "max", "any", "all"} or dtype.kind in "biu" + exact = exact and name not in {"mean", "var", "std"} + try: + if exact: + assert_array_equal(out, np.asarray(expected)) + elif name == "prod" and not np.isfinite(dense).all(): + # (a product of infinities and zeros: NaN or not depends on the order) + pass + else: + finite = dense[np.isfinite(dense)] if dense.dtype.kind in "fc" else dense + scale = float(np.abs(finite.astype(np.complex128)).max(initial=1.0)) + assert_close(out, expected, scale * dense.size) + except AssertionError as exc: + raise AssertionError( + f"{name} of {dtype} fill {x.fill_value!r} axis={axis}: {exc}" + ) from None + + +@pytest.mark.parametrize("name", ["sum", "prod"]) +@pytest.mark.parametrize("dtype", ["int8", "uint8", "int32", "float32", "bool"]) +def test_sum_prod_dtype(name, dtype): + func = getattr(np, name) + for dense, x in inputs(np.dtype(dtype), (3, 5), seed=2): + for out_dtype in [None, np.dtype("int64"), np.dtype("float64"), np.dtype("uint64")]: + if out_dtype is not None and dtype == "float32" and out_dtype.kind in "iu": + continue + expected = np_call(func, dense, axis=1, dtype=out_dtype) + out = getattr(xp, name)(x, axis=1, dtype=out_dtype) + assert_close(out, expected) + + +def test_sum_wraps_like_numpy(): + a = np.array([[100, 100, 0, 0], [0, 0, 0, 50]], dtype=np.int8) + x = make(a, 0) + assert_array_equal(xp.sum(x, axis=1, dtype=xp.int8), np.sum(a, axis=1, dtype=np.int8)) + # a nonzero fill counted many times over + x = make(np.full((2, 300), 7, dtype=np.int8), 0) + 0 + y = xp.full((2, 300), 7, dtype=xp.int8) + assert_array_equal( + xp.sum(y, axis=1, dtype=xp.int8), np.full(2, 7 * 300, np.int64).astype(np.int8) + ) + assert_array_equal( + xp.prod(y, axis=1, dtype=xp.int8), + np.prod(np.full((2, 300), 7, np.int8), axis=1, dtype=np.int8), + ) + assert_array_equal(xp.prod(y, axis=1), np.prod(np.full((2, 300), 7, np.int8), axis=1)) + del x + + +def test_min_max_nan_propagates(): + a = np.array([[1.0, np.nan, 0.0, 0.0], [0.0, 0.0, 2.0, 0.0], [0, 0, 0, 0]]) + x = make(a, 0.0) + for name in ["min", "max"]: + for axis in [None, 0, 1]: + assert_array_equal(getattr(xp, name)(x, axis=axis), getattr(np, name)(a, axis=axis)) + x = make(np.array([[1.0, np.nan], [np.nan, np.nan]]), np.nan) + assert_array_equal(xp.max(x, axis=1), np.array([np.nan, np.nan])) + + +def test_empty_axis(): + x = xp.zeros((0, 3)) + assert_array_equal(xp.sum(x, axis=0), np.zeros(3)) + assert_array_equal(xp.prod(x, axis=0), np.ones(3)) + assert_array_equal(xp.any(x, axis=0), np.zeros(3, bool)) + assert_array_equal(xp.all(x, axis=0), np.ones(3, bool)) + assert_array_equal(xp.sum(x, axis=1), np.zeros(0)) + with pytest.raises(ValueError, match="empty"): + xp.max(x, axis=0) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + assert_array_equal(xp.mean(x, axis=0), np.mean(np.zeros((0, 3)), axis=0)) + assert_array_equal(xp.var(x, axis=0), np.var(np.zeros((0, 3)), axis=0)) + assert_array_equal(xp.max(x, axis=1), np.zeros(0)) + + +@pytest.mark.parametrize("correction", [0, 1, 2.5, 4, 10]) +def test_var_correction(correction): + rng = np.random.default_rng(3) + a = np.where(rng.random((3, 4)) < 0.5, 0.0, np.round(rng.normal(size=(3, 4)), 2)) + for fill in [0.0, 1.5]: + dense = np.where(a == 0, fill, a) + x = make(dense, fill) + for axis in [None, 0, 1, ()]: + for name in ["var", "std"]: + expected = np_call(getattr(np, name), dense, axis=axis, ddof=correction) + out = np_call(getattr(xp, name), x, axis=axis, correction=correction) + assert_close(out, expected) + + +def test_reductions_on_full_arrays(): + a = np.arange(1.0, 13.0).reshape(3, 4) + x = make(a, 0.0) + assert x.nnz == 12 + for name, func in REDUCTIONS: + for axis in [None, 0, 1]: + assert_close(getattr(xp, name)(x, axis=axis), func(a, axis=axis)) + + +def test_reductions_of_fill_only_arrays(): + x = xp.full((4, 5), 2.5) + assert_close(xp.sum(x, axis=1), np.full(4, 12.5)) + assert xp.sum(x, axis=1).nnz == 0 + assert_close(xp.prod(x), np.prod(np.full((4, 5), 2.5))) + assert_close(xp.var(x, axis=0), np.zeros(5)) + x = xp.full((3, 3), np.inf) + assert_array_equal(np_call(xp.var, x, axis=0), np_call(np.var, np.full((3, 3), np.inf), axis=0)) + + +@pytest.mark.parametrize("name", ["cumulative_sum", "cumulative_prod"]) +@pytest.mark.parametrize("include_initial", [False, True]) +def test_cumulative(name, include_initial): + func = np.cumsum if name == "cumulative_sum" else np.cumprod + for dtype in [np.dtype(t) for t in ("bool", "int8", "uint16", "int64", "float32", "float64")]: + for shape, axes in [((7,), [None, 0, -1]), ((3, 4), [0, 1, -1]), ((2, 3, 2), [0, 2])]: + for dense, x in inputs(dtype, shape, seed=5): + for axis in axes: + np_axis = 0 if axis is None else axis + expected = np_call(func, dense, axis=np_axis) + if include_initial: + initial = np.zeros if name == "cumulative_sum" else np.ones + shp = list(expected.shape) + shp[np_axis] = 1 + expected = np.concatenate( + [initial(shp, expected.dtype), expected], axis=np_axis + ) + out = getattr(xp, name)(x, axis=axis, include_initial=include_initial) + assert_close(out, expected) + x = make(np.ones((2, 3)), 0.0) + with pytest.raises(ValueError, match="axis"): + getattr(xp, name)(x) + out = getattr(xp, name)(make(np.array([1, 2, 3], np.int8), 0), dtype=xp.float32) + assert out.dtype == xp.float32 + + +def test_cumulative_is_guarded(): + with gb.config.set(array_api_densify_limit=10), pytest.raises(xp.DensifyError): + xp.cumulative_sum(xp.zeros(100)) + + +@pytest.mark.parametrize("n", [0, 1, 2, 5]) +def test_diff(n): + for dtype in [np.dtype(t) for t in ("bool", "int16", "uint8", "float64", "complex64")]: + for shape, axes in [((6,), [0, -1]), ((3, 4), [0, 1])]: + for dense, x in inputs(dtype, shape, seed=6): + for axis in axes: + expected = np_call(np.diff, dense, n=n, axis=axis) + assert_array_equal(xp.diff(x, n=n, axis=axis), expected) + + +def test_diff_prepend_append(): + a = np.array([[1.0, 0.0, 3.0], [0.0, 0.0, -1.0]]) + x = make(a, 0.0) + p = np.array([[5.0], [0.0]]) + q = np.array([[0.0, 2.0, 0.0]]) + assert_array_equal(xp.diff(x, axis=1, prepend=make(p, 0.0)), np.diff(a, axis=1, prepend=p)) + assert_array_equal(xp.diff(x, axis=0, append=make(q, 1.0)), np.diff(a, axis=0, append=q)) + assert_array_equal(xp.diff(x, axis=1, append=2.0), np.diff(a, axis=1, append=2.0)) + assert_array_equal( + xp.diff(x, axis=1, n=2, prepend=make(p, 0.0), append=make(p, 7.0)), + np.diff(a, axis=1, n=2, prepend=p, append=p), + ) + + +def test_axis_errors(): + x = make(np.ones((2, 3)), 0.0) + with pytest.raises(IndexError): + xp.sum(x, axis=2) + with pytest.raises(ValueError, match="repeated"): + xp.sum(x, axis=(0, -2)) From 6bbbb70fa74733bf9f40e6e7be505d52aa0bd4a9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:35:12 -0600 Subject: [PATCH 04/77] array_api: test creation, data type functions, the Array object and basic indexing isdtype takes only this namespace's dtypes (or numpy scalar types), as numpy's does, not strings. --- graphblas/array_api/_dtypes.py | 7 +- graphblas/tests/test_array_api_creation.py | 529 +++++++++++++++++++++ 2 files changed, 532 insertions(+), 4 deletions(-) create mode 100644 graphblas/tests/test_array_api_creation.py diff --git a/graphblas/array_api/_dtypes.py b/graphblas/array_api/_dtypes.py index f107a3afa..176121836 100644 --- a/graphblas/array_api/_dtypes.py +++ b/graphblas/array_api/_dtypes.py @@ -71,11 +71,10 @@ def _dtype_of(x): def isdtype(dtype, kind): + if isinstance(dtype, type) and issubclass(dtype, np.generic): + dtype = np.dtype(dtype) if not isinstance(dtype, np.dtype) or dtype not in _SUPPORTED: - try: - dtype = as_dtype(dtype) - except TypeError: - raise TypeError(f"isdtype needs a data type, not {dtype!r}") from None + raise TypeError(f"isdtype needs a data type of this namespace, not {dtype!r}") if isinstance(kind, tuple): return builtins.any(isdtype(dtype, k) for k in kind) if isinstance(kind, str): diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py new file mode 100644 index 000000000..a4af7c1d2 --- /dev/null +++ b/graphblas/tests/test_array_api_creation.py @@ -0,0 +1,529 @@ +"""graphblas.array_api: creation, data types, the Array object, basic indexing and the shape +functions, against numpy. +""" + +import math +import pickle + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp +from graphblas import Matrix, Scalar, Tensor, Vector + +from .test_array_api_elementwise import ( + DTYPES, + assert_array_equal, + check_invariants, + inputs, + make, + np_call, +) + + +def test_namespace(): + assert xp.__array_api_version__ == "2025.12" + for name in xp.__all__: + assert hasattr(xp, name), name + assert xp.float64 is np.dtype("float64") + assert xp.bool == np.dtype(bool) + x = xp.asarray([1, 2]) + assert x.__array_namespace__() is xp + assert x.__array_namespace__(api_version="2024.12") is xp + with pytest.raises(ValueError, match="version"): + x.__array_namespace__(api_version="2019.12") + + +# asarray +def test_asarray_stores_what_differs_from_the_fill(): + x = xp.asarray([1.0, 0.0, -0.0, np.nan, 2.0]) + assert x.fill_value == 0.0 + assert x.nnz == 4 # -0.0 and NaN are stored + assert_array_equal(x, np.array([1.0, 0.0, -0.0, np.nan, 2.0])) + y = xp.asarray([1.0, np.nan, np.nan], fill_value=np.nan) + assert y.nnz == 1 + assert_array_equal(y, np.array([1.0, np.nan, np.nan])) + z = xp.asarray([[1.0, 2.0], [-0.0, 0.0]], fill_value=-0.0) + assert z.nnz == 3 + assert_array_equal(z, np.array([[1.0, 2.0], [-0.0, 0.0]])) + c = xp.asarray([complex(0, -0.0), 0j, complex(np.nan, 0)]) + assert c.nnz == 2 + assert_array_equal(c, np.array([complex(0, -0.0), 0j, complex(np.nan, 0)])) + + +@pytest.mark.parametrize( + ("obj", "dtype"), + [ + (True, "bool"), + (3, "int64"), + (2.5, "float64"), + (1j, "complex128"), + ([1, 2], "int64"), + ([1.0, 2], "float64"), + ([True, False], "bool"), + ([[1, 0], [0, 3]], "int64"), + ([], "float64"), + (np.float32(1.5), "float32"), + (np.arange(4, dtype=np.uint16), "uint16"), + ], +) +def test_asarray_dtypes(obj, dtype): + x = xp.asarray(obj) + assert x.dtype == np.dtype(dtype) + assert_array_equal(x, np.asarray(obj, dtype=dtype)) + + +def test_asarray_dtype_and_copy(): + x = xp.asarray([1, 2, 3], dtype=xp.float32) + assert x.dtype == xp.float32 + assert xp.asarray(x) is x + assert xp.asarray(x, copy=False) is x + y = xp.asarray(x, copy=True) + assert y is not x + y += 1 + assert_array_equal(x, np.array([1, 2, 3], np.float32)) + assert_array_equal(xp.asarray(x, dtype=xp.int8), np.array([1, 2, 3], np.int8)) + with pytest.raises(ValueError, match="copy"): + xp.asarray(x, dtype=xp.int8, copy=False) + with pytest.raises(ValueError, match="copy"): + xp.asarray([1, 2], copy=False) + with pytest.raises(TypeError): + xp.asarray([1, 2], dtype=np.float16) + with pytest.raises(TypeError): + xp.asarray(np.array(["a"])) + with pytest.raises(OverflowError): + xp.asarray([300], dtype=xp.int8) + with pytest.raises(ValueError, match="device"): + xp.asarray([1], device="gpu") + + +def test_asarray_nested_arrays(): + x = xp.asarray([xp.asarray(1.0), xp.asarray(2.0)]) + assert_array_equal(x, np.array([1.0, 2.0])) + x = xp.asarray([[1, 2], xp.asarray([3, 4])]) + assert_array_equal(x, np.array([[1, 2], [3, 4]])) + + +def test_asarray_graphblas_objects(): + v = Vector.from_coo([0, 3], [1.5, 2.5], size=5) + x = xp.asarray(v) + assert x.nnz == 2 + assert_array_equal(x, np.array([1.5, 0, 0, 2.5, 0])) + x = xp.asarray(v, fill_value=np.inf) + assert_array_equal(x, np.array([1.5, np.inf, np.inf, 2.5, np.inf])) + A = Matrix.from_coo([0, 1], [1, 0], [5, 6], nrows=2, ncols=3) + assert_array_equal(xp.asarray(A), np.array([[0, 5, 0], [6, 0, 0]])) + assert_array_equal(xp.asarray(A.T), np.array([[0, 6], [5, 0], [0, 0]])) + assert_array_equal( + xp.asarray(A, dtype=xp.float32), np.array([[0, 5, 0], [6, 0, 0]], np.float32) + ) + T = Tensor.from_coo([[0, 1], [1, 0], [2, 3]], [1.0, 2.0], shape=(2, 2, 4)) + x = xp.asarray(T, fill_value=-1.0) + expected = np.full((2, 2, 4), -1.0) + expected[0, 1, 2], expected[1, 0, 3] = 1.0, 2.0 + assert_array_equal(x, expected) + T[0, 0, 0] = 9.0 # the Array has its own data + assert_array_equal(x, expected) + assert_array_equal(xp.asarray(Scalar.from_value(3)), np.array(3)) + assert_array_equal(xp.asarray(Scalar(float)), np.array(0.0)) + assert isinstance(x.tensor, Tensor) + assert x.tensor.nvals == 2 + + +def test_asarray_scipy_and_pydata_sparse(): + sp = pytest.importorskip("scipy.sparse") + A = sp.csr_array(np.array([[0.0, 1.5], [2.0, 0.0]])) + assert_array_equal(xp.asarray(A), np.array([[0.0, 1.5], [2.0, 0.0]])) + sparse = pytest.importorskip("sparse") + s = sparse.COO.from_numpy(np.array([[1.0, 2.0], [2.0, 3.0]]), fill_value=2.0) + x = xp.asarray(s) + assert x.fill_value == 2.0 + assert x.nnz == 2 + assert_array_equal(x, np.array([[1.0, 2.0], [2.0, 3.0]])) + + +def test_from_coo(): + x = xp.from_coo(([0, 2], [1, 0]), [1.0, 2.0], (3, 2), fill_value=5.0) + assert_array_equal(x, np.array([[5.0, 1.0], [5.0, 5.0], [2.0, 5.0]])) + x = xp.from_coo(np.array([[0, 2], [1, 0]]), 7, (3, 2)) + assert x.dtype == np.dtype(int) + assert_array_equal(x, np.array([[0, 7], [0, 0], [7, 0]])) + x = xp.from_coo(([],), [], (4,), dtype=xp.int8) + assert_array_equal(x, np.zeros(4, np.int8)) + coords, values = xp.from_coo(([3, 1],), [1.0, 2.0], (5,)).to_coo() + np.testing.assert_array_equal(coords[0], [1, 3]) + np.testing.assert_array_equal(values, [2.0, 1.0]) + + +# Creation functions +@pytest.mark.parametrize("shape", [(), 3, (0,), (2, 3), (2, 0, 4)]) +@pytest.mark.parametrize("dtype", [None, "bool", "int16", "uint8", "float32", "complex64"]) +def test_full_family(shape, dtype): + dt = None if dtype is None else np.dtype(dtype) + for name, value in [("zeros", 0), ("ones", 1), ("empty", 0)]: + x = getattr(xp, name)(shape, dtype=dt) + assert x.nnz == 0 + assert_array_equal(x, getattr(np, name if name != "empty" else "zeros")(shape, dtype=dt)) + like = getattr(xp, f"{name}_like")(x) + assert like.nnz == 0 + assert_array_equal(like, np.full(x.shape, value, dtype=x.dtype)) + if dt is None: + for value in [True, 3, 2.5, 1j, -0.0, np.nan]: + assert_array_equal(xp.full(shape, value), np.full(shape, value)) + else: + assert_array_equal(xp.full(shape, 1, dtype=dt), np.full(shape, 1, dtype=dt)) + assert_array_equal(xp.full_like(xp.zeros(shape), 2, dtype=dt), np.full(shape, 2, dtype=dt)) + with pytest.raises(OverflowError): + xp.full(shape, 1000, dtype=xp.int8) + + +@pytest.mark.parametrize("dtype", [None, "int32", "bool", "complex128"]) +def test_eye(dtype): + for args in [(3,), (3, 5), (5, 3), (0,), (4, 4)]: + for k in [-6, -2, -1, 0, 1, 3, 5]: + assert_array_equal(xp.eye(*args, k=k, dtype=dtype), np.eye(*args, k=k, dtype=dtype)) + + +@pytest.mark.parametrize( + "args", + [(5,), (2, 9), (2, 9, 3), (9, 2, -2), (0,), (1.5, 4.0, 0.5), (10, 0, -3), (0, 1, 0.1)], +) +@pytest.mark.parametrize("dtype", [None, "int16", "float32"]) +def test_arange(args, dtype): + if dtype == "int16" and any(isinstance(a, float) for a in args): + return + expected = np.arange(*args, dtype=dtype) + if dtype is None and all(isinstance(a, int) for a in args): + expected = expected.astype(np.int64) + assert_array_equal(xp.arange(*args, dtype=dtype), expected) + with pytest.raises(ValueError, match="step"): + xp.arange(0, 3, 0) + + +@pytest.mark.parametrize( + ("start", "stop", "num"), [(0, 1, 5), (2.5, -1, 4), (0, 1, 1), (0, 1, 0), (1j, 2, 3)] +) +@pytest.mark.parametrize("endpoint", [True, False]) +def test_linspace(start, stop, num, endpoint): + expected = np.linspace(start, stop, num, endpoint=endpoint) + assert_array_equal(xp.linspace(start, stop, num, endpoint=endpoint), expected) + if not isinstance(start, complex): + assert_array_equal( + xp.linspace(start, stop, num, endpoint=endpoint, dtype=xp.float32), + np.linspace(start, stop, num, endpoint=endpoint, dtype=np.float32), + ) + + +def test_tril_triu(): + rng = np.random.default_rng(0) + for shape in [(4, 4), (3, 5), (2, 3, 4), (0, 3)]: + a = np.where(rng.random(shape) < 0.5, 0.0, rng.normal(size=shape)) + for fill in [0.0, 2.0]: + dense = np.where(a == 0, fill, a) + x = make(dense, fill) + for k in [-3, -1, 0, 1, 2, 5]: + assert_array_equal(xp.tril(x, k=k), np.tril(dense, k=k)) + assert_array_equal(xp.triu(x, k=k), np.triu(dense, k=k)) + with pytest.raises(ValueError, match="2 dimensions"): + xp.tril(xp.zeros(3)) + + +@pytest.mark.parametrize("indexing", ["xy", "ij"]) +def test_meshgrid(indexing): + a, b, c = np.array([1.0, 0.0, 2.0]), np.array([0.0, 5.0]), np.array([7.0, 0.0, 0.0, 1.0]) + for arrays in [(a,), (a, b), (a, b, c)]: + out = xp.meshgrid(*[make(v, 0.0) for v in arrays], indexing=indexing) + expected = np.meshgrid(*arrays, indexing=indexing) + assert len(out) == len(expected) + for o, e in zip(out, expected, strict=True): + assert_array_equal(o, e) + assert xp.meshgrid() == [] + + +def test_from_dlpack(): + a = np.array([[1.0, 0.0], [0.0, 2.0]]) + assert_array_equal(xp.from_dlpack(a), a) + x = make(a, 0.0) + assert_array_equal(xp.from_dlpack(x), a) + np.testing.assert_array_equal(np.from_dlpack(x), a) + + +# Data type functions +def test_astype_all_pairs(): + for src in DTYPES: + for dense, x in inputs(src, (3, 4), seed=1): + for dst in DTYPES: + # (numpy's astype from complex to real keeps the real part, with a warning) + real = dense.real if src.kind == "c" and dst.kind not in "cb" else dense + expected = np_call(real.astype, dst) + if src.kind == "f" and dst.kind in "iu" and not np.isfinite(dense).all(): + # numpy's cast of NaN and inf to an integer is undefined; any value will do + out = xp.astype(x, dst) + assert out.dtype == dst + finite = np.isfinite(dense) + np.testing.assert_array_equal(out.todense()[finite], expected[finite]) + continue + if src.kind == "f" and dst.kind in "iu": + big = np.abs(dense) >= 2.0 ** (8 * dst.itemsize - 1) + out = xp.astype(x, dst) + np.testing.assert_array_equal(out.todense()[~big], expected[~big]) + continue + assert_array_equal(xp.astype(x, dst), expected) + x = xp.asarray([1, 2]) + assert xp.astype(x, xp.int64, copy=False) is x + assert xp.astype(x, xp.int64) is not x + + +def test_can_cast(): + assert xp.can_cast(xp.int8, xp.int16) + assert not xp.can_cast(xp.int16, xp.int8) + assert xp.can_cast(xp.asarray([1], dtype=xp.float32), xp.float64) + assert not xp.can_cast(xp.float64, xp.float32) + assert xp.can_cast(xp.bool, xp.uint8) + + +def test_finfo_iinfo(): + for dt in [xp.float32, xp.float64, xp.complex64, xp.complex128]: + info = xp.finfo(dt) + ref = np.finfo(dt) + assert info.bits == ref.bits + assert info.eps == float(ref.eps) + assert info.max == float(ref.max) + assert info.min == float(ref.min) + assert info.smallest_normal == float(ref.smallest_normal) + assert isinstance(info.eps, float) + assert xp.finfo(xp.asarray(1, dtype=dt)) == info + assert xp.finfo(xp.complex64).dtype == xp.float32 + for dt in [xp.int8, xp.uint64, xp.int64]: + info = xp.iinfo(dt) + assert (info.bits, info.min, info.max) == ( + np.iinfo(dt).bits, + int(np.iinfo(dt).min), + int(np.iinfo(dt).max), + ) + assert info.dtype == dt + with pytest.raises(TypeError): + xp.finfo(xp.int8) + with pytest.raises(TypeError): + xp.iinfo(xp.float32) + + +def test_isdtype(): + assert xp.isdtype(xp.int8, "signed integer") + assert xp.isdtype(xp.uint8, "integral") + assert not xp.isdtype(xp.bool, "numeric") + assert xp.isdtype(xp.complex64, ("real floating", "complex floating")) + assert xp.isdtype(xp.float32, xp.float32) + assert not xp.isdtype(xp.float32, xp.float64) + with pytest.raises(ValueError, match="kind"): + xp.isdtype(xp.int8, "integer") + with pytest.raises(TypeError): + xp.isdtype("int8", "integral") + + +def test_result_type(): + assert xp.result_type(xp.int8, xp.uint8) == xp.int16 + assert xp.result_type(xp.asarray([1], dtype=xp.float32), 1.5) == xp.float32 + assert xp.result_type(xp.int8, 1) == xp.int8 + assert xp.result_type(xp.float32, xp.int8, xp.uint8) == np.result_type( + np.result_type(np.float32, np.int8), np.uint8 + ) + with pytest.raises(ValueError, match="at least one"): + xp.result_type(1, 2.0) + + +def test_info(): + info = xp.__array_namespace_info__() + assert info.default_device() == "cpu" + assert info.devices() == ["cpu"] + assert info.default_dtypes()["real floating"] == xp.float64 + assert info.default_dtypes(device="cpu")["indexing"] == xp.int64 + assert set(info.dtypes()) == {str(dt) for dt in DTYPES} + assert set(info.dtypes(kind="integral")) == { + "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64" + } # fmt: skip + assert set(info.dtypes(kind=("bool", "complex floating"))) == { + "bool", + "complex64", + "complex128", + } + assert info.capabilities()["data-dependent shapes"] + + +# The Array object +def test_array_attributes(): + a = np.arange(24.0).reshape(2, 3, 4) + x = make(a, 0.0) + assert x.shape == (2, 3, 4) + assert x.ndim == 3 + assert x.size == 24 + assert x.dtype == xp.float64 + assert x.device == "cpu" + assert x.to_device("cpu") is x + with pytest.raises(ValueError, match="device"): + x.to_device("gpu") + assert_array_equal(x.T, a.T) + assert_array_equal(x.mT, np.swapaxes(a, -1, -2)) + with pytest.raises(ValueError, match="2 dimensions"): + xp.zeros(3).mT + assert x.nnz == 23 + assert x.fill_value == 0.0 + assert repr(x) == "Array(shape=(2, 3, 4), dtype=float64, nnz=23, fill_value=0.0)" + assert repr(xp.asarray(1.5)) == "Array(1.5, dtype=float64)" + with pytest.raises(TypeError, match="asarray"): + xp.Array() + assert len(x) == 2 + assert [y.shape for y in x] == [(3, 4), (3, 4)] + with pytest.raises(TypeError, match="len"): + _ = len(xp.asarray(1)) + with pytest.raises(TypeError): + _ = iter(xp.asarray(1)) + with pytest.raises(TypeError): + hash(x) + + +def test_python_scalars(): + assert bool(xp.asarray(True)) is True + assert int(xp.asarray(-3)) == -3 + assert float(xp.asarray(2.5)) == 2.5 + assert complex(xp.asarray(1 + 2j)) == 1 + 2j + assert [10, 20, 30][xp.asarray(1)] == 20 + assert math.isnan(float(xp.asarray(np.nan))) + with pytest.raises(TypeError): + [1, 2][xp.asarray(1.0)] + with pytest.raises(TypeError): + bool(xp.asarray([1, 2])) + with pytest.raises(TypeError): + float(xp.asarray(1j)) + with pytest.raises(ValueError, match="NaN"): + int(xp.asarray(np.nan)) + + +def test_array_protocol_is_guarded(): + x = make(np.array([1.0, 0.0, 2.0]), 0.0) + np.testing.assert_array_equal(np.asarray(x), [1.0, 0.0, 2.0]) + assert np.asarray(x, dtype=np.float32).dtype == np.float32 + with pytest.raises(ValueError, match="copies"): + np.asarray(x, copy=False) + big = xp.zeros((10**5, 10**5)) + with pytest.raises(xp.DensifyError, match="array_api_densify_limit"): + np.asarray(big) + with gb.config.set(array_api_densify_limit=2), pytest.raises(xp.DensifyError): + x.todense() + assert big.nnz == 0 + assert xp.sum(big) == 0.0 + + +def test_pickle_not_supported_yet_or_roundtrips(): + x = make(np.array([1.0, 0.0, 2.0]), 0.0) + try: + y = pickle.loads(pickle.dumps(x)) + except (TypeError, AttributeError, pickle.PicklingError): + return + assert_array_equal(y, x.todense()) + + +# Indexing and shape functions +@pytest.mark.parametrize( + "key", + [ + 0, + -1, + (1, 2), + (1, 2, 3), + slice(None), + slice(1, None), + slice(None, None, -1), + slice(None, None, 2), + (slice(None), 1), + (Ellipsis, 1), + (1, Ellipsis), + (Ellipsis,), + (None, 0), + (0, None, slice(None, None, -2)), + (Ellipsis, None), + (None, Ellipsis, None), + (slice(5, 1, -1), slice(0, 0)), + (), + (1, slice(0, 2), -1), + ], +) +def test_getitem_basic(key): + rng = np.random.default_rng(2) + a = np.where(rng.random((3, 4, 5)) < 0.5, 0.0, np.round(rng.normal(size=(3, 4, 5)), 1)) + for fill in [0.0, -1.5]: + dense = np.where(a == 0, fill, a) + x = make(dense, fill, rng) + assert_array_equal(x[key], dense[key]) + + +def test_getitem_errors_and_0d(): + x = make(np.arange(6.0).reshape(2, 3), 0.0) + with pytest.raises(IndexError): + x[2] + with pytest.raises(IndexError): + x[0, -4] + with pytest.raises(IndexError): + x[0, 0, 0] + with pytest.raises(IndexError): + x[..., 0, ...] + z = xp.asarray(2.5) + assert_array_equal(z[()], np.array(2.5)) + assert_array_equal(z[...], np.array(2.5)) + assert_array_equal(z[None], np.array([2.5])) + assert_array_equal(x[xp.asarray(1)], np.arange(3.0, 6.0)) + + +@pytest.mark.parametrize( + ("shape", "new"), + [ + ((2, 3, 4), (6, 4)), + ((2, 3, 4), (4, -1)), + ((2, 3, 4), (24,)), + ((2, 3, 4), (2, 3, 4)), + ((6,), (1, 6, 1)), + ((1,), ()), + ((), (1, 1)), + ((), -1), + ((0, 3), (3, 0)), + ], +) +def test_reshape(shape, new): + for dtype in [np.dtype("int16"), np.dtype("float64")]: + for dense, x in inputs(dtype, shape, seed=3): + assert_array_equal(xp.reshape(x, new), np.reshape(dense, new)) + assert_array_equal(xp.reshape(x, new, copy=True), np.reshape(dense, new)) + with pytest.raises(ValueError, match="reshape"): + xp.reshape(xp.zeros((2, 3)), (4, 2)) + with pytest.raises(ValueError, match="-1"): + xp.reshape(xp.zeros((2, 3)), (-1, -1)) + + +def test_reshape_copy_false(): + x = make(np.arange(24.0).reshape(2, 3, 4), 0.0) + out = xp.reshape(x, (2, 12), copy=False) + assert_array_equal(out, np.arange(24.0).reshape(2, 12)) + + +def test_permute_expand_squeeze(): + rng = np.random.default_rng(4) + a = np.where(rng.random((2, 1, 3)) < 0.5, 0, rng.integers(1, 9, (2, 1, 3))) + x = make(a, 0) + for axes in [(0, 1, 2), (2, 0, 1), (-1, -2, -3)]: + assert_array_equal(xp.permute_dims(x, axes), np.transpose(a, axes)) + with pytest.raises(ValueError, match="permutation"): + xp.permute_dims(x, (0, 0, 1)) + for axis in [0, 1, 3, -1, -4, (0, 2), (0, -1)]: + assert_array_equal(xp.expand_dims(x, axis=axis), np.expand_dims(a, axis)) + with pytest.raises(IndexError): + xp.expand_dims(x, axis=5) + assert_array_equal(xp.squeeze(x, axis=1), np.squeeze(a, axis=1)) + assert_array_equal(xp.squeeze(x, axis=(1,)), np.squeeze(a, axis=1)) + with pytest.raises(ValueError, match="squeeze"): + xp.squeeze(x, axis=0) + one = make(np.array([[7]]), 0) + assert_array_equal(xp.squeeze(one, axis=(0, 1)), np.array(7)) + assert_array_equal(xp.expand_dims(xp.asarray(3), axis=0), np.array([3])) + assert_array_equal(xp.permute_dims(xp.asarray(3), ()), np.array(3)) + check_invariants(xp.squeeze(one, axis=(0, 1))) From 6b30d8ca3384c285ad410771edf0ecba66747e6c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:37:17 -0600 Subject: [PATCH 05/77] array_api: devices() and meshgrid return tuples, as the 2025.12 standard says --- graphblas/array_api/_creation.py | 2 +- graphblas/array_api/_info.py | 2 +- graphblas/tests/test_array_api_creation.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/graphblas/array_api/_creation.py b/graphblas/array_api/_creation.py index dc7ee3732..3908de60e 100644 --- a/graphblas/array_api/_creation.py +++ b/graphblas/array_api/_creation.py @@ -354,7 +354,7 @@ def meshgrid(*arrays, indexing="xy"): s = [1] * n s[axis] = x.shape[0] out.append(_broadcast_to(reshape(x, tuple(s)), tuple(shape))) - return out + return tuple(out) def _tri(x, k, lower): diff --git a/graphblas/array_api/_info.py b/graphblas/array_api/_info.py index aab48b339..d155d409f 100644 --- a/graphblas/array_api/_info.py +++ b/graphblas/array_api/_info.py @@ -20,7 +20,7 @@ def default_device(self): return CPU_DEVICE def devices(self): - return [CPU_DEVICE] + return (CPU_DEVICE,) def default_dtypes(self, *, device=None): _check(device) diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py index a4af7c1d2..d67356301 100644 --- a/graphblas/tests/test_array_api_creation.py +++ b/graphblas/tests/test_array_api_creation.py @@ -238,7 +238,7 @@ def test_meshgrid(indexing): assert len(out) == len(expected) for o, e in zip(out, expected, strict=True): assert_array_equal(o, e) - assert xp.meshgrid() == [] + assert xp.meshgrid() == () def test_from_dlpack(): @@ -336,7 +336,7 @@ def test_result_type(): def test_info(): info = xp.__array_namespace_info__() assert info.default_device() == "cpu" - assert info.devices() == ["cpu"] + assert info.devices() == ("cpu",) assert info.default_dtypes()["real floating"] == xp.float64 assert info.default_dtypes(device="cpu")["indexing"] == xp.int64 assert set(info.dtypes()) == {str(dt) for dt in DTYPES} From df280d48dd128e46416bde0d5011053d673189e9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:40:32 -0600 Subject: [PATCH 06/77] array_api: multiply floats on the intersection when that is exact With both fills 0.0, the product at a position only one operand stores is value * 0.0, which is 0.0 (the fill) only for a finite value with its sign bit clear: -1.0 * 0.0 is -0.0 and inf * 0.0 is NaN. Each Array caches whether its stored values are all such, and multiply takes the intersection (ewise_mult) when both are, the union otherwise. --- graphblas/array_api/_array.py | 4 +- graphblas/array_api/_elementwise.py | 48 +++++++++++++++---- graphblas/tests/test_array_api_elementwise.py | 15 ++++++ 3 files changed, 58 insertions(+), 9 deletions(-) diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py index 6b1b55ed7..ff72da66b 100644 --- a/graphblas/array_api/_array.py +++ b/graphblas/array_api/_array.py @@ -65,7 +65,9 @@ def _replace(self, other): self._facts = None def _fact(self, name, compute): - """A cached property of the stored values; assignment clears the cache.""" + """A cached property of the stored values. Code that changes the stored entries in + place must clear the cache (``_facts = None``), as ``_replace`` does. + """ facts = self._facts if facts is None: facts = self._facts = {} diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index b7b71db26..2b6ee5df7 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -381,8 +381,9 @@ class _Binary: ``kind`` is "arith" (the op works in the result dtype), "cmp" (in the promoted dtype of the operands, with a bool result) or "logical" (in bool). ``gb(dtype)`` gives the GraphBLAS operator that equals ``ufunc`` on operands of the dtype it works in, or None. - ``annihilator(dtype)`` gives the fill value that, held by both operands, makes the union - equal to the intersection (``ewise_mult``), or None. + ``annihilator(a, b, fill_a, fill_b, dtype)`` says whether the union equals the + intersection (``ewise_mult``): whether each operand's fill makes the result the result's + fill whatever the other operand holds. """ def __init__(self, name, ufunc, kind="arith", gb=None, annihilator=None): @@ -471,8 +472,7 @@ def _binary_arrays(spec, a, b, op, ct, rdt, fill): if op is not None: fa = _utils.scalar(a._fill, ct) fb = _utils.scalar(b._fill, ct) - zero = spec.annihilator(ct) if spec.annihilator is not None else None - if zero is not None and _utils.same_value(fa, zero) and _utils.same_value(fb, zero): + if spec.annihilator is not None and spec.annihilator(a, b, fa, fb, ct): # Where either operand holds its fill the result is the fill: the intersection R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) return _result(R, fill) @@ -509,6 +509,37 @@ def _zero(dtype): return _utils.scalar(0, dtype) +def _both_fills_zero(a, b, fa, fb, ct): + zero = _zero(ct) + return _utils.same_value(fa, zero) and _utils.same_value(fb, zero) + + +def _multiply_intersects(a, b, fa, fb, ct): + """Whether ``a * b`` stores only the intersection of entries: both fills are 0 and, for + floats, ``value * 0.0`` is ``0.0`` for every stored value (not ``-0.0``, nor NaN for inf). + """ + if not _both_fills_zero(a, b, fa, fb, ct): + return False + if ct.kind != "f": + return True + return _finite_nonnegative(a) and _finite_nonnegative(b) + + +def _finite_nonnegative(x): + """Whether every stored value is finite with its sign bit clear (cached per Array).""" + + def compute(x): + T = x._T() + if x.dtype.kind == "b" or T.nvals == 0: + return True + values = _utils.tensor_values(T) + if x.dtype.kind in _INTEGRAL: + return bool((values >= 0).all()) + return bool(np.isfinite(values).all() and not np.signbit(values).any()) + + return x._fact("finite_nonnegative", compute) + + _SPECS = { "add": _Binary( "add", np.add, gb=_typed_for({"b": binary.lor, **dict.fromkeys("iufc", binary.plus)}) @@ -517,12 +548,13 @@ def _zero(dtype): "subtract", np.subtract, gb=_typed_for(dict.fromkeys("iufc", binary.minus)) ), # Complex products and quotients are computed by numpy: C's formulas round otherwise. - # A float zero fill is not an annihilator: 0.0 * -1.0 is -0.0, and 0.0 * inf is NaN. + # A float zero fill annihilates only nonnegative finite values: 0.0 * -1.0 is -0.0, and + # 0.0 * inf is NaN. "multiply": _Binary( "multiply", np.multiply, gb=_typed_for({"b": binary.land, **dict.fromkeys("iuf", binary.times)}), - annihilator=lambda dt: _zero(dt) if dt.kind in "biu" else None, + annihilator=_multiply_intersects, ), "divide": _Binary("divide", np.divide, gb=_typed_for({"f": binary.truediv})), "equal": _Binary("equal", np.equal, "cmp", gb=_typed_for(dict.fromkeys("biufc", binary.eq))), @@ -544,7 +576,7 @@ def _zero(dtype): np.logical_and, "logical", gb=lambda dt: binary.land[gb_dtype(dt)], - annihilator=_zero, + annihilator=_both_fills_zero, ), "logical_or": _Binary( "logical_or", np.logical_or, "logical", gb=lambda dt: binary.lor[gb_dtype(dt)] @@ -556,7 +588,7 @@ def _zero(dtype): "bitwise_and", np.bitwise_and, gb=_typed_for({"b": binary.land, "i": binary.band, "u": binary.band}), - annihilator=_zero, + annihilator=_both_fills_zero, ), "bitwise_or": _Binary( "bitwise_or", diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index d3a0d0c5c..2b3cacbf1 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -468,6 +468,21 @@ def test_integer_multiply_intersects(): assert_array_equal(out, np.array([0, 0, 10, 0, 18])) +def test_float_multiply_of_nonnegative_values_intersects(): + x = make(np.array([1.0, 0.0, 2.0, 0.0, 3.0]), 0.0) + y = make(np.array([0.0, 4.0, 5.0, 0.0, 6.0]), 0.0) + out = x * y + assert out.nnz == 2 + assert_array_equal(out, x.todense() * y.todense()) + # a negative value makes -0.0 where the other operand is unstored + x *= -1 + out = x * y + assert_array_equal(out, x.todense() * y.todense()) + assert out.nnz > 2 + z = make(np.array([0.0, np.inf, 1.0, 0.0, 0.0]), 0.0) + assert_array_equal(z * y, np_call(np.multiply, z.todense(), y.todense())) + + def test_float_multiply_keeps_union(): x = make(np.array([-1.0, 0.0, 2.0, np.inf]), 0.0) y = make(np.array([0.0, 4.0, 5.0, 0.0]), 0.0) From 23bd854c9cecea4a42707d1afcfd9ac5257c4762 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:42:32 -0600 Subject: [PATCH 07/77] array_api: maximum and minimum of floats in GraphBLAS when no NaN or -0.0 is present GraphBLAS's max and min of floats are C's fmax and fmin, which skip NaN and may return either of 0.0 and -0.0; numpy's maximum and minimum propagate NaN and order -0.0 below 0.0. Without NaN or -0.0 in either operand they agree, so an Array now caches a summary of its stored values (finite, nonnegative, NaN, -0.0), computed once, and maximum and minimum use GraphBLAS when the summaries and fills allow. min and max reductions and multiply read the same summary. --- graphblas/array_api/_elementwise.py | 49 +++++++++++++------ graphblas/array_api/_statistical.py | 8 +-- graphblas/array_api/_utils.py | 32 ++++++++++++ graphblas/tests/test_array_api_elementwise.py | 25 ++++++++++ 4 files changed, 93 insertions(+), 21 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 2b6ee5df7..14fe5106b 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -381,17 +381,19 @@ class _Binary: ``kind`` is "arith" (the op works in the result dtype), "cmp" (in the promoted dtype of the operands, with a bool result) or "logical" (in bool). ``gb(dtype)`` gives the GraphBLAS operator that equals ``ufunc`` on operands of the dtype it works in, or None. + ``gb_values(a, b, dtype)``, in place of ``gb``, may look at the operands' values too. ``annihilator(a, b, fill_a, fill_b, dtype)`` says whether the union equals the intersection (``ewise_mult``): whether each operand's fill makes the result the result's fill whatever the other operand holds. """ - def __init__(self, name, ufunc, kind="arith", gb=None, annihilator=None): + def __init__(self, name, ufunc, kind="arith", gb=None, annihilator=None, gb_values=None): self.name = name self.ufunc = ufunc self.kind = kind self.gb = gb self.annihilator = annihilator + self.gb_values = gb_values def _binary(spec, x1, x2): @@ -448,6 +450,10 @@ def _gb_op(spec, a, b, ct): """The typed GraphBLAS op for ``spec`` working in ``ct``, when it gives numpy's answer; else None. """ + if spec.gb_values is not None: + if any(not isinstance(x, Array) and not _python_scalar_fits(x, ct) for x in (a, b)): + return None + return spec.gb_values(a, b, ct) if spec.gb is None: return None for x in (a, b): @@ -526,18 +532,34 @@ def _multiply_intersects(a, b, fa, fb, ct): def _finite_nonnegative(x): - """Whether every stored value is finite with its sign bit clear (cached per Array).""" + facts = _utils.summary(x) + return facts["finite"] and facts["nonnegative"] + - def compute(x): - T = x._T() - if x.dtype.kind == "b" or T.nvals == 0: - return True - values = _utils.tensor_values(T) - if x.dtype.kind in _INTEGRAL: - return bool((values >= 0).all()) - return bool(np.isfinite(values).all() and not np.signbit(values).any()) +def _max_min_gb(op): + """GraphBLAS's max and min of floats are C's fmax and fmin, which skip NaN and may take + either zero of 0.0 and -0.0; numpy's propagate NaN and order -0.0 below 0.0. They agree + when neither operand holds NaN or -0.0 (cached per Array), fills included. + """ - return x._fact("finite_nonnegative", compute) + def get(a, b, ct): + if ct.kind in "biu": + return op[gb_dtype(ct)] + if ct.kind != "f": + return None + for x in (a, b): + if isinstance(x, Array): + fill = np.asarray(x._fill) + if np.isnan(fill) or (fill == 0 and np.signbit(fill)): + return None + facts = _utils.summary(x) + if facts["nan"] or facts["negzero"]: + return None + elif isinstance(x, float) and (np.isnan(x) or (x == 0 and np.signbit(x))): + return None + return op[gb_dtype(ct)] + + return get _SPECS = { @@ -600,9 +622,8 @@ def compute(x): np.bitwise_xor, gb=_typed_for({"b": binary.lxor, "i": binary.bxor, "u": binary.bxor}), ), - # GraphBLAS's max and min of floats drop NaN (C's fmax); numpy's propagate it - "maximum": _Binary("maximum", np.maximum, gb=_typed_for(dict.fromkeys("biu", binary.max))), - "minimum": _Binary("minimum", np.minimum, gb=_typed_for(dict.fromkeys("biu", binary.min))), + "maximum": _Binary("maximum", np.maximum, gb_values=_max_min_gb(binary.max)), + "minimum": _Binary("minimum", np.minimum, gb_values=_max_min_gb(binary.min)), "copysign": _Binary("copysign", np.copysign, gb=_typed_for({"f": binary.copysign})), # numpy's own definitions: floored division and remainder (with its results for division # by zero), integer powers, shifts past the width, and libm-free special values diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index bb897021b..8c9734e77 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -244,13 +244,7 @@ def _reduce_no_axes(name, x, ct): def _has_nan(x): - def compute(x): - T = x._T() - if T.nvals == 0: - return False - return bool(T.apply(unary.isnan).new().reduce_scalar(monoid.lor).new().value) - - return x._fact("has_nan", compute) + return _utils.summary(x)["nan"] def sum(x, /, *, axis=None, dtype=None, keepdims=False): # noqa: A001 (the Array API's name) diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index 6736def17..726a43453 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -105,6 +105,38 @@ def tensor_values(T): return T._M.to_coo(rows=False, columns=False, sort=False)[2] +def summary(x): + """Facts about the stored values of the Array ``x``, computed once (cached on ``x``): + ``finite`` (all finite), ``nonnegative`` (no sign bit set: no negative value, no -0.0), + ``nan`` (some NaN) and ``negzero`` (some -0.0). + """ + + def compute(x): + T = x._T() + kind = x.dtype.kind + if kind == "b" or T.nvals == 0: + return {"finite": True, "nonnegative": True, "nan": False, "negzero": False} + values = tensor_values(T) + if kind in "iu": + nonneg = kind == "u" or bool((values >= 0).all()) + return {"finite": True, "nonnegative": nonneg, "nan": False, "negzero": False} + if kind == "c": + parts = [values.real, values.imag] + else: + parts = [values] + signbit = [np.signbit(p) for p in parts] + return { + "finite": all(bool(np.isfinite(p).all()) for p in parts), + "nonnegative": not any(bool(sb.any()) for sb in signbit), + "nan": any(bool(np.isnan(p).any()) for p in parts), + "negzero": any( + bool((sb & (p == 0)).any()) for p, sb in zip(parts, signbit, strict=True) + ), + } + + return x._fact("summary", compute) + + def _x_array(cont, dtype, stack): """The value array of an unloaded Container as a numpy view (None when it is empty).""" return stack.enter_context(tl._VectorArray(cont.x, np.dtype(dtype))).array diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 2b3cacbf1..24887ea93 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -483,6 +483,31 @@ def test_float_multiply_of_nonnegative_values_intersects(): assert_array_equal(z * y, np_call(np.multiply, z.todense(), y.todense())) +@pytest.mark.parametrize("name", ["maximum", "minimum"]) +def test_float_maximum_minimum(name): + # Without NaN or -0.0 GraphBLAS's max and min are numpy's; with them, numpy computes + rng = np.random.default_rng(8) + func = getattr(np, name) + for fill1, fill2 in [(0.0, 0.0), (2.5, -1.0), (np.inf, 0.0), (0.0, -0.0), (np.nan, 1.0)]: + a = np.where(rng.random((4, 5)) < 0.5, fill1, np.round(rng.normal(size=(4, 5)), 1)) + b = np.where(rng.random(5) < 0.5, fill2, np.round(rng.normal(size=5), 1)) + x, y = make(a, fill1), make(b, fill2) + assert_array_equal(getattr(xp, name)(x, y), func(a, b)) + for s in [0.0, -0.0, 1.5, np.nan, -np.inf, 3]: + assert_array_equal(getattr(xp, name)(x, s), func(a, s)) + assert_array_equal(getattr(xp, name)(s, y), func(s, b)) + a[0, 0] = -0.0 + a[1, 1] = np.nan + x = make(a, fill1) + assert_array_equal(getattr(xp, name)(x, y), func(a, b)) + assert_array_equal(getattr(xp, name)(y, x), func(b, a)) + x = make(np.array([1.0, 0.0, 2.0]), 0.0) + y = make(np.array([0.5, 0.5, 0.5]), 0.0) + assert_array_equal(getattr(xp, name)(x, y), func(x.todense(), y.todense())) + x *= -0.0 # (the cached facts about x's values must not survive) + assert_array_equal(getattr(xp, name)(x, y), func(x.todense(), y.todense())) + + def test_float_multiply_keeps_union(): x = make(np.array([-1.0, 0.0, 2.0, np.inf]), 0.0) y = make(np.array([0.0, 4.0, 5.0, 0.0]), 0.0) From de48d91b7fabf863f220354b2a29dbf7f6de73b3 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:43:52 -0600 Subject: [PATCH 08/77] array_api: less overhead per call Reuse the GraphBLAS Scalars that carry fills and Python scalars (a new one costs about as much as a small union), skip the broadcast checks for operands of one shape, and reduce over every axis to a Scalar, whose count of unstored positions is known without a count reduction. A 10x10 x + y now costs about 24 us (pydata/sparse: about 110 us; a bare Tensor union: about 20 us). --- graphblas/array_api/_elementwise.py | 2 +- graphblas/array_api/_statistical.py | 15 ++++++++++++++- graphblas/array_api/_utils.py | 19 +++++++++++++++++-- 3 files changed, 32 insertions(+), 4 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 14fe5106b..0baf858cd 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -468,7 +468,7 @@ def _gb_op(spec, a, b, ct): def _binary_arrays(spec, a, b, op, ct, rdt, fill): Ta, Tb = a._T(), b._T() try: - shape = np.broadcast_shapes(Ta.shape, Tb.shape) + shape = Ta.shape if Ta.shape == Tb.shape else np.broadcast_shapes(Ta.shape, Tb.shape) except ValueError: raise ValueError( f"{spec.name}: shapes {Ta.shape} and {Tb.shape} do not broadcast" diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index 8c9734e77..73f47f83a 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -209,9 +209,22 @@ def _reduce(name, x, axis, keepdims, dtype=None): fill_out = spec.repeat(fill, L, ct) if T.nvals == 0 or math.prod(shape) == 0: return _keepdims(Array._new(Tensor(gb_dtype(ct), shape), fill_out), x, axes, keepdims) - R = _reduce_tensor(T, spec.gb[gb_dtype(ct)], axes) nan_fix = name in {"min", "max"} and ct.kind == "f" and _has_nan(x) full = T.nvals == x.size + if len(axes) == x.ndim: + # Over every axis: one Scalar, and the count is the number of entries + r = T.reduce_scalar(spec.gb[gb_dtype(ct)]).new().value + if nan_fix: + value = np.nan + elif full or _utils.same_value(fill, spec.identity(ct)): + value = r + else: + with np.errstate(all="ignore"): + count = np.array([T.nvals], dtype=np.int64) + value = spec.combine(L, fill, ct)(np.array([r], dtype=ct), count)[0] + out = Array._new(Tensor(gb_dtype(ct), ()), value) + return _keepdims(out, x, axes, keepdims) + R = _reduce_tensor(T, spec.gb[gb_dtype(ct)], axes) if not nan_fix and (full or _utils.same_value(fill, spec.identity(ct))): out = Array._new(R, fill if full else fill_out) else: diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index 726a43453..849727626 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -45,9 +45,22 @@ def scalar(value, dtype): return np.asarray(value).astype(dtype, copy=False)[()] +_SCALARS = {} + + def gb_scalar(value, dtype): - """A GraphBLAS Scalar of the numpy ``dtype`` holding ``value`` cast as numpy casts it.""" - return Scalar.from_value(scalar(value, dtype), gb_dtype(dtype)) + """A GraphBLAS Scalar of the numpy ``dtype`` holding ``value`` cast as numpy casts it. + + The Scalars are only read (as operands), so they are cached by their bytes: a new one + costs about as much as a small element-wise operation. + """ + value = scalar(value, dtype) + key = (value.dtype.str, value.tobytes()) + if (s := _SCALARS.get(key)) is None: + if len(_SCALARS) > 1024: + _SCALARS.clear() + s = _SCALARS[key] = Scalar.from_value(value, gb_dtype(dtype)) + return s def same_value(a, b): @@ -80,6 +93,8 @@ def expansion_check(Ta, Tb, shape): The union replicates each broadcast operand's entries over the axes it broadcasts along. """ + if Ta.shape == Tb.shape: + return replicated = 0 for T in (Ta, Tb): factor = _broadcast_factor(T.shape, shape) From 9026d241ef1aeeea56d08c967ec82d31217f5027 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:49:09 -0600 Subject: [PATCH 09/77] array_api: declare the graphblas.array_api package in pyproject.toml --- graphblas/array_api/_indexing.py | 2 ++ pyproject.toml | 1 + 2 files changed, 3 insertions(+) diff --git a/graphblas/array_api/_indexing.py b/graphblas/array_api/_indexing.py index 49bfef70b..a3c155b9e 100644 --- a/graphblas/array_api/_indexing.py +++ b/graphblas/array_api/_indexing.py @@ -96,4 +96,6 @@ def _advanced_getitem(x, key): def setitem(x, key, value): + # (an implementation that changes x's Tensor in place must also clear x._facts, the cached + # summary of its stored values; Array._replace does both) raise NotImplementedError("assignment by index is not supported yet") diff --git a/pyproject.toml b/pyproject.toml index c46d012eb..23e0cfda7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -115,6 +115,7 @@ all = [ packages = [ "graphblas", "graphblas.agg", + "graphblas.array_api", "graphblas.binary", "graphblas.core", "graphblas.core.operator", From f4f84b0fc2e1b071d7bd60f3ab9069992717c4a6 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:50:09 -0600 Subject: [PATCH 10/77] array_api: a fill that is never seen does not raise An Array that stores every position never shows its fill. Where the fills of an operation raise (an integer to a negative power) and an operand is such an Array, the result's fill is never seen either: 0 stands in for it instead of raising an error numpy would not raise. --- graphblas/array_api/_elementwise.py | 19 +++++++++++++++++-- graphblas/tests/test_array_api_elementwise.py | 9 +++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 0baf858cd..e95d179e3 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -129,6 +129,21 @@ def _fill_operand(x): return x +def _fill_of(ufunc, rdt, operands): + """``ufunc`` of the operands' fills: the result's fill. + + An Array that stores every position never shows its fill, so where the fills raise + (integers to a negative power) and one operand is such an Array, the result's fill is + never seen either, and 0 stands in for it. + """ + try: + return ufunc(*[_fill_operand(x) for x in operands]) + except (ValueError, ArithmeticError): + if any(isinstance(x, Array) and x.ndim and x._T().nvals == x.size for x in operands): + return _utils.scalar(0, rdt) + raise + + def _result(T, fill): return Array._new(T, fill) @@ -152,7 +167,7 @@ def _unary(name, ufunc, x, gb=None, *, const=None): x = asarray(x) with np.errstate(all="ignore"): rdt = ufunc(np.empty(0, x.dtype)).dtype - fill = ufunc(np.asarray(x._fill)) + fill = _fill_of(ufunc, rdt, (x,)) check_supported(rdt, f"{name} of {x.dtype}") T = x._T() if T.ndim == 0: @@ -406,7 +421,7 @@ def _binary(spec, x1, x2): a = asarray(a) with np.errstate(all="ignore"): rdt = ufunc(_probe(a), _probe(b)).dtype - fill = ufunc(_fill_operand(a), _fill_operand(b)) + fill = _fill_of(ufunc, rdt, (a, b)) check_supported(rdt, f"{spec.name} of {_dtype_name(a)} and {_dtype_name(b)}") a_arr = isinstance(a, Array) and a.ndim > 0 b_arr = isinstance(b, Array) and b.ndim > 0 diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 24887ea93..128c2c708 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -553,6 +553,15 @@ def test_values_mapped_on_every_storage_format(): del gb +def test_unseen_fill_does_not_raise(): + # x stores every position, so its fill (-1) never meets y's fill as an exponent + x = xp.from_coo(([0, 1, 2],), [2, 3, 4], (3,), fill_value=-1) + y = xp.full((3,), -2) + assert_array_equal(xp.pow(y, x), np.power(np.full(3, -2), [2, 3, 4])) + with pytest.raises(ValueError, match="negative"): + xp.pow(make(np.array([2, 0, 3]), 0), -1) + + def test_zero_d_results_store_nothing(): x = xp.asarray(2.5) y = xp.add(x, 1) From 5ff28226b4277cf3cac3ceea8115eb9a2d0e2dda Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:54:04 -0600 Subject: [PATCH 11/77] array_api: test argument checking and rare storage cases; unload a shared Matrix once map_values unloaded each operand's Matrix through a GxB_Container. Two Tensors that share one Matrix (copy-on-write) made it unload the same Container twice, and the second read found no values. It now unloads each Matrix once. Also drop branches no input reaches: numpy's type resolution already refuses what they checked. --- graphblas/array_api/_elementwise.py | 19 +- graphblas/array_api/_statistical.py | 9 +- graphblas/array_api/_utils.py | 17 +- graphblas/tests/test_array_api_errors.py | 216 +++++++++++++++++++++++ 4 files changed, 231 insertions(+), 30 deletions(-) create mode 100644 graphblas/tests/test_array_api_errors.py diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index e95d179e3..c31d97b66 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -283,13 +283,7 @@ def logical_not(x, /): def bitwise_invert(x, /): - def gb(dtype): - if dtype.kind == "b": - return unary.lnot[gb_dtype(dtype)] - if dtype.kind in _INTEGRAL: - return unary.bnot[gb_dtype(dtype)] - return None - + gb = _first_of(_typed(unary.lnot, "b"), _typed(unary.bnot, _INTEGRAL)) return _unary("bitwise_invert", np.invert, x, gb) @@ -430,13 +424,12 @@ def _binary(spec, x1, x2): if spec.kind == "arith": ct = rdt elif spec.kind == "cmp": - try: - ct = np.result_type(_probe(a), _probe(b)) - except (OverflowError, TypeError): - ct = None # (a Python int out of range: numpy compares it exactly) + # (a Python int out of range of an integer dtype keeps that dtype here; _gb_op then + # leaves the comparison to numpy, which compares such an int exactly) + ct = np.result_type(_probe(a), _probe(b)) else: ct = np.dtype(bool) - op = None if ct is None else _gb_op(spec, a, b, ct) + op = _gb_op(spec, a, b, ct) if a_arr and b_arr: return _binary_arrays(spec, a, b, op, ct, rdt, fill) if a_arr: @@ -466,8 +459,6 @@ def _gb_op(spec, a, b, ct): else None. """ if spec.gb_values is not None: - if any(not isinstance(x, Array) and not _python_scalar_fits(x, ct) for x in (a, b)): - return None return spec.gb_values(a, b, ct) if spec.gb is None: return None diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index 73f47f83a..0e2ff9669 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -79,15 +79,8 @@ def _reduce_tensor(T, op, axes): def _combine(func, tensors, dtype, fill): """An Array whose entries are ``func`` of the aligned values of ``tensors`` (one - pattern), and ``fill`` elsewhere; 0-d tensors give a 0-d Array. + pattern), and ``fill`` elsewhere. """ - T0 = tensors[0] - if T0.ndim == 0: - if T0.nvals == 0: - return Array._new(Tensor(gb_dtype(dtype), ()), fill) - values = [T.to_coo()[1] for T in tensors] - with np.errstate(all="ignore"): - return Array._new(Tensor(gb_dtype(dtype), ()), np.asarray(func(*values))[0]) return Array._new(_utils.map_values(func, tensors, dtype), fill) diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index 849727626..0c33bf227 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -179,16 +179,17 @@ def map_values(func, tensors, dtype): W = tl.new_matrix(gbdt, template._nrows, template._ncols) W << template with contextlib.ExitStack() as stack: - unloaded = {} - for M in Ms: - if id(M) not in unloaded: - unloaded[id(M)] = stack.enter_context(tl._Unloaded(M)) - conts = [unloaded[id(M)] for M in Ms] + # (tensors sharing one Matrix, copy-on-write, unload it once) + unique = {id(M): M for M in Ms} + conts = {key: stack.enter_context(tl._Unloaded(M)) for key, M in unique.items()} wcont = stack.enter_context(tl._Unloaded(W)) - formats = {c.format for c in conts} | {wcont.format} - nvals = {c.nvals for c in conts} | {wcont.nvals} + formats = {c.format for c in conts.values()} | {wcont.format} + nvals = {c.nvals for c in conts.values()} | {wcont.nvals} if len(formats) == 1 and len(nvals) == 1: - arrays = [_x_array(c, np_dtype(M.dtype), stack) for c, M in zip(conts, Ms, strict=True)] + values = { + key: _x_array(conts[key], np_dtype(M.dtype), stack) for key, M in unique.items() + } + arrays = [values[id(M)] for M in Ms] out = _x_array(wcont, dtype, stack) with np.errstate(all="ignore"): result = func(*arrays) diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py new file mode 100644 index 000000000..9ca0aa975 --- /dev/null +++ b/graphblas/tests/test_array_api_errors.py @@ -0,0 +1,216 @@ +"""graphblas.array_api: argument checking, rare storage cases and extension paths.""" + +import numpy as np +import pytest + +import graphblas.array_api as xp +from graphblas.array_api import _utils +from graphblas.core.tensor import Tensor + +from .test_array_api_elementwise import assert_array_equal, make + + +def test_array_arguments_are_checked(): + a = np.array([1.0, 2.0]) + for func in [xp.reshape, xp.squeeze, xp.permute_dims, xp.expand_dims]: + with pytest.raises(TypeError, match="needs an Array"): + func(a, 0) + for func in [xp.sum, xp.diff, xp.cumulative_sum, xp.tril, xp.zeros_like, xp.clip]: + with pytest.raises(TypeError, match="needs an Array"): + func(a) + with pytest.raises(TypeError, match="needs an Array"): + xp.astype(a, xp.float32) + with pytest.raises(TypeError, match="arrays or Python scalars"): + xp.add(make(a, 0.0), [1, 2]) + + +def test_scalars_alone(): + assert_array_equal(xp.sin(1.0), np.sin(np.float64(1.0))) + assert_array_equal(xp.add(1, 2), np.array(3)) + assert_array_equal(xp.multiply(np.float32(2), 3), np.array(6, np.float32)) + + +def test_shape_errors(): + x = make(np.arange(6.0).reshape(2, 3), 0.0) + with pytest.raises(ValueError, match="negative"): + xp.zeros((2, -1)) + with pytest.raises(ValueError, match="invalid shape"): + xp.reshape(x, (-2, 3)) + with pytest.raises(ValueError, match="cannot reshape"): + xp.reshape(x, (4, -1)) + with pytest.raises(ValueError, match="repeated"): + xp.expand_dims(x, axis=(0, 0)) + with pytest.raises(ValueError, match="repeated"): + xp.squeeze(xp.zeros((1, 1)), axis=(0, -2)) + assert_array_equal(xp.squeeze(x, axis=()), np.arange(6.0).reshape(2, 3)) + with pytest.raises(ValueError, match="do not broadcast"): + x + make(np.ones(4), 0.0) + with pytest.raises(ValueError, match="1-d"): + xp.meshgrid(x) + with pytest.raises(ValueError, match="one data type"): + xp.meshgrid(xp.zeros(2), xp.zeros(2, dtype=xp.int8)) + with pytest.raises(ValueError, match="indexing"): + xp.meshgrid(xp.zeros(2), indexing="yx") + + +def test_creation_errors(): + with pytest.raises(TypeError, match="Python scalar"): + xp.full(3, "a") + with pytest.raises(ValueError, match="copy"): + xp.from_dlpack(np.ones(2), copy=False) + with pytest.raises(ValueError, match="keeps its fill"): + xp.asarray(xp.zeros(2), fill_value=1.0) + assert_array_equal(xp.asarray(np.array([1, 2], np.float16)), np.array([1, 2], np.float32)) + assert_array_equal(xp.full(2, np.float32(1.5)), np.full(2, 1.5, np.float32)) + assert_array_equal(xp.full(2, xp.asarray(3, dtype=xp.int8)), np.full(2, 3, np.int8)) + assert_array_equal(xp.full_like(xp.zeros(2), xp.asarray(4.0)), np.full(2, 4.0)) + assert_array_equal(xp.from_coo(([1, 3],), [1, 2]), np.array([0, 1, 0, 2])) + assert_array_equal(xp.from_coo((), 5, ()), np.array(5)) + assert_array_equal(xp.from_coo((), [], (), dtype=xp.int8), np.array(0, np.int8)) + + +def test_sparse_input_errors(): + sp = pytest.importorskip("scipy.sparse") + A = sp.coo_array((np.array([1.0, 2.0]), (np.array([0, 0]), np.array([1, 1]))), shape=(2, 2)) + assert_array_equal(xp.asarray(A), np.array([[0.0, 3.0], [0.0, 0.0]])) # duplicates summed + with pytest.raises(ValueError, match="fill value 0"): + xp.asarray(A, fill_value=1.0) + sparse = pytest.importorskip("sparse") + s = sparse.GCXS.from_numpy(np.array([[1.0, 0.0], [0.0, 3.0]])) + assert_array_equal(xp.asarray(s), np.array([[1.0, 0.0], [0.0, 3.0]])) + with pytest.raises(ValueError, match="differs"): + xp.asarray(s, fill_value=1.0) + + +def test_dtype_errors(): + with pytest.raises(TypeError, match="not a data type"): + xp.asarray([1], dtype="not a dtype") + assert xp.isdtype(np.float32, "real floating") + with pytest.raises(ValueError, match="at least one"): + xp.result_type() + info = xp.__array_namespace_info__() + with pytest.raises(ValueError, match="kind"): + info.dtypes(kind="integer") + with pytest.raises(ValueError, match="device"): + info.default_dtypes(device="gpu") + + +def test_object_protocol_errors(): + x = make(np.array([1.0, 0.0]), 0.0) + with pytest.raises(BufferError, match="copies"): + x.__dlpack__(copy=False) + assert x.__dlpack_device__() == np.empty(0).__dlpack_device__() + np.testing.assert_array_equal(np.from_dlpack(x), [1.0, 0.0]) + with pytest.raises(TypeError, match="complex"): + int(xp.asarray(1j)) + with pytest.raises(TypeError): + x @ "a" + with pytest.raises(TypeError): + "a" @ x + with pytest.raises(TypeError): + x + "a" + with pytest.raises(TypeError): + "a" + x + np.testing.assert_array_equal(np.from_dlpack(xp.from_dlpack(x)), [1.0, 0.0]) + capsule = x.__dlpack__(dl_device=x.__dlpack_device__(), max_version=(1, 0)) + assert capsule is not None + with pytest.raises(TypeError, match="None"): + xp.astype(x, None) + + +def test_indexing_not_supported_yet(): + x = make(np.arange(4.0), 0.0) + with pytest.raises(NotImplementedError): + x[True] + with pytest.raises(NotImplementedError): + x[[0, 1]] + with pytest.raises(NotImplementedError): + x[0] = 1.0 + + +def test_clip_mixed_bounds(): + a = np.array([1, 5, 0, 9], np.int16) + x = make(a, 0) + lo = make(np.array([2, 2, 2, 2], np.int32), 2) + out = xp.clip(x, min=lo, max=np.int16(6)) + assert out.dtype == xp.int16 + assert_array_equal(out, np.minimum(np.maximum(a, 2), 6).astype(np.int16)) + assert_array_equal(xp.clip(x, max=4.5), np.clip(a, None, 4.5).astype(np.int16)) + assert_array_equal(xp.clip(x, min=xp.asarray(np.int16(3))), np.clip(a, 3, None)) + hi = make(np.array([4, 4, 4, 4], np.int16), 4) + assert_array_equal(xp.clip(x, max=hi), np.minimum(a, 4)) + + +def test_complex_comparisons_mixed_precision(): + a = np.array([1 + 1j, 0, 2], np.complex64) + b = np.array([1 + 1j, 3, 0], np.complex128) + assert_array_equal(xp.equal(make(a, 0), make(b, 0)), a == b) + assert_array_equal(xp.not_equal(make(a, 0), make(b, 0)), a != b) + + +def test_bool_products(): + a = np.array([[True, False, True], [True, True, True]]) + x = make(a, False) + for axis in [None, 0, 1]: + assert_array_equal(xp.prod(x, axis=axis, dtype=xp.bool), np.prod(a, axis=axis, dtype=bool)) + y = make(a, True) + assert_array_equal(xp.prod(y, axis=1, dtype=xp.bool), np.prod(a, axis=1, dtype=bool)) + + +def test_statistical_errors(): + with pytest.raises(ValueError, match="at least one dimension"): + xp.cumulative_sum(xp.asarray(1.0)) + with pytest.raises(ValueError, match="at least one dimension"): + xp.diff(xp.asarray(1.0)) + with pytest.raises(ValueError, match="nonnegative"): + xp.diff(xp.zeros(3), n=-1) + + +def test_concat_helper(): + from graphblas.array_api._manipulation import _broadcast_to, _concat + + a = make(np.array([[1.0, 0.0]]), 0.0) + b = make(np.array([[2.0, 2.0]]), 2.0) + assert_array_equal(_concat([a, b], 0), np.array([[1.0, 0.0], [2.0, 2.0]])) + assert_array_equal(_concat([xp.zeros((1, 2)), xp.zeros((2, 2))], 0), np.zeros((3, 2))) + with pytest.raises(ValueError, match="at least one"): + _concat([], 0) + with pytest.raises(ValueError, match="positive number"): + _concat([xp.asarray(1.0)], 0) + with pytest.raises(ValueError, match="do not concatenate"): + _concat([a, xp.zeros((1, 3))], 0) + with pytest.raises(ValueError, match="broadcast"): + _broadcast_to(a, (3, 3)) + with pytest.raises(ValueError, match="cannot broadcast"): + _broadcast_to(a, (2,)) + + +def test_map_values_storage_cases(): + T = Tensor.from_coo([[0, 1], [1, 2]], [1.0, 2.0], shape=(2, 3)) + U = T.relayout(((1,), (0,))) + out = _utils.map_values(np.add, [T, U], np.dtype(float)) + np.testing.assert_array_equal(out.to_dense(fill_value=0), [[0.0, 2.0, 0.0], [0.0, 0.0, 4.0]]) + out = _utils.map_values(np.add, [T, T.dup()], np.dtype(float)) # (one Matrix, shared) + np.testing.assert_array_equal(out.to_coo()[1], [2.0, 4.0]) + E = Tensor(T.dtype, (2, 3)) + assert _utils.map_values(np.sin, [E], np.dtype(float)).nvals == 0 + out = _utils._map_values_coo(np.add, [T, T], np.dtype(float)) + np.testing.assert_array_equal(out.to_coo()[1], [2.0, 4.0]) + out = _utils._map_values_coo(np.negative, [T], np.dtype(float)) + np.testing.assert_array_equal(out.to_coo()[1], [-1.0, -2.0]) + out = _utils._map_values_coo(lambda v: np.float64(7.0), [T], np.dtype(float)) + np.testing.assert_array_equal(out.to_coo()[1], [7.0, 7.0]) + assert _utils.full_tensor(np.dtype(float), (0, 3), 1.0).nvals == 0 + + +def test_scalar_cache_is_bounded(): + for i in range(1100): + _utils.gb_scalar(i, np.dtype(np.int64)) + assert len(_utils._SCALARS) <= 1025 + assert _utils.gb_scalar(-0.0, np.dtype(float)).value == 0.0 + + +def test_summary_of_complex_values(): + x = make(np.array([1 + 0j, complex(-0.0, 1), complex(np.nan, 0)]), 0j) + facts = _utils.summary(x) + assert facts == {"finite": False, "nonnegative": False, "nan": True, "negzero": True} From 7b2d3adbd2fca118f3dc896cdecf40cf4952237a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 00:58:32 -0600 Subject: [PATCH 12/77] array_api: var and std broadcast the stored means instead of storing every mean The means of a reduction store exactly the fibers that hold entries, and every entry lies in one of them, so the deviations are one broadcast ewise_mult of the entries with the means as they are. Before, the mean of every fiber was stored densely first (and refused past the densify limit): var(x, axis=0) of a 1e9 x 1e9 array with a few entries now takes about a millisecond. --- graphblas/array_api/_statistical.py | 27 +++++++++++++------ graphblas/tests/test_array_api_statistical.py | 15 +++++++++++ 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index 0e2ff9669..779f2f924 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -322,17 +322,11 @@ def var(x, /, *, axis=None, correction=0.0, keepdims=False): return out if out.dtype == real_dt else cast(out, real_dt) T = x._T() m = mean(x, axis=axes) - keep_shape = _out_shape(x.shape, axes, keepdims=True) if T.nvals: - _utils.check_dense(math.prod(keep_shape), "var or std (the means)") - M = _utils.full_tensor(rdt, keep_shape, m._fill) - Tm = m._T() - if Tm.nvals: - M(accum=binary.second) << Tm.expand_dims(axes) - D = T.ewise_mult(M, binary.minus[gb_dtype(rdt)]).new() + C = _reduce_tensor(T, agg.count, axes) + D = _deviations(T, m, C, axes, rdt) S2 = _utils.map_values(lambda d: np.real(d * np.conj(d)).astype(real_dt), [D], real_dt) S2 = _reduce_tensor(S2, monoid.plus[gb_dtype(real_dt)], axes) - C = _reduce_tensor(T, agg.count, axes) else: S2 = Tensor(gb_dtype(real_dt), m.shape) C = Tensor(gb_dtype(np.int64), m.shape) @@ -346,6 +340,23 @@ def var(x, /, *, axis=None, correction=0.0, keepdims=False): return _keepdims(out, x, axes, keepdims) +def _deviations(T, m, C, axes, rdt): + """Each entry of ``T`` minus its fiber's mean (``m``, ``C`` the entries per fiber).""" + minus = binary.minus[gb_dtype(rdt)] + if m.ndim == 0: + return T.apply(minus, right=_utils.gb_scalar(m._fill, rdt)).new() + Tm = m._T() + if Tm.nvals != C.nvals: # pragma: no cover (safety: store the mean of every fiber) + keep_shape = _out_shape(T.shape, axes, keepdims=True) + _utils.check_dense(math.prod(keep_shape), "var or std (the means)") + M = _utils.full_tensor(rdt, keep_shape, m._fill) + M(accum=binary.second) << Tm.expand_dims(axes) + return T.ewise_mult(M, minus).new() + # The means store exactly the fibers that hold entries (the pattern of the sums), and every + # entry lies in one of them: broadcast them as they are, with no dense copy + return T.ewise_mult(Tm.expand_dims(axes), minus).new() + + def _times_count(k, d): return np.where(k == 0, d.dtype.type(0), k * d) diff --git a/graphblas/tests/test_array_api_statistical.py b/graphblas/tests/test_array_api_statistical.py index f8c2ed90a..dbc634bf8 100644 --- a/graphblas/tests/test_array_api_statistical.py +++ b/graphblas/tests/test_array_api_statistical.py @@ -246,3 +246,18 @@ def test_axis_errors(): xp.sum(x, axis=2) with pytest.raises(ValueError, match="repeated"): xp.sum(x, axis=(0, -2)) + + +def test_var_of_a_huge_sparse_array(): + # The means are not stored densely: a 1e9 x 1e9 array with 4 entries + n = 10**9 + coords = (np.array([0, 5, 5, 10**8]), np.array([3, 3, 7, 3])) + x = xp.from_coo(coords, np.array([1.0, 2.0, 4.0, 8.0]), (n, n)) + v = xp.var(x, axis=0) + assert v.nnz == 2 + assert v.fill_value == 0.0 + for col, values in [(3, [1.0, 2.0, 8.0]), (7, [4.0])]: + values = np.array(values) + mean = values.sum() / n + expected = (((values - mean) ** 2).sum() + (n - values.size) * mean**2) / n + assert float(v[col]) == pytest.approx(expected, rel=1e-12) From ec6c5a3c054cd5e04eb0f802f88557f05dcb4b47 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:16:26 -0600 Subject: [PATCH 13/77] array_api: where, concat and stack where takes the union of the entries of its three operands (broadcast) and numpy's where over their aligned values. concat joins the stored entries in O(entries); the result's fill is the one held by the most positions, and a part with another fill stores its fill explicitly. stack is concat of expand_dims. --- graphblas/array_api/_manipulation.py | 58 ++++++++++++---- graphblas/array_api/_searching.py | 67 ++++++++++++++++++- graphblas/array_api/_statistical.py | 4 +- graphblas/array_api/_utils.py | 53 ++++++++++++++- graphblas/tests/test_array_api_errors.py | 16 ++--- .../tests/test_array_api_manipulation.py | 54 +++++++++++++++ graphblas/tests/test_array_api_searching.py | 56 ++++++++++++++++ 7 files changed, 282 insertions(+), 26 deletions(-) create mode 100644 graphblas/tests/test_array_api_manipulation.py create mode 100644 graphblas/tests/test_array_api_searching.py diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py index 4b6d8d8f2..8a33bdd2e 100644 --- a/graphblas/array_api/_manipulation.py +++ b/graphblas/array_api/_manipulation.py @@ -10,7 +10,7 @@ from ._array import Array from ._utils import gb_dtype -__all__ = ["expand_dims", "permute_dims", "reshape", "squeeze"] +__all__ = ["concat", "expand_dims", "permute_dims", "reshape", "squeeze", "stack"] def _check_array(x, name): @@ -116,46 +116,76 @@ def _broadcast_to(x, shape): return Array._new(T.ewise_add(E, first, allow_expansion=True).new(), x._fill) -def _concat(arrays, axis): - """The arrays joined along ``axis``. Parts whose fill differs from the first part's store - their fill explicitly (densely, under the limit). +def concat(arrays, /, *, axis=0): + """Join arrays along ``axis`` (None: flattened first), O(entries). + + The result's fill is the one held by the most positions; a part with another fill stores + its fill explicitly (O(that part's size), under the densify limit). """ - from ._creation import _from_numpy from ._dtypes import result_type from ._elementwise import cast arrays = list(arrays) if not arrays: - raise ValueError("need at least one array to concatenate") + raise ValueError("concat needs at least one array") + for a in arrays: + _check_array(a, "concat") + if axis is None: + arrays = [reshape(a, (-1,)) for a in arrays] + axis = 0 ndim = arrays[0].ndim if ndim == 0 or any(a.ndim != ndim for a in arrays): - raise ValueError("concatenation needs arrays of one positive number of dimensions") + raise ValueError("concat needs arrays of one number of dimensions, at least 1") axis = _normalize_axis(axis, ndim) shape = list(arrays[0].shape) for a in arrays[1:]: if any(a.shape[i] != n for i, n in enumerate(shape) if i != axis): raise ValueError(f"shapes {arrays[0].shape} and {a.shape} do not concatenate") shape[axis] = sum(a.shape[axis] for a in arrays) + shape = tuple(shape) dtype = result_type(*arrays) arrays = [a if a.dtype == dtype else cast(a, dtype) for a in arrays] - fill = arrays[0]._fill + fill = _common_fill(arrays) coords, values = [], [] offset = 0 for a in arrays: - if not _utils.same_value(a._fill, fill): - _utils.check_dense(a.size, "concatenating arrays with different fill values") - a = _from_numpy(a.todense(), fill) - T = a._T() + T = _utils.with_fill(a, fill)._T() if T.nvals: - idx, vals = T.to_coo() + # (stored order: no relayout to C order; the build sorts) + idx, vals = T.to_coo(sort=False) idx = list(idx) idx[axis] = idx[axis] + offset coords.append(idx) values.append(vals) offset += a.shape[axis] - shape = tuple(shape) if not values: return Array._new(Tensor(gb_dtype(dtype), shape), fill) coords = tuple(np.concatenate([c[i] for c in coords]) for i in range(ndim)) T = Tensor.from_coo(coords, np.concatenate(values), gb_dtype(dtype), shape=shape) return Array._new(T, fill) + + +def _common_fill(arrays): + """The fill held by the most positions of ``arrays`` (the first such, on ties).""" + totals = [] + for a in arrays: + for i, (fill, _) in enumerate(totals): + if _utils.same_value(fill, a._fill): + totals[i] = (fill, totals[i][1] + a.size) + break + else: + totals.append((a._fill, a.size)) + best = max(n for _, n in totals) + return next(fill for fill, n in totals if n == best) + + +def stack(arrays, /, *, axis=0): + arrays = list(arrays) + if not arrays: + raise ValueError("stack needs at least one array") + for a in arrays: + _check_array(a, "stack") + if len({a.shape for a in arrays}) > 1: + raise ValueError(f"stack needs arrays of one shape, not {[a.shape for a in arrays]}") + axis = _normalize_axis(axis, arrays[0].ndim + 1) + return concat([expand_dims(a, axis=axis) for a in arrays], axis=axis) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index cc6f15e5b..f58a07f28 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -1,3 +1,68 @@ """Searching functions: argmax, argmin, count_nonzero, nonzero, searchsorted, where.""" -__all__ = [] +import builtins + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._dtypes import check_supported +from ._utils import gb_dtype + +__all__ = ["where"] + + +def _operand(x, name): + if isinstance(x, Array): + return x + if isinstance(x, (np.generic, np.ndarray)): + from ._creation import asarray + + return asarray(x) + if isinstance(x, (builtins.bool, int, float, complex)): + return x + raise TypeError(f"{name} needs arrays or Python scalars, not {type(x).__name__}") + + +def _probe(x): + return np.empty(0, x.dtype) if isinstance(x, Array) else x + + +def _value(x): + return np.asarray(x._fill) if isinstance(x, Array) else x + + +def where(condition, x1, x2, /): + """``x1`` where ``condition`` is true, else ``x2`` (broadcast); numpy's ``where``. + + The entries are the union of the three arrays' entries, O(entries) (with the broadcast + operands' entries replicated, as for any union). + """ + cond = _operand(condition, "where") + if not isinstance(cond, Array): + from ._creation import asarray + + cond = asarray(cond) + a = _operand(x1, "where") + b = _operand(x2, "where") + rdt = check_supported(np.result_type(_probe(a), _probe(b)), "where") + fill = np.where(np.asarray(cond._fill, dtype=bool), _value(a), _value(b)).astype(rdt)[()] + shapes = [x.shape for x in (cond, a, b) if isinstance(x, Array)] + shape = np.broadcast_shapes(*shapes) + arrays = [x for x in (cond, a, b) if isinstance(x, Array) and x.ndim] + if not shape or all(x._T().nvals == 0 for x in arrays): + return Array._new(Tensor(gb_dtype(rdt), shape), fill) + tensors = iter(_utils.aligned(arrays, shape)) + slots = [next(tensors) if isinstance(x, Array) and x.ndim else None for x in (cond, a, b)] + constants = [_value(x) for x in (cond, a, b)] + + def func(*values): + it = iter(values) + c, v1, v2 = ( + next(it) if t is not None else k for t, k in zip(slots, constants, strict=True) + ) + return np.where(np.asarray(c, dtype=bool), v1, v2).astype(rdt, copy=False) + + R = _utils.map_values(func, [t for t in slots if t is not None], rdt) + return Array._new(R, fill) diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index 779f2f924..85cd28d5e 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -404,7 +404,7 @@ def cumulative_prod(x, /, *, axis=None, dtype=None, include_initial=False): def diff(x, /, *, axis=-1, n=1, prepend=None, append=None): from ._elementwise import not_equal, subtract - from ._manipulation import _concat + from ._manipulation import concat _check_array(x, "diff") if x.ndim == 0: @@ -415,7 +415,7 @@ def diff(x, /, *, axis=-1, n=1, prepend=None, append=None): raise ValueError(f"diff's n must be nonnegative, not {n}") parts = [_diff_part(p, x, axis) for p in (prepend, x, append) if p is not None] if len(parts) > 1: - x = _concat(parts, axis) + x = concat(parts, axis=axis) if n == 0: return x._copy() # (numpy's diff of booleans is not_equal: subtracting booleans is an error) diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index 0c33bf227..e6ca9ab80 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -5,7 +5,7 @@ import numpy as np -from .. import config +from .. import binary, config from ..core import tensor_layout as tl from ..core.scalar import Scalar from ..core.tensor import Tensor @@ -115,6 +115,57 @@ def full_tensor(dtype, shape, value): return T +def union_pattern(tensors, shape): + """A BOOL Tensor of ``shape`` with an entry wherever one of ``tensors``, broadcast to + ``shape``, has one. + """ + pair = binary.pair[lookup_dtype(bool)] + P = None + for T in tensors: + if P is None: + if T.shape == shape: + P = T.apply(pair, right=True).new() + continue + P = Tensor(lookup_dtype(bool), shape) + expansion_check(P, T, shape) + P = P.ewise_add(T, pair, allow_expansion=True).new() + return P + + +def aligned(arrays, shape): + """Tensors of ``shape`` with one pattern, the union of the arrays' entries (broadcast), + holding each array's values there: its entry, or its fill where it has none. + """ + tensors = [x._T() for x in arrays] + if len(tensors) == 1 and tensors[0].shape == shape: + return tensors + P = union_pattern(tensors, shape) + out = [] + for x, T in zip(arrays, tensors, strict=True): + dt = x.dtype + second = binary.second[gb_dtype(dt)] + V = P.ewise_union(T, second, gb_scalar(0, dt), gb_scalar(x._fill, dt), allow_expansion=True) + out.append(V.new()) + return out + + +def with_fill(x, fill): + """``x`` with the fill ``fill``: where it stores nothing, its old fill is stored (every + position it does not store is written: O(size), under the densify limit). + """ + from ._array import Array + + fill = scalar(fill, x.dtype) + if same_value(fill, x._fill): + return x + check_dense(x.size, "giving an array another fill value") + T = full_tensor(x.dtype, x.shape, x._fill) + Tx = x._T() + if Tx.nvals: + T(accum=binary.second[gb_dtype(x.dtype)]) << Tx + return Array._new(T, fill) + + def tensor_values(T): """The stored values of ``T`` in storage order (a copy).""" return T._M.to_coo(rows=False, columns=False, sort=False)[2] diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 9ca0aa975..758772fde 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -166,19 +166,19 @@ def test_statistical_errors(): xp.diff(xp.zeros(3), n=-1) -def test_concat_helper(): - from graphblas.array_api._manipulation import _broadcast_to, _concat +def test_concat_errors(): + from graphblas.array_api._manipulation import _broadcast_to a = make(np.array([[1.0, 0.0]]), 0.0) b = make(np.array([[2.0, 2.0]]), 2.0) - assert_array_equal(_concat([a, b], 0), np.array([[1.0, 0.0], [2.0, 2.0]])) - assert_array_equal(_concat([xp.zeros((1, 2)), xp.zeros((2, 2))], 0), np.zeros((3, 2))) + assert_array_equal(xp.concat([a, b]), np.array([[1.0, 0.0], [2.0, 2.0]])) + assert_array_equal(xp.concat([xp.zeros((1, 2)), xp.zeros((2, 2))]), np.zeros((3, 2))) with pytest.raises(ValueError, match="at least one"): - _concat([], 0) - with pytest.raises(ValueError, match="positive number"): - _concat([xp.asarray(1.0)], 0) + xp.concat([]) + with pytest.raises(ValueError, match="at least 1"): + xp.concat([xp.asarray(1.0)]) with pytest.raises(ValueError, match="do not concatenate"): - _concat([a, xp.zeros((1, 3))], 0) + xp.concat([a, xp.zeros((1, 3))]) with pytest.raises(ValueError, match="broadcast"): _broadcast_to(a, (3, 3)) with pytest.raises(ValueError, match="cannot broadcast"): diff --git a/graphblas/tests/test_array_api_manipulation.py b/graphblas/tests/test_array_api_manipulation.py new file mode 100644 index 000000000..1c5ae4d93 --- /dev/null +++ b/graphblas/tests/test_array_api_manipulation.py @@ -0,0 +1,54 @@ +"""graphblas.array_api manipulation functions against numpy.""" + +import numpy as np +import pytest + +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, inputs, make + + +@pytest.mark.parametrize( + ("shapes", "axis"), + [ + ([(2, 3), (4, 3)], 0), + ([(2, 3), (2, 1), (2, 4)], 1), + ([(2, 3), (2, 3)], -1), + ([(2, 3, 2), (2, 1, 2)], 1), + ([(3,), (0,), (2,)], 0), + ([(2, 3), (3,)], None), + ], +) +def test_concat(shapes, axis): + for dtype in [np.dtype("float64"), np.dtype("int16"), np.dtype("bool")]: + parts = [list(inputs(dtype, shape, seed=i)) for i, shape in enumerate(shapes)] + for i in range(len(parts[0])): + dense = [p[i % len(p)][0] for p in parts] + arrays = [p[i % len(p)][1] for p in parts] + assert_array_equal(xp.concat(arrays, axis=axis), np.concatenate(dense, axis=axis)) + mixed = [make(np.array([1, 0, 2], np.int8), 0), make(np.array([1.5, 2.5]), 2.5)] + expected = np.concatenate([np.array([1, 0, 2], np.int8), np.array([1.5, 2.5])]) + assert_array_equal(xp.concat(mixed), expected) + + +def test_concat_keeps_the_most_common_fill(): + big = xp.full((10,), 7.0) + small = make(np.array([0.0, 1.0]), 0.0) + out = xp.concat([small, big]) + assert out.fill_value == 7.0 + assert out.nnz == 2 + assert_array_equal(out, np.concatenate([[0.0, 1.0], np.full(10, 7.0)])) + + +@pytest.mark.parametrize("axis", [0, 1, 2, -1, -3]) +def test_stack(axis): + for dtype in [np.dtype("float32"), np.dtype("uint8")]: + a = list(inputs(dtype, (2, 3), seed=1)) + b = list(inputs(dtype, (2, 3), seed=2)) + for (da, xa), (db, xb) in zip(a, b, strict=False): + assert_array_equal(xp.stack([xa, xb, xa], axis=axis), np.stack([da, db, da], axis=axis)) + assert_array_equal(xp.stack([xp.asarray(1), xp.asarray(2)]), np.array([1, 2])) + with pytest.raises(ValueError, match="one shape"): + xp.stack([xp.zeros(2), xp.zeros(3)]) + with pytest.raises(ValueError, match="at least one"): + xp.stack([]) diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py new file mode 100644 index 000000000..38a67c589 --- /dev/null +++ b/graphblas/tests/test_array_api_searching.py @@ -0,0 +1,56 @@ +"""graphblas.array_api searching functions against numpy.""" + +import numpy as np +import pytest + +import graphblas.array_api as xp + +from .test_array_api_elementwise import DTYPES, assert_array_equal, inputs, make, np_call + +WHERE_SHAPES = [ + ((3, 4), (3, 4), (3, 4)), + ((3, 4), (4,), (3, 1)), + ((3, 1), (1, 4), ()), + ((), (2, 3), (2, 3)), + ((2, 1, 4), (3, 1), (4,)), + ((), (), ()), +] + + +@pytest.mark.parametrize("shapes", WHERE_SHAPES, ids=str) +def test_where(shapes): + rng = np.random.default_rng(1) + for dt1, dt2 in [ + ("float64", "float64"), + ("int8", "uint8"), + ("float32", "int16"), + ("bool", "complex64"), + ]: + dt1, dt2 = np.dtype(dt1), np.dtype(dt2) + for (c, xc), (a, xa), (b, xb) in zip( + inputs(np.dtype(bool), shapes[0], seed=2), + inputs(dt1, shapes[1], seed=3), + inputs(dt2, shapes[2], seed=4), + strict=False, + ): + assert_array_equal(xp.where(xc, xa, xb), np.where(c, a, b)) + for s in [0, 1.5, True]: + assert_array_equal(xp.where(xc, xa, s), np.where(c, a, s)) + assert_array_equal(xp.where(xc, s, xb), np.where(c, s, b)) + # a condition that is not bool counts by truth value + a = np.where(rng.random((3, 4)) < 0.5, 0.0, rng.normal(size=(3, 4))) + x = make(a, 0.0) + assert_array_equal(xp.where(x, x, -1.0), np.where(a, a, -1.0)) + assert_array_equal(xp.where(x > 0, x, 0.0), np.where(a > 0, a, 0.0)) + with pytest.raises(ValueError, match="broadcast"): + xp.where(make(np.ones(3, bool), False), x, x) + + +def test_where_dtypes(): + for dt in DTYPES: + for dense, x in inputs(dt, (5,), seed=6): + c = dense.astype(bool) if dt.kind != "c" else dense.real.astype(bool) + assert_array_equal(xp.where(xp.asarray(c), x, x), np.where(c, dense, dense)) + assert_array_equal( + np_call(xp.where, xp.asarray(c), x, 2), np_call(np.where, c, dense, 2) + ) From 83cd44b123f1c34199ff11c350629cbc11aa025a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:17:38 -0600 Subject: [PATCH 14/77] array_api: broadcast_*, flip, moveaxis, repeat, roll, tile and unstack All cost O(entries of the result): flip slices, moveaxis permutes (O(1)), and repeat, roll and tile compute the new coordinates of the stored entries. broadcast_to replicates the entries over the broadcast axes, as a union does, under the densify limit. --- graphblas/array_api/_manipulation.py | 171 +++++++++++++++++- .../tests/test_array_api_manipulation.py | 96 ++++++++++ 2 files changed, 264 insertions(+), 3 deletions(-) diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py index 8a33bdd2e..86c5007dc 100644 --- a/graphblas/array_api/_manipulation.py +++ b/graphblas/array_api/_manipulation.py @@ -10,7 +10,23 @@ from ._array import Array from ._utils import gb_dtype -__all__ = ["concat", "expand_dims", "permute_dims", "reshape", "squeeze", "stack"] +__all__ = [ + "broadcast_arrays", + "broadcast_shapes", + "broadcast_to", + "concat", + "expand_dims", + "flip", + "moveaxis", + "permute_dims", + "repeat", + "reshape", + "roll", + "squeeze", + "stack", + "tile", + "unstack", +] def _check_array(x, name): @@ -100,9 +116,27 @@ def squeeze(x, /, axis): return Array._new(T.squeeze(tuple(norm)), x._fill) +def broadcast_to(x, /, shape): + """``x`` broadcast to ``shape``: its entries replicated over the broadcast axes (the + result's entries, under the densify limit for the replication). + """ + _check_array(x, "broadcast_to") + return _broadcast_to(x, shape) + + +def broadcast_shapes(*shapes): + return tuple(np.broadcast_shapes(*shapes)) + + +def broadcast_arrays(*arrays): + for x in arrays: + _check_array(x, "broadcast_arrays") + shape = np.broadcast_shapes(*[x.shape for x in arrays]) + return tuple(_broadcast_to(x, shape) for x in arrays) + + def _broadcast_to(x, shape): - """``x`` broadcast to ``shape``: its entries replicated over the broadcast axes.""" - shape = tuple(shape) + shape = tuple(int(n) for n in shape) if np.broadcast_shapes(x.shape, shape) != shape: raise ValueError(f"cannot broadcast shape {x.shape} to {shape}") if x.shape == shape: @@ -189,3 +223,134 @@ def stack(arrays, /, *, axis=0): raise ValueError(f"stack needs arrays of one shape, not {[a.shape for a in arrays]}") axis = _normalize_axis(axis, arrays[0].ndim + 1) return concat([expand_dims(a, axis=axis) for a in arrays], axis=axis) + + +def unstack(x, /, *, axis=0): + _check_array(x, "unstack") + if x.ndim == 0: + raise ValueError("unstack needs an array of at least one dimension") + axis = _normalize_axis(axis, x.ndim) + before = (slice(None),) * axis + return tuple(x[(*before, i)] for i in range(x.shape[axis])) + + +def _axes(axis, ndim): + if axis is None: + return tuple(range(ndim)) + axes = axis if isinstance(axis, tuple) else (axis,) + norm = tuple(_normalize_axis(a, ndim) for a in axes) + if len(set(norm)) != len(norm): + raise ValueError(f"repeated axis in {axis}") + return norm + + +def flip(x, /, *, axis=None): + _check_array(x, "flip") + axes = _axes(axis, x.ndim) + if x.ndim == 0: + return x._copy() + return x[tuple(slice(None, None, -1) if a in axes else slice(None) for a in range(x.ndim))] + + +def moveaxis(x, source, destination, /): + _check_array(x, "moveaxis") + src = source if isinstance(source, tuple) else (source,) + dst = destination if isinstance(destination, tuple) else (destination,) + if len(src) != len(dst): + raise ValueError("moveaxis needs as many destinations as sources") + src = _axes(src, x.ndim) + dst = _axes(dst, x.ndim) + order = [a for a in range(x.ndim) if a not in src] + for d, s in sorted(zip(dst, src, strict=True)): + order.insert(d, s) + return permute_dims(x, tuple(order)) + + +def _entries(x): + """The stored entries of ``x`` as (coordinates, values), in stored order.""" + T = x._T() + if T.nvals == 0: + return [np.empty(0, np.int64) for _ in range(x.ndim)], np.empty(0, x.dtype) + idx, values = T.to_coo(sort=False) + return list(idx), values + + +def _from_entries(coords, values, shape, x): + """An Array of ``x``'s dtype and fill storing ``values`` at ``coords`` in ``shape``.""" + shape = tuple(shape) + if len(values) == 0: + return Array._new(Tensor(x._T().dtype, shape), x._fill) + T = Tensor.from_coo(tuple(coords), values, x._T().dtype, shape=shape) + return Array._new(T, x._fill) + + +def repeat(x, repeats, /, *, axis=None): + """Each element repeated along ``axis``, O(entries of the result).""" + _check_array(x, "repeat") + if axis is None: + x = reshape(x, (-1,)) + axis = 0 + axis = _normalize_axis(axis, x.ndim) + n = x.shape[axis] + if isinstance(repeats, Array): + if repeats.dtype.kind not in "iu" or repeats.ndim != 1: + raise TypeError("repeats must be an int or a 1-d integer array") + counts = repeats.todense().astype(np.int64) + else: + counts = np.asarray(int(repeats), dtype=np.int64) + counts = np.broadcast_to(counts, (n,)) + if (counts < 0).any(): + raise ValueError("repeats must be nonnegative") + starts = np.zeros(n + 1, np.int64) + np.cumsum(counts, out=starts[1:]) + shape = list(x.shape) + shape[axis] = int(starts[-1]) + coords, values = _entries(x) + k = counts[coords[axis]] + which = np.repeat(np.arange(len(values)), k) + within = np.arange(which.size) - np.repeat(np.cumsum(k) - k, k) + coords = [c[which] for c in coords] + coords[axis] = starts[:-1][coords[axis]] + within + return _from_entries(coords, values[which], shape, x) + + +def roll(x, /, shift, *, axis=None): + _check_array(x, "roll") + if axis is None: + return reshape(roll(reshape(x, (-1,)), shift, axis=0), x.shape) + shifts = shift if isinstance(shift, tuple) else (shift,) + axes = axis if isinstance(axis, tuple) else (axis,) + if len(shifts) == 1 and len(axes) > 1: + shifts = shifts * len(axes) + if len(shifts) != len(axes): + raise ValueError("roll needs one shift per axis, or one shift") + total = [0] * x.ndim + for s, a in zip(shifts, axes, strict=True): + total[_normalize_axis(a, x.ndim)] += int(s) + coords, values = _entries(x) + for a, s in enumerate(total): + if s and x.shape[a]: + coords[a] = (coords[a] + s) % x.shape[a] + return _from_entries(coords, values, x.shape, x) + + +def tile(x, repetitions, /): + """``x`` repeated ``repetitions`` times along each axis, O(entries of the result).""" + _check_array(x, "tile") + reps = tuple(int(r) for r in repetitions) + if any(r < 0 for r in reps): + raise ValueError("repetitions must be nonnegative") + ndim = max(len(reps), x.ndim) + reps = (1,) * (ndim - len(reps)) + reps + if x.ndim < ndim: + x = reshape(x, (1,) * (ndim - x.ndim) + x.shape) + shape = tuple(n * r for n, r in zip(x.shape, reps, strict=True)) + coords, values = _entries(x) + ntiles = math.prod(reps) + if ntiles == 0 or len(values) == 0: + return Array._new(Tensor(x._T().dtype, shape), x._fill) + tiles = np.unravel_index(np.arange(ntiles), reps) + out = [] + for c, t, n in zip(coords, tiles, x.shape, strict=True): + out.append((c[None, :] + (t * n)[:, None]).reshape(-1)) + return _from_entries(out, np.tile(values, ntiles), shape, x) diff --git a/graphblas/tests/test_array_api_manipulation.py b/graphblas/tests/test_array_api_manipulation.py index 1c5ae4d93..8fc56c488 100644 --- a/graphblas/tests/test_array_api_manipulation.py +++ b/graphblas/tests/test_array_api_manipulation.py @@ -52,3 +52,99 @@ def test_stack(axis): xp.stack([xp.zeros(2), xp.zeros(3)]) with pytest.raises(ValueError, match="at least one"): xp.stack([]) + + +DTYPES3 = [np.dtype("float64"), np.dtype("int8"), np.dtype("bool")] + + +def _cases(shape, seed=0): + for dtype in DTYPES3: + yield from inputs(dtype, shape, seed=seed) + + +def test_broadcast(): + assert xp.broadcast_shapes((2, 1), (3,), ()) == (2, 3) + assert xp.broadcast_shapes() == () + with pytest.raises(ValueError, match="broadcast"): + xp.broadcast_shapes((2,), (3,)) + for dense, x in _cases((3, 1)): + for shape in [(3, 1), (3, 4), (2, 3, 4), (0, 3, 1)]: + assert_array_equal(xp.broadcast_to(x, shape), np.broadcast_to(dense, shape)) + with pytest.raises(ValueError, match="broadcast"): + xp.broadcast_to(x, (3,)) + y = make(np.array([1, 0, 2, 0], dtype=dense.dtype), np.zeros((), dense.dtype)[()]) + out = xp.broadcast_arrays(x, y, xp.asarray(1)) + expected = np.broadcast_arrays(dense, y.todense(), np.array(1)) + assert isinstance(out, tuple) + for o, e in zip(out, expected, strict=True): + assert_array_equal(o, e) + + +@pytest.mark.parametrize("axis", [0, 1, -1, 2]) +def test_unstack(axis): + for dense, x in _cases((2, 3, 4)): + out = xp.unstack(x, axis=axis) + assert isinstance(out, tuple) + expected = [np.take(dense, i, axis=axis) for i in range(dense.shape[axis])] + assert len(out) == len(expected) + for o, e in zip(out, expected, strict=True): + assert_array_equal(o, e) + with pytest.raises(ValueError, match="dimension"): + xp.unstack(xp.asarray(1.0)) + + +@pytest.mark.parametrize("axis", [None, 0, 1, -1, (0, 2), ()]) +def test_flip(axis): + for dense, x in _cases((2, 3, 4)): + assert_array_equal(xp.flip(x, axis=axis), np.flip(dense, axis=axis)) + assert_array_equal(xp.flip(xp.asarray(3)), np.array(3)) + + +@pytest.mark.parametrize( + ("source", "destination"), [(0, 2), (-1, 0), ((0, 1), (2, 0)), ((), ()), ((0, 1, 2), (2, 1, 0))] +) +def test_moveaxis(source, destination): + for dense, x in _cases((2, 3, 4)): + assert_array_equal( + xp.moveaxis(x, source, destination), np.moveaxis(dense, source, destination) + ) + with pytest.raises(ValueError, match="as many"): + xp.moveaxis(xp.zeros((2, 3)), (0, 1), 0) + + +@pytest.mark.parametrize("axis", [None, 0, 1, -1]) +def test_repeat(axis): + for dense, x in _cases((2, 3)): + n = dense.size if axis is None else dense.shape[axis] + for repeats in [0, 1, 3, np.arange(n) % 3, np.array([2])]: + if isinstance(repeats, np.ndarray): + r = xp.asarray(repeats) + expected = np.repeat(dense, np.broadcast_to(repeats, (n,)), axis=axis) + else: + r = repeats + expected = np.repeat(dense, repeats, axis=axis) + assert_array_equal(xp.repeat(x, r, axis=axis), expected) + with pytest.raises(ValueError, match="nonnegative"): + xp.repeat(xp.zeros(3), -1) + with pytest.raises(TypeError, match="integer"): + xp.repeat(xp.zeros(3), xp.asarray([1.0, 2.0, 3.0])) + + +@pytest.mark.parametrize( + ("shift", "axis"), + [(1, None), (-7, None), (2, 0), (-1, 1), ((1, 2), (0, 1)), (3, (0, 1)), (1, (1, 1))], +) +def test_roll(shift, axis): + for dense, x in _cases((3, 4)): + assert_array_equal(xp.roll(x, shift, axis=axis), np.roll(dense, shift, axis=axis)) + with pytest.raises(ValueError, match="one shift"): + xp.roll(xp.zeros((2, 2)), (1, 2, 3), axis=(0, 1)) + + +@pytest.mark.parametrize("reps", [(2,), (2, 3), (1, 1), (0, 2), (2, 1, 3), (3,)]) +def test_tile(reps): + for dense, x in _cases((2, 3)): + assert_array_equal(xp.tile(x, reps), np.tile(dense, reps)) + assert_array_equal(xp.tile(xp.asarray(2.0), (3,)), np.tile(np.array(2.0), 3)) + with pytest.raises(ValueError, match="nonnegative"): + xp.tile(xp.zeros(2), (-1,)) From 16d48782746a6fe8c97fe1c947c713dad2c0bd59 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:24:11 -0600 Subject: [PATCH 15/77] array_api: advanced indexing, assignment by index, take and take_along_axis x[key] takes numpy's keys: integer and boolean arrays (as Arrays, numpy arrays or lists) with ints, slices, ... and None. A boolean index whose fill is False costs O(its entries). The values are gathered from the stored entries (Tensor.take or vindex), so a position the array does not store holds the fill in the result too. numpy's axis order applies: advanced axes stay in place when adjacent, else go first. x[key] = value assigns a region by GraphBLAS assign (a basic key, or one 1-d index array among slices, where outer and vectorized indexing agree); a value whose fill is x's carries its structure, so its unstored positions delete. Several index arrays scatter: the selected positions are written, the last of repeated ones winning, as in numpy. Assignment clears the cached summary of the stored values. --- graphblas/array_api/_indexing.py | 480 ++++++++++++++++++--- graphblas/tests/test_array_api_errors.py | 10 - graphblas/tests/test_array_api_indexing.py | 222 ++++++++++ 3 files changed, 645 insertions(+), 67 deletions(-) create mode 100644 graphblas/tests/test_array_api_indexing.py diff --git a/graphblas/array_api/_indexing.py b/graphblas/array_api/_indexing.py index a3c155b9e..7876379bd 100644 --- a/graphblas/array_api/_indexing.py +++ b/graphblas/array_api/_indexing.py @@ -1,101 +1,467 @@ -"""Indexing: ``x[key]`` (and later ``x[key] = value``, ``take``, ``take_along_axis``). +"""Indexing: ``x[key]``, ``x[key] = value``, ``take`` and ``take_along_axis``. -Basic keys so far: integers, slices, ``...`` and ``None``, as numpy and the Array API index. +Keys follow numpy: integers, slices, ``...`` and ``None`` (basic indexing, extracted in +GraphBLAS), and integer and boolean arrays (advanced indexing, gathered from the stored +entries, so a position the array does not store is the fill in the result too). The Array API +standard specifies an index of integers and integer arrays only (vectorized indexing), or one +boolean array alone; other combinations follow numpy. + +Costs: basic indexing is O(entries in the region). An integer index array is read densely +(its size is the result's); a boolean index array costs O(its entries) when its fill is +False, O(its size) when its fill is True. Assigning a scalar other than the fill, or an array +whose fill is not ``x``'s, to a region stores every position of the region; assigning by +several index arrays writes every selected position (O(selection)). """ +import builtins +import math import operator import numpy as np from ..core.tensor import Tensor +from . import _utils from ._array import Array +from ._utils import gb_dtype + +__all__ = ["take", "take_along_axis"] + + +# Keys +class _Adv: + """An integer index array (advanced indexing) for one axis, nonnegative.""" + + __slots__ = ("index",) + + def __init__(self, index): + self.index = index + + +class _NewAxis: + """``None``, or a 0-d boolean index: a new axis of length 1 (0 for False).""" + + __slots__ = ("length",) + + def __init__(self, length=1): + self.length = length -__all__ = [] +def _is_int(k): + return isinstance(k, (int, np.integer)) and not isinstance(k, (bool, np.bool_)) -def _basic_key(x, key): - """``key`` as a tuple of ints and slices for every axis, and the positions of the new - axes (None) in the result; None when the key is not basic. + +def _index_array(k): + """An index of array type: a boolean Array as it is, else a numpy array; None when ``k`` + is not an array. + """ + if isinstance(k, Array): + if k.dtype.kind == "b": + return k + if k.dtype.kind not in "iu": + raise IndexError( + f"arrays used as indices must be of integer or boolean type: {k.dtype}" + ) + return k.todense() + if isinstance(k, (list, np.ndarray)): + arr = np.asarray(k) + if arr.dtype.kind not in "biu": + if arr.size == 0: + return arr.astype(np.int64) + raise IndexError( + f"arrays used as indices must be of integer or boolean type: {arr.dtype}" + ) + return arr + return None + + +def _mask_coords(mask, shape): + """The coordinates, in C order, where the boolean ``mask`` (an Array or ndarray) is True; + ``shape`` is that of the axes it covers. + """ + if mask.shape != tuple(shape): + # (the standard lets a dimension of the mask be 0 instead of the array's) + if any(m not in (n, 0) for m, n in zip(mask.shape, shape, strict=True)): + raise IndexError(f"boolean index of shape {mask.shape} does not match {shape}") + return tuple(np.empty(0, np.int64) for _ in mask.shape) + if isinstance(mask, np.ndarray): + return np.nonzero(mask) + if not mask._fill: + T = mask._T() + if T.nvals == 0: + return tuple(np.empty(0, np.int64) for _ in mask.shape) + # (stored entries may be False, equal to the fill) + return T.select("==", True).new().to_coo(values=False)[0] + _utils.check_dense(mask.size, "a boolean index whose fill is True") + return np.nonzero(mask.todense()) + + +def _normalize(x, key): + """``key`` as a list of items: an int, slice or _Adv for each axis of ``x``, and _NewAxis + items for new axes. A boolean array of M dimensions becomes M _Adv items. """ if not isinstance(key, tuple): key = (key,) items = [] for k in key: - if k is None or k is Ellipsis or isinstance(k, slice): + if k is None: + items.append(_NewAxis()) + elif k is Ellipsis or isinstance(k, slice): items.append(k) elif isinstance(k, (bool, np.bool_)): - return None # (a boolean index is a mask) - elif isinstance(k, (int, np.integer)): + items.append(_NewAxis(1 if k else 0)) + elif _is_int(k): items.append(int(k)) - elif isinstance(k, Array) and k.ndim == 0 and k.dtype.kind in "iu": - items.append(operator.index(k)) + elif isinstance(k, Array) and k.ndim == 0: + if k.dtype.kind == "b": + items.append(_NewAxis(1 if k._fill else 0)) + elif k.dtype.kind in "iu": + items.append(operator.index(k)) + else: + raise IndexError(f"a 0-d index must be an integer or boolean, not {k.dtype}") else: - return None - if sum(k is Ellipsis for k in items) > 1: + arr = _index_array(k) + if arr is None: + raise IndexError(f"invalid index {k!r}") + items.append(arr) + if sum(k is Ellipsis for k in items if not isinstance(k, (np.ndarray, Array))) > 1: raise IndexError("an index can only have a single ellipsis ('...')") - nindexed = sum(k is not None and k is not Ellipsis for k in items) - if nindexed > x.ndim: - raise IndexError(f"too many indices: {nindexed} for an array of {x.ndim} dimensions") - if Ellipsis in items: - i = items.index(Ellipsis) - items[i : i + 1] = [slice(None)] * (x.ndim - nindexed) + + def width(k): + if isinstance(k, (np.ndarray, Array)) and k.dtype.kind == "b": + return k.ndim + return 0 if k is Ellipsis or isinstance(k, _NewAxis) else 1 + + consumed = sum(width(k) for k in items) + if consumed > x.ndim: + raise IndexError(f"too many indices: {consumed} for an array of {x.ndim} dimensions") + i = next((i for i, k in enumerate(items) if k is Ellipsis), None) + if i is None: + items.extend([slice(None)] * (x.ndim - consumed)) else: - items.extend([slice(None)] * (x.ndim - nindexed)) - axes_key = [] - new_axes = [] - out_axis = 0 + items[i : i + 1] = [slice(None)] * (x.ndim - consumed) + out = [] axis = 0 for k in items: - if k is None: - new_axes.append(out_axis) - out_axis += 1 + if isinstance(k, _NewAxis): + out.append(k) + continue + if isinstance(k, (np.ndarray, Array)) and k.dtype.kind == "b": + if k.ndim == 0: # (a 0-d boolean numpy array) + out.append(_NewAxis(1 if k[()] else 0)) + continue + coords = _mask_coords(k, x.shape[axis : axis + k.ndim]) + out.extend(_Adv(np.asarray(c, dtype=np.int64)) for c in coords) + axis += k.ndim continue n = x.shape[axis] - if isinstance(k, int): + if isinstance(k, np.ndarray): + idx = k.astype(np.int64, copy=False) + if idx.size and (idx.min() < -n or idx.max() >= n): + raise IndexError(f"index out of bounds for axis {axis} with size {n}") + out.append(_Adv(np.where(idx < 0, idx + n, idx))) + elif isinstance(k, int): if not -n <= k < n: raise IndexError(f"index {k} is out of bounds for axis {axis} with size {n}") - axes_key.append(k % n) + out.append(k % n) else: - axes_key.append(k) - out_axis += 1 + out.append(k) axis += 1 - return tuple(axes_key), new_axes + return out + + +def _has_adv(items): + return any(isinstance(k, _Adv) for k in items) + + +def _adjacent(items): + """Whether the advanced items (index arrays, with the ints among them) are next to each + other; numpy then keeps their axes in place, else puts them first. + """ + positions = [i for i, k in enumerate(items) if isinstance(k, (_Adv, int))] + return positions == list(range(positions[0], positions[-1] + 1)) + + +def _finish(out, items, nadv): + """``out`` holds the advanced axes (``nadv`` of them) first, then the sliced axes in + order: move the advanced ones where numpy puts them and insert the new axes. + """ + from ._manipulation import expand_dims, permute_dims + + comps = [] + placed = True + if nadv or _has_adv(items): + placed = not _adjacent(items) + if placed: + comps.append("adv") + for k in items: + if isinstance(k, (_NewAxis, slice)): + comps.append(k) + elif not placed: + comps.append("adv") + placed = True + order = [] + new = [] + rest = nadv + for c in comps: + if isinstance(c, str): + order.extend(range(nadv)) + elif isinstance(c, slice): + order.append(rest) + rest += 1 + else: + new.append((len(order) + len(new), c.length)) + if order != list(range(len(order))): + out = permute_dims(out, tuple(order)) + if new: + out = expand_dims(out, tuple(p for p, _ in new)) + if any(length == 0 for _, length in new): + empty = {p for p, length in new if length == 0} + out = out[tuple(slice(0, 0) if a in empty else slice(None) for a in range(out.ndim))] + return out + + +def _basic_shape(shape, axes_key): + return tuple( + len(range(*k.indices(n))) + for k, n in zip(axes_key, shape, strict=True) + if isinstance(k, slice) + ) +# getitem def getitem(x, key): - basic = _basic_key(x, key) - if basic is None: - return _advanced_getitem(x, key) - axes_key, new_axes = basic + items = _normalize(x, key) + axes_key = [k for k in items if not isinstance(k, _NewAxis)] + if _has_adv(items): + out, nadv = _advanced_extract(x, axes_key) + else: + out, nadv = _basic_extract(x, tuple(axes_key)), 0 + return _finish(out, items, nadv) + + +def _basic_extract(x, axes_key): T = x._T() + if x.ndim == 0 or all(k == slice(None) for k in axes_key): + return x._copy() + shape = _basic_shape(T.shape, axes_key) + if T.nvals == 0 or 0 in shape: + return Array._new(Tensor(T.dtype, shape), x._fill) + return Array._new(T[axes_key].new(), x._fill) + + +def _advanced_extract(x, axes_key): + """``x`` indexed by ``axes_key`` (one item per axis) with the advanced axes first, then the + sliced ones in order, and the number of advanced axes. + """ + sliced = tuple(k if isinstance(k, slice) else slice(None) for k in axes_key) + y = _basic_extract(x, sliced) if any(k != slice(None) for k in sliced) else x + adv_axes = [a for a, k in enumerate(axes_key) if not isinstance(k, slice)] + rest = [a for a, k in enumerate(axes_key) if isinstance(k, slice)] + indices = np.broadcast_arrays( + *[ + k.index if isinstance(k, _Adv) else np.asarray(k, np.int64) + for k in axes_key + if not isinstance(k, slice) + ] + ) + S2 = indices[0].shape + adv_shape = tuple(y.shape[a] for a in adv_axes) + rest_shape = tuple(y.shape[a] for a in rest) + T = y._T() + if T.nvals == 0 or math.prod(S2) == 0 or math.prod(rest_shape) == 0: + return Array._new(Tensor(T.dtype, S2 + rest_shape), x._fill), len(S2) + if not rest: + # Vectorized (pointwise) indexing of every axis + return Array._new(T.vindex[tuple(indices)], x._fill), len(S2) + Tp = T.permute_dims(tuple(adv_axes) + tuple(rest)) + if len(adv_axes) > 1: + flat = np.ravel_multi_index(tuple(indices), adv_shape) + Tp = Tp.reshape((math.prod(adv_shape), *rest_shape)) + else: + flat = indices[0] + return Array._new(Tp.take(flat.astype(np.int64, copy=False), axis=0), x._fill), len(S2) + + +# setitem +def setitem(x, key, value): + if not isinstance(x, Array): # pragma: no cover (safety) + raise TypeError(type(x)) + value = _setitem_value(x, value) + items = _normalize(x, key) + axes_key = [k for k in items if not isinstance(k, _NewAxis)] + if any(isinstance(k, _NewAxis) and k.length == 0 for k in items): + return # (a False index selects nothing) if x.ndim == 0: - out = x._copy() - elif all(k == slice(None) for k in axes_key): - out = Array._new(T.dup(), x._fill) + if isinstance(value, Array): + raise ValueError(f"cannot assign an array of shape {value.shape} to a 0-d array") + x._fill = value else: - shape = tuple( - len(range(*k.indices(n))) - for k, n in zip(axes_key, T.shape, strict=True) - if isinstance(k, slice) - ) - if T.nvals == 0 or 0 in shape: - out = Array._new(Tensor(T.dtype, shape), x._fill) + adv = [k for k in axes_key if isinstance(k, _Adv)] + if ( + len(adv) == 1 + and adv[0].index.ndim == 1 + and not any(isinstance(k, int) for k in axes_key) + ): + # One 1-d index array among slices: outer (GraphBLAS) and vectorized indexing agree + _region_setitem( + x, tuple(k.index if isinstance(k, _Adv) else k for k in axes_key), value + ) + elif adv: + _scatter_setitem(x, axes_key, value) else: - out = Array._new(T[axes_key].new(), x._fill) - if new_axes: - from ._manipulation import expand_dims + _region_setitem(x, tuple(axes_key), value) + x._facts = None - out = expand_dims(out, tuple(new_axes)) - return out +def _setitem_value(x, value): + """The value cast to ``x``'s dtype: a numpy scalar, or an Array of at least 1 dimension.""" + from ._creation import asarray + from ._elementwise import cast + + dt = x.dtype + if isinstance(value, (builtins.bool, int, float, complex, np.generic)): + return _scalar_value(value, dt) + if not isinstance(value, Array): + value = asarray(value) + if value.ndim == 0: + return _scalar_value(value._fill, dt) + return value if value.dtype == dt else cast(value, dt) -def _advanced_getitem(x, key): - raise NotImplementedError( - f"indexing with {key!r}: only integers, slices, ... and None are supported so far" + +def _scalar_value(value, dt): + if isinstance(value, (complex, np.complexfloating)) and dt.kind != "c": + raise TypeError(f"cannot assign a complex value to an array of {dt}") + with np.errstate(all="ignore"): + return np.array(value, dtype=dt)[()] + + +def _region_setitem(x, key, value): + """Assign to the region an outer index selects (GraphBLAS assign).""" + T = x._T() + region = tuple( + len(range(*k.indices(n))) if isinstance(k, slice) else len(k) + for k, n in zip(key, x.shape, strict=True) + if not isinstance(k, int) ) + if not isinstance(value, Array): + if math.prod(region) == 0: + return + if _utils.same_value(value, x._fill): + if T.nvals: + del T[key] + else: + T[key] = _utils.gb_scalar(value, x.dtype) + return + try: + ok = np.broadcast_shapes(value.shape, region) == region + except ValueError: + ok = False + if not ok: + raise ValueError( + f"cannot assign an array of shape {value.shape} to a region of shape {region}" + ) + if math.prod(region) == 0: + return + # The region takes the value's structure: an unstored position of the value deletes + T[key] = _utils.with_fill(value, x._fill)._T() -def setitem(x, key, value): - # (an implementation that changes x's Tensor in place must also clear x._facts, the cached - # summary of its stored values; Array._replace does both) - raise NotImplementedError("assignment by index is not supported yet") +def _scatter_setitem(x, axes_key, value): + """Assign by several index arrays: write every selected position.""" + # numpy lays out the selection; index broadcast views of the coordinates with the key + key = tuple(k.index if isinstance(k, _Adv) else k for k in axes_key) + coords = [] + for a, n in enumerate(x.shape): + grid = np.arange(n, dtype=np.int64).reshape([n if b == a else 1 for b in range(x.ndim)]) + coords.append(np.asarray(np.broadcast_to(grid, x.shape)[key])) + shape = coords[0].shape + size = math.prod(shape) + if size == 0: + return + _utils.check_dense(size, "assigning to the selected positions") + if isinstance(value, Array): + try: + values = np.broadcast_to(value.todense(), shape) + except ValueError: + raise ValueError( + f"cannot assign an array of shape {value.shape} to a selection of shape {shape}" + ) from None + else: + values = np.broadcast_to(np.asarray(value, dtype=x.dtype), shape) + coords = [c.reshape(-1) for c in coords] + values = values.reshape(-1) + # The last of repeated positions wins, as in numpy + flat = np.ravel_multi_index(coords, x.shape) + _, last = np.unique(flat[::-1], return_index=True) + keep = np.sort(flat.size - 1 - last) + coords = [c[keep] for c in coords] + values = values[keep] + stored = _utils.differs(values, x._fill) + T = x._T() + gbdt = gb_dtype(x.dtype) + if stored.any(): + U = Tensor.from_coo([c[stored] for c in coords], values[stored], gbdt, shape=x.shape) + T(U.S) << U + if not stored.all() and T.nvals: + D = Tensor.from_coo([c[~stored] for c in coords], True, shape=x.shape) + T(D.S) << Tensor(gbdt, x.shape) + + +# take, take_along_axis +def _index_operand(indices, name): + if isinstance(indices, Array): + if indices.dtype.kind not in "iu": + raise TypeError(f"{name} needs integer indices, not {indices.dtype}") + return indices.todense().astype(np.int64, copy=False) + arr = np.asarray(indices) + if arr.dtype.kind not in "iu" and arr.size: + raise TypeError(f"{name} needs integer indices, not {arr.dtype}") + return arr.astype(np.int64, copy=False) + + +def _check_bounds(idx, n, axis): + if idx.size and (idx.min() < -n or idx.max() >= n): + raise IndexError(f"index out of bounds for axis {axis} with size {n}") + + +def take(x, indices, /, *, axis=None): + """``x``'s elements at ``indices`` along ``axis`` (the flattened array if None).""" + from ._manipulation import _normalize_axis + + if not isinstance(x, Array): + raise TypeError(f"take needs an Array, not {type(x).__name__}") + idx = _index_operand(indices, "take") + if x.ndim == 0: + raise ValueError("take needs an array of at least one dimension") + if axis is None: + if x.ndim > 1: + raise ValueError("take needs an axis for an array of more than one dimension") + axis = 0 + axis = _normalize_axis(axis, x.ndim) + _check_bounds(idx, x.shape[axis], axis) + T = x._T() + shape = x.shape[:axis] + idx.shape + x.shape[axis + 1 :] + if T.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + return Array._new(T.take(idx, axis=axis), x._fill) + + +def take_along_axis(x, indices, /, *, axis=-1): + from ._manipulation import _normalize_axis + + if not isinstance(x, Array): + raise TypeError(f"take_along_axis needs an Array, not {type(x).__name__}") + idx = _index_operand(indices, "take_along_axis") + if idx.ndim != x.ndim: + raise ValueError("take_along_axis needs indices with as many dimensions as x") + axis = _normalize_axis(axis, x.ndim) + _check_bounds(idx, x.shape[axis], axis) + others = [s for a, s in enumerate(x.shape) if a != axis] + iothers = [s for a, s in enumerate(idx.shape) if a != axis] + bshape = list(np.broadcast_shapes(tuple(others), tuple(iothers))) + bshape.insert(axis, idx.shape[axis]) + shape = tuple(bshape) + T = x._T() + if T.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + return Array._new(T.take_along_axis(idx, axis=axis), x._fill) diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 758772fde..529c9ef0c 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -118,16 +118,6 @@ def test_object_protocol_errors(): xp.astype(x, None) -def test_indexing_not_supported_yet(): - x = make(np.arange(4.0), 0.0) - with pytest.raises(NotImplementedError): - x[True] - with pytest.raises(NotImplementedError): - x[[0, 1]] - with pytest.raises(NotImplementedError): - x[0] = 1.0 - - def test_clip_mixed_bounds(): a = np.array([1, 5, 0, 9], np.int16) x = make(a, 0) diff --git a/graphblas/tests/test_array_api_indexing.py b/graphblas/tests/test_array_api_indexing.py new file mode 100644 index 000000000..243fa41d4 --- /dev/null +++ b/graphblas/tests/test_array_api_indexing.py @@ -0,0 +1,222 @@ +"""graphblas.array_api indexing (get and set), take and take_along_axis against numpy.""" + +import numpy as np +import pytest + +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, inputs, make + +rng = np.random.default_rng(0) +A = np.arange(60).reshape(3, 4, 5) +MASK = (A % 7) < 3 + + +def _keys(): + return [ + [0, 2], + np.array([2, -1, 0, 0]), + ([1, 0], [3, 3]), + ([[0], [2]], [1, 0, 1], [4, 4, 0]), + (2, [1, 3]), + ([1], 2, [0, 4]), + (0, Ellipsis, [4, 1]), + ([0, 1], slice(None), [1, 3]), + (slice(None), [2, 0], [1, 1]), + (slice(1, None), [3, 0, 3]), + (slice(None, None, -1), slice(1, 3), np.array([[0, 1], [2, 3]])), + (None, [1], slice(None), 2), + ([2, 1], None, 0), + (Ellipsis, [3, 0]), + MASK, + MASK[:, :, 0], + (slice(None), MASK[0]), + (MASK[:, 0, :].any(axis=1),), + (MASK[:, 0, 0], 1), + True, + False, + (np.array([], dtype=np.int64),), + (np.zeros((3, 4, 5), bool),), + (Ellipsis, None), + ] + + +@pytest.mark.parametrize("key", range(len(_keys()))) +def test_getitem_advanced(key): + key = _keys()[key] + for dtype in [np.dtype("float64"), np.dtype("int16"), np.dtype("bool"), np.dtype("complex64")]: + for dense, x in inputs(dtype, (3, 4, 5), seed=3): + assert_array_equal(x[key], dense[key]) + # index arrays given as Arrays + if ( + isinstance(key, np.ndarray) + or isinstance(key, tuple) + and any(isinstance(k, np.ndarray) for k in key) + ): + k2 = key if not isinstance(key, np.ndarray) else (key,) + k2 = tuple(xp.asarray(k) if isinstance(k, np.ndarray) else k for k in k2) + assert_array_equal(x[k2], dense[key]) + + +def test_boolean_mask_from_comparison(): + for fill in [0.0, 1.5]: + a = np.where(rng.random((4, 5)) < 0.5, fill, np.round(rng.normal(size=(4, 5)), 1)) + x = make(a, fill) + for mask in [x > 0, x < 0.5, x == fill, xp.isnan(x)]: + assert_array_equal(x[mask], a[mask.todense()]) + assert_array_equal(x[x > 100], a[a > 100]) + big = xp.zeros((10**5, 10**5)) + mask = big != 0 + assert big[mask].shape == (0,) + + +def test_getitem_errors(): + x = make(A.astype(float), 0.0) + with pytest.raises(IndexError): + x[[0, 3]] + with pytest.raises(IndexError): + x[[0], [-5]] + with pytest.raises(IndexError): + x[MASK[:2]] + with pytest.raises(IndexError): + xp.zeros((0, 0))[np.zeros((0, 1), bool)] + assert xp.zeros((2, 3))[np.zeros((0, 3), bool)].shape == (0,) + with pytest.raises(IndexError): + x[np.array([0.5])] + with pytest.raises(IndexError): + x[xp.asarray([0.5])] + with pytest.raises(IndexError): + x["a"] + with pytest.raises(IndexError): + x[..., 0, ...] + with pytest.raises(IndexError): + x[MASK, 0] + with pytest.raises(IndexError): + x[xp.asarray(1.5)] + + +SET_CASES = [ + ((0, 1), 9.0), + ((0, [1, 2]), 9.0), + (([0, 1], [0, 2]), 7.0), + (([0, 1, 0], [0, 2, 0], [1, 1, 1]), np.array([1.0, 2.0, 3.0])), + ((slice(None), [0, 0], 1), np.array([5.0, 6.0])), + (([1], slice(1, 3)), 0.0), + ((Ellipsis, 1), -1.0), + ((1,), np.arange(20.0).reshape(4, 5)), + ((slice(None), slice(1, 3)), np.array([1.0, 0.0, 2.0, 0.0, 3.0])), + ((slice(None, None, -1), 2), 0.0), + (([2, 0, 2],), np.arange(20.0).reshape(4, 5) * 0), + (MASK, 0.0), + (MASK, 4.0), + (MASK[:, :, 0], np.array([1.0, 0.0, 2.0, 0.0, 3.0])), + (([0, 2], [1, 3]), np.array([[1.0], [2.0]])), + ((), 3.5), + ((Ellipsis,), np.zeros(5)), + ((None, 0), 1.0), +] + + +@pytest.mark.parametrize("case", range(len(SET_CASES))) +def test_setitem(case): + key, value = SET_CASES[case] + for fill in [0.0, 2.5, -0.0]: + dense = np.where(A % 3 == 0, fill, A.astype(float)) + x = make(dense, fill) + expected = dense.copy() + expected[key] = value + for v in [value, make(np.asarray(value), 0.0), make(np.asarray(value), fill)]: + y = make(dense, fill) + y[key] = v + assert_array_equal(y, expected) + assert y.fill_value == x.fill_value + y = make(dense, fill) + y[MASK.tolist()] = 1.0 # (a nested list is an index array too) + expected = dense.copy() + expected[MASK] = 1.0 + assert_array_equal(y, expected) + + +def test_setitem_dtypes_and_errors(): + x = make(np.array([1, 0, 3], np.int8), 0) + x[0] = 1.9 + assert_array_equal(x, np.array([1, 0, 3], np.int8)) + with pytest.raises(OverflowError): + x[1] = 300 + with pytest.raises(TypeError, match="complex"): + x[1] = 1j + with pytest.raises(ValueError, match="cannot assign"): + x[:2] = xp.asarray([1, 2, 3]) + with pytest.raises(ValueError, match="cannot assign"): + x[[0, 1]] = xp.asarray([1, 2, 3]) + y = make(np.zeros((2, 2)), 0.0) + with pytest.raises(ValueError, match="cannot assign"): + y[[0, 1], [0, 1]] = np.ones(3) + z = xp.asarray(1.5) + z[False] = 3.0 + assert_array_equal(z, np.array(1.5)) + z[True] = 2.5 + assert_array_equal(z, np.array(2.5)) + q = make(np.array([1.0, 2.0]), 0.0) + q[False] = 3.0 + assert_array_equal(q, np.array([1.0, 2.0])) + z[()] = 2.0 + assert_array_equal(z, np.array(2.0)) + with pytest.raises(ValueError, match="0-d"): + z[...] = xp.asarray([1.0]) + # the cached summary of the values does not survive an assignment + w = make(np.array([1.0, 0.0, 2.0]), 0.0) + assert_array_equal(w * make(np.array([0.0, 3.0, 0.0]), 0.0), np.zeros(3)) + w[0] = -1.0 + assert_array_equal(w * make(np.array([0.0, 3.0, 0.0]), 0.0), np.array([-0.0, 0.0, 0.0])) + # an assignment does not change another array that shares its data + a = make(np.array([1.0, 0.0, 2.0]), 0.0) + b = xp.asarray(a, copy=True) + c = a[:] + a[[0, 2]] = 5.0 + assert_array_equal(b, np.array([1.0, 0.0, 2.0])) + assert_array_equal(c, np.array([1.0, 0.0, 2.0])) + + +@pytest.mark.parametrize("axis", [0, 1, -1, None]) +def test_take(axis): + for dtype in [np.dtype("float64"), np.dtype("uint8"), np.dtype("bool")]: + shape = (7,) if axis is None else (3, 4) + for dense, x in inputs(dtype, shape, seed=2): + n = dense.shape[0 if axis is None else axis] + for idx in [np.array([0, n - 1, 1, 1]), np.array([-1, 0]), np.array([], np.int64)]: + expected = np.take(dense, idx, axis=axis) + assert_array_equal(xp.take(x, xp.asarray(idx), axis=axis), expected) + assert_array_equal(xp.take(x, idx, axis=axis), expected) + with pytest.raises(IndexError): + xp.take(x, xp.asarray([n]), axis=axis) + with pytest.raises(ValueError, match="axis"): + xp.take(xp.zeros((2, 2)), xp.asarray([0])) + with pytest.raises(TypeError, match="integer"): + xp.take(xp.zeros(2), xp.asarray([0.0])) + + +@pytest.mark.parametrize("axis", [0, 1, -1]) +def test_take_along_axis(axis): + for dense, x in inputs(np.dtype("float32"), (3, 4), seed=5): + n = dense.shape[axis] + idx = np.argsort(dense, axis=axis) + assert_array_equal( + xp.take_along_axis(x, xp.asarray(idx), axis=axis), + np.take_along_axis(dense, idx, axis=axis), + ) + shape = [3, 4] + shape[axis] = 2 + idx = rng.integers(-n, n, shape) + assert_array_equal( + xp.take_along_axis(x, xp.asarray(idx), axis=axis), + np.take_along_axis(dense, idx, axis=axis), + ) + bshape = [1, 1] + bshape[axis] = 5 + idx = rng.integers(0, n, bshape) + assert_array_equal( + xp.take_along_axis(x, idx, axis=axis), np.take_along_axis(dense, idx, axis=axis) + ) + with pytest.raises(ValueError, match="dimensions"): + xp.take_along_axis(xp.zeros((2, 2)), xp.asarray([0])) From a5da9a8a565572c0b93b19e615d2d0fbe5bc4057 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:26:21 -0600 Subject: [PATCH 16/77] array_api: argmax, argmin, count_nonzero, nonzero and searchsorted argmax and argmin rank each fiber's stored entries (NaN first, then the earliest position) and let the fill compete once where the fiber has an unstored position, at the first such position: numpy's answer in O(entries log entries). nonzero reads the stored entries when the fill is zero (every unstored position is nonzero otherwise, so it is dense). searchsorted searches the stored values of the sorted array and counts the run of the fill where it sorts. --- graphblas/array_api/_searching.py | 169 +++++++++++++++++++- graphblas/tests/test_array_api_searching.py | 86 ++++++++++ 2 files changed, 254 insertions(+), 1 deletion(-) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index f58a07f28..e0847cb78 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -1,6 +1,7 @@ """Searching functions: argmax, argmin, count_nonzero, nonzero, searchsorted, where.""" import builtins +import math import numpy as np @@ -10,7 +11,7 @@ from ._dtypes import check_supported from ._utils import gb_dtype -__all__ = ["where"] +__all__ = ["argmax", "argmin", "count_nonzero", "nonzero", "searchsorted", "where"] def _operand(x, name): @@ -66,3 +67,169 @@ def func(*values): R = _utils.map_values(func, [t for t in slots if t is not None], rdt) return Array._new(R, fill) + + +def _check_array(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + + +def _arg(name, x, axis, keepdims): + """The first position of the extreme value of each fiber (argmax, argmin), NaN counting + as the extreme (numpy's rule). O(entries log entries): the stored entries of each fiber + are ranked, and the fill competes once if the fiber has an unstored position. + """ + from ._manipulation import _normalize_axis, expand_dims, permute_dims, reshape + + _check_array(x, name) + if x.dtype.kind == "c": + raise TypeError(f"{name} needs a real-valued array, not {x.dtype}") + if axis is None: + out = _arg(name, reshape(x, (-1,)), 0, False) + if keepdims: + out = reshape(out, (1,) * x.ndim) + return out + axis = _normalize_axis(axis, x.ndim) + L = x.shape[axis] + if L == 0: + raise ValueError(f"{name} of an empty axis") + out_shape = tuple(n for a, n in enumerate(x.shape) if a != axis) + others = tuple(a for a in range(x.ndim) if a != axis) + T = permute_dims(x, (*others, axis))._T() + nfibers = math.prod(out_shape) + if T.nvals == 0 or nfibers == 0: + out = Array._new(Tensor(gb_dtype(np.int64), out_shape), 0) + return expand_dims(out, axis) if keepdims else out + coords, values = T.to_coo() # (C order: by fiber, then by position along the axis) + fiber = ( + np.ravel_multi_index(coords[:-1], out_shape) + if out_shape + else np.zeros(len(values), np.int64) + ) + pos = coords[-1] + starts = np.flatnonzero(np.r_[True, fiber[1:] != fiber[:-1]]) + counts = np.diff(np.r_[starts, len(fiber)]) + # Rank the values, NaN first: then the lowest key, earliest position, wins + isnan = np.isnan(values) if values.dtype.kind == "f" else np.zeros(len(values), bool) + _, rank = np.unique(values, return_inverse=True) + rank = rank.reshape(-1) + if name == "argmax": + rank = rank.max(initial=0) - rank + key = np.where(isnan, -1, rank) + order = np.lexsort((pos, key, fiber)) + best = order[starts] + best_pos, best_val = pos[best], values[best] + # The first position each fiber does not store (L if it stores them all) + within = np.arange(len(pos)) - np.repeat(starts, counts) + gap = pos != within + first_gap = np.where( + np.add.reduceat(gap, starts) > 0, + _first_true(gap, starts, counts, within), + counts, + ) + fill = x._fill + with np.errstate(invalid="ignore"): + if fill != fill: # (NaN) + better = ~np.isnan(best_val) if best_val.dtype.kind == "f" else np.ones(len(best), bool) + tie = ~better + else: + best_is_nan = ( + np.isnan(best_val) if best_val.dtype.kind == "f" else np.zeros(len(best), bool) + ) + beats = (fill > best_val) if name == "argmax" else (fill < best_val) + better = beats & ~best_is_nan + tie = (fill == best_val) & ~best_is_nan + unstored = first_gap < L + result = np.where( + unstored & better, + first_gap, + np.where(unstored & tie, np.minimum(first_gap, best_pos), best_pos), + ) + fibers = fiber[starts] + if out_shape: + R = Tensor.from_coo( + np.unravel_index(fibers, out_shape), result.astype(np.int64), gb_dtype(np.int64), + shape=out_shape, + ) # fmt: skip + out = Array._new(R, 0) + else: + out = Array._new(Tensor(gb_dtype(np.int64), ()), result[0]) + return expand_dims(out, axis) if keepdims else out + + +def _first_true(flags, starts, counts, within): + """Per group (``starts``, ``counts``): ``within`` at the group's first True flag.""" + idx = np.where(flags, np.arange(len(flags)), len(flags)) + first = np.minimum.reduceat(idx, starts) + first = np.minimum(first, len(flags) - 1) + return within[first] + + +def argmax(x, /, *, axis=None, keepdims=False): + return _arg("argmax", x, axis, keepdims) + + +def argmin(x, /, *, axis=None, keepdims=False): + return _arg("argmin", x, axis, keepdims) + + +def count_nonzero(x, /, *, axis=None, keepdims=False): + from ._elementwise import not_equal + from ._statistical import sum as sum_ + + _check_array(x, "count_nonzero") + return sum_(not_equal(x, 0), axis=axis, dtype=np.dtype(np.int64), keepdims=keepdims) + + +def nonzero(x, /): + """The coordinates (C order) of the nonzero positions: O(entries) when the fill is zero, + O(size) otherwise (every unstored position is nonzero). + """ + from ._creation import asarray + from ._elementwise import not_equal + + _check_array(x, "nonzero") + if x.ndim == 0: + raise ValueError("nonzero needs an array of at least one dimension") + nz = not_equal(x, 0) + if nz._fill: + _utils.check_dense(x.size, "nonzero of an array whose fill is not zero") + coords = np.nonzero(nz.todense()) + else: + T = nz._T() + if T.nvals == 0: + coords = tuple(np.empty(0, np.int64) for _ in range(x.ndim)) + else: + coords = T.select("==", True).new().to_coo(values=False)[0] + return tuple(asarray(np.asarray(c, dtype=np.int64)) for c in coords) + + +def searchsorted(x1, x2, /, *, side="left", sorter=None): + """Insertion positions in the sorted 1-d ``x1``: the stored values are searched, and the + run of the fill counted where it sorts. O(entries of x1 + entries of x2). + """ + from ._creation import asarray + from ._indexing import take + + _check_array(x1, "searchsorted") + if x1.ndim != 1: + raise ValueError("searchsorted needs a 1-d x1") + if side not in {"left", "right"}: + raise ValueError(f"side must be 'left' or 'right', not {side!r}") + if sorter is not None: + x1 = take(x1, sorter) + T = x1._T() + stored = T.to_coo()[1] if T.nvals else np.empty(0, x1.dtype) + k = x1.size - T.nvals + fill = np.asarray([x1._fill]) + + def count(v): + n = np.searchsorted(stored, v, side=side) + return (n + k * np.searchsorted(fill, v, side=side)).astype(np.int64) + + if not isinstance(x2, Array): + x2 = asarray(x2) + out_fill = count(np.asarray(x2._fill)) + if x2.ndim == 0 or x2._T().nvals == 0: + return Array._new(Tensor(gb_dtype(np.int64), x2.shape), out_fill) + return Array._new(_utils.map_values(count, [x2._T()], np.dtype(np.int64)), out_fill) diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py index 38a67c589..578b2ccbf 100644 --- a/graphblas/tests/test_array_api_searching.py +++ b/graphblas/tests/test_array_api_searching.py @@ -54,3 +54,89 @@ def test_where_dtypes(): assert_array_equal( np_call(xp.where, xp.asarray(c), x, 2), np_call(np.where, c, dense, 2) ) + + +SEARCH_DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint16", "int64", "float32", "float64")] + + +@pytest.mark.parametrize("name", ["argmax", "argmin"]) +@pytest.mark.parametrize("keepdims", [False, True]) +def test_argmax_argmin(name, keepdims): + func = getattr(np, name) + for dtype in SEARCH_DTYPES: + for shape, axes in [((7,), [None, 0, -1]), ((3, 5), [None, 0, 1]), ((2, 3, 4), [0, 1, 2])]: + for dense, x in inputs(dtype, shape, seed=8): + for axis in axes: + expected = func(dense, axis=axis, keepdims=keepdims) + out = getattr(xp, name)(x, axis=axis, keepdims=keepdims) + assert_array_equal(out, np.asarray(expected, dtype=np.int64)) + with pytest.raises(ValueError, match="empty"): + getattr(xp, name)(xp.zeros((0, 3)), axis=0) + with pytest.raises(TypeError, match="real"): + getattr(xp, name)(xp.zeros(3, dtype=xp.complex64)) + assert_array_equal(getattr(xp, name)(xp.zeros((2, 0)), axis=0), np.zeros(0, np.int64)) + + +def test_argmax_ties_with_the_fill(): + # the fill ties with a stored value: the first of either position + for a, fill in [ + ([0.0, 3.0, 0.0, 3.0], 3.0), + ([3.0, 3.0, 1.0, 0.0], 3.0), + ([-0.0, 0.0, -1.0], 0.0), + ([np.nan, 1.0, 1.0, np.nan], np.nan), + ([1.0, 1.0, np.nan, 1.0], 1.0), + ]: + a = np.array(a) + for x in [make(a, fill), make(a, fill, np.random.default_rng(0), explicit=0.7)]: + assert int(xp.argmax(x)) == np.argmax(a) + assert int(xp.argmin(x)) == np.argmin(a) + + +@pytest.mark.parametrize("axis", [None, 0, 1, (0, 1), ()]) +def test_count_nonzero(axis): + for dtype in [*SEARCH_DTYPES, np.dtype("complex64")]: + for dense, x in inputs(dtype, (3, 4), seed=9): + for keepdims in [False, True]: + expected = np.count_nonzero(dense, axis=axis, keepdims=keepdims) + out = xp.count_nonzero(x, axis=axis, keepdims=keepdims) + assert_array_equal(out, np.asarray(expected, dtype=np.int64)) + + +def test_nonzero(): + for dtype in [*SEARCH_DTYPES, np.dtype("complex128")]: + for shape in [(7,), (3, 4), (2, 3, 2)]: + for dense, x in inputs(dtype, shape, seed=10): + out = xp.nonzero(x) + expected = np.nonzero(dense) + assert isinstance(out, tuple) + assert len(out) == len(expected) + for o, e in zip(out, expected, strict=True): + assert_array_equal(o, e.astype(np.int64)) + with pytest.raises(ValueError, match="dimension"): + xp.nonzero(xp.asarray(1)) + big = xp.zeros((10**6, 10**6)) + big[3, 4] = 1.0 + out = xp.nonzero(big) + assert [int(c[0]) for c in out] == [3, 4] + + +@pytest.mark.parametrize("side", ["left", "right"]) +def test_searchsorted(side): + rng = np.random.default_rng(11) + for fill in [0.0, 2.0, -np.inf, np.nan]: + for _ in range(20): + a = np.sort(np.where(rng.random(9) < 0.5, fill, np.round(rng.normal(size=9) * 2))) + x1 = make(a, fill) + v = np.array([-10.0, -1.0, 0.0, fill, 1.0, 2.0, 2.5, np.inf, np.nan]) + for x2, vv in [(make(v, 0.0), v), (make(v, fill), v), (xp.asarray(1.0), np.array(1.0))]: + out = xp.searchsorted(x1, x2, side=side) + assert_array_equal(out, np.searchsorted(a, vv, side=side).astype(np.int64)) + assert int(xp.searchsorted(x1, 2.0, side=side)) == np.searchsorted(a, 2.0, side=side) + a = np.array([3, 0, 1, 0, 5]) + sorter = np.argsort(a, kind="stable") + out = xp.searchsorted(make(a, 0), xp.asarray([0, 4]), side=side, sorter=xp.asarray(sorter)) + assert_array_equal(out, np.searchsorted(a, [0, 4], side=side, sorter=sorter).astype(np.int64)) + with pytest.raises(ValueError, match="1-d"): + xp.searchsorted(xp.zeros((2, 2)), 1.0) + with pytest.raises(ValueError, match="side"): + xp.searchsorted(xp.zeros(2), 1.0, side="middle") From 35a4b17d18046dd5d615474728031ed92ad06377 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:27:32 -0600 Subject: [PATCH 17/77] array_api: sort, which keeps sparsity, and argsort sort places each fiber's stored values below the fill first, then the run of the fill (stored values equal to it keep their places in it), then those above it, NaN last: numpy's stable sort, signed zeros included, in O(entries log entries). Descending order is flip(sort(flip(x))): values reversed, equal values in their order, NaN first. argsort's result is a permutation of every position, so it goes through numpy under the densify limit. --- graphblas/array_api/_sorting.py | 124 +++++++++++++++++++++- graphblas/tests/test_array_api_sorting.py | 68 ++++++++++++ 2 files changed, 190 insertions(+), 2 deletions(-) create mode 100644 graphblas/tests/test_array_api_sorting.py diff --git a/graphblas/array_api/_sorting.py b/graphblas/array_api/_sorting.py index 156e958c6..d14489f49 100644 --- a/graphblas/array_api/_sorting.py +++ b/graphblas/array_api/_sorting.py @@ -1,3 +1,123 @@ -"""Sorting functions: argsort, sort.""" +"""Sorting functions: sort and argsort. -__all__ = [] +``sort`` keeps the result as sparse as its input: along each fiber, the values below the fill +come first, then the run of values equal to the fill (the unstored positions among them), +then the values above it, NaN last; equal values keep their order (stable). Descending order +reverses the values but keeps equal values in their order, NaN first (as +``flip(sort(flip(x)))``). O(entries log entries). + +``argsort``'s result is a permutation of every position, dense: it goes through numpy, under +the densify limit. +""" + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._array import Array + +__all__ = ["argsort", "sort"] + + +def _check(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + if x.ndim == 0: + raise ValueError(f"{name} needs an array of at least one dimension") + if x.dtype.kind == "c": + raise TypeError(f"{name} needs a real-valued array, not {x.dtype}") + + +def sort(x, /, *, axis=-1, descending=False, stable=True): + from ._manipulation import _normalize_axis, permute_dims + + _check(x, "sort") + axis = _normalize_axis(axis, x.ndim) + L = x.shape[axis] + T = x._T() + if T.nvals == 0 or L <= 1: + return x._copy() + others = tuple(a for a in range(x.ndim) if a != axis) + perm = (*others, axis) + coords, values = permute_dims(x, perm)._T().to_coo() # (by fiber, then position) + fshape = tuple(x.shape[a] for a in others) + fiber = np.ravel_multi_index(coords[:-1], fshape) if fshape else np.zeros(len(values), np.int64) + pos = coords[-1] + fill = x._fill + with np.errstate(invalid="ignore"): + if values.dtype.kind == "f": + nan_v, nan_f = np.isnan(values), bool(np.isnan(fill)) + equal = (values == fill) | (nan_v & nan_f) + below = (values < fill) | (~nan_v & nan_f) + else: + equal = values == fill + below = values < fill + above = ~(equal | below) + if descending: + below, above = above, below + # Rank the values (np.unique sorts NaN last and makes them one value) + _, rank = np.unique(values, return_inverse=True) + rank = rank.reshape(-1) + if descending: + rank = rank.max(initial=0) - rank + # Groups per fiber: 0 below the fill, 1 equal to it, 2 above it + group = np.where(below, 0, np.where(equal, 1, 2)) + order = np.lexsort((pos, rank, group, fiber)) + f_o, g_o, p_o = fiber[order], group[order], pos[order] + # Rank of each entry within its (fiber, group), and the sizes of the groups per fiber + key = f_o * 3 + g_o + starts = np.flatnonzero(np.r_[True, key[1:] != key[:-1]]) + sizes = np.diff(np.r_[starts, len(key)]) + within = np.arange(len(key)) - np.repeat(starts, sizes) + fibers, fiber_index = np.unique(f_o, return_inverse=True) + counts = np.zeros((len(fibers), 3), np.int64) + np.add.at(counts, (fiber_index, g_o), 1) + n_below = counts[fiber_index, 0] + n_equal = counts[fiber_index, 1] + unstored = L - counts.sum(axis=1)[fiber_index] + # An entry equal to the fill keeps its place among the run of the fill (stable): the + # positions before it in the run are every position before it but the other groups' + others_before = _count_before(f_o, p_o, g_o != 1) + new_pos = np.where( + g_o == 0, + within, + np.where(g_o == 1, n_below + p_o - others_before, n_below + n_equal + unstored + within), + ) + vals = values[order] + keep = (g_o != 1) | _utils.differs(vals, fill) + if not keep.any(): + return Array._new(Tensor(T.dtype, x.shape), fill) + out_coords = list(np.unravel_index(f_o[keep], fshape)) if fshape else [] + out_coords.append(new_pos[keep]) + R = Tensor.from_coo(out_coords, vals[keep], T.dtype, shape=tuple(x.shape[a] for a in perm)) + inverse = tuple(int(i) for i in np.argsort(perm)) + return permute_dims(Array._new(R, fill), inverse) + + +def _count_before(fiber, pos, flags): + """For each entry: how many flagged entries of its fiber have a smaller position.""" + order = np.lexsort((pos, fiber)) + f, fl = fiber[order], flags[order].astype(np.int64) + starts = np.flatnonzero(np.r_[True, f[1:] != f[:-1]]) + csum = np.cumsum(fl) - fl # (flagged before, over all fibers) + base = np.repeat(csum[starts], np.diff(np.r_[starts, len(f)])) + out = np.empty(len(fiber), np.int64) + out[order] = csum - base + return out + + +def argsort(x, /, *, axis=-1, descending=False, stable=True): + from ._creation import _from_numpy + from ._manipulation import _normalize_axis + + _check(x, "argsort") + axis = _normalize_axis(axis, x.ndim) + _utils.check_dense(x.size, "argsort (every position's index)") + dense = x.todense() + if descending: + # Descending and stable: equal values keep their order, as flip(sort(flip(x))) + flipped = np.flip(dense, axis=axis) + out = x.shape[axis] - 1 - np.flip(np.argsort(flipped, axis=axis, kind="stable"), axis=axis) + else: + out = np.argsort(dense, axis=axis, kind="stable") + return _from_numpy(out.astype(np.int64, copy=False)) diff --git a/graphblas/tests/test_array_api_sorting.py b/graphblas/tests/test_array_api_sorting.py new file mode 100644 index 000000000..85cba8519 --- /dev/null +++ b/graphblas/tests/test_array_api_sorting.py @@ -0,0 +1,68 @@ +"""graphblas.array_api sort and argsort against numpy (stable sorts).""" + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, inputs, make + +DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint16", "int64", "float32", "float64")] + + +def _np_sort(a, axis, descending): + if descending: + return np.flip(np.sort(np.flip(a, axis=axis), axis=axis, kind="stable"), axis=axis) + return np.sort(a, axis=axis, kind="stable") + + +def _np_argsort(a, axis, descending): + if descending: + n = a.shape[axis] + flipped = np.argsort(np.flip(a, axis=axis), axis=axis, kind="stable") + return n - 1 - np.flip(flipped, axis=axis) + return np.argsort(a, axis=axis, kind="stable") + + +@pytest.mark.parametrize("descending", [False, True]) +def test_sort(descending): + for dtype in DTYPES: + for shape, axes in [((9,), [0, -1]), ((3, 6), [0, 1]), ((2, 3, 4), [0, 1, 2])]: + for dense, x in inputs(dtype, shape, seed=12): + for axis in axes: + out = xp.sort(x, axis=axis, descending=descending) + assert_array_equal(out, _np_sort(dense, axis, descending)) + assert out.fill_value == x.fill_value or np.isnan(out.fill_value) + out = xp.argsort(x, axis=axis, descending=descending) + assert_array_equal(out, _np_argsort(dense, axis, descending).astype(np.int64)) + + +def test_sort_keeps_sparsity(): + x = xp.zeros((1000, 1000)) + x[5, [3, 7, 9]] = np.array([2.0, -1.0, 4.0]) + x[8, 0] = -0.0 + out = xp.sort(x, axis=1) + assert out.nnz == 4 + row = out[5].todense() + assert (row[[0, 998, 999]] == [-1.0, 2.0, 4.0]).all() + assert np.signbit(out[8].todense()[0]) # -0.0 keeps its place among the zeros: first + with gb.config.set(array_api_densify_limit=10**5), pytest.raises(xp.DensifyError): + xp.argsort(x) + + +def test_sort_signed_zeros_and_nan_stable(): + a = np.array([0.0, -0.0, np.nan, 1.0, -0.0, 0.0, np.nan, -1.0]) + for fill in [0.0, -0.0, np.nan, 1.0]: + for x in [make(a, fill), make(a, fill, np.random.default_rng(1), explicit=0.6)]: + for descending in [False, True]: + assert_array_equal(xp.sort(x, descending=descending), _np_sort(a, 0, descending)) + + +def test_sort_errors(): + with pytest.raises(ValueError, match="dimension"): + xp.sort(xp.asarray(1.0)) + with pytest.raises(TypeError, match="real"): + xp.argsort(xp.zeros(2, dtype=xp.complex64)) + with pytest.raises(TypeError, match="Array"): + xp.sort(np.zeros(3)) From cc388a77f9d909634f7a9e69957fe6c83b4f08f2 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:29:06 -0600 Subject: [PATCH 18/77] array_api: isin and unique_* The unique values are those of the stored entries plus the fill, counted once per unstored position, so unique_all costs O(entries log entries) and its inverse_indices keep the input's sparsity (the fill's index is their fill). NaNs are distinct, as the standard says: a NaN fill with unstored positions goes through numpy densely, under the limit. isin checks x1's stored values and fill against x2's. --- graphblas/array_api/_set.py | 165 +++++++++++++++++++++++++- graphblas/tests/test_array_api_set.py | 86 ++++++++++++++ 2 files changed, 249 insertions(+), 2 deletions(-) create mode 100644 graphblas/tests/test_array_api_set.py diff --git a/graphblas/array_api/_set.py b/graphblas/array_api/_set.py index 48c4f080d..9e2d27faa 100644 --- a/graphblas/array_api/_set.py +++ b/graphblas/array_api/_set.py @@ -1,3 +1,164 @@ -"""Set functions: isin, unique_all, unique_counts, unique_inverse, unique_values.""" +"""Set functions: isin, unique_all, unique_counts, unique_inverse, unique_values. -__all__ = [] +The unique values of an array are those of its stored entries plus the fill, counted once per +unstored position: O(entries log entries). NaNs are distinct (each is its own value), so a +fill of NaN with unstored positions makes as many distinct values: that case goes through +numpy densely, under the densify limit. Values come sorted, as numpy's. +""" + +from typing import NamedTuple + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._utils import gb_dtype + +__all__ = ["isin", "unique_all", "unique_counts", "unique_inverse", "unique_values"] + + +class UniqueAllResult(NamedTuple): + values: Array + indices: Array + inverse_indices: Array + counts: Array + + +class UniqueCountsResult(NamedTuple): + values: Array + counts: Array + + +class UniqueInverseResult(NamedTuple): + values: Array + inverse_indices: Array + + +def _vector(a): + from ._creation import _from_numpy + + return _from_numpy(np.asarray(a)) + + +def _unique(x, *, index=False, inverse=False, counts=False): + """The sorted unique values of ``x`` (as numpy's unique with equal_nan=False), and the + first flat index, the inverse (an Array of x's shape) and the counts when asked. + """ + if not isinstance(x, Array): + raise TypeError(f"needs an Array, not {type(x).__name__}") + T = x._T() + fill = x._fill + k = x.size - T.nvals + if k and x.dtype.kind in "fc" and np.isnan(fill): + _utils.check_dense(x.size, "unique of an array whose fill is NaN (each NaN is distinct)") + return _unique_dense(x, index, inverse, counts) + if T.nvals: + coords, values = T.to_coo() + flat = np.ravel_multi_index(coords, x.shape) if x.ndim else np.zeros(1, np.int64) + else: + coords = tuple(np.empty(0, np.int64) for _ in range(x.ndim)) + values, flat = np.empty(0, x.dtype), np.empty(0, np.int64) + if k: + # The fill once, at its first unstored position, with the other k - 1 added later + first = _first_unstored(flat, x.size) + at = np.searchsorted(flat, first) + values = np.insert(values, at, fill) + flat = np.insert(flat, at, first) + out = np.unique( + values, return_index=True, return_inverse=True, return_counts=True, equal_nan=False + ) + uvalues, uindex, uinverse, ucounts = out + uinverse = uinverse.reshape(-1) + result = {"values": _vector(uvalues)} + if k: + fill_slot = uinverse[at] + ucounts = ucounts.copy() + ucounts[fill_slot] += k - 1 + if index: + result["indices"] = _vector(flat[uindex].astype(np.int64)) + if counts: + result["counts"] = _vector(ucounts.astype(np.int64)) + if inverse: + inv_fill = int(uinverse[at]) if k else 0 + if k: + uinverse = np.delete(uinverse, at) + if x.ndim == 0: + inv = Array._new( + Tensor(gb_dtype(np.int64), ()), uinverse[0] if len(uinverse) else inv_fill + ) + elif T.nvals: + R = Tensor.from_coo( + coords, uinverse.astype(np.int64), gb_dtype(np.int64), shape=x.shape + ) + inv = Array._new(R, inv_fill) + else: + inv = Array._new(Tensor(gb_dtype(np.int64), x.shape), inv_fill) + result["inverse_indices"] = inv + return result + + +def _first_unstored(flat, size): + """The smallest flat position in ``range(size)`` not in the sorted ``flat``.""" + gaps = np.flatnonzero(flat != np.arange(len(flat))) + return int(gaps[0]) if len(gaps) else len(flat) + + +def _unique_dense(x, index, inverse, counts): + from ._creation import _from_numpy + + dense = x.todense() + uvalues, uindex, uinverse, ucounts = np.unique( + dense.reshape(-1), return_index=True, return_inverse=True, return_counts=True, + equal_nan=False, + ) # fmt: skip + result = {"values": _vector(uvalues)} + if index: + result["indices"] = _vector(uindex.astype(np.int64)) + if counts: + result["counts"] = _vector(ucounts.astype(np.int64)) + if inverse: + result["inverse_indices"] = _from_numpy(uinverse.reshape(x.shape).astype(np.int64)) + return result + + +def unique_values(x, /): + return _unique(x)["values"] + + +def unique_counts(x, /): + r = _unique(x, counts=True) + return UniqueCountsResult(r["values"], r["counts"]) + + +def unique_inverse(x, /): + r = _unique(x, inverse=True) + return UniqueInverseResult(r["values"], r["inverse_indices"]) + + +def unique_all(x, /): + r = _unique(x, index=True, inverse=True, counts=True) + return UniqueAllResult(r["values"], r["indices"], r["inverse_indices"], r["counts"]) + + +def isin(x1, x2, /, *, invert=False): + """Whether each element of ``x1`` is among ``x2``'s values (its stored values and its + fill): O(entries of both). + """ + from ._creation import asarray + + x1 = x1 if isinstance(x1, Array) else asarray(x1) + x2 = x2 if isinstance(x2, Array) else asarray(x2) + T2 = x2._T() + values = T2.to_coo(indices=False)[1] if T2.nvals else np.empty(0, x2.dtype) + if x2.size > T2.nvals: + values = np.append(values, x2._fill) + + def member(v): + return np.isin(v, values, invert=invert) + + fill = member(np.asarray(x1._fill)) + T1 = x1._T() + if x1.ndim == 0 or T1.nvals == 0: + return Array._new(Tensor(gb_dtype(bool), x1.shape), fill) + return Array._new(_utils.map_values(member, [T1], np.dtype(bool)), fill) diff --git a/graphblas/tests/test_array_api_set.py b/graphblas/tests/test_array_api_set.py new file mode 100644 index 000000000..6bcd36084 --- /dev/null +++ b/graphblas/tests/test_array_api_set.py @@ -0,0 +1,86 @@ +"""graphblas.array_api set functions against numpy.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, check_invariants, inputs, make + +DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint32", "int64", "float32", "float64")] + + +def _check_values(out, expected): + """Unique values: as numpy's, except that which of 0.0 and -0.0 stands for both is free.""" + check_invariants(out) + got = out.todense() + assert got.dtype == expected.dtype + np.testing.assert_array_equal(got, expected) # (0.0 == -0.0, NaN matches NaN) + + +@pytest.mark.parametrize("shape", [(9,), (3, 4), (2, 3, 2), ()]) +def test_unique(shape): + for dtype in DTYPES: + for dense, x in inputs(dtype, shape, seed=13): + # (numpy's unique_* need not sort; np.unique does, as graphblas.array_api does) + values, indices, inverse, counts = np.unique( + dense, return_index=True, return_inverse=True, return_counts=True, equal_nan=False + ) + e = SimpleNamespace( + values=values, + indices=indices, + inverse_indices=inverse.reshape(dense.shape), + counts=counts, + ) + r = xp.unique_all(x) + _check_values(r.values, e.values) + assert_array_equal(r.counts, e.counts.astype(np.int64)) + # the first index of each value (of either zero, for 0.0 and -0.0) + assert_array_equal(r.indices, e.indices.astype(np.int64)) + assert_array_equal(r.inverse_indices, e.inverse_indices.astype(np.int64)) + _check_values(xp.unique_values(x), e.values) + c = xp.unique_counts(x) + _check_values(c.values, e.values) + assert_array_equal(c.counts, e.counts.astype(np.int64)) + i = xp.unique_inverse(x) + _check_values(i.values, e.values) + assert_array_equal(i.inverse_indices, e.inverse_indices.astype(np.int64)) + + +def test_unique_is_sparse(): + x = xp.zeros((10**6, 10**6)) + x[3, [1, 5]] = np.array([2.0, -1.0]) + r = xp.unique_all(x) + assert_array_equal(r.values, np.array([-1.0, 0.0, 2.0])) + assert_array_equal(r.indices, np.array([3 * 10**6 + 5, 0, 3 * 10**6 + 1])) + assert_array_equal(r.counts, np.array([1, 10**12 - 2, 1])) + assert r.inverse_indices.nnz == 2 + assert r.inverse_indices.fill_value == 1 + + +def test_unique_nan_fill(): + a = np.array([np.nan, 1.0, np.nan, 2.0]) + x = make(a, np.nan) + _check_values(xp.unique_values(x), np.unique(a, equal_nan=False)) + assert_array_equal(xp.unique_counts(x).counts, np.ones(4, np.int64)) + with gb.config.set(array_api_densify_limit=2), pytest.raises(xp.DensifyError): + xp.unique_values(x) + + +@pytest.mark.parametrize("invert", [False, True]) +def test_isin(invert): + for dtype in DTYPES: + for (d1, x1), (d2, x2) in zip( + inputs(dtype, (3, 4), seed=14), inputs(dtype, (5,), seed=15), strict=False + ): + expected = np.isin(d1, d2, invert=invert) + assert_array_equal(xp.isin(x1, x2, invert=invert), expected) + x = make(np.array([1, 0, 5, 2]), 0) + assert_array_equal(xp.isin(x, 2), np.isin([1, 0, 5, 2], 2)) + assert_array_equal(xp.isin(3, x), np.array(False)) + assert_array_equal(xp.isin(xp.asarray(1.0), x), np.array(True)) + with pytest.raises(TypeError, match="Array"): + xp.unique_values(np.zeros(3)) From ea0fdccf8710233d235e2e947c19459b00a72551 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:33:48 -0600 Subject: [PATCH 19/77] array_api: matmul, tensordot, vecdot, matrix_transpose and the linalg extension Contractions run as GraphBLAS semirings (plus_times, or lor_land for bool) and treat a zero fill as a structural zero, as scipy.sparse does: 0 * inf at an unstored position does not make NaN. A nonzero finite fill is the fill times ones plus an array with fill zero, which becomes sums of the operands broadcast over the result. A fill that is not finite, or True, goes through numpy densely, under the densify limit. xp.linalg adds diagonal, trace, vector_norm, the Frobenius, 1 and inf matrix norms, outer, cross and matrix_power (n >= 0) on the stored entries, and the decompositions and solvers through numpy densely. pyproject.toml's codespell list learns 'fro', the standard's name for the Frobenius norm. --- graphblas/array_api/__init__.py | 1 + graphblas/array_api/_linalg.py | 248 +++++++++++++++- graphblas/array_api/linalg.py | 342 +++++++++++++++++++++++ graphblas/tests/test_array_api_linalg.py | 245 ++++++++++++++++ pyproject.toml | 2 +- 5 files changed, 835 insertions(+), 3 deletions(-) create mode 100644 graphblas/array_api/linalg.py create mode 100644 graphblas/tests/test_array_api_linalg.py diff --git a/graphblas/array_api/__init__.py b/graphblas/array_api/__init__.py index e9db10cf7..b35c16924 100644 --- a/graphblas/array_api/__init__.py +++ b/graphblas/array_api/__init__.py @@ -20,6 +20,7 @@ _set, _sorting, _statistical, + linalg, ) from ._array import Array from ._constants import * # noqa: F403 diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index a22e74f0a..5d1957618 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -1,3 +1,247 @@ -"""Linear algebra in the main namespace: matmul, matrix_transpose, tensordot, vecdot.""" +"""Linear algebra in the main namespace: matmul, matrix_transpose, tensordot and vecdot. -__all__ = [] +Contractions treat a fill of zero as a structural zero (DESIGN D6), as scipy.sparse and +pydata/sparse do: a product with an unstored position is skipped, so 0 * inf there is not NaN, +unlike numpy. Sums may round differently from numpy's (the order of the terms differs). A +nonzero fill is the fill times a matrix of ones plus an array with fill zero, which turns into +sums of the operands broadcast over the result: the result is then dense in general (under the +densify limit). A fill that is not finite, or True, makes the result go through numpy densely. +""" + +import math + +import numpy as np + +from .. import semiring +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._dtypes import check_supported +from ._utils import gb_dtype + +__all__ = ["matmul", "matrix_transpose", "tensordot", "vecdot"] + + +def _check(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs Arrays, not {type(x).__name__}") + + +def _semiring(dt): + return (semiring.lor_land if dt.kind == "b" else semiring.plus_times)[gb_dtype(dt)] + + +def _is_zero(fill): + return not bool(fill) + + +def _dense_ok(*arrays): + """Whether the decomposition A = A0 + fill works: finite fills, not bool.""" + for x in arrays: + if _is_zero(x._fill): + continue + if x.dtype.kind == "b" or (x.dtype.kind in "fc" and not np.isfinite(x._fill)): + return False + return True + + +def _dense(func, *arrays): + from ._creation import _from_numpy + + for x in arrays: + _utils.check_dense(x.size, "a contraction with a fill that is not finite or is True") + out = np.asarray(func(*[x.todense() for x in arrays])) + return _from_numpy(out) + + +def _zero_fill(x): + """``x`` minus its fill (exact where the fill is 0), and the fill.""" + from ._elementwise import subtract + + if _is_zero(x._fill): + return x, None + return subtract(x, x._fill), x._fill + + +def _rdt(name, func, x1, x2): + probe = func( + np.ones((1,) * max(x1.ndim, 1), x1.dtype), np.ones((1,) * max(x2.ndim, 1), x2.dtype) + ) + return check_supported(np.asarray(probe).dtype, name) + + +def matmul(x1, x2, /): + from ._elementwise import add, multiply + from ._manipulation import expand_dims + from ._statistical import sum as sum_ + + _check(x1, "matmul") + _check(x2, "matmul") + if x1.ndim == 0 or x2.ndim == 0: + raise ValueError("matmul needs arrays of at least one dimension") + K = x1.shape[-1] + if x2.shape[-2 if x2.ndim > 1 else 0] != K: + raise ValueError(f"matmul: shapes {x1.shape} and {x2.shape} do not align") + b1 = x1.shape[:-2] + b2 = x2.shape[:-2] + try: + np.broadcast_shapes(b1, b2) + except ValueError: + raise ValueError(f"matmul: shapes {x1.shape} and {x2.shape} do not broadcast") from None + rdt = _rdt("matmul", np.matmul, x1, x2) + if not _dense_ok(x1, x2): + return _dense(np.matmul, x1, x2) + a0, fa = _zero_fill(x1) + b0, fb = _zero_fill(x2) + out = _product(a0, b0, rdt) + if fa is None and fb is None: + return out + terms = [] + if fa is not None: + # fa * (ones @ b0): b0 summed over its rows, one row of the result + s = sum_(b0, axis=-2 if x2.ndim > 1 else 0, dtype=rdt) + if x1.ndim > 1 and x2.ndim > 1: + s = expand_dims(s, axis=-2) + terms.append(multiply(fa, s)) + if fb is not None: + r = sum_(a0, axis=-1, dtype=rdt) + if x2.ndim > 1 and x1.ndim > 1: + r = expand_dims(r, axis=-1) + terms.append(multiply(fb, r)) + if fa is not None and fb is not None: + terms.append(multiply(multiply(fa, fb), K)) + for t in terms: + out = add(out, t) + from ._elementwise import cast + + return out if out.dtype == rdt else cast(out, rdt) + + +def _product(a, b, rdt): + """``a @ b`` of arrays with fill 0, in GraphBLAS.""" + Ta, Tb = a._T(), b._T() + shape = np.broadcast_shapes(a.shape[:-2], b.shape[:-2]) + shape = shape + (a.shape[-2:-1] if a.ndim > 1 else ()) + (b.shape[-1:] if b.ndim > 1 else ()) + if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(gb_dtype(rdt), shape), 0) + return Array._new(Ta.matmul(Tb, _semiring(rdt)).new(dtype=gb_dtype(rdt)), 0) + + +def matrix_transpose(x, /): + _check(x, "matrix_transpose") + return x.mT + + +def _tensordot_axes(x1, x2, axes): + if isinstance(axes, int): + if axes < 0: + raise ValueError("tensordot's axes must be nonnegative") + if axes > min(x1.ndim, x2.ndim): + raise ValueError("tensordot: more axes than the arrays have") + return tuple(range(x1.ndim - axes, x1.ndim)), tuple(range(axes)) + a, b = axes + a = tuple(int(i) % x1.ndim for i in (a if isinstance(a, (list, tuple)) else (a,))) + b = tuple(int(i) % x2.ndim for i in (b if isinstance(b, (list, tuple)) else (b,))) + if len(a) != len(b): + raise ValueError("tensordot needs as many axes of each array") + return a, b + + +def tensordot(x1, x2, /, *, axes=2): + from ._elementwise import add, cast, multiply + from ._manipulation import reshape + from ._statistical import sum as sum_ + + _check(x1, "tensordot") + _check(x2, "tensordot") + a_axes, b_axes = _tensordot_axes(x1, x2, axes) + for a, b in zip(a_axes, b_axes, strict=True): + if x1.shape[a] != x2.shape[b]: + raise ValueError(f"tensordot: axis {a} of x1 and axis {b} of x2 differ in size") + rdt = check_supported(np.result_type(x1.dtype, x2.dtype), "tensordot") + if not _dense_ok(x1, x2): + return _dense(lambda a, b: np.tensordot(a, b, axes=(a_axes, b_axes)), x1, x2) + free1 = tuple(n for i, n in enumerate(x1.shape) if i not in a_axes) + free2 = tuple(n for i, n in enumerate(x2.shape) if i not in b_axes) + shape = free1 + free2 + a0, fa = _zero_fill(x1) + b0, fb = _zero_fill(x2) + Ta, Tb = a0._T(), b0._T() + if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: + out = Array._new(Tensor(gb_dtype(rdt), shape), 0) + elif not shape: + value = Ta.tensordot(Tb, (list(a_axes), list(b_axes)), _semiring(rdt)).new() + out = Array._new(value, 0) + else: + R = Ta.tensordot(Tb, (list(a_axes), list(b_axes)), _semiring(rdt)).new(dtype=gb_dtype(rdt)) + out = Array._new(R, 0) + if out.dtype != rdt: + out = cast(out, rdt) + if fa is not None: + s = sum_(b0, axis=b_axes, dtype=rdt) if b_axes else cast(b0, rdt) + out = add(out, multiply(fa, reshape(s, (1,) * len(free1) + free2))) + if fb is not None: + r = sum_(a0, axis=a_axes, dtype=rdt) if a_axes else cast(a0, rdt) + out = add(out, multiply(fb, reshape(r, free1 + (1,) * len(free2)))) + if fa is not None and fb is not None: + K = math.prod(x1.shape[a] for a in a_axes) + out = add(out, multiply(multiply(fa, fb), K)) + return out if out.dtype == rdt else cast(out, rdt) + + +def vecdot(x1, x2, /, *, axis=-1): + """Sum over ``axis`` of ``conj(x1) * x2`` (broadcast), with zero fills structural.""" + from ._elementwise import add, cast, conj, multiply + from ._statistical import sum as sum_ + + _check(x1, "vecdot") + _check(x2, "vecdot") + ndim = min(x1.ndim, x2.ndim) + if ndim == 0: + raise ValueError("vecdot needs arrays of at least one dimension") + if not -ndim <= axis < 0: + if 0 <= axis < ndim and x1.ndim == x2.ndim: + axis -= x1.ndim + else: + raise ValueError(f"vecdot's axis must be in [-{ndim}, -1], not {axis}") + if x1.shape[axis] != x2.shape[axis]: + raise ValueError("vecdot: the axis differs in size") + shape = np.broadcast_shapes(x1.shape, x2.shape) + rdt = check_supported(np.result_type(x1.dtype, x2.dtype), "vecdot") + if not _dense_ok(x1, x2): + return _dense(lambda a, b: np.vecdot(a, b, axis=axis), x1, x2) + if x1.dtype.kind == "c": + x1 = conj(x1) + a0, fa = _zero_fill(x1) + b0, fb = _zero_fill(x2) + out_shape = tuple(n for i, n in enumerate(shape) if i != len(shape) + axis) + Ta, Tb = a0._T(), b0._T() + gbdt = gb_dtype(rdt) + if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: + out = Array._new(Tensor(gbdt, out_shape), 0) + else: + mult = _semiring(rdt).binaryop + if _both_broadcast(Ta.shape, Tb.shape, shape): + # (an outer product: ewise_mult replicates one operand) + _utils.expansion_check(Ta, Tb, shape) + P = Array._new(Ta.ewise_mult(Tb, mult).new(dtype=gbdt), 0) + out = sum_(P, axis=axis, dtype=rdt) if rdt.kind != "b" else _any(P, axis) + if fa is not None: + out = add(out, multiply(fa, sum_(b0, axis=axis, dtype=rdt))) + if fb is not None: + out = add(out, multiply(fb, sum_(a0, axis=axis, dtype=rdt))) + if fa is not None and fb is not None: + out = add(out, multiply(multiply(fa, fb), shape[axis])) + return out if out.dtype == rdt else cast(out, rdt) + + +def _any(x, axis): + from ._statistical import any as any_ + + return any_(x, axis=axis) + + +def _both_broadcast(sa, sb, shape): + from ..core.tensor_ewise import _broadcast_factor + + return _broadcast_factor(sa, shape) > 1 and _broadcast_factor(sb, shape) > 1 diff --git a/graphblas/array_api/linalg.py b/graphblas/array_api/linalg.py new file mode 100644 index 000000000..3f7634209 --- /dev/null +++ b/graphblas/array_api/linalg.py @@ -0,0 +1,342 @@ +"""The Array API's ``linalg`` extension (``xp.linalg``). + +matmul, matrix_transpose, tensordot, vecdot, diagonal, trace, vector_norm, outer, cross, the +Frobenius, 1 and inf matrix norms, and matrix_power with n >= 0 work on the stored entries. +The decompositions and solvers (cholesky, det, eig, eigh, eigvals, eigvalsh, inv, pinv, qr, +slogdet, solve, svd, svdvals, matrix_rank, and matrix_norm of order 2, -2 or "nuc") are dense +algebra: they go through numpy densely, under the densify limit. +""" + +import math +from typing import NamedTuple + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._linalg import matmul, matrix_transpose, tensordot, vecdot + +__all__ = [ + "cholesky", + "cross", + "det", + "diagonal", + "eig", + "eigh", + "eigvals", + "eigvalsh", + "inv", + "matmul", + "matrix_norm", + "matrix_power", + "matrix_rank", + "matrix_transpose", + "outer", + "pinv", + "qr", + "slogdet", + "solve", + "svd", + "svdvals", + "tensordot", + "trace", + "vecdot", + "vector_norm", +] + + +class EighResult(NamedTuple): + eigenvalues: Array + eigenvectors: Array + + +class EigResult(NamedTuple): + eigenvalues: Array + eigenvectors: Array + + +class QRResult(NamedTuple): + Q: Array + R: Array + + +class SlogdetResult(NamedTuple): + sign: Array + logabsdet: Array + + +class SVDResult(NamedTuple): + U: Array + S: Array + Vh: Array + + +def _check(x, name, ndim=2): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + if x.ndim < ndim: + raise ValueError(f"{name} needs an array of at least {ndim} dimensions") + + +def _float(x): + """``x`` in a floating-point dtype (numpy's linalg computes integers in float64).""" + from ._elementwise import cast + + return x if x.dtype.kind in "fc" else cast(x, np.dtype(np.float64)) + + +def _dense_call(name, func, *arrays): + from ._creation import _from_numpy + + for x in arrays: + _utils.check_dense(x.size, f"linalg.{name} (dense algebra)") + out = func(*[x.todense() for x in arrays]) + if isinstance(out, tuple): + return tuple(_from_numpy(np.asarray(o)) for o in out) + return _from_numpy(np.asarray(out)) + + +# Sparse +def diagonal(x, /, *, offset=0): + """The diagonals of the last two axes: O(entries).""" + _check(x, "diagonal") + n, m = x.shape[-2:] + offset = int(offset) + length = max(0, min(n, m - offset) if offset >= 0 else min(n + offset, m)) + shape = (*x.shape[:-2], length) + T = x._T() + if T.nvals == 0 or length == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + coords, values = T.to_coo(sort=False) + i, j = coords[-2], coords[-1] + keep = j - i == offset + k = i if offset >= 0 else j + out = [c[keep] for c in coords[:-2]] + [k[keep]] + if not keep.any(): + return Array._new(Tensor(T.dtype, shape), x._fill) + return Array._new(Tensor.from_coo(out, values[keep], T.dtype, shape=shape), x._fill) + + +def trace(x, /, *, offset=0, dtype=None): + from ._statistical import sum as sum_ + + return sum_(diagonal(x, offset=offset), axis=-1, dtype=dtype) + + +def outer(x1, x2, /): + from ._elementwise import multiply + from ._manipulation import expand_dims + + for x in (x1, x2): + _check(x, "outer", ndim=1) + if x.ndim != 1: + raise ValueError("outer needs 1-d arrays") + return multiply(expand_dims(x1, axis=1), expand_dims(x2, axis=0)) + + +def cross(x1, x2, /, *, axis=-1): + """The cross product of 3-vectors along ``axis`` (numpy's products and differences).""" + from ._elementwise import multiply, subtract + from ._manipulation import stack + + for x in (x1, x2): + _check(x, "cross", ndim=1) + if not -min(x1.ndim, x2.ndim) <= axis < 0: + raise ValueError(f"cross's axis must be negative and within both arrays, not {axis}") + if x1.shape[axis] != 3 or x2.shape[axis] != 3: + raise ValueError("cross needs vectors of size 3 along the axis") + np.broadcast_shapes(x1.shape, x2.shape) + + def part(x, i): + return x[(Ellipsis, i) + (slice(None),) * (-1 - axis)] + + a = [part(x1, i) for i in range(3)] + b = [part(x2, i) for i in range(3)] + c = [ + subtract(multiply(a[1], b[2]), multiply(a[2], b[1])), + subtract(multiply(a[2], b[0]), multiply(a[0], b[2])), + subtract(multiply(a[0], b[1]), multiply(a[1], b[0])), + ] + return stack(c, axis=axis) + + +def vector_norm(x, /, *, axis=None, keepdims=False, ord=2): # noqa: A002 (the Array API's name) + """O(entries): the norms are reductions of ``abs(x)``.""" + from ._elementwise import abs as abs_ + from ._elementwise import pow as pow_ + from ._elementwise import sqrt, square + from ._searching import count_nonzero + from ._statistical import max as max_ + from ._statistical import min as min_ + from ._statistical import sum as sum_ + + if not isinstance(x, Array): + raise TypeError(f"vector_norm needs an Array, not {type(x).__name__}") + a = abs_(_float(x)) + if axis is None and x.ndim: + axis = tuple(range(x.ndim)) + if ord == 2: + return sqrt(sum_(square(a), axis=axis, keepdims=keepdims)) + if ord == 1: + return sum_(a, axis=axis, keepdims=keepdims) + if ord == math.inf: + return max_(a, axis=axis, keepdims=keepdims) + if ord == -math.inf: + return min_(a, axis=axis, keepdims=keepdims) + if ord == 0: + from ._elementwise import cast + + return cast(count_nonzero(a, axis=axis, keepdims=keepdims), a.dtype) + p = float(ord) + with np.errstate(divide="ignore"): + inv = 1.0 / p + return pow_(sum_(pow_(a, p), axis=axis, keepdims=keepdims), inv) + + +def matrix_norm(x, /, *, keepdims=False, ord="fro"): # noqa: A002 (the Array API's name) + from ._elementwise import abs as abs_ + from ._elementwise import sqrt, square + from ._manipulation import expand_dims + from ._statistical import max as max_ + from ._statistical import min as min_ + from ._statistical import sum as sum_ + + _check(x, "matrix_norm") + a = abs_(_float(x)) + if ord == "fro": + out = sqrt(sum_(square(a), axis=(-2, -1))) + elif ord in (1, -1): + cols = sum_(a, axis=-2) + out = max_(cols, axis=-1) if ord == 1 else min_(cols, axis=-1) + elif ord in (math.inf, -math.inf): + rows = sum_(a, axis=-1) + out = max_(rows, axis=-1) if ord == math.inf else min_(rows, axis=-1) + elif ord in (2, -2, "nuc"): + s = svdvals(x) + if ord == 2: + out = max_(s, axis=-1) + elif ord == -2: + out = min_(s, axis=-1) + else: + out = sum_(s, axis=-1) + else: + raise ValueError(f"invalid matrix_norm order {ord!r}") + if keepdims: + out = expand_dims(expand_dims(out, axis=-1), axis=-1) + return out + + +def matrix_power(x, n, /): + """``x`` to the integer power ``n``: repeated sparse products for n >= 0; inv for n < 0.""" + _check(x, "matrix_power") + if x.shape[-1] != x.shape[-2]: + raise ValueError("matrix_power needs square matrices") + n = int(n) + if n < 0: + return matrix_power(inv(x), -n) + if n == 0: + from ._creation import eye + from ._manipulation import broadcast_to + + return broadcast_to(eye(x.shape[-1], dtype=x.dtype), x.shape) + result = None + base = x + while n: + if n & 1: + result = base if result is None else matmul(result, base) + n >>= 1 + if n: + base = matmul(base, base) + return result + + +# Dense algebra (numpy under the densify limit) +def cholesky(x, /, *, upper=False): + _check(x, "cholesky") + + def func(a): + L = np.linalg.cholesky(a) + return np.conj(np.swapaxes(L, -1, -2)) if upper else L + + return _dense_call("cholesky", func, _float(x)) + + +def det(x, /): + _check(x, "det") + return _dense_call("det", np.linalg.det, _float(x)) + + +def eigh(x, /): + _check(x, "eigh") + return EighResult(*_dense_call("eigh", lambda a: tuple(np.linalg.eigh(a)), _float(x))) + + +def eigvalsh(x, /): + _check(x, "eigvalsh") + return _dense_call("eigvalsh", np.linalg.eigvalsh, _float(x)) + + +def eig(x, /): + _check(x, "eig") + return EigResult(*_dense_call("eig", lambda a: tuple(np.linalg.eig(a)), _float(x))) + + +def eigvals(x, /): + _check(x, "eigvals") + return _dense_call("eigvals", np.linalg.eigvals, _float(x)) + + +def inv(x, /): + _check(x, "inv") + return _dense_call("inv", np.linalg.inv, _float(x)) + + +def pinv(x, /, *, rtol=None): + _check(x, "pinv") + if isinstance(rtol, Array): + rtol = rtol.todense() + return _dense_call("pinv", lambda a: np.linalg.pinv(a, rtol=rtol), _float(x)) + + +def matrix_rank(x, /, *, rtol=None): + _check(x, "matrix_rank") + if isinstance(rtol, Array): + rtol = rtol.todense() + return _dense_call("matrix_rank", lambda a: np.linalg.matrix_rank(a, rtol=rtol), _float(x)) + + +def qr(x, /, *, mode="reduced"): + _check(x, "qr") + return QRResult(*_dense_call("qr", lambda a: tuple(np.linalg.qr(a, mode=mode)), _float(x))) + + +def slogdet(x, /): + _check(x, "slogdet") + return SlogdetResult(*_dense_call("slogdet", lambda a: tuple(np.linalg.slogdet(a)), _float(x))) + + +def solve(x1, x2, /): + _check(x1, "solve") + _check(x2, "solve", ndim=1) + a, b = _float(x1), _float(x2) + + def func(a, b): + if b.ndim == 1: # (the standard's rule: a vector only when x2 is 1-d) + return np.linalg.solve(a, b[..., None])[..., 0] + return np.linalg.solve(a, b) + + return _dense_call("solve", func, a, b) + + +def svd(x, /, *, full_matrices=True): + _check(x, "svd") + out = _dense_call( + "svd", lambda a: tuple(np.linalg.svd(a, full_matrices=full_matrices)), _float(x) + ) + return SVDResult(*out) + + +def svdvals(x, /): + _check(x, "svdvals") + return _dense_call("svdvals", lambda a: np.linalg.svd(a, compute_uv=False), _float(x)) diff --git a/graphblas/tests/test_array_api_linalg.py b/graphblas/tests/test_array_api_linalg.py new file mode 100644 index 000000000..f8d545bb3 --- /dev/null +++ b/graphblas/tests/test_array_api_linalg.py @@ -0,0 +1,245 @@ +"""graphblas.array_api linear algebra (main namespace and the linalg extension) against numpy.""" + +import math + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, check_invariants, make, np_call + +rng = np.random.default_rng(21) + + +def _rand(shape, fill, dtype=np.float64, density=0.5): + if np.dtype(dtype).kind == "b": + a = rng.random(shape) < 0.5 + return np.where(rng.random(shape) < density, a, bool(fill)) + if np.dtype(dtype).kind in "iu": + a = rng.integers(0, 5, shape) + else: + a = np.round(rng.normal(size=shape), 2) + return np.where(rng.random(shape) < density, a, fill).astype(dtype) + + +def assert_close(x, expected): + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape + got = x.todense() + if got.dtype.kind in "fc": + tol = 1e-5 if got.dtype.itemsize * (2 if got.dtype.kind == "f" else 1) == 8 else 1e-10 + np.testing.assert_allclose(got, expected, rtol=tol, atol=tol) + else: + np.testing.assert_array_equal(got, expected) + + +MATMUL_SHAPES = [ + ((3, 4), (4, 5)), + ((4,), (4, 5)), + ((3, 4), (4,)), + ((4,), (4,)), + ((2, 3, 4), (4, 5)), + ((2, 1, 3, 4), (3, 4, 2)), + ((0, 4), (4, 2)), +] + + +@pytest.mark.parametrize("shapes", MATMUL_SHAPES, ids=str) +def test_matmul(shapes): + for dtype in [np.float64, np.int64, np.float32, np.complex128, np.uint8, np.bool_]: + for fa, fb in [(0, 0), (1, 0), (0, 2), (1, 3)]: + if np.dtype(dtype).kind == "b" and (fa or fb) and (fa, fb) != (1, 3): + continue + a, b = _rand(shapes[0], fa, dtype), _rand(shapes[1], fb, dtype) + x1, x2 = make(a, fa), make(b, fb) + expected = np.matmul(a, b) + assert_close(xp.matmul(x1, x2), expected) + assert_close(x1 @ x2, expected) + with pytest.raises(ValueError, match="align"): + xp.matmul(xp.zeros((2, 3)), xp.zeros((2, 3))) + with pytest.raises(ValueError, match="at least one"): + xp.matmul(xp.asarray(1.0), xp.zeros(2)) + with pytest.raises(ValueError, match="broadcast"): + xp.matmul(xp.zeros((2, 2, 3)), xp.zeros((3, 3, 2))) + + +def test_matmul_zero_fill_is_structural(): + # (DESIGN D6: as scipy.sparse; numpy would give NaN for 0 * inf) + a = make(np.array([[np.inf, 0.0], [0.0, 1.0]]), 0.0) + b = make(np.array([[0.0, 1.0], [1.0, 0.0]]), 0.0) + assert_close(a @ b, np.array([[0.0, np.inf], [1.0, 0.0]])) + # a fill that is not finite: numpy, densely + c = make(np.array([[np.inf, 1.0], [np.inf, np.inf]]), np.inf) + assert_array_equal(np_call(xp.matmul, c, b), np_call(np.matmul, c.todense(), b.todense())) + with gb.config.set(array_api_densify_limit=2), pytest.raises(xp.DensifyError): + c @ b + + +def test_matmul_inplace(): + a = make(np.eye(3), 0.0) + b = make(np.arange(9.0).reshape(3, 3), 0.0) + a @= b + assert_close(a, np.arange(9.0).reshape(3, 3)) + + +@pytest.mark.parametrize( + "axes", [2, 1, 0, ([1, 2], [1, 0]), ([2], [0]), ([0, 1, 2], [2, 1, 0]), ((), ())] +) +def test_tensordot(axes): + for dtype in [np.float64, np.int32, np.complex64]: + for fa, fb in [(0, 0), (2, 0), (0, 1), (1, 1)]: + a = _rand((3, 4, 3), fa, dtype) + b = _rand((3, 4, 3), fb, dtype) if axes != 1 else _rand((3, 2), fb, dtype) + if axes == 2: + b = _rand((4, 3, 2), fb, dtype) + if axes == 0: + b = _rand((2,), fb, dtype) + expected = np.tensordot(a, b, axes=axes) + out = xp.tensordot(make(a, fa), make(b, fb), axes=axes) + if np.dtype(dtype) == np.complex64: + np.testing.assert_allclose(out.todense(), expected, rtol=1e-5, atol=1e-5) + assert out.dtype == expected.dtype + else: + assert_close(out, expected) + with pytest.raises(ValueError, match="differ"): + xp.tensordot(xp.zeros((2, 3)), xp.zeros((2, 3)), axes=1) + + +@pytest.mark.parametrize("axis", [-1, -2]) +def test_vecdot(axis): + for dtype in [np.float64, np.int16, np.complex128, np.bool_]: + for fa, fb in [(0, 0), (1, 0), (0, 2), (1, 2)]: + if np.dtype(dtype).kind == "b" and (fa or fb) and (fa, fb) != (1, 2): + continue + for sa, sb in [((3, 4), (3, 4)), ((2, 3, 4), (3, 4)), ((3, 1), (1, 3, 1))]: + a, b = _rand(sa, fa, dtype), _rand(sb, fb, dtype) + if -axis > min(a.ndim, b.ndim) or a.shape[axis] != b.shape[axis]: + continue + out = xp.vecdot(make(a, fa), make(b, fb), axis=axis) + assert_close(out, np.vecdot(a, b, axis=axis)) + with pytest.raises(ValueError, match="axis"): + xp.vecdot(xp.zeros(3), xp.zeros(3), axis=-2) + a = np.array([1 + 2j, np.inf, 1j]) + b = np.array([2j, 1, 1 - 1j]) + assert_array_equal(np_call(xp.vecdot, make(a, np.inf), make(b, 0)), np_call(np.vecdot, a, b)) + + +def test_matrix_transpose(): + a = _rand((2, 3, 4), 0.0) + assert_array_equal(xp.matrix_transpose(make(a, 0.0)), np.swapaxes(a, -1, -2)) + assert_array_equal(xp.linalg.matrix_transpose(make(a, 1.0)), np.swapaxes(a, -1, -2)) + with pytest.raises(ValueError, match="2 dimensions"): + xp.matrix_transpose(xp.zeros(3)) + + +@pytest.mark.parametrize("offset", [-3, -1, 0, 1, 2, 5]) +def test_diagonal_trace(offset): + for fill in [0.0, 1.5]: + a = _rand((2, 3, 4), fill) + x = make(a, fill) + assert_array_equal( + xp.linalg.diagonal(x, offset=offset), np.linalg.diagonal(a, offset=offset) + ) + assert_close(xp.linalg.trace(x, offset=offset), np.linalg.trace(a, offset=offset)) + b = _rand((4, 4), fill, np.int8) + assert_close( + xp.linalg.trace(make(b, int(fill)), offset=offset, dtype=xp.int8), + np.linalg.trace(b, offset=offset, dtype=np.int8), + ) + + +def test_outer_cross(): + a, b = _rand((4,), 0.0), _rand((3,), 1.0) + assert_array_equal(xp.linalg.outer(make(a, 0.0), make(b, 1.0)), np.outer(a, b)) + u, v = _rand((2, 3), 0.0, density=0.7), _rand((3,), 2.0, density=0.7) + assert_array_equal(xp.linalg.cross(make(u, 0.0), make(v, 2.0)), np.cross(u, v)) + u, v = _rand((3, 2), 0.0, density=0.7), _rand((3, 2), 0.0, density=0.7) + assert_array_equal( + xp.linalg.cross(make(u, 0.0), make(v, 0.0), axis=-2), np.cross(u, v, axis=-2) + ) + with pytest.raises(ValueError, match="size 3"): + xp.linalg.cross(xp.zeros(2), xp.zeros(2)) + with pytest.raises(ValueError, match="1-d"): + xp.linalg.outer(xp.zeros((2, 2)), xp.zeros(2)) + + +@pytest.mark.parametrize("order", [2, 1, 0, math.inf, -math.inf, 3, -1, 0.5]) +def test_vector_norm(order): + for dtype, fill in [(np.float64, 0.0), (np.float64, -2.0), (np.complex128, 0), (np.int32, 0)]: + a = _rand((3, 4), fill, dtype) + x = make(a, fill) + for axis in [None, 0, -1, (0, 1)]: + for keepdims in [False, True]: + expected = np_call( + np.linalg.vector_norm, a, axis=axis, keepdims=keepdims, ord=order + ) + assert_close( + np_call(xp.linalg.vector_norm, x, axis=axis, keepdims=keepdims, ord=order), + expected, + ) + + +@pytest.mark.parametrize("order", ["fro", 1, -1, math.inf, -math.inf, 2, -2, "nuc"]) +def test_matrix_norm(order): + for fill in [0.0, 1.0]: + a = _rand((2, 3, 4), fill) + for keepdims in [False, True]: + expected = np.linalg.matrix_norm(a, ord=order, keepdims=keepdims) + assert_close( + xp.linalg.matrix_norm(make(a, fill), ord=order, keepdims=keepdims), expected + ) + with pytest.raises(ValueError, match="order"): + xp.linalg.matrix_norm(xp.zeros((2, 2)), ord=3) + + +@pytest.mark.parametrize("n", [0, 1, 2, 5, -1, -2]) +def test_matrix_power(n): + a = np.array([[2.0, 1.0, 0.0], [0.0, 1.0, 0.0], [1.0, 0.0, 3.0]]) + for fill in [0.0, 1.0]: + dense = np.where(a == 0, fill, a) + expected = np.linalg.matrix_power(dense, n) + assert_close(xp.linalg.matrix_power(make(dense, fill), n), expected) + b = np.stack([a, a.T]) + assert_close(xp.linalg.matrix_power(make(b, 0.0), 3), np.linalg.matrix_power(b, 3)) + with pytest.raises(ValueError, match="square"): + xp.linalg.matrix_power(xp.zeros((2, 3)), 2) + + +def test_dense_algebra(): + a = np.array([[4.0, 1.0, 0.0], [1.0, 3.0, 0.0], [0.0, 0.0, 2.0]]) + x = make(a, 0.0) + la = np.linalg + assert_close(xp.linalg.cholesky(x), la.cholesky(a)) + assert_close(xp.linalg.cholesky(x, upper=True), la.cholesky(a).T) + assert_close(xp.linalg.det(x), la.det(a)) + w, _ = xp.linalg.eigh(x) + assert_close(w, la.eigh(a).eigenvalues) + assert_close(xp.linalg.eigvalsh(x), la.eigvalsh(a)) + assert_close(xp.linalg.eigvals(x), la.eigvals(a)) + assert xp.linalg.eig(x).eigenvectors.shape == (3, 3) + assert_close(xp.linalg.inv(x), la.inv(a)) + assert_close(xp.linalg.pinv(x), la.pinv(a)) + assert_close(xp.linalg.matrix_rank(x), la.matrix_rank(a)) + q, r = xp.linalg.qr(x) + assert_close(q @ r, a) + sign, logdet = xp.linalg.slogdet(x) + assert_close(sign, la.slogdet(a).sign) + assert_close(logdet, la.slogdet(a).logabsdet) + b = np.array([1.0, 0.0, 2.0]) + assert_close(xp.linalg.solve(x, make(b, 0.0)), la.solve(a, b)) + B = np.stack([a, a]) + assert_close(xp.linalg.solve(make(B, 0.0), make(B, 0.0)), la.solve(B, B)) + s = xp.linalg.svd(x, full_matrices=False).S + assert_close(s, la.svd(a).S) + assert_close(xp.linalg.svdvals(x), la.svdvals(a)) + assert_close(xp.linalg.inv(make(np.eye(2, dtype=np.int64), 0)), np.eye(2)) + with gb.config.set(array_api_densify_limit=4), pytest.raises(xp.DensifyError, match="dense"): + xp.linalg.inv(x) + with pytest.raises(ValueError, match="2 dimensions"): + xp.linalg.det(xp.zeros(3)) + with pytest.raises(TypeError, match="Array"): + xp.linalg.det(a) diff --git a/pyproject.toml b/pyproject.toml index 23e0cfda7..21dd78bf8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -230,7 +230,7 @@ skip_empty = true exclude_lines = ["pragma: no cover", "raise AssertionError", "raise NotImplementedError"] [tool.codespell] -ignore-words-list = "bu,coo,ba" +ignore-words-list = "bu,coo,ba,fro" [tool.ruff] # https://github.com/charliermarsh/ruff/ From b4d0938d09644e279093688fd9041fa987738ca3 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:36:48 -0600 Subject: [PATCH 20/77] array_api: test thread 3's argument checks and edge cases; ndarray @ Array matmul takes a numpy array operand as an Array, so ndarray @ Array (which numpy hands to Array.__rmatmul__) works as Array @ Array does. --- graphblas/array_api/_linalg.py | 15 ++-- graphblas/tests/test_array_api_errors.py | 94 ++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 5 deletions(-) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 5d1957618..057879a70 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -27,6 +27,15 @@ def _check(x, name): raise TypeError(f"{name} needs Arrays, not {type(x).__name__}") +def _operand(x): + """A numpy array as an Array (``ndarray @ Array`` comes here through ``__rmatmul__``).""" + if isinstance(x, np.ndarray): + from ._creation import asarray + + return asarray(x) + return x + + def _semiring(dt): return (semiring.lor_land if dt.kind == "b" else semiring.plus_times)[gb_dtype(dt)] @@ -75,6 +84,7 @@ def matmul(x1, x2, /): from ._manipulation import expand_dims from ._statistical import sum as sum_ + x1, x2 = _operand(x1), _operand(x2) _check(x1, "matmul") _check(x2, "matmul") if x1.ndim == 0 or x2.ndim == 0: @@ -169,14 +179,9 @@ def tensordot(x1, x2, /, *, axes=2): Ta, Tb = a0._T(), b0._T() if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: out = Array._new(Tensor(gb_dtype(rdt), shape), 0) - elif not shape: - value = Ta.tensordot(Tb, (list(a_axes), list(b_axes)), _semiring(rdt)).new() - out = Array._new(value, 0) else: R = Ta.tensordot(Tb, (list(a_axes), list(b_axes)), _semiring(rdt)).new(dtype=gb_dtype(rdt)) out = Array._new(R, 0) - if out.dtype != rdt: - out = cast(out, rdt) if fa is not None: s = sum_(b0, axis=b_axes, dtype=rdt) if b_axes else cast(b0, rdt) out = add(out, multiply(fa, reshape(s, (1,) * len(free1) + free2))) diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 529c9ef0c..3ea41bb1f 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -3,6 +3,7 @@ import numpy as np import pytest +import graphblas as gb import graphblas.array_api as xp from graphblas.array_api import _utils from graphblas.core.tensor import Tensor @@ -204,3 +205,96 @@ def test_summary_of_complex_values(): x = make(np.array([1 + 0j, complex(-0.0, 1), complex(np.nan, 0)]), 0j) facts = _utils.summary(x) assert facts == {"finite": False, "nonnegative": False, "nan": True, "negzero": True} + + +def test_thread3_argument_checks(): + x = make(np.arange(6.0).reshape(2, 3), 0.0) + for func in [xp.take, xp.take_along_axis]: + with pytest.raises(TypeError, match="Array"): + func(np.zeros(3), np.array([0])) + with pytest.raises(TypeError, match="integer"): + xp.take(x, np.array([0.5]), axis=0) + with pytest.raises(ValueError, match="dimension"): + xp.take(xp.asarray(1.0), np.array([0])) + with pytest.raises(TypeError, match="Arrays"): + xp.matmul(x, [[1.0], [2.0], [3.0]]) + with pytest.raises(ValueError, match="nonnegative"): + xp.tensordot(x, x, axes=-1) + with pytest.raises(ValueError, match="more axes"): + xp.tensordot(x, x, axes=3) + with pytest.raises(ValueError, match="as many"): + xp.tensordot(x, x, axes=([0], [0, 1])) + with pytest.raises(ValueError, match="dimension"): + xp.vecdot(xp.asarray(1.0), x) + with pytest.raises(ValueError, match="differs"): + xp.vecdot(x, make(np.ones((2, 2)), 0.0)) + assert_array_equal(xp.vecdot(x, x, axis=1), np.vecdot(x.todense(), x.todense(), axis=1)) + with pytest.raises(ValueError, match="repeated"): + xp.flip(x, axis=(0, 0)) + with pytest.raises(TypeError, match="Python scalars"): + xp.where(x > 0, x, "a") + with pytest.raises(TypeError, match="Array"): + xp.argmax(np.zeros(3)) + with pytest.raises(ValueError, match="negative"): + xp.linalg.cross(xp.zeros(3), xp.zeros(3), axis=0) + with pytest.raises(TypeError, match="Array"): + xp.linalg.vector_norm(np.zeros(3)) + a = np.array([[1.0, 2.0], [3.0, 4.0]]) + rtol = xp.asarray(1e-3) + assert_array_equal(xp.linalg.pinv(make(a, 0.0), rtol=rtol), np.linalg.pinv(a, rtol=1e-3)) + assert int(xp.linalg.matrix_rank(make(a, 0.0), rtol=rtol)) == 2 + assert_array_equal(x @ x.mT, x.todense() @ x.todense().T) + + +def test_thread3_edge_cases(): + x = make(np.arange(6.0).reshape(2, 3), 0.0) + assert_array_equal(x[np.array([], dtype=float)], np.zeros((0, 3))) + assert_array_equal(x[xp.asarray(True)], x.todense()[None]) + assert_array_equal(x[np.array(False)], x.todense()[np.array(False)]) + y = make(np.arange(6.0).reshape(2, 3), 0.0) + y[:, 0:0] = 5.0 + y[[0, 1], 0:0] = xp.zeros((2, 0)) + y[[], 1] = 3.0 + assert_array_equal(y, x.todense()) + e = xp.zeros((2, 3)) + e[0] = 0.0 # (deleting from an array that stores nothing) + assert e.nnz == 0 + assert_array_equal( + xp.take_along_axis(xp.zeros((2, 3)), np.zeros((2, 1), int), axis=1), np.zeros((2, 1)) + ) + assert_array_equal(xp.nonzero(xp.zeros(4))[0], np.zeros(0, np.int64)) + assert_array_equal(xp.sort(xp.zeros(3)), np.zeros(3)) + assert_array_equal(xp.sort(make(np.array([0.0, 2.0, 0.0]), 2.0)), np.array([0.0, 0.0, 2.0])) + z = make(np.array([[2.0, 2.0], [2.0, 2.0]]), 2.0, np.random.default_rng(0), explicit=1.0) + assert_array_equal(xp.sort(z, axis=0), np.full((2, 2), 2.0)) + full = make(np.array([[1.0, 2.0], [3.0, 1.0]]), 0.0) + assert_array_equal(xp.unique_inverse(full).inverse_indices, np.array([[0, 1], [2, 0]])) + assert_array_equal(xp.unique_values(xp.zeros((2, 0))), np.zeros(0)) + assert_array_equal( + xp.unique_inverse(xp.zeros((2, 0))).inverse_indices, np.zeros((2, 0), np.int64) + ) + assert_array_equal(xp.isin(full, full), np.ones((2, 2), bool)) + with np.errstate(all="ignore"): + assert_array_equal(xp.var(xp.zeros((0, 3)), axis=0, keepdims=True), np.full((1, 3), np.nan)) + # tensordot with a nonzero fill on a 0-d result, and int32 * int8 promotion + a = make(np.array([1.0, 0.0, 2.0]), 1.0) + assert_array_equal(xp.tensordot(a, a, axes=1), np.tensordot(a.todense(), a.todense(), axes=1)) + i32 = make(np.array([[1, 2], [3, 4]], np.int32), 0) + i8 = make(np.array([[1, 0], [0, 1]], np.int8), 0) + assert_array_equal( + xp.tensordot(i32, i8, axes=1), np.tensordot(i32.todense(), i8.todense(), axes=1) + ) + inf = make(np.array([[np.inf, 1.0], [2.0, np.inf]]), np.inf) + assert_array_equal( + xp.tensordot(inf, i32, axes=1), np.tensordot(inf.todense(), i32.todense(), axes=1) + ) + # both operands broadcast: an outer product of the entries, under the limit + u, v = make(np.ones((3, 1, 4)), 0.0), make(np.ones((5, 4)), 0.0) + assert_array_equal(xp.vecdot(u, v), np.vecdot(u.todense(), v.todense())) + with gb.config.set(array_api_densify_limit=4), pytest.raises(xp.DensifyError): + xp.vecdot(u, v) + # numpy operands + c = np.array([True, False, True]) + assert_array_equal(xp.where(c, np.array([1, 2, 3]), 0), np.where(c, [1, 2, 3], 0)) + assert_array_equal(xp.where(True, x, 0.0), x.todense()) + assert_array_equal(np.eye(3) @ x.mT, np.eye(3) @ x.todense().T) From 4f6e29ebd205db1a0a0abbf908c82f7619bd14a5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:51:39 -0600 Subject: [PATCH 21/77] Tensor: dup(dtype=...) keeps an n x 1 stored Matrix held by row Matrix.dup holds an n x 1 Matrix by column, as SuiteSparse holds a Vector, which broke the Tensor's invariant that its Matrix is held by row (graphblas.array_api's astype of a (2, 1) array failed the check). The copy is now made into a Matrix created by row. --- graphblas/core/tensor.py | 5 ++++- graphblas/tests/test_tensor.py | 4 ++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/graphblas/core/tensor.py b/graphblas/core/tensor.py index 399ebc82c..817e12dae 100644 --- a/graphblas/core/tensor.py +++ b/graphblas/core/tensor.py @@ -803,7 +803,10 @@ def dup(self, dtype=None, *, layout=None, clear=False, name=None): if layout is not None: M, lay = tl.relayout(M, self._shape, lay, tl.as_layout(layout, self._shape)) if dtype != self.dtype: - M = M.dup(dtype) + # (Matrix.dup would hold an n x 1 Matrix by column, as SuiteSparse holds a Vector) + D = tl.new_matrix(dtype, M._nrows, M._ncols) + D << M + M = D if M is self._M: T = Tensor._from_M(M, self._shape, lay, name=name, store=self._store) else: diff --git a/graphblas/tests/test_tensor.py b/graphblas/tests/test_tensor.py index 4713b119c..950fc207f 100644 --- a/graphblas/tests/test_tensor.py +++ b/graphblas/tests/test_tensor.py @@ -383,6 +383,10 @@ def test_dup_and_clear(): D.clear() assert D.nvals == 0 check(T, values, present) + # A new dtype keeps an n x 1 stored Matrix held by row (Matrix.dup holds it by column) + for shape in [(2, 1), (3, 1, 1)]: + T, values, present = random_tensor(shape, rng, density=1.0) + check(T.dup(dtypes.FP32), values.astype(np.float32), present) def test_vanilla_check(monkeypatch): From f1f199277a9403a4ad716215edeac485c49ed54c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:51:43 -0600 Subject: [PATCH 22/77] array_api: signed zeros give way to sparsity by default; guard outer products; numpy's clip By default (graphblas.config array_api_signed_zeros False), a product with an unstored zero reads 0.0 whatever its sign: multiply of finite floats with zero fills stores only the intersection, and divide of a zero-filled numerator by an array with no 0 or NaN stores only the numerator's entries. A row scaling d[:, None] * A with signed d stored every position to keep the -0.0s; it now stores A's entries. NaN and inf stay exact, as does everything else; array_api_signed_zeros True keeps every signed zero, as numpy. An intersection of two broadcast operands (an outer product) is now checked against the densify limit, as a union is. clip with array bounds is numpy's clip over the union of the entries (numpy gives the bound on a tie, which minimum(maximum()) did not for signed zeros). The tests that compare with numpy bit for bit run in the strict mode. --- graphblas/array_api/_elementwise.py | 83 ++++++++++++++----- graphblas/array_api/_utils.py | 37 +++++++-- graphblas/graphblas.yaml | 4 + graphblas/tests/test_array_api_elementwise.py | 83 +++++++++++++++++++ graphblas/tests/test_array_api_errors.py | 5 +- graphblas/tests/test_array_api_indexing.py | 1 + graphblas/tests/test_array_api_linalg.py | 1 + .../tests/test_array_api_manipulation.py | 1 + graphblas/tests/test_array_api_searching.py | 1 + graphblas/tests/test_array_api_set.py | 1 + graphblas/tests/test_array_api_sorting.py | 1 + 11 files changed, 192 insertions(+), 26 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index c31d97b66..14e3d6fc7 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -14,7 +14,7 @@ import numpy as np -from .. import binary, unary +from .. import binary, config, unary from ..core.tensor import Tensor from . import _utils from ._array import Array @@ -486,12 +486,17 @@ def _binary_arrays(spec, a, b, op, ct, rdt, fill): fb = _utils.scalar(b._fill, ct) if spec.annihilator is not None and spec.annihilator(a, b, fa, fb, ct): # Where either operand holds its fill the result is the fill: the intersection + _utils.mult_expansion_check(Ta, Tb, shape) R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) return _result(R, fill) _utils.expansion_check(Ta, Tb, shape) - R = Ta.ewise_union(Tb, op, gb_scalar(fa, ct), gb_scalar(fb, ct), allow_expansion=True).new( - dtype=gb_dtype(rdt) - ) + expr = Ta.ewise_union(Tb, op, gb_scalar(fa, ct), gb_scalar(fb, ct), allow_expansion=True) + if spec.name == "divide" and Ta.shape == shape and _divide_left_joins(b, fa, fb): + # The result's entries lie within the numerator's: compute only there + R = Tensor(gb_dtype(rdt), shape) + R(Ta.S) << expr + else: + R = expr.new(dtype=gb_dtype(rdt)) return _result(R, fill) # numpy over the values: both operands' values on the union of entries, aligned _utils.expansion_check(Ta, Tb, shape) @@ -529,14 +534,34 @@ def _both_fills_zero(a, b, fa, fb, ct): def _multiply_intersects(a, b, fa, fb, ct): """Whether ``a * b`` stores only the intersection of entries: both fills are 0 and, for floats, ``value * 0.0`` is ``0.0`` for every stored value (not ``-0.0``, nor NaN for inf). + + Unless ``graphblas.config["array_api_signed_zeros"]`` is True (strict), a ``-0.0`` + product counts as the fill: finite values suffice (the union of a row scaling + ``d[:, None] * A`` with signed ``d`` would store every position, to keep the signs of its + zeros; DESIGN D3a). """ if not _both_fills_zero(a, b, fa, fb, ct): return False if ct.kind != "f": return True + if not config.get("array_api_signed_zeros", False): + return _utils.summary(a)["finite"] and _utils.summary(b)["finite"] return _finite_nonnegative(a) and _finite_nonnegative(b) +def _divide_left_joins(b, fa, fb): + """Whether ``a / b`` stores only ``a``'s entries: ``a``'s fill is 0, and 0 divided by + any value of ``b`` is a zero (b holds no 0 or NaN). The zero's sign then reads as the + fill's, which numpy's may not be: only when signed zeros give way (the default). + """ + if config.get("array_api_signed_zeros", False) or fa != 0: + return False + if fb == 0 or np.isnan(fb): + return False + facts = _utils.summary(b) + return not facts["zero"] and not facts["nan"] + + def _finite_nonnegative(x): facts = _utils.summary(x) return facts["finite"] and facts["nonnegative"] @@ -694,30 +719,50 @@ def func(x1, x2, /): def clip(x, /, min=None, max=None): - """``x`` clipped to ``[min, max]`` in the dtype of ``x``; NaN propagates. - - With scalar bounds this is numpy's ``clip`` (which keeps ``-0.0`` against a bound of - ``0.0``); with array bounds it is ``minimum(maximum(x, min), max)``, as numpy's ``clip`` - with array bounds is. + """``x`` clipped as numpy's ``clip`` does (NaN propagates; a tie gives the bound), in + the dtype of ``x``, over the union of the entries of ``x`` and of array bounds. """ if not isinstance(x, Array): raise TypeError(f"clip needs an Array, not {type(x).__name__}") lo = _clip_bound(min, x) hi = _clip_bound(max, x) if lo is None and hi is None: - out = x._copy() - elif all(b is None or not isinstance(b, Array) or b.ndim == 0 for b in (lo, hi)): + return x._copy() + operands = [x] + [b for b in (lo, hi) if isinstance(b, Array) and b.ndim] + if len(operands) == 1: lo_v, hi_v = (b._fill if isinstance(b, Array) else b for b in (lo, hi)) out = _unary("clip", lambda v: np.clip(v, lo_v, hi_v), x) else: - out = x - if lo is not None: - out = maximum(out, lo) - if hi is not None: - out = minimum(out, hi) - if out.dtype != x.dtype: - out = cast(out, x.dtype) - return out + out = _clip_arrays(x, lo, hi, operands) + return out if out.dtype == x.dtype else cast(out, x.dtype) + + +def _clip_arrays(x, lo, hi, operands): + """Clip with array bounds as numpy does, over the union of the entries (aligned values).""" + shape = np.broadcast_shapes(*[b.shape for b in operands]) + if shape != x.shape: + raise ValueError(f"clip's bounds do not broadcast to the shape {x.shape}") + is_array = [isinstance(b, Array) and b.ndim > 0 for b in (lo, hi)] + + def value(b): # (a bound's value where it stores nothing, typed as numpy types it) + return np.asarray(b._fill) if isinstance(b, Array) else b + + def probe(b): + return np.empty(0, b.dtype) if isinstance(b, Array) else b + + with np.errstate(all="ignore"): + rdt = np.clip(np.empty(0, x.dtype), probe(lo), probe(hi)).dtype + fill = np.clip(np.asarray(x._fill), value(lo), value(hi)) + + def func(v, *arrays): + it = iter(arrays) + lo_v, hi_v = ( + next(it) if is_nd else value(b) for b, is_nd in zip((lo, hi), is_array, strict=True) + ) + return np.clip(v, lo_v, hi_v) + + tensors = _utils.aligned(operands, shape) + return _result(_utils.map_values(func, tensors, rdt), fill) def _clip_bound(bound, x): diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index e6ca9ab80..c7c2c898e 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -103,6 +103,17 @@ def expansion_check(Ta, Tb, shape): check_dense(replicated, f"broadcasting shapes {Ta.shape} and {Tb.shape} to {shape}") +def mult_expansion_check(Ta, Tb, shape): + """Before an intersection (ewise_mult) of ``Ta`` and ``Tb`` broadcast to ``shape``: when + both broadcast, one is replicated (an outer product); raise past the limit. + """ + fa = _broadcast_factor(Ta.shape, shape) + fb = _broadcast_factor(Tb.shape, shape) + if fa > 1 and fb > 1: + replicated = min(fa * Ta.nvals, fb * Tb.nvals) + check_dense(replicated, f"an outer product of shapes {Ta.shape} and {Tb.shape}") + + def empty_tensor(dtype, shape): return Tensor(gb_dtype(dtype), tuple(shape)) @@ -174,18 +185,31 @@ def tensor_values(T): def summary(x): """Facts about the stored values of the Array ``x``, computed once (cached on ``x``): ``finite`` (all finite), ``nonnegative`` (no sign bit set: no negative value, no -0.0), - ``nan`` (some NaN) and ``negzero`` (some -0.0). + ``nan`` (some NaN), ``negzero`` (some -0.0) and ``zero`` (some 0.0 or -0.0). """ def compute(x): T = x._T() kind = x.dtype.kind - if kind == "b" or T.nvals == 0: - return {"finite": True, "nonnegative": True, "nan": False, "negzero": False} + if T.nvals == 0: + return { + "finite": True, + "nonnegative": True, + "nan": False, + "negzero": False, + "zero": False, + } values = tensor_values(T) - if kind in "iu": - nonneg = kind == "u" or bool((values >= 0).all()) - return {"finite": True, "nonnegative": nonneg, "nan": False, "negzero": False} + if kind in "biu": + nonneg = kind in "bu" or bool((values >= 0).all()) + zero = not bool(values.all()) + return { + "finite": True, + "nonnegative": nonneg, + "nan": False, + "negzero": False, + "zero": zero, + } if kind == "c": parts = [values.real, values.imag] else: @@ -198,6 +222,7 @@ def compute(x): "negzero": any( bool((sb & (p == 0)).any()) for p, sb in zip(parts, signbit, strict=True) ), + "zero": not bool(values.all()), } return x._fact("summary", compute) diff --git a/graphblas/graphblas.yaml b/graphblas/graphblas.yaml index 4711f66b2..5bc7615fd 100644 --- a/graphblas/graphblas.yaml +++ b/graphblas/graphblas.yaml @@ -10,3 +10,7 @@ jit_c_control: null tensor_max_expansion: 1000000000 # graphblas.array_api: most elements an array may hold densely (np.asarray, cumulative_sum, ...) array_api_densify_limit: 100000000 +# graphblas.array_api: keep the sign of every zero a product or quotient makes with an unstored +# zero, as numpy (True), even where that stores every position (a row scaling d[:, None] * A with +# negative d); by default (False) those positions read 0.0, and multiply stores the intersection +array_api_signed_zeros: false diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 128c2c708..884a3fcd9 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -11,6 +11,7 @@ import numpy as np import pytest +import graphblas as gb import graphblas.array_api as xp from graphblas import backend from graphblas.core.tensor import _has_container as tensor_has_container @@ -21,6 +22,15 @@ pytest.skip("graphblas.array_api needs SuiteSparse:GraphBLAS 10", allow_module_level=True) +@pytest.fixture(autouse=True) +def _strict_signed_zeros(): + """These tests compare with numpy bit for bit: keep every signed zero (DESIGN D3a's strict + mode). The default mode, where products with unstored zeros read 0.0, has its own tests. + """ + with gb.config.set(array_api_signed_zeros=True): + yield + + DTYPES = [ np.dtype(name) for name in ( @@ -615,3 +625,76 @@ def test_unsupported_operands(): assert (x == "a") is False with pytest.raises(TypeError): xp.sin(make(np.array([True, False]), False)) # numpy would give float16 + + +def test_intersection_of_broadcast_operands_is_guarded(): + # an intersection of two broadcast operands is an outer product (review finding M4) + import graphblas as gb + + x = make(np.arange(1, 101).reshape(100, 1), 0) + y = make(np.arange(1, 101).reshape(1, 100), 0) + b = x > 0 + for func, args in [ + (xp.multiply, (x, y)), + (xp.bitwise_and, (x, y)), + (xp.logical_and, (b, b.mT)), + ]: + with gb.config.set(array_api_densify_limit=1000), pytest.raises(xp.DensifyError): + func(*args) + assert func(*args).nnz == 10000 + + +def test_clip_with_array_bounds_ties(): + # numpy's clip gives the bound on a tie, so the sign of a zero is the bound's (review M6) + x = make(np.array([0.0, -0.0, 0.0, 1.0, np.nan]), 0.0) + lo = make(np.array([-0.0, 0.0, 0.0, -0.0, 0.0]), 0.0) + hi = make(np.array([-0.0, 0.0, 2.0, 0.5, 1.0]), 0.0) + a, b, c = x.todense(), lo.todense(), hi.todense() + assert_array_equal(xp.clip(x, min=lo, max=hi), np_call(np.clip, a, b, c)) + assert_array_equal(xp.clip(x, min=lo), np_call(np.clip, a, b, None)) + assert_array_equal(xp.clip(x, max=hi), np_call(np.clip, a, None, c)) + assert_array_equal(xp.clip(x, min=lo, max=0.5), np_call(np.clip, a, b, 0.5)) + with pytest.raises(ValueError, match="broadcast"): + xp.clip(x, min=make(np.zeros((2, 5)), 0.0)) + + +def test_signed_zeros_give_way_by_default(): + # (review finding H1, DESIGN D3a) by default a product with an unstored zero is the fill + # whatever its sign: a row scaling stores only A's entries + rng = np.random.default_rng(0) + A = make(np.where(rng.random((50, 50)) < 0.05, rng.normal(size=(50, 50)), 0.0), 0.0) + d = make(rng.normal(size=(50, 1)), 0.0) + exact = d * A + assert_array_equal(exact, d.todense() * A.todense()) + assert exact.nnz == 2500 + with gb.config.set(array_api_signed_zeros=False): + loose = d * A + assert loose.nnz == A.nnz + np.testing.assert_array_equal(loose.todense(), d.todense() * A.todense()) # (0.0 == -0.0) + # inf and NaN stay exact + inf = make(np.array([[np.inf, 0.0]]), 0.0) + assert_array_equal(inf * make(np.array([[0.0, 1.0]]), 0.0), np.array([[np.nan, 0.0]])) + # divide: the numerator's entries, when no value of the denominator is 0 or NaN + a = np.where(rng.random((4, 5)) < 0.5, 0.0, rng.normal(size=(4, 5))) + b = rng.normal(size=5) + 3.0 + for fb in [1.0, -2.0]: + q = make(a, 0.0) / make(b, fb) + assert q.nnz == np.count_nonzero(a) + np.testing.assert_array_equal(q.todense(), a / b) + q = make(a, 0.0) / make(np.r_[b[:4], 0.0], 1.0) + np.testing.assert_array_equal( + np_call(np.isnan, q.todense()), np.isnan(np_call(np.divide, a, np.r_[b[:4], 0.0])) + ) + + +@pytest.mark.parametrize("name", ["multiply", "divide"]) +def test_default_mode_matches_numpy_but_for_zero_signs(name): + func = getattr(np, name) + with gb.config.set(array_api_signed_zeros=False): + for shapes in [((5,), (5,)), ((3, 4), (4,)), ((3, 1), (1, 4))]: + for dt1, dt2 in [(np.dtype("float64"),) * 2, (np.dtype("float32"), np.dtype("int8"))]: + for dense1, x1, dense2, x2 in _binary_cases(dt1, dt2, shapes, 5): + out = getattr(xp, name)(x1, x2).todense() + expected = np_call(func, dense1, dense2) + assert out.dtype == expected.dtype + np.testing.assert_array_equal(out, expected) # (NaN == NaN, 0.0 == -0.0) diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 3ea41bb1f..9a9e3948e 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -8,6 +8,7 @@ from graphblas.array_api import _utils from graphblas.core.tensor import Tensor +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, make @@ -204,7 +205,9 @@ def test_scalar_cache_is_bounded(): def test_summary_of_complex_values(): x = make(np.array([1 + 0j, complex(-0.0, 1), complex(np.nan, 0)]), 0j) facts = _utils.summary(x) - assert facts == {"finite": False, "nonnegative": False, "nan": True, "negzero": True} + assert facts == { + "finite": False, "nonnegative": False, "nan": True, "negzero": True, "zero": False + } # fmt: skip def test_thread3_argument_checks(): diff --git a/graphblas/tests/test_array_api_indexing.py b/graphblas/tests/test_array_api_indexing.py index 243fa41d4..f55c0c028 100644 --- a/graphblas/tests/test_array_api_indexing.py +++ b/graphblas/tests/test_array_api_indexing.py @@ -5,6 +5,7 @@ import graphblas.array_api as xp +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, inputs, make rng = np.random.default_rng(0) diff --git a/graphblas/tests/test_array_api_linalg.py b/graphblas/tests/test_array_api_linalg.py index f8d545bb3..53e1792ef 100644 --- a/graphblas/tests/test_array_api_linalg.py +++ b/graphblas/tests/test_array_api_linalg.py @@ -8,6 +8,7 @@ import graphblas as gb import graphblas.array_api as xp +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, check_invariants, make, np_call rng = np.random.default_rng(21) diff --git a/graphblas/tests/test_array_api_manipulation.py b/graphblas/tests/test_array_api_manipulation.py index 8fc56c488..4b002e55b 100644 --- a/graphblas/tests/test_array_api_manipulation.py +++ b/graphblas/tests/test_array_api_manipulation.py @@ -5,6 +5,7 @@ import graphblas.array_api as xp +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, inputs, make diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py index 578b2ccbf..2fb53d07a 100644 --- a/graphblas/tests/test_array_api_searching.py +++ b/graphblas/tests/test_array_api_searching.py @@ -5,6 +5,7 @@ import graphblas.array_api as xp +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import DTYPES, assert_array_equal, inputs, make, np_call WHERE_SHAPES = [ diff --git a/graphblas/tests/test_array_api_set.py b/graphblas/tests/test_array_api_set.py index 6bcd36084..da41fb1ee 100644 --- a/graphblas/tests/test_array_api_set.py +++ b/graphblas/tests/test_array_api_set.py @@ -8,6 +8,7 @@ import graphblas as gb import graphblas.array_api as xp +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, check_invariants, inputs, make DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint32", "int64", "float32", "float64")] diff --git a/graphblas/tests/test_array_api_sorting.py b/graphblas/tests/test_array_api_sorting.py index 85cba8519..3a3f8c85e 100644 --- a/graphblas/tests/test_array_api_sorting.py +++ b/graphblas/tests/test_array_api_sorting.py @@ -6,6 +6,7 @@ import graphblas as gb import graphblas.array_api as xp +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, inputs, make DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint16", "int64", "float32", "float64")] From 2a5fba9b2a82618db34d20b061be86c35fb3584d Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:51:46 -0600 Subject: [PATCH 23/77] array_api: reductions as numpy for signed zeros, complex infinities and bool sums - sum starts from +0.0 as numpy's does, so a sum of -0.0 is 0.0. - min and max of arrays holding -0.0 (or with a -0.0 fill) reduce with numpy's minimum and maximum, which order -0.0 below 0.0; GraphBLAS's fmin and fmax take either zero. - k copies of a complex fill add part by part: inf+0j summed is inf+0j, not inf+nanj. - sum and prod with dtype=bool are logical or and and. - the counts of unstored positions no longer overflow for fibers of 2**63 positions or more. - mean of an empty axis is 0/0 as numpy's: nan, or nan+nanj for complex. - diff with n=0 ignores prepend and append, as numpy's does. - asarray of a list of Arrays stacks them (sparse) instead of densifying each, and takes numpy arrays of either byte order. --- graphblas/array_api/_creation.py | 25 ++++ graphblas/array_api/_statistical.py | 123 +++++++++++++++--- graphblas/tests/test_array_api_creation.py | 29 +++++ graphblas/tests/test_array_api_statistical.py | 103 +++++++++++++++ 4 files changed, 259 insertions(+), 21 deletions(-) diff --git a/graphblas/array_api/_creation.py b/graphblas/array_api/_creation.py index 3908de60e..6441c51ee 100644 --- a/graphblas/array_api/_creation.py +++ b/graphblas/array_api/_creation.py @@ -156,17 +156,42 @@ def asarray(obj, /, *, dtype=None, device=None, copy=None, fill_value=None): return _from_sparse(obj, dtype, fill_value) dt = None if dtype is None else as_dtype(dtype) if isinstance(obj, (list, tuple)): + stacked = _stack_arrays(obj) + if stacked is not None: + return asarray(stacked, dtype=dtype, fill_value=fill_value) obj = _unwrap_nested(obj) if dt is None and isinstance(obj, builtins.bool): dt = np.dtype(bool) elif dt is None and isinstance(obj, int) and not isinstance(obj, np.generic): dt = np.dtype("int64") a = np.asarray(obj, dtype=dt) + if not a.dtype.isnative: + a = a.astype(a.dtype.newbyteorder("=")) if dt is None and a.dtype.kind == "f" and a.dtype.itemsize < 4: a = a.astype(np.float32) return _from_numpy(a, fill_value) +def _stack_arrays(obj): + """A sequence of n-d Arrays of one shape (or such sequences), stacked sparsely; None for + anything else. + """ + from ._manipulation import stack + + if not obj: + return None + parts = [] + for o in obj: + if isinstance(o, (list, tuple)): + o = _stack_arrays(o) + if not isinstance(o, Array) or o.ndim == 0: + return None + parts.append(o) + if len({p.shape for p in parts}) != 1: + return None + return stack(parts) + + def _unwrap_nested(obj): """A (nested) sequence with its 0-d Arrays replaced by their values (numpy would densify every Array through ``__array__``, which works too, but slower). diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index 85cd28d5e..8dc2cdc7b 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -93,6 +93,35 @@ def _int_power(base, exponent, dtype): return np.asarray(value).astype(dtype)[()] +def _unstored(L, c, ct): + """``L - c`` (counts of unstored positions), in a dtype that does not overflow: uint64 + for integer sums (which wrap anyway), float64 otherwise (a fiber may hold 2**63 or more + positions). + """ + if ct.kind in "iub": + return np.uint64(L % (1 << 64)) - np.asarray(c).astype(np.uint64) + return np.float64(L) - c + + +def _scaled(k, fill, ct): + """``fill`` added ``k`` times (k >= 0), as repeated additions give it: wrapping for + integers, logical or for bool, and part by part for complex (0 * inf in one part must not + make NaN in the other). + """ + with np.errstate(all="ignore"): + if ct.kind in "iu": + f = np.asarray(fill).astype(np.int64 if ct.kind == "i" else np.uint64) + return (np.asarray(k, dtype=np.uint64) * f.astype(np.uint64)).astype(ct) + if ct.kind == "b": + return np.logical_and(np.asarray(k) > 0, fill) + if ct.kind == "c": + out = np.empty(np.shape(k), ct) + out.real = k * fill.real + out.imag = k * fill.imag + return out + return np.asarray(np.multiply(k, fill)).astype(ct) + + class _Monoid: def __init__(self, name, gb, npfunc, identity): self.name = name @@ -102,17 +131,16 @@ def __init__(self, name, gb, npfunc, identity): def repeat(self, fill, L, ct): """The monoid over ``L`` copies of ``fill``.""" - with np.errstate(all="ignore"): - if self.name == "sum": - if ct.kind in "iu": - return np.multiply(fill, np.asarray(L % (1 << 64)).astype(ct), dtype=ct) - return np.multiply(fill, L, dtype=ct) - if self.name == "prod": + if self.name == "sum": + k = np.uint64(L % (1 << 64)) if ct.kind in "iub" else np.float64(L) + return _scaled(k, fill, ct)[()] + if self.name == "prod": + with np.errstate(all="ignore"): if ct.kind in "iu": return _int_power(fill, L, ct) if ct.kind == "b": return fill - return np.power(fill, ct.type(L), dtype=ct) + return np.power(fill, np.float64(L), dtype=ct) return fill def combine(self, L, fill, ct): @@ -121,17 +149,14 @@ def combine(self, L, fill, ct): """ name = self.name if name == "sum": - if ct.kind in "iu": - return lambda r, c: r + (L - c).astype(ct) * fill # (a fiber with no unstored position must not see 0 * inf) - return lambda r, c: np.where(c < L, r + (L - c) * fill, r) + return lambda r, c: np.where(c < L, r + _scaled(_unstored(L, c, ct), fill, ct), r) if name == "prod": if ct.kind in "iu": - big = np.dtype("uint64" if ct.kind == "u" else "int64") - return lambda r, c: r * np.power(big.type(fill), (L - c).astype(big)).astype(ct) + return lambda r, c: r * _int_powers(fill, _unstored(L, c, ct), ct) if ct.kind == "b": return lambda r, c: np.where(c < L, r & fill, r) - return lambda r, c: np.where(c < L, r * np.power(fill, (L - c).astype(ct)), r) + return lambda r, c: np.where(c < L, r * np.power(fill, _unstored(L, c, ct)), r) func = { "min": np.minimum, "max": np.maximum, @@ -141,6 +166,13 @@ def combine(self, L, fill, ct): return lambda r, c: np.where(c < L, func(r, fill), r) +def _int_powers(base, exponents, dtype): + """``base ** e`` for each count ``e`` (uint64), wrapped to the integer ``dtype``.""" + big = np.dtype("uint64" if dtype.kind == "u" else "int64") + with np.errstate(all="ignore"): + return np.power(big.type(base), exponents.astype(big)).astype(dtype) + + def _min_identity(dt): if dt.kind == "f": return np.inf @@ -168,6 +200,16 @@ def _max_identity(dt): def _reduce(name, x, axis, keepdims, dtype=None): + out = _reduce_main(name, x, axis, keepdims, dtype) + if name == "sum" and out.dtype.kind in "fc": + # numpy's sum starts from +0.0: a sum of -0.0 is 0.0 (GraphBLAS starts from the first) + from ._elementwise import add + + out = add(out, 0.0) + return out + + +def _reduce_main(name, x, axis, keepdims, dtype): _check_array(x, name) spec = _MONOIDS[name] if name in {"min", "max"} and x.dtype.kind == "c": @@ -202,11 +244,15 @@ def _reduce(name, x, axis, keepdims, dtype=None): fill_out = spec.repeat(fill, L, ct) if T.nvals == 0 or math.prod(shape) == 0: return _keepdims(Array._new(Tensor(gb_dtype(ct), shape), fill_out), x, axes, keepdims) + if name in {"min", "max"} and ct.kind == "f" and _signed_zeros(x, fill): + # GraphBLAS's min and max are C's fmin and fmax, which may take either of 0.0 and -0.0 + out = _reduce_numpy(x, axes, np.minimum if name == "min" else np.maximum, fill, ct) + return _keepdims(out, x, axes, keepdims) nan_fix = name in {"min", "max"} and ct.kind == "f" and _has_nan(x) full = T.nvals == x.size if len(axes) == x.ndim: # Over every axis: one Scalar, and the count is the number of entries - r = T.reduce_scalar(spec.gb[gb_dtype(ct)]).new().value + r = T.reduce_scalar(_gb_monoid(spec, ct)).new().value if nan_fix: value = np.nan elif full or _utils.same_value(fill, spec.identity(ct)): @@ -217,7 +263,7 @@ def _reduce(name, x, axis, keepdims, dtype=None): value = spec.combine(L, fill, ct)(np.array([r], dtype=ct), count)[0] out = Array._new(Tensor(gb_dtype(ct), ()), value) return _keepdims(out, x, axes, keepdims) - R = _reduce_tensor(T, spec.gb[gb_dtype(ct)], axes) + R = _reduce_tensor(T, _gb_monoid(spec, ct), axes) if not nan_fix and (full or _utils.same_value(fill, spec.identity(ct))): out = Array._new(R, fill if full else fill_out) else: @@ -249,6 +295,43 @@ def _reduce_no_axes(name, x, ct): return cast(x, ct) if x.dtype != ct else x._copy() +def _gb_monoid(spec, ct): + # (a sum of booleans in bool is their logical or, a product their logical and) + if ct.kind == "b" and spec.name in {"sum", "prod"}: + return (monoid.lor if spec.name == "sum" else monoid.land)[gb_dtype(ct)] + return spec.gb[gb_dtype(ct)] + + +def _signed_zeros(x, fill): + """Whether a min or max might meet both zeros: a stored -0.0, or a fill of -0.0.""" + return _utils.summary(x)["negzero"] or bool(fill == 0 and np.signbit(fill)) + + +def _reduce_numpy(x, axes, ufunc, fill, ct): + """``ufunc.reduce`` over ``axes`` by numpy on the stored entries of each fiber, then with + the fill once where a fiber has an unstored position: O(entries log entries). + """ + from ._manipulation import permute_dims + + others = tuple(a for a in range(x.ndim) if a not in axes) + shape = tuple(x.shape[a] for a in others) + L = math.prod(x.shape[a] for a in axes) + coords, values = permute_dims(x, others + tuple(axes))._T().to_coo() + if shape: + fiber = np.ravel_multi_index(coords[: len(others)], shape) + else: + fiber = np.zeros(len(values), np.int64) + starts = np.flatnonzero(np.r_[True, fiber[1:] != fiber[:-1]]) + counts = np.diff(np.r_[starts, len(fiber)]) + with np.errstate(all="ignore"): + r = ufunc.reduceat(values, starts) + r = np.where(counts < L, ufunc(r, fill), r).astype(ct) + if not shape: + return Array._new(Tensor(gb_dtype(ct), ()), r[0]) + idx = np.unravel_index(fiber[starts], shape) + return Array._new(Tensor.from_coo(idx, r, gb_dtype(ct), shape=shape), fill) + + def _has_nan(x): return _utils.summary(x)["nan"] @@ -287,10 +370,8 @@ def mean(x, /, *, axis=None, keepdims=False): _check_array(x, "mean") rdt = _mean_dtype(x, "mean") L = math.prod(x.shape[a] for a in _axes(x, axis)) - total = _reduce("sum", x, axis, keepdims, rdt) - if L == 0: - return Array._new(Tensor(total._T().dtype, total.shape), np.nan) - return divide(total, L) + # (an empty axis divides 0 by 0, as numpy does: nan, or nan+nanj for complex) + return divide(_reduce("sum", x, axis, keepdims, rdt), L) def _squared_magnitude(x): @@ -413,11 +494,11 @@ def diff(x, /, *, axis=-1, n=1, prepend=None, append=None): n = int(n) if n < 0: raise ValueError(f"diff's n must be nonnegative, not {n}") + if n == 0: + return x._copy() # (numpy's diff ignores prepend and append then) parts = [_diff_part(p, x, axis) for p in (prepend, x, append) if p is not None] if len(parts) > 1: x = concat(parts, axis=axis) - if n == 0: - return x._copy() # (numpy's diff of booleans is not_equal: subtracting booleans is an error) op = not_equal if x.dtype.kind == "b" else subtract before = (slice(None),) * axis diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py index d67356301..8a14cec04 100644 --- a/graphblas/tests/test_array_api_creation.py +++ b/graphblas/tests/test_array_api_creation.py @@ -12,6 +12,7 @@ import graphblas.array_api as xp from graphblas import Matrix, Scalar, Tensor, Vector +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import ( DTYPES, assert_array_equal, @@ -527,3 +528,31 @@ def test_permute_expand_squeeze(): assert_array_equal(xp.expand_dims(xp.asarray(3), axis=0), np.array([3])) assert_array_equal(xp.permute_dims(xp.asarray(3), ()), np.array(3)) check_invariants(xp.squeeze(one, axis=(0, 1))) + + +def test_asarray_of_arrays_stays_sparse(): + # a list of Arrays is stacked, not densified (review finding L7) + import graphblas as gb + + x = xp.zeros((10**5, 10**5)) + x[1, 2] = 3.0 + s = xp.asarray([x, x]) + assert s.shape == (2, 10**5, 10**5) + assert s.nnz == 2 + assert_array_equal( + xp.asarray([[make(np.array([1, 0]), 0), make(np.array([0, 2]), 0)]]), + np.array([[[1, 0], [0, 2]]]), + ) + assert_array_equal(xp.asarray([xp.asarray([1, 2]), [3, 4]]), np.array([[1, 2], [3, 4]])) + assert_array_equal( + xp.asarray([make(np.array([1.0]), 0.0)], dtype=xp.float32), np.array([[1.0]], np.float32) + ) + del gb + + +def test_asarray_byte_order(): + # (review finding L9) + assert_array_equal(xp.asarray(np.array([1, 2], dtype=">i8")), np.array([1, 2])) + assert_array_equal( + xp.asarray(np.array([1.5, 0.0], dtype=">f4")), np.array([1.5, 0.0], np.float32) + ) diff --git a/graphblas/tests/test_array_api_statistical.py b/graphblas/tests/test_array_api_statistical.py index dbc634bf8..d9c1a82e2 100644 --- a/graphblas/tests/test_array_api_statistical.py +++ b/graphblas/tests/test_array_api_statistical.py @@ -8,6 +8,7 @@ import graphblas as gb import graphblas.array_api as xp +from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import ( DTYPES, assert_array_equal, @@ -261,3 +262,105 @@ def test_var_of_a_huge_sparse_array(): mean = values.sum() / n expected = (((values - mean) ** 2).sum() + (n - values.size) * mean**2) / n assert float(v[col]) == pytest.approx(expected, rel=1e-12) + + +def test_sum_starts_from_positive_zero(): + # numpy's sum starts from +0.0: a sum of -0.0 is 0.0 (review finding M1) + for a in [ + np.array([-0.0, -0.0]), + np.array([[-0.0, 0.0], [-0.0, 1.0]]), + np.array([-0.0], np.float32), + ]: + x = make(a, 0.0) + for axis in [None, -1]: + assert_array_equal(xp.sum(x, axis=axis), np.sum(a, axis=axis)) + assert_array_equal(xp.mean(x, axis=axis), np.mean(a, axis=axis)) + c = np.array([complex(-0.0, -0.0)]) + assert_array_equal(xp.sum(make(c, 0j)), np.sum(c)) + x = make(np.array([1.0, -0.0]), -0.0) + assert_array_equal(xp.sum(x, axis=0), np.sum(np.array([1.0, -0.0]))) + assert_array_equal(xp.sum(xp.full((3,), -0.0)), np.sum(np.full(3, -0.0))) + + +def test_min_max_order_signed_zeros(): + # GraphBLAS's fmin/fmax take either zero; numpy orders -0.0 below 0.0 (review finding M2) + cases = [ + (np.array([-0.0, 0.0]), 1.0), + (np.array([0.0, -0.0]), 1.0), + (np.array([[0.0, 2.0], [-0.0, 3.0]]), 5.0), + (np.array([[np.nan, -0.0], [np.nan, 0.0]]), np.nan), + (np.array([[0.0, 1.0], [2.0, 0.0]]), -0.0), + (np.array([[-0.0, 1.0], [2.0, -0.0]]), 0.0), + ] + for a, fill in cases: + for x in [make(a, fill), make(np.where(a == 0, fill, a), fill)]: + dense = x.todense() + for name in ["min", "max"]: + for axis in [None, 0, -1]: + assert_array_equal( + getattr(xp, name)(x, axis=axis), getattr(np, name)(dense, axis=axis) + ) + + +def test_complex_sum_of_infinite_fill(): + # k copies of inf+0j sum to inf+0j, not inf+nanj (review finding M3) + for fill in [complex(np.inf, 0), complex(np.nan, 0), complex(0, -np.inf)]: + a = np.array([[1 + 1j, fill, fill], [fill, fill, fill]]) + x = make(a, fill) + for axis in [None, 0, 1]: + for name in ["sum", "mean"]: + assert_array_equal( + np_call(getattr(xp, name), x, axis=axis), + np_call(getattr(np, name), a, axis=axis), + ) + + +def test_sum_in_bool(): + # a sum of booleans in bool is their logical or (review finding M5) + a = np.array([[1, 0, 3], [0, 0, 2]]) + for fill in [0, 2]: + x = make(np.where(a == 0, fill, a), fill) + dense = x.todense() + for axis in [None, 0, 1]: + assert_array_equal( + xp.sum(x, axis=axis, dtype=xp.bool), np.sum(dense, axis=axis, dtype=bool) + ) + assert_array_equal( + xp.prod(x, axis=axis, dtype=xp.bool), np.prod(dense, axis=axis, dtype=bool) + ) + + +def test_mean_of_empty_complex(): + # 0j / 0 is nan+nanj, as numpy's mean of nothing (review finding L1) + for dt in [np.complex128, np.complex64, np.float32]: + assert_array_equal( + np_call(xp.mean, xp.zeros((0,), dtype=dt)), np_call(np.mean, np.zeros(0, dt)) + ) + assert_array_equal( + np_call(xp.mean, xp.zeros((2, 0), dtype=dt), axis=1), + np_call(np.mean, np.zeros((2, 0), dt), axis=1), + ) + + +def test_diff_n0_ignores_prepend_append(): + # (review finding L4) + a = np.array([1.0, 0.0, 2.0]) + x = make(a, 0.0) + assert_array_equal( + xp.diff(x, n=0, prepend=make(np.array([5.0]), 0.0)), np.diff(a, n=0, prepend=5.0) + ) + assert_array_equal(xp.diff(x, n=0, append=7.0), np.diff(a, n=0, append=7.0)) + + +def test_reductions_of_fibers_past_2_63_positions(): + # (review finding L6) a fiber of 2**63 positions: counts do not overflow + shape = (2**21,) * 3 + y = xp.full(shape, 1.0) + y[0, 0, 0] = 3.0 + assert float(xp.sum(y)) == float(2**63 + 2) + assert float(xp.sum(y, axis=(1, 2))[0]) == 2.0**42 + 2 + assert float(xp.max(y)) == 3.0 + z = xp.full(shape, 1, dtype=xp.int64) + z[0, 0, 0] = 3 + assert int(xp.sum(z)) == 3 + (2**63 - 1) - 2**64 # (wraps, as integer sums do) + assert int(xp.prod(z)) == 3 From 084b647968b1038d6b455cbfd0319a8c2b3f0876 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:53:55 -0600 Subject: [PATCH 24/77] array_api: cumulative_sum and cumulative_prod stay sparse when the fill is the identity With fill 0 (sum) or 1 (product), each fiber holds the fill up to its first entry, so only the positions from there on are stored. numpy accumulates one position at a time, so a run of unstored positions adds the fill once more (0.0 + -0.0 is 0.0): each run becomes one element of the fiber's sequence, and numpy accumulates the sequences in order, for numpy's rounding and signed zeros. Other fills still go through numpy densely. --- graphblas/array_api/_statistical.py | 96 +++++++++++++++++++ graphblas/tests/test_array_api_statistical.py | 37 ++++++- 2 files changed, 132 insertions(+), 1 deletion(-) diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index 8dc2cdc7b..a44446e4b 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -10,6 +10,7 @@ """ import builtins +import itertools import math import numpy as np @@ -462,6 +463,11 @@ def _cumulative(name, npfunc, x, axis, dtype, include_initial): ct = check_supported(_MONOIDS[reduce_name].npfunc(np.zeros(1, x.dtype)).dtype, name) else: ct = as_dtype(dtype) + with np.errstate(all="ignore"): + fill = np.asarray(x._fill).astype(ct) + identity = 0 if name == "cumulative_sum" else 1 + if x.size and fill == identity and not (include_initial and np.signbit(fill)): + return _cumulative_sparse(npfunc, x, axis, ct, fill, include_initial) _utils.check_dense(x.size + (x.size // builtins.max(x.shape[axis], 1)), name) with np.errstate(all="ignore"): out = npfunc(x.todense(), axis=axis, dtype=ct) @@ -475,6 +481,96 @@ def _cumulative(name, npfunc, x, axis, dtype, include_initial): return _from_numpy(out) +def _cumulative_sparse(npfunc, x, axis, ct, fill, include_initial): + """A cumulative sum (product) along ``axis`` when the fill is the identity, 0 (1): each + fiber is the fill up to its first entry, so only the positions from there on are stored, + O(sum over fibers of the positions from the first entry on). + + numpy accumulates one position at a time, so each run of unstored positions adds the fill + once more (0.0 + -0.0 is 0.0): each run becomes one element of the fiber's sequence, and + the sequences are accumulated by numpy itself, in order, for numpy's rounding. + """ + from ._manipulation import permute_dims + + T = x._T() + L = x.shape[axis] + others = tuple(a for a in range(x.ndim) if a != axis) + shape = tuple(x.shape[a] for a in others) + out_L = L + 1 if include_initial else L + perm = (*others, axis) + out_shape = (*shape, out_L) + if T.nvals == 0: + out = Array._new(Tensor(gb_dtype(ct), out_shape), fill) + return permute_dims(out, tuple(int(i) for i in np.argsort(perm))) + coords, values = permute_dims(x, perm)._T().to_coo() + with np.errstate(all="ignore"): + values = values.astype(ct) + fiber = np.ravel_multi_index(coords[:-1], shape) if shape else np.zeros(len(values), np.int64) + pos = coords[-1] + first = np.r_[True, fiber[1:] != fiber[:-1]] + last = np.r_[fiber[1:] != fiber[:-1], True] + prev = np.where(first, -1, np.r_[0, pos[:-1]]) + # The sequence per fiber: [fill if a run precedes the entry] + entry, ..., [fill if a run + # follows the last entry] + gap = pos > prev + 1 + tail = last & (pos < L - 1) + n_elem = 1 + gap + tail + seq_index = np.cumsum(n_elem) - n_elem # (where each entry's elements start) + seq = np.full(int(n_elem.sum()), fill, dtype=ct) + entry_at = seq_index + gap + seq[entry_at] = values + starts = np.flatnonzero(first) + acc = _grouped_accumulate(npfunc, seq, seq_index[starts], ct) + # Positions and values: each entry; the run before it (not before the first entry: there + # the values are the fill); the run after the last entry + out_pos, out_val, out_fib = [pos], [acc[entry_at]], [fiber] + inner = gap & ~first + if inner.any(): + lengths = (pos - prev - 1)[inner] + out_pos.append(_runs(prev[inner] + 1, lengths)) + out_val.append(np.repeat(acc[seq_index[inner]], lengths)) + out_fib.append(np.repeat(fiber[inner], lengths)) + if tail.any(): + lengths = (L - 1 - pos)[tail] + out_pos.append(_runs(pos[tail] + 1, lengths)) + out_val.append(np.repeat(acc[entry_at[tail] + 1], lengths)) + out_fib.append(np.repeat(fiber[tail], lengths)) + out_pos = np.concatenate(out_pos) + (1 if include_initial else 0) + _utils.check_dense(len(out_pos), "a cumulative sum or product (the positions after an entry)") + idx = list(np.unravel_index(np.concatenate(out_fib), shape)) if shape else [] + idx.append(out_pos) + R = Tensor.from_coo(idx, np.concatenate(out_val), gb_dtype(ct), shape=out_shape) + return permute_dims(Array._new(R, fill), tuple(int(i) for i in np.argsort(perm))) + + +def _runs(starts, lengths): + """Concatenated ranges ``start, start + 1, ..., start + length - 1``.""" + offsets = np.arange(int(lengths.sum())) - np.repeat(np.cumsum(lengths) - lengths, lengths) + return np.repeat(starts, lengths) + offsets + + +def _grouped_accumulate(npfunc, seq, group_starts, ct): + """The cumsum (cumprod) of each group of ``seq`` by numpy (groups start at ``group_starts``), + each in order, as numpy accumulates one fiber. + """ + bounds = np.r_[group_starts, len(seq)] + lengths = np.diff(bounds) + longest = int(lengths.max()) + out = np.empty_like(seq) + with np.errstate(all="ignore"): + if longest * len(lengths) <= 4 * len(seq) + 1024: + # One 2-d accumulation, the groups padded at their ends (padding comes after) + padded = np.zeros((len(lengths), longest), dtype=ct) + col = np.arange(len(seq)) - np.repeat(group_starts, lengths) + row = np.repeat(np.arange(len(lengths)), lengths) + padded[row, col] = seq + out[:] = npfunc(padded, axis=1, dtype=ct)[row, col] + else: + for a, b in itertools.pairwise(bounds): + out[a:b] = npfunc(seq[a:b], dtype=ct) + return out + + def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False): return _cumulative("cumulative_sum", np.cumsum, x, axis, dtype, include_initial) diff --git a/graphblas/tests/test_array_api_statistical.py b/graphblas/tests/test_array_api_statistical.py index d9c1a82e2..ea3b4c7d8 100644 --- a/graphblas/tests/test_array_api_statistical.py +++ b/graphblas/tests/test_array_api_statistical.py @@ -214,7 +214,42 @@ def test_cumulative(name, include_initial): def test_cumulative_is_guarded(): with gb.config.set(array_api_densify_limit=10), pytest.raises(xp.DensifyError): - xp.cumulative_sum(xp.zeros(100)) + xp.cumulative_sum(xp.full(100, 1.0)) + # with the identity as fill only the positions from each fiber's first entry are stored + # (review finding L8) + y = xp.zeros((20000, 20000)) + y[[1, 5, 9], [3, 19000, 7]] = 1.0 + assert xp.cumulative_sum(y, axis=1).nnz == 19997 + 1000 + 19993 + assert xp.cumulative_sum(xp.zeros(10**12)).nnz == 0 + with gb.config.set(array_api_densify_limit=100), pytest.raises(xp.DensifyError): + xp.cumulative_sum(y, axis=1) + + +@pytest.mark.parametrize("name", ["cumulative_sum", "cumulative_prod"]) +def test_cumulative_sparse_exact(name): + # runs of unstored positions add (multiply by) the fill: 0.0 + -0.0 is 0.0, as in numpy + func = np.cumsum if name == "cumulative_sum" else np.cumprod + identity = 0 if name == "cumulative_sum" else 1 + rng = np.random.default_rng(17) + specials = np.array([-0.0, 1.5, -2.0, 3.0, np.inf, np.nan, 0.5]) + for dtype in [np.float64, np.float32, np.int16, np.uint8, np.bool_]: + for shape in [(9,), (3, 7), (2, 3, 5), (1, 40)]: + if np.dtype(dtype).kind == "f": + vals = rng.choice(specials, shape) + else: + vals = rng.integers(0, 4, shape) + a = np.where(rng.random(shape) < 0.6, identity, vals).astype(dtype) + x = make(a, identity, rng) + for axis in range(len(shape)): + for include_initial in [False, True]: + expected = np_call(func, a, axis=axis) + if include_initial: + s = list(expected.shape) + s[axis] = 1 + initial = np.full(s, identity, expected.dtype) + expected = np.concatenate([initial, expected], axis=axis) + out = np_call(getattr(xp, name), x, axis=axis, include_initial=include_initial) + assert_array_equal(out, expected) @pytest.mark.parametrize("n", [0, 1, 2, 5]) From 4e5c3dfb42b531dc5312f55fb19c9a7e4467c8ce Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:57:28 -0600 Subject: [PATCH 25/77] Tensor: ewise_mult of broadcast operands into a shape with no positions The diagonal scaling of a broadcast ewise_mult multiplied matrices of mismatched shapes when the result had a dimension of size 0 (GrB_DIMENSION_MISMATCH, e.g. (2, 1) times (1, 0)). Such a result has no positions, so it now takes the plain path, as ewise_add does. --- graphblas/core/tensor_ewise.py | 5 ++++- graphblas/tests/test_tensor.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/graphblas/core/tensor_ewise.py b/graphblas/core/tensor_ewise.py index 6102f1dd6..c40dd7e09 100644 --- a/graphblas/core/tensor_ewise.py +++ b/graphblas/core/tensor_ewise.py @@ -190,7 +190,10 @@ def _bcast_recipe(updater, method, a, b, op, untyped_op, extra, shape): # result, so "covers the result" must also mean "has its axes". a = _lift_rank(a, len(shape)) b = _lift_rank(b, len(shape)) - if method == "ewise_mult": + if method == "ewise_mult" and math.prod(shape) == 0: + # (no positions: the diagonal scaling below would multiply mismatched shapes) + _ewise_recipe(updater, method, _expand(a, shape), _expand(b, shape), op, extra) + elif method == "ewise_mult": _bcast_mult(updater, a, b, op, untyped_op, shape) else: if census.on: diff --git a/graphblas/tests/test_tensor.py b/graphblas/tests/test_tensor.py index 950fc207f..0b7e4ded4 100644 --- a/graphblas/tests/test_tensor.py +++ b/graphblas/tests/test_tensor.py @@ -2977,6 +2977,21 @@ def test_broadcast(xshape, yshape): check((X > Y).new(), *np_bcast("mult", xv, xp, yv, yp, np.greater)) +@pytest.mark.parametrize( + ("xshape", "yshape"), [((2, 1), (1, 0)), ((1, 0), (2, 1)), ((0, 1), (1, 3)), ((2, 0, 1), (4,))] +) +def test_broadcast_to_no_positions(xshape, yshape): + # A broadcast result with no positions: ewise_mult's diagonal scaling multiplied + # mismatched shapes (GrB_DIMENSION_MISMATCH) + X = Tensor.from_dense(np.ones(xshape)) + Y = Tensor.from_dense(np.ones(yshape)) + shape = np.broadcast_shapes(xshape, yshape) + for method in ["ewise_mult", "ewise_add"]: + R = getattr(X, method)(Y).new() + assert R.shape == shape + assert R.nvals == 0 + + def test_broadcast_operands_and_dtypes(): rng = np.random.default_rng(42) X, xv, xp = random_tensor((3, 4, 5), rng, dtype=dtypes.INT64) From 8af35623d0356fc19401ebdddd26d921f67626de Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 01:59:46 -0600 Subject: [PATCH 26/77] array_api: fixes from the conformance suite: clip broadcasts, from_dlpack of an Array, complex cumsum - clip with array bounds takes the shape x and its bounds broadcast to, as numpy's does (a 0-d x with a bound of shape (0,) is empty). - from_dlpack of an Array of this namespace needs no copy: copy=False returns it, as asarray does. - cumulative_sum of a complex array with include_initial checked the sign of a complex fill with signbit, which takes no complex values. --- graphblas/array_api/_creation.py | 6 ++-- graphblas/array_api/_elementwise.py | 35 ++++++++++--------- graphblas/array_api/_statistical.py | 8 ++++- graphblas/tests/test_array_api_creation.py | 7 ++++ graphblas/tests/test_array_api_elementwise.py | 8 +++-- graphblas/tests/test_array_api_errors.py | 2 -- graphblas/tests/test_array_api_statistical.py | 11 ++++++ 7 files changed, 53 insertions(+), 24 deletions(-) diff --git a/graphblas/array_api/_creation.py b/graphblas/array_api/_creation.py index 6441c51ee..25144966e 100644 --- a/graphblas/array_api/_creation.py +++ b/graphblas/array_api/_creation.py @@ -350,10 +350,10 @@ def linspace(start, stop, /, num, *, dtype=None, device=None, endpoint=True): def from_dlpack(x, /, *, device=None, copy=None): check_device(device) - if copy is False: - raise ValueError("from_dlpack stores its own copy; copy=False") if isinstance(x, Array): - return x._copy() + return x._copy() if copy else x + if copy is False: + raise ValueError("from_dlpack stores its own copy of anything but an Array; copy=False") return _from_numpy(np.from_dlpack(x)) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 14e3d6fc7..ae35aa79c 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -11,6 +11,7 @@ """ import builtins +import math import numpy as np @@ -728,40 +729,42 @@ def clip(x, /, min=None, max=None): hi = _clip_bound(max, x) if lo is None and hi is None: return x._copy() - operands = [x] + [b for b in (lo, hi) if isinstance(b, Array) and b.ndim] - if len(operands) == 1: + if all(not isinstance(b, Array) or b.ndim == 0 for b in (lo, hi)): lo_v, hi_v = (b._fill if isinstance(b, Array) else b for b in (lo, hi)) out = _unary("clip", lambda v: np.clip(v, lo_v, hi_v), x) else: - out = _clip_arrays(x, lo, hi, operands) + out = _clip_arrays(x, lo, hi) return out if out.dtype == x.dtype else cast(out, x.dtype) -def _clip_arrays(x, lo, hi, operands): - """Clip with array bounds as numpy does, over the union of the entries (aligned values).""" - shape = np.broadcast_shapes(*[b.shape for b in operands]) - if shape != x.shape: - raise ValueError(f"clip's bounds do not broadcast to the shape {x.shape}") - is_array = [isinstance(b, Array) and b.ndim > 0 for b in (lo, hi)] +def _clip_arrays(x, lo, hi): + """Clip with array bounds as numpy does, over the union of the entries (aligned values); + the result has the shape the three broadcast to. + """ + operands = (x, lo, hi) + shape = np.broadcast_shapes(*[b.shape for b in operands if isinstance(b, Array)]) + is_array = [isinstance(b, Array) and b.ndim > 0 for b in operands] - def value(b): # (a bound's value where it stores nothing, typed as numpy types it) + def value(b): # (a value where nothing is stored, typed as numpy types it) return np.asarray(b._fill) if isinstance(b, Array) else b def probe(b): return np.empty(0, b.dtype) if isinstance(b, Array) else b with np.errstate(all="ignore"): - rdt = np.clip(np.empty(0, x.dtype), probe(lo), probe(hi)).dtype - fill = np.clip(np.asarray(x._fill), value(lo), value(hi)) + rdt = np.clip(probe(x), probe(lo), probe(hi)).dtype + fill = np.clip(value(x), value(lo), value(hi)) + if math.prod(shape) == 0: + return _result(_utils.empty_tensor(rdt, shape), fill) - def func(v, *arrays): + def func(*arrays): it = iter(arrays) - lo_v, hi_v = ( - next(it) if is_nd else value(b) for b, is_nd in zip((lo, hi), is_array, strict=True) + v, lo_v, hi_v = ( + next(it) if arr else value(b) for b, arr in zip(operands, is_array, strict=True) ) return np.clip(v, lo_v, hi_v) - tensors = _utils.aligned(operands, shape) + tensors = _utils.aligned([b for b, arr in zip(operands, is_array, strict=True) if arr], shape) return _result(_utils.map_values(func, tensors, rdt), fill) diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index a44446e4b..b78ec5737 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -466,7 +466,7 @@ def _cumulative(name, npfunc, x, axis, dtype, include_initial): with np.errstate(all="ignore"): fill = np.asarray(x._fill).astype(ct) identity = 0 if name == "cumulative_sum" else 1 - if x.size and fill == identity and not (include_initial and np.signbit(fill)): + if x.size and fill == identity and not (include_initial and _negative_zero(fill)): return _cumulative_sparse(npfunc, x, axis, ct, fill, include_initial) _utils.check_dense(x.size + (x.size // builtins.max(x.shape[axis], 1)), name) with np.errstate(all="ignore"): @@ -481,6 +481,12 @@ def _cumulative(name, npfunc, x, axis, dtype, include_initial): return _from_numpy(out) +def _negative_zero(fill): + """Whether a zero ``fill`` has a sign bit (an included initial 0 would then be stored).""" + parts = [fill.real, fill.imag] if np.iscomplexobj(fill) else [fill] + return builtins.any(bool(np.signbit(p)) for p in parts if np.asarray(p).dtype.kind == "f") + + def _cumulative_sparse(npfunc, x, axis, ct, fill, include_initial): """A cumulative sum (product) along ``axis`` when the fill is the identity, 0 (1): each fiber is the fill up to its first entry, so only the positions from there on are stored, diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py index 8a14cec04..e7ff27ed7 100644 --- a/graphblas/tests/test_array_api_creation.py +++ b/graphblas/tests/test_array_api_creation.py @@ -248,6 +248,10 @@ def test_from_dlpack(): x = make(a, 0.0) assert_array_equal(xp.from_dlpack(x), a) np.testing.assert_array_equal(np.from_dlpack(x), a) + assert xp.from_dlpack(x, copy=False) is x + assert xp.from_dlpack(x, copy=True) is not x + with pytest.raises(ValueError, match="copy"): + xp.from_dlpack(a, copy=False) # Data type functions @@ -544,6 +548,9 @@ def test_asarray_of_arrays_stays_sparse(): np.array([[[1, 0], [0, 2]]]), ) assert_array_equal(xp.asarray([xp.asarray([1, 2]), [3, 4]]), np.array([[1, 2], [3, 4]])) + ragged = [make(np.array([1, 2]), 0), make(np.array([3]), 0)] + with pytest.raises(ValueError, match="inhomogeneous"): # (as numpy refuses it) + xp.asarray(ragged) assert_array_equal( xp.asarray([make(np.array([1.0]), 0.0)], dtype=xp.float32), np.array([[1.0]], np.float32) ) diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 884a3fcd9..297ff483e 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -654,8 +654,12 @@ def test_clip_with_array_bounds_ties(): assert_array_equal(xp.clip(x, min=lo), np_call(np.clip, a, b, None)) assert_array_equal(xp.clip(x, max=hi), np_call(np.clip, a, None, c)) assert_array_equal(xp.clip(x, min=lo, max=0.5), np_call(np.clip, a, b, 0.5)) - with pytest.raises(ValueError, match="broadcast"): - xp.clip(x, min=make(np.zeros((2, 5)), 0.0)) + # the bounds may broadcast x, as in numpy + lo2 = make(np.array([[0.5] * 5, [-1.0] * 5]), 0.0) + assert_array_equal(xp.clip(x, min=lo2), np_call(np.clip, a, lo2.todense(), None)) + s = xp.asarray(np.uint8(0)) + assert_array_equal(xp.clip(s, max=xp.zeros((0,), dtype=xp.uint8)), np.zeros(0, np.uint8)) + assert_array_equal(xp.clip(xp.asarray(2.0), min=make(np.array([1.0, 3.0]), 0.0)), [2.0, 3.0]) def test_signed_zeros_give_way_by_default(): diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 9a9e3948e..d01a0d4c7 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -58,8 +58,6 @@ def test_shape_errors(): def test_creation_errors(): with pytest.raises(TypeError, match="Python scalar"): xp.full(3, "a") - with pytest.raises(ValueError, match="copy"): - xp.from_dlpack(np.ones(2), copy=False) with pytest.raises(ValueError, match="keeps its fill"): xp.asarray(xp.zeros(2), fill_value=1.0) assert_array_equal(xp.asarray(np.array([1, 2], np.float16)), np.array([1, 2], np.float32)) diff --git a/graphblas/tests/test_array_api_statistical.py b/graphblas/tests/test_array_api_statistical.py index ea3b4c7d8..1182b16fb 100644 --- a/graphblas/tests/test_array_api_statistical.py +++ b/graphblas/tests/test_array_api_statistical.py @@ -225,6 +225,17 @@ def test_cumulative_is_guarded(): xp.cumulative_sum(y, axis=1) +def test_cumulative_sparse_uneven_fibers(): + # one long fiber among many short ones: the fibers accumulate one by one + a = np.zeros((200, 300)) + a[0] = np.arange(1.0, 301.0) + a[np.arange(1, 200), np.arange(1, 200)] = -0.5 + x = make(a, 0.0) + assert_array_equal(xp.cumulative_sum(x, axis=1), np.cumsum(a, axis=1)) + c = xp.zeros(1, dtype=xp.complex64) + assert_array_equal(xp.cumulative_sum(c, include_initial=True), np.zeros(2, np.complex64)) + + @pytest.mark.parametrize("name", ["cumulative_sum", "cumulative_prod"]) def test_cumulative_sparse_exact(name): # runs of unstored positions add (multiply by) the fill: 0.0 + -0.0 is 0.0, as in numpy From 85da672704814f4ce0fe0b059a35de470956565a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 02:16:26 -0600 Subject: [PATCH 27/77] array_api: an operand that stores every position keeps products sparse An Array that stores every position never shows its fill, so the intersection rules now take its fill as 0: A / d with a dense d (fill 0) and A * (1 / deg)[:, None] (fill inf) store A's entries instead of every position. divide by such an array is one broadcast ewise_mult. Fills are computed with numpy's loop for arrays (1-element arrays), which numpy's 0-d clip does not share: it broke a tie between 0.0 and -0.0 the other way. --- graphblas/array_api/_elementwise.py | 54 ++++++++++++++----- graphblas/tests/test_array_api_elementwise.py | 36 +++++++++++++ 2 files changed, 76 insertions(+), 14 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index ae35aa79c..3e328de8e 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -137,8 +137,11 @@ def _fill_of(ufunc, rdt, operands): (integers to a negative power) and one operand is such an Array, the result's fill is never seen either, and 0 stands in for it. """ + # (1-element arrays: numpy's loop for arrays, which a few functions, clip among them, do + # not share with their 0-d path; a Python scalar stays weak) + args = [np.reshape(x._fill, (1,)) if isinstance(x, Array) else x for x in operands] try: - return ufunc(*[_fill_operand(x) for x in operands]) + return np.asarray(ufunc(*args)).reshape(-1)[0] except (ValueError, ArithmeticError): if any(isinstance(x, Array) and x.ndim and x._T().nvals == x.size for x in operands): return _utils.scalar(0, rdt) @@ -485,20 +488,35 @@ def _binary_arrays(spec, a, b, op, ct, rdt, fill): if op is not None: fa = _utils.scalar(a._fill, ct) fb = _utils.scalar(b._fill, ct) - if spec.annihilator is not None and spec.annihilator(a, b, fa, fb, ct): - # Where either operand holds its fill the result is the fill: the intersection - _utils.mult_expansion_check(Ta, Tb, shape) - R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) - return _result(R, fill) + if spec.annihilator is not None: + # An operand that stores every position never shows its fill: take it as 0 + zero = _zero(ct) + fa_seen = zero if _stores_all(a) else fa + fb_seen = zero if _stores_all(b) else fb + if spec.annihilator(a, b, fa_seen, fb_seen, ct): + # Where either operand holds its fill the result is the fill: the intersection + _utils.mult_expansion_check(Ta, Tb, shape) + R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) + with np.errstate(all="ignore"): + return _result(R, spec.ufunc(fa_seen, fb_seen)) + if spec.name == "divide" and Ta.shape == shape: + fb_seen = _utils.scalar(1, ct) if _stores_all(b) else fb + if _divide_left_joins(b, fa, fb_seen): + # The result's entries lie within the numerator's: compute only there + with np.errstate(all="ignore"): + fill = spec.ufunc(fa, fb_seen) + if _stores_all(b): + _utils.mult_expansion_check(Ta, Tb, shape) + return _result(Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)), fill) + _utils.expansion_check(Ta, Tb, shape) + R = Tensor(gb_dtype(rdt), shape) + R(Ta.S) << Ta.ewise_union( + Tb, op, gb_scalar(fa, ct), gb_scalar(fb, ct), allow_expansion=True + ) + return _result(R, fill) _utils.expansion_check(Ta, Tb, shape) expr = Ta.ewise_union(Tb, op, gb_scalar(fa, ct), gb_scalar(fb, ct), allow_expansion=True) - if spec.name == "divide" and Ta.shape == shape and _divide_left_joins(b, fa, fb): - # The result's entries lie within the numerator's: compute only there - R = Tensor(gb_dtype(rdt), shape) - R(Ta.S) << expr - else: - R = expr.new(dtype=gb_dtype(rdt)) - return _result(R, fill) + return _result(expr.new(dtype=gb_dtype(rdt)), fill) # numpy over the values: both operands' values on the union of entries, aligned _utils.expansion_check(Ta, Tb, shape) da, db = a.dtype, b.dtype @@ -527,6 +545,11 @@ def _zero(dtype): return _utils.scalar(0, dtype) +def _stores_all(x): + """Whether ``x`` stores every position, so that its fill is never seen.""" + return x._T().nvals == x.size + + def _both_fills_zero(a, b, fa, fb, ct): zero = _zero(ct) return _utils.same_value(fa, zero) and _utils.same_value(fb, zero) @@ -751,9 +774,12 @@ def value(b): # (a value where nothing is stored, typed as numpy types it) def probe(b): return np.empty(0, b.dtype) if isinstance(b, Array) else b + def value1(b): # (as a 1-element array: numpy's 0-d clip breaks ties otherwise) + return np.reshape(b._fill, (1,)) if isinstance(b, Array) else b + with np.errstate(all="ignore"): rdt = np.clip(probe(x), probe(lo), probe(hi)).dtype - fill = np.clip(value(x), value(lo), value(hi)) + fill = np.clip(value1(x), value1(lo), value1(hi))[0] if math.prod(shape) == 0: return _result(_utils.empty_tensor(rdt, shape), fill) diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 297ff483e..ab3adecf2 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -702,3 +702,39 @@ def test_default_mode_matches_numpy_but_for_zero_signs(name): expected = np_call(func, dense1, dense2) assert out.dtype == expected.dtype np.testing.assert_array_equal(out, expected) # (NaN == NaN, 0.0 == -0.0) + + +def test_operands_storing_every_position_keep_products_sparse(): + # An operand that stores every position never shows its fill (review finding N1): A / d + # and A * (1 / deg)[:, None] store A's entries, whatever d's fill + rng = np.random.default_rng(3) + A = make(np.where(rng.random((40, 40)) < 0.05, rng.normal(size=(40, 40)), 0.0), 0.0) + d = xp.asarray(rng.normal(size=40) + 5.0) + deg = xp.asarray(rng.integers(1, 9, 40).astype(float)) + with gb.config.set(array_api_signed_zeros=False): + q = A / d + assert q.nnz == A.nnz + np.testing.assert_array_equal(q.todense(), A.todense() / d.todense()) + inv = 1 / deg + assert inv.fill_value == np.inf + p = A * xp.reshape(inv, (40, 1)) + assert p.nnz == A.nnz + assert p.fill_value == 0.0 + np.testing.assert_array_equal(p.todense(), A.todense() * inv.todense()[:, None]) + # integers too: a full operand's fill takes no part + i = make(np.where(rng.random((5, 6)) < 0.5, 0, rng.integers(1, 5, (5, 6))), 0) + full = xp.asarray(rng.integers(1, 5, 6), fill_value=7) + assert (i * full).nnz == i.nnz + assert_array_equal(i * full, i.todense() * full.todense()) + # strict mode: the intersection only for values whose product with 0.0 is 0.0 + p = A * xp.reshape(inv, (40, 1)) + assert_array_equal(p, A.todense() * inv.todense()[:, None]) + + +def test_clip_fill_uses_numpys_array_loop(): + # numpy's 0-d clip breaks a tie differently from its array loop (review finding N3) + x = xp.full((3,), 0.0) + lo = xp.full((3,), -0.0) + a, b = np.zeros(3), np.full(3, -0.0) + assert_array_equal(xp.clip(x, min=lo, max=0.5), np.clip(a, b, 0.5)) + assert_array_equal(xp.clip(x, min=-0.0, max=0.5), np.clip(a, -0.0, 0.5)) From 44b074482f70c64f497cca00837a5274ff49d1d5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 02:26:38 -0600 Subject: [PATCH 28/77] array_api: one-sided products stay sparse; fills take numpy's loop for their operands multiply and divide where one operand's fill is 0 and every value of the other makes 0 times (or divided by) it a zero now store only the first operand's entries, whatever the other's fill: A * exp(v)[None, :] with a sparse v (fill 1) and A / d with d of fill 1. A broadcast second operand smaller than the first's entries is stored in full first (O(its size)) so that the product is an intersection, not a replication. The strict mode takes this only where the zeros keep their sign (nonnegative values), which now includes divide by a positive array. A fill is computed as numpy computes the dense array: an n-d operand's fill as a 1-element array and a 0-d operand as 0-d. numpy picks its loop by those shapes: an array to a 0-d power of 0.5 is sqrt (pow(-inf, 0.5) is NaN), between two arrays it is pow (inf). The tests seed their inputs with crc32 of the function's name (hash() of a str varies between runs), which had made this case show up only in some runs. --- graphblas/array_api/_elementwise.py | 98 +++++++++++++------ graphblas/tests/test_array_api_elementwise.py | 52 +++++++++- 2 files changed, 119 insertions(+), 31 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 3e328de8e..ae134630d 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -137,9 +137,18 @@ def _fill_of(ufunc, rdt, operands): (integers to a negative power) and one operand is such an Array, the result's fill is never seen either, and 0 stands in for it. """ - # (1-element arrays: numpy's loop for arrays, which a few functions, clip among them, do - # not share with their 0-d path; a Python scalar stays weak) - args = [np.reshape(x._fill, (1,)) if isinstance(x, Array) else x for x in operands] + # numpy picks its loop by its operands' shapes (an array and a 0-d exponent of 0.5 is + # sqrt, so pow(-inf, 0.5) is NaN there and inf between two arrays; clip breaks ties by + # its 0-d path otherwise): a fill stands in for an array as a 1-element array, a 0-d + # Array stays 0-d, and a Python scalar stays weak + args = [ + ( + (np.reshape(x._fill, (1,)) if x.ndim else np.asarray(x._fill)) + if isinstance(x, Array) + else x + ) + for x in operands + ] try: return np.asarray(ufunc(*args)).reshape(-1)[0] except (ValueError, ArithmeticError): @@ -499,21 +508,10 @@ def _binary_arrays(spec, a, b, op, ct, rdt, fill): R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) with np.errstate(all="ignore"): return _result(R, spec.ufunc(fa_seen, fb_seen)) - if spec.name == "divide" and Ta.shape == shape: - fb_seen = _utils.scalar(1, ct) if _stores_all(b) else fb - if _divide_left_joins(b, fa, fb_seen): - # The result's entries lie within the numerator's: compute only there - with np.errstate(all="ignore"): - fill = spec.ufunc(fa, fb_seen) - if _stores_all(b): - _utils.mult_expansion_check(Ta, Tb, shape) - return _result(Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)), fill) - _utils.expansion_check(Ta, Tb, shape) - R = Tensor(gb_dtype(rdt), shape) - R(Ta.S) << Ta.ewise_union( - Tb, op, gb_scalar(fa, ct), gb_scalar(fb, ct), allow_expansion=True - ) - return _result(R, fill) + if spec.name in {"multiply", "divide"} and ct.kind == "f": + out = _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape) + if out is not None: + return out _utils.expansion_check(Ta, Tb, shape) expr = Ta.ewise_union(Tb, op, gb_scalar(fa, ct), gb_scalar(fb, ct), allow_expansion=True) return _result(expr.new(dtype=gb_dtype(rdt)), fill) @@ -573,17 +571,56 @@ def _multiply_intersects(a, b, fa, fb, ct): return _finite_nonnegative(a) and _finite_nonnegative(b) -def _divide_left_joins(b, fa, fb): - """Whether ``a / b`` stores only ``a``'s entries: ``a``'s fill is 0, and 0 divided by - any value of ``b`` is a zero (b holds no 0 or NaN). The zero's sign then reads as the - fill's, which numpy's may not be: only when signed zeros give way (the default). +def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): + """``a * b`` or ``a / b`` on the entries of one operand only, when that operand's fill + is 0 and every value of the other makes 0 times (or divided by) it a zero: a row scaling + ``A * d[:, None]`` or ``A / d`` stores A's entries. The other operand's fill is computed + at the first operand's entries (a broadcast operand smaller than them is stored in full + first, so that it is not replicated). None when that does not hold. """ - if config.get("array_api_signed_zeros", False) or fa != 0: - return False - if fb == 0 or np.isnan(fb): + strict = config.get("array_api_signed_zeros", False) + sides = [(a, b)] if spec.name == "divide" else [(a, b), (b, a)] + for X, Y in sides: + fx = fa if X is a else fb + fy = fb if X is a else fa + if _stores_all(X) or X._T().shape != shape or fx != 0: + continue + if _stores_all(Y): + fy = _utils.scalar(1, ct) # (never seen) + if not _zero_times(spec, Y, fy, strict) or (strict and np.signbit(fx)): + continue + Tx, Ty = X._T(), Y._T() + with np.errstate(all="ignore"): + fill = spec.ufunc(fx, fy) if X is a else spec.ufunc(fy, fx) + if not _stores_all(Y) and Y.size <= Tx.nvals + 1024: + # (O(size of Y): every position of Y stored, so the product is an intersection) + other = 0 if Y._fill != 0 else 1 + Ty = _utils.with_fill(Y, other)._T() + Ta, Tb = (Tx, Ty) if X is a else (Ty, Tx) + if Ty.nvals == Y.size: + _utils.mult_expansion_check(Ta, Tb, shape) + return _result(Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)), fill) + _utils.expansion_check(Ta, Tb, shape) + R = Tensor(gb_dtype(rdt), shape) + fa_s, fb_s = gb_scalar(fa, ct), gb_scalar(fb, ct) + R(Tx.S) << Ta.ewise_union(Tb, op, fa_s, fb_s, allow_expansion=True) + return _result(R, fill) + return None + + +def _zero_times(spec, y, fy, strict): + """Whether 0 * v (or 0 / v) is a zero for every value v of ``y`` (its fill ``fy`` too): + finite values (for divide also no 0), and in the strict mode nonnegative ones (so that + the zero is +0.0, the fill). NaN never qualifies. + """ + facts = _utils.summary(y) + if spec.name == "divide": + ok = not facts["zero"] and not facts["nan"] and fy != 0 and not np.isnan(fy) + # (strict: 0.0 / -2.0 is -0.0, which an unstored position would read as 0.0) + return ok and (not strict or (facts["nonnegative"] and fy > 0)) + if not facts["finite"] or not np.isfinite(fy): return False - facts = _utils.summary(b) - return not facts["zero"] and not facts["nan"] + return not strict or (facts["nonnegative"] and not np.signbit(fy)) def _finite_nonnegative(x): @@ -774,8 +811,11 @@ def value(b): # (a value where nothing is stored, typed as numpy types it) def probe(b): return np.empty(0, b.dtype) if isinstance(b, Array) else b - def value1(b): # (as a 1-element array: numpy's 0-d clip breaks ties otherwise) - return np.reshape(b._fill, (1,)) if isinstance(b, Array) else b + def value1(b): # (an array's fill as a 1-element array: numpy's 0-d clip breaks ties + # otherwise) + if isinstance(b, Array): + return np.reshape(b._fill, (1,)) if b.ndim else np.asarray(b._fill) + return b with np.errstate(all="ignore"): rdt = np.clip(probe(x), probe(lo), probe(hi)).dtype diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index ab3adecf2..bbf73f2c7 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -7,6 +7,7 @@ import math import operator import warnings +import zlib import numpy as np import pytest @@ -274,7 +275,7 @@ def test_all_elementwise_functions_are_tested(): @pytest.mark.parametrize("shape", [(), (7,), (3, 4), (2, 3, 4)]) def test_unary(name, func, shape): for dtype in DTYPES: - for dense, x in inputs(dtype, shape, seed=hash(name) % 1000): + for dense, x in inputs(dtype, shape, seed=zlib.crc32(name.encode()) % 1000): expected = expected_or_error(func, dense) if expected is None: with pytest.raises(TypeError): @@ -324,7 +325,9 @@ def _binary_cases(dt1, dt2, shapes, seed): @pytest.mark.parametrize("shapes", BINARY_SHAPES, ids=str) def test_binary(name, func, shapes): for dt1, dt2 in DTYPE_PAIRS: - for dense1, x1, dense2, x2 in _binary_cases(dt1, dt2, shapes, hash(name) % 1000): + for dense1, x1, dense2, x2 in _binary_cases( + dt1, dt2, shapes, zlib.crc32(name.encode()) % 1000 + ): expected = expected_or_error(func, dense1, dense2) if expected is None: with pytest.raises(TypeError): @@ -728,7 +731,42 @@ def test_operands_storing_every_position_keep_products_sparse(): assert_array_equal(i * full, i.todense() * full.todense()) # strict mode: the intersection only for values whose product with 0.0 is 0.0 p = A * xp.reshape(inv, (40, 1)) + assert p.nnz == A.nnz assert_array_equal(p, A.todense() * inv.todense()[:, None]) + assert (A / xp.exp(d)).nnz == A.nnz + assert_array_equal(A / xp.exp(d), A.todense() / np.exp(d.todense())) + signed = xp.asarray(rng.normal(size=40)) + assert_array_equal(A / signed, A.todense() / signed.todense()) + + +def test_one_sided_products(): + # one operand's fill is 0, the other's is not: the result lies within the first's entries + # (a broadcast second operand is stored in full first, not replicated) + rng = np.random.default_rng(4) + A = make(np.where(rng.random((30, 30)) < 0.05, rng.normal(size=(30, 30)), 0.0), 0.0) + v = make(np.where(rng.random(30) < 0.5, 0.0, rng.normal(size=30)), 0.0) + e = xp.exp(v) # (fill 1, some positions unstored) + w = make(np.where(rng.random(30) < 0.3, 1.0, rng.normal(size=30) + 4), 1.0) + for mode in [False, True]: + with gb.config.set(array_api_signed_zeros=mode): + for out, expected in [ + (A * xp.reshape(e, (1, 30)), A.todense() * e.todense()[None, :]), + (xp.reshape(e, (30, 1)) * A, e.todense()[:, None] * A.todense()), + (A / xp.reshape(w, (30, 1)), A.todense() / w.todense()[:, None]), + (A / e, A.todense() / e.todense()), + ]: + assert out.nnz <= A.nnz + if mode: + assert_array_equal(out, expected) + else: + np.testing.assert_array_equal(out.todense(), expected) + # a second operand as large as the result: the masked union, no full copy + B = make(np.where(rng.random((60, 60)) < 0.05, rng.normal(size=(60, 60)), 0.0), 0.0) + big = make(np.where(rng.random((60, 60)) < 0.5, 2.0, rng.normal(size=(60, 60)) + 5), 2.0) + with gb.config.set(array_api_signed_zeros=False): + out = B / big + assert out.nnz == B.nnz + np.testing.assert_array_equal(out.todense(), B.todense() / big.todense()) def test_clip_fill_uses_numpys_array_loop(): @@ -738,3 +776,13 @@ def test_clip_fill_uses_numpys_array_loop(): a, b = np.zeros(3), np.full(3, -0.0) assert_array_equal(xp.clip(x, min=lo, max=0.5), np.clip(a, b, 0.5)) assert_array_equal(xp.clip(x, min=-0.0, max=0.5), np.clip(a, -0.0, 0.5)) + + +def test_fill_follows_numpys_loop_choice(): + # numpy's power of an array by a 0-d 0.5 is sqrt (pow(-inf, 0.5) is NaN), and between two + # arrays it is pow (inf): an array's fill takes the loop the dense array would + x = make(np.array([[-np.inf, 2.0], [-np.inf, -np.inf]], np.float32), np.float32(-np.inf)) + half = xp.asarray(np.float32(0.5)) + halves = make(np.full((2, 2), 0.5, np.float32), np.float32(0.5)) + assert_array_equal(np_call(xp.pow, x, half), np_call(np.power, x.todense(), np.float32(0.5))) + assert_array_equal(np_call(xp.pow, x, halves), np_call(np.power, x.todense(), halves.todense())) From a037a29274d2e6b5fba645a91d0439853260fc04 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 03:07:50 -0600 Subject: [PATCH 29/77] array_api: semiring= on matmul, tensordot and vecdot, and einsum matmul, tensordot and vecdot (and their xp.linalg names) take a keyword-only semiring=, a graphblas Semiring or its name ("min_plus", "max_min", ...), and the new xp.einsum(subscripts, *operands, semiring=None) runs numpy subscripts over the Tensor's einsum (plus_times by default, lor_land for bool). A contraction skips unstored positions, so a semiring other than the default needs every operand's fill to be the identity of its monoid, and that identity must annihilate the multiply from both sides (inf for min_plus, -inf for max_plus, False for lor_land); the result's fill is the identity. The annihilator is checked once per typed semiring on finite probe values, so max_times on floats, plus_pair, min_first and integer min_plus raise ValueError, while unsigned max_times and plus_min pass. The default semiring keeps the existing rules for nonzero fills; einsum takes numpy's dense path for them, under the densify limit. Tensor contractions under a mask now pass an operand stored as the transpose of the layout the product needs through the transpose descriptor instead of a transposed copy, so SuiteSparse can use its masked dot product: an SDDMM-shaped A @ B.T under a mask of 2e5 entries (n = 2e4, k = 64) went from 2.0 s to 1 ms. --- graphblas/array_api/_linalg.py | 275 ++++++++++++++++-- graphblas/core/tensor_contract.py | 29 +- graphblas/tests/test_array_api_semiring.py | 317 +++++++++++++++++++++ graphblas/tests/test_tensor.py | 36 +++ 4 files changed, 626 insertions(+), 31 deletions(-) create mode 100644 graphblas/tests/test_array_api_semiring.py diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 057879a70..0080be1f5 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -1,4 +1,5 @@ -"""Linear algebra in the main namespace: matmul, matrix_transpose, tensordot and vecdot. +"""Linear algebra in the main namespace: matmul, matrix_transpose, tensordot and vecdot, and +the extension einsum. Contractions treat a fill of zero as a structural zero (DESIGN D6), as scipy.sparse and pydata/sparse do: a product with an unstored position is skipped, so 0 * inf there is not NaN, @@ -6,20 +7,34 @@ nonzero fill is the fill times a matrix of ones plus an array with fill zero, which turns into sums of the operands broadcast over the result: the result is then dense in general (under the densify limit). A fill that is not finite, or True, makes the result go through numpy densely. + +Semirings (an extension, DESIGN D11): ``semiring=`` takes a ``graphblas`` Semiring or its +name ("min_plus", "max_min", ...), whose monoid takes the place of the sum and whose binary +operator that of the product. A contraction skips the unstored positions, which is right when +every operand's fill is the identity ``e`` of the semiring's monoid and ``e`` annihilates the +multiply from both sides, ``mult(e, y) == mult(y, e) == e``: inf for min_plus, -inf for +max_plus, False for lor_land, 0 for plus_times. The result's fill is then ``e``; other fills +raise ValueError. As in D6 the skipping is structural: an infinity meeting an unstored +position makes no NaN (inf + -inf in min_plus). The annihilator is checked on a set of finite +probe values of each operand type, once per typed semiring. The default (None) is plus_times, +or lor_land for bool, under the rules of the first paragraph. """ import math import numpy as np -from .. import semiring +from .. import semiring as semirings +from ..core.operator.semiring import Semiring, TypedBuiltinSemiring, TypedUserSemiring +from ..core.scalar import Scalar from ..core.tensor import Tensor +from ..core.vector import Vector from . import _utils from ._array import Array from ._dtypes import check_supported -from ._utils import gb_dtype +from ._utils import gb_dtype, np_dtype -__all__ = ["matmul", "matrix_transpose", "tensordot", "vecdot"] +__all__ = ["einsum", "matmul", "matrix_transpose", "tensordot", "vecdot"] def _check(x, name): @@ -37,7 +52,131 @@ def _operand(x): def _semiring(dt): - return (semiring.lor_land if dt.kind == "b" else semiring.plus_times)[gb_dtype(dt)] + return (semirings.lor_land if dt.kind == "b" else semirings.plus_times)[gb_dtype(dt)] + + +def _resolve(semiring, dt, name): + """``(typed semiring, whether it is the default)`` for the ``semiring=`` argument, typed + for ``dt``, the operands' promoted dtype. + """ + default = _semiring(dt) + if semiring is None: + return default, True + if isinstance(semiring, str): + try: + sr = semirings.from_string(semiring) + except ValueError: + raise ValueError(f"{name}: unknown semiring {semiring!r}") from None + elif isinstance(semiring, (Semiring, TypedBuiltinSemiring, TypedUserSemiring)): + sr = semiring + else: + raise TypeError( + f"{name}'s semiring must be a graphblas Semiring or its name, not " + f"{type(semiring).__name__}" + ) + if isinstance(sr, Semiring): + if dt.kind == "b" and sr.name == "plus_times": + return default, True # (numpy's sum of bools is their lor) + try: + sr = sr[gb_dtype(dt)] + except KeyError: + raise TypeError(f"{name}: the semiring {sr.name} does not work with {dt}") from None + same = sr is default or ( + sr.name == default.name and sr.type == default.type and sr.type2 == default.type2 + ) + return sr, same + + +def _probes(dt): + """Values of the numpy dtype ``dt`` to check an annihilator on: both signs, zeros and the + extremes, all finite (an infinity meets an unstored position only structurally, as in D6). + """ + if dt.kind == "b": + return np.array([False, True]) + if dt.kind in "iu": + info = np.iinfo(dt) + values = [0, 1, 2, 7, info.max, info.max // 3] + if dt.kind == "i": + values += [-1, -5, info.min, info.min // 3] + return np.array(values, dtype=dt) + big = np.finfo(dt).max + values = [0.0, -0.0, 1.0, -1.0, 0.5, -2.5, 3e9, big, -big, np.finfo(dt).tiny] + if dt.kind == "c": + values += [1j, -1j, 1 + 2j, -3.5 + 0.5j, complex(0, big)] + return np.array(values, dtype=dt) + + +_ANNIHILATES = {} + + +def _annihilates(typed): + """Whether the identity of the typed semiring's monoid annihilates its multiply from both + sides on the probe values (cached). + """ + key = (typed.name, typed.type.name, typed.type2.name) + if (ok := _ANNIHILATES.get(key)) is None: + ok = _ANNIHILATES[key] = _check_annihilates(typed) + return ok + + +def _check_annihilates(typed): + if typed.is_positional: + return False + mult = typed.binaryop + e = typed.monoid.identity + for left in (True, False): + e_type, y_type = (typed.type, typed.type2) if left else (typed.type2, typed.type) + ys = _probes(np_dtype(y_type)) + v = Vector.from_coo(np.arange(ys.size), ys, dtype=y_type, size=ys.size) + s = Scalar.from_value(_utils.scalar(e, np_dtype(e_type)), dtype=e_type) + r = (v.apply(mult, left=s) if left else v.apply(mult, right=s)).new() + if r.nvals != ys.size or not bool(np.all(r.to_coo()[1] == e)): + return False + return True + + +def _operand_name(name, i): + return f"operand {i}" if name == "einsum" else f"x{i + 1}" + + +def _spell(value): + """``value`` as the namespace spells it (``xp.inf`` for inf).""" + value = np.asarray(value).item() + if isinstance(value, float) and np.isinf(value): + return "xp.inf" if value > 0 else "-xp.inf" + return repr(value) + + +def _sound(typed, arrays, name, *, multiplies=True): + """The result dtype and fill of a contraction of ``arrays`` with a typed semiring other + than the default, after checking that it may skip their unstored positions (ValueError + if not). ``multiplies``: whether values meet the multiply (einsum of one operand only + sums). + """ + e = typed.monoid.identity + out_dt = check_supported(np_dtype(typed.return_type), name) + if multiplies and not _annihilates(typed): + mult = typed.binaryop.name + raise ValueError( + f"{name}: the semiring {typed.name} on {np_dtype(typed.type)} cannot skip unstored " + f"positions: the identity {e!r} of its monoid {typed.monoid.name} does not " + f"annihilate {mult} ({mult}(e, y) and {mult}(y, e) must be e for every y), so no " + "fill works. Use it on Tensors (x.tensor), which skip missing entries." + ) + for i, x in enumerate(arrays): + in_dt = np_dtype(typed.type if i == 0 else typed.type2) + fill = _utils.scalar(x._fill, in_dt) + if fill != e: # (a NaN fill too) + spelled = _spell(e) + raise ValueError( + f"{name} with the semiring {typed.name} needs operands whose fill is {e!r}, " + f"the identity of {typed.monoid.name}" + + (f", which annihilates {typed.binaryop.name}" if multiplies else "") + + f"; {_operand_name(name, i)} has the fill {x._fill.item()!r}. Make the array " + f"with fill_value={spelled} (xp.asarray(x.tensor, fill_value={spelled}) reads " + f"x's unstored positions as {e!r})." + ) + return out_dt, _utils.scalar(e, out_dt) def _is_zero(fill): @@ -79,7 +218,10 @@ def _rdt(name, func, x1, x2): return check_supported(np.asarray(probe).dtype, name) -def matmul(x1, x2, /): +def matmul(x1, x2, /, *, semiring=None): + """The matrix product (numpy's matmul rules). ``semiring`` (an extension): a graphblas + Semiring or its name, under the module's rules for fills. + """ from ._elementwise import add, multiply from ._manipulation import expand_dims from ._statistical import sum as sum_ @@ -99,11 +241,15 @@ def matmul(x1, x2, /): except ValueError: raise ValueError(f"matmul: shapes {x1.shape} and {x2.shape} do not broadcast") from None rdt = _rdt("matmul", np.matmul, x1, x2) + typed, default = _resolve(semiring, rdt, "matmul") + if not default: + out_dt, fill = _sound(typed, (x1, x2), "matmul") + return _product(x1, x2, typed, out_dt, fill) if not _dense_ok(x1, x2): return _dense(np.matmul, x1, x2) a0, fa = _zero_fill(x1) b0, fb = _zero_fill(x2) - out = _product(a0, b0, rdt) + out = _product(a0, b0, typed, rdt, _utils.scalar(0, rdt)) if fa is None and fb is None: return out terms = [] @@ -127,14 +273,17 @@ def matmul(x1, x2, /): return out if out.dtype == rdt else cast(out, rdt) -def _product(a, b, rdt): - """``a @ b`` of arrays with fill 0, in GraphBLAS.""" +def _product(a, b, typed, out_dt, fill): + """``a @ b`` with the typed semiring, skipping unstored positions; the result's fill is + ``fill``. + """ Ta, Tb = a._T(), b._T() shape = np.broadcast_shapes(a.shape[:-2], b.shape[:-2]) shape = shape + (a.shape[-2:-1] if a.ndim > 1 else ()) + (b.shape[-1:] if b.ndim > 1 else ()) + gbdt = gb_dtype(out_dt) if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: - return Array._new(Tensor(gb_dtype(rdt), shape), 0) - return Array._new(Ta.matmul(Tb, _semiring(rdt)).new(dtype=gb_dtype(rdt)), 0) + return Array._new(Tensor(gbdt, shape), fill) + return Array._new(Ta.matmul(Tb, typed).new(dtype=gbdt), fill) def matrix_transpose(x, /): @@ -157,7 +306,10 @@ def _tensordot_axes(x1, x2, axes): return a, b -def tensordot(x1, x2, /, *, axes=2): +def tensordot(x1, x2, /, *, axes=2, semiring=None): + """The tensor dot product (numpy's tensordot). ``semiring`` (an extension): a graphblas + Semiring or its name, under the module's rules for fills. + """ from ._elementwise import add, cast, multiply from ._manipulation import reshape from ._statistical import sum as sum_ @@ -169,19 +321,18 @@ def tensordot(x1, x2, /, *, axes=2): if x1.shape[a] != x2.shape[b]: raise ValueError(f"tensordot: axis {a} of x1 and axis {b} of x2 differ in size") rdt = check_supported(np.result_type(x1.dtype, x2.dtype), "tensordot") - if not _dense_ok(x1, x2): - return _dense(lambda a, b: np.tensordot(a, b, axes=(a_axes, b_axes)), x1, x2) + typed, default = _resolve(semiring, rdt, "tensordot") free1 = tuple(n for i, n in enumerate(x1.shape) if i not in a_axes) free2 = tuple(n for i, n in enumerate(x2.shape) if i not in b_axes) shape = free1 + free2 + if not default: + out_dt, fill = _sound(typed, (x1, x2), "tensordot") + return _tensordot(x1, x2, a_axes, b_axes, shape, typed, out_dt, fill) + if not _dense_ok(x1, x2): + return _dense(lambda a, b: np.tensordot(a, b, axes=(a_axes, b_axes)), x1, x2) a0, fa = _zero_fill(x1) b0, fb = _zero_fill(x2) - Ta, Tb = a0._T(), b0._T() - if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: - out = Array._new(Tensor(gb_dtype(rdt), shape), 0) - else: - R = Ta.tensordot(Tb, (list(a_axes), list(b_axes)), _semiring(rdt)).new(dtype=gb_dtype(rdt)) - out = Array._new(R, 0) + out = _tensordot(a0, b0, a_axes, b_axes, shape, typed, rdt, _utils.scalar(0, rdt)) if fa is not None: s = sum_(b0, axis=b_axes, dtype=rdt) if b_axes else cast(b0, rdt) out = add(out, multiply(fa, reshape(s, (1,) * len(free1) + free2))) @@ -194,8 +345,21 @@ def tensordot(x1, x2, /, *, axes=2): return out if out.dtype == rdt else cast(out, rdt) -def vecdot(x1, x2, /, *, axis=-1): - """Sum over ``axis`` of ``conj(x1) * x2`` (broadcast), with zero fills structural.""" +def _tensordot(a, b, a_axes, b_axes, shape, typed, out_dt, fill): + Ta, Tb = a._T(), b._T() + gbdt = gb_dtype(out_dt) + if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(gbdt, shape), fill) + R = Ta.tensordot(Tb, (list(a_axes), list(b_axes)), typed).new(dtype=gbdt) + return Array._new(R, fill) + + +def vecdot(x1, x2, /, *, axis=-1, semiring=None): + """Sum over ``axis`` of ``conj(x1) * x2`` (broadcast), with zero fills structural. + + ``semiring`` (an extension): a graphblas Semiring or its name, whose monoid reduces + ``mult(x1, x2)`` over ``axis`` (no conjugate), under the module's rules for fills. + """ from ._elementwise import add, cast, conj, multiply from ._statistical import sum as sum_ @@ -213,23 +377,23 @@ def vecdot(x1, x2, /, *, axis=-1): raise ValueError("vecdot: the axis differs in size") shape = np.broadcast_shapes(x1.shape, x2.shape) rdt = check_supported(np.result_type(x1.dtype, x2.dtype), "vecdot") + typed, default = _resolve(semiring, rdt, "vecdot") + out_shape = tuple(n for i, n in enumerate(shape) if i != len(shape) + axis) + if not default: + out_dt, fill = _sound(typed, (x1, x2), "vecdot") + return _vecdot(x1, x2, axis, shape, out_shape, typed, out_dt, fill) if not _dense_ok(x1, x2): return _dense(lambda a, b: np.vecdot(a, b, axis=axis), x1, x2) if x1.dtype.kind == "c": x1 = conj(x1) a0, fa = _zero_fill(x1) b0, fb = _zero_fill(x2) - out_shape = tuple(n for i, n in enumerate(shape) if i != len(shape) + axis) Ta, Tb = a0._T(), b0._T() gbdt = gb_dtype(rdt) if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: out = Array._new(Tensor(gbdt, out_shape), 0) else: - mult = _semiring(rdt).binaryop - if _both_broadcast(Ta.shape, Tb.shape, shape): - # (an outer product: ewise_mult replicates one operand) - _utils.expansion_check(Ta, Tb, shape) - P = Array._new(Ta.ewise_mult(Tb, mult).new(dtype=gbdt), 0) + P = Array._new(_vecdot_product(Ta, Tb, shape, typed.binaryop, gbdt), 0) out = sum_(P, axis=axis, dtype=rdt) if rdt.kind != "b" else _any(P, axis) if fa is not None: out = add(out, multiply(fa, sum_(b0, axis=axis, dtype=rdt))) @@ -240,6 +404,22 @@ def vecdot(x1, x2, /, *, axis=-1): return out if out.dtype == rdt else cast(out, rdt) +def _vecdot_product(Ta, Tb, shape, mult, gbdt=None): + if _both_broadcast(Ta.shape, Tb.shape, shape): + # (an outer product: ewise_mult replicates one operand) + _utils.expansion_check(Ta, Tb, shape) + return Ta.ewise_mult(Tb, mult).new(dtype=gbdt) + + +def _vecdot(a, b, axis, shape, out_shape, typed, out_dt, fill): + Ta, Tb = a._T(), b._T() + gbdt = gb_dtype(out_dt) + if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(gbdt, out_shape), fill) + P = _vecdot_product(Ta, Tb, shape, typed.binaryop) + return Array._new(P.reduce(typed.monoid, axis=len(shape) + axis).new(dtype=gbdt), fill) + + def _any(x, axis): from ._statistical import any as any_ @@ -250,3 +430,42 @@ def _both_broadcast(sa, sb, shape): from ..core.tensor_ewise import _broadcast_factor return _broadcast_factor(sa, shape) > 1 and _broadcast_factor(sb, shape) > 1 + + +def einsum(subscripts, /, *operands, semiring=None): + """Einstein summation (an extension), with numpy's subscripts: ``einsum("ij,jk->ik", a, + b)`` is ``a @ b``. Each index appears at most once per operand (no diagonals), and + ``...`` is not supported. + + ``semiring``: a graphblas Semiring or its name; the default is plus_times, or lor_land + for bool. Its multiply combines the operands of each term, and its monoid sums the terms + over the indices not in the output. Fills follow the module's rules; with the default + semiring an operand whose fill is not zero makes the computation go through numpy + densely (under the densify limit). + """ + from ..core.tensor_einsum import _einsum_plan + from ..core.tensor_einsum import einsum as tensor_einsum + + if not isinstance(subscripts, str): + raise TypeError(f"einsum's subscripts must be a str, not {type(subscripts).__name__}") + if not operands: + raise ValueError("einsum needs at least one operand") + ops = [_operand(x) for x in operands] + for x in ops: + _check(x, "einsum") + rdt = check_supported(np.result_type(*(x.dtype for x in ops)), "einsum") + typed, default = _resolve(semiring, rdt, "einsum") + shapes = tuple(x.shape for x in ops) + shape = _einsum_plan(subscripts, shapes, "greedy")[2] + if default: + if not all(_is_zero(x._fill) for x in ops): + _utils.check_dense(math.prod(shape), "einsum with a nonzero fill") + return _dense(lambda *a: np.einsum(subscripts, *a), *ops) + out_dt, fill = rdt, _utils.scalar(0, rdt) + else: + out_dt, fill = _sound(typed, ops, "einsum", multiplies=len(ops) > 1) + tensors = [x._T() for x in ops] + gbdt = gb_dtype(out_dt) + if math.prod(shape) == 0 or any(T.nvals == 0 for T in tensors): + return Array._new(Tensor(gbdt, shape), fill) + return Array._new(tensor_einsum(subscripts, *tensors, op=typed).new(dtype=gbdt), fill) diff --git a/graphblas/core/tensor_contract.py b/graphblas/core/tensor_contract.py index b10ef17ba..38100e5af 100644 --- a/graphblas/core/tensor_contract.py +++ b/graphblas/core/tensor_contract.py @@ -137,8 +137,19 @@ def _contract(updater, a, b, op, a_axes, b_axes, a_batch=(), b_batch=()): if best is None: raise ValueError("the contracted or free axes do not fit in a Matrix dimension (2**60)") _, sa, sb, beta = best - A = _in_layout(a, Layout.from_state(sa), opt_nthreads) - B = _in_layout(b, Layout.from_state(sb), opt_nthreads) + A = B = None + if not nb and updater.kwargs.get("mask") is not None: + # Under a mask, an operand stored as the transpose of what the product needs goes in + # through the transpose descriptor, not a transposed copy: SuiteSparse can then take + # its masked dot product (C = A*B' with both held by row: an SDDMM). Given the + # copy, it ran a masked saxpy over the whole product (n = 2e4, k = 64, a mask of 2e5 + # entries: 2.0 s against 0.9 ms) + A = _transposed(a, sa) + B = _transposed(b, sb) + if A is None: + A = _in_layout(a, Layout.from_state(sa), opt_nthreads) + if B is None: + B = _in_layout(b, Layout.from_state(sb), opt_nthreads) out_batch = {x: a_batch.index(x) for x in a_batch} out_a = {x: nb + i for i, x in enumerate(free_a)} out_b = {y: nb + len(free_a) + i for i, y in enumerate(free_b)} @@ -148,7 +159,8 @@ def _contract(updater, a, b, op, a_axes, b_axes, a_batch=(), b_batch=()): layout = Layout(rows, tuple(out_b[y] for y in sb[0][nb + k :])) if census.on: census.note("contract", "batched, block-diagonal" if nb else "one mxm", - moved=[A is not a._M, B is not b._M], cost=best[0]) # fmt: skip + moved=[getattr(A, "_matrix", A) is not a._M, + getattr(B, "_matrix", B) is not b._M], cost=best[0]) # fmt: skip if not nb: return A.mxm(B, op), layout, shape nblocks = tl.prod(a._shape, a_batch) @@ -172,6 +184,17 @@ def _contract(updater, a, b, op, a_axes, b_axes, a_batch=(), b_batch=()): return A.mxm(B, op), layout, shape +def _transposed(t, state): + """``t``'s Matrix transposed by descriptor (no copy) when that puts it in the layout + ``state``, else None. + """ + rows, cols = t._layout + seq, p = state + if rows and cols and seq == cols + rows and p == len(cols): + return t._M.T + return None + + def _matmul_expr(self, other, op): if (tensor := _as_tensor(other, "matmul")) is None: self._expect_type(other, Tensor, within="matmul", argname="other", op=op) diff --git a/graphblas/tests/test_array_api_semiring.py b/graphblas/tests/test_array_api_semiring.py new file mode 100644 index 000000000..c53179b6b --- /dev/null +++ b/graphblas/tests/test_array_api_semiring.py @@ -0,0 +1,317 @@ +"""graphblas.array_api: the semiring= extension of matmul, tensordot and vecdot, and einsum.""" + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import check_invariants, make + +rng = np.random.default_rng(58) + + +def _rand(shape, fill, dtype=np.float64, density=0.5): + dtype = np.dtype(dtype) + if dtype.kind == "b": + a = rng.random(shape) < 0.5 + elif dtype.kind in "iu": + a = rng.integers(0, 9, shape) + else: + a = np.round(rng.normal(size=shape) * 4, 1) + return np.where(rng.random(shape) < density, a, fill).astype(dtype) + + +def assert_close(x, expected, fill=None): + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape, (x.shape, expected.shape) + if fill is not None and x.ndim: + assert x.fill_value == fill or (np.isnan(fill) and np.isnan(x.fill_value)) + got = x.todense() + if got.dtype.kind in "fc": + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + else: + np.testing.assert_array_equal(got, expected) + + +# (name, numpy add-monoid reduction, numpy multiply, identity) for the oracles +SEMIRINGS = { + "min_plus": (np.min, np.add, np.inf), + "max_plus": (np.max, np.add, -np.inf), + "min_max": (np.min, np.maximum, np.inf), + "max_min": (np.max, np.minimum, -np.inf), +} + + +def semiring_einsum(subscripts, *arrays, reduce, mult): + """The numpy oracle: each term's operands combined by ``mult``, then ``reduce`` over the + summed indices (every index spelled out, so each operand broadcasts over all of them). + """ + inputs, output = subscripts.replace(" ", "").split("->") + terms = inputs.split(",") + letters = sorted(set("".join(terms))) + total = None + for term, a in zip(terms, arrays, strict=True): + a = np.asarray(a) + perm = [term.index(c) for c in letters if c in term] + a = np.transpose(a, perm) + shape = [a.shape[sorted(term, key=letters.index).index(c)] if c in term else 1 + for c in letters] # fmt: skip + a = a.reshape(shape) + total = a if total is None else mult(total, a) + with np.errstate(invalid="ignore"): + summed = tuple(i for i, c in enumerate(letters) if c not in output) + out = reduce(total, axis=summed) if summed else total + kept = [c for c in letters if c in output] + return np.transpose(out, [kept.index(c) for c in output]) + + +@pytest.mark.parametrize("name", list(SEMIRINGS)) +@pytest.mark.parametrize( + "shapes", + [((3, 4), (4, 5)), ((4,), (4, 5)), ((3, 4), (4,)), ((4,), (4,)), ((2, 1, 3, 4), (3, 4, 2))], + ids=str, +) +def test_matmul_semirings(name, shapes): + reduce, mult, e = SEMIRINGS[name] + for dtype in [np.float64, np.float32]: + a, b = _rand(shapes[0], e, dtype), _rand(shapes[1], e, dtype) + x1, x2 = make(a, e, rng), make(b, e) + a2 = a[None] if a.ndim == 1 else a + b2 = b[:, None] if b.ndim == 1 else b + expected = reduce(mult(a2[..., :, :, None], b2[..., None, :, :]), axis=-2) + if a.ndim == 1: + expected = expected[..., 0, :] + if b.ndim == 1: + expected = expected[..., 0] + expected = np.asarray(expected, dtype) + for sr in [name, gb.semiring.from_string(name)]: + assert_close(xp.matmul(x1, x2, semiring=sr), expected, e) + assert_close(xp.linalg.matmul(x1, x2, semiring=sr), expected, e) + # Stored values equal to the fill do not change the answer + assert_close(xp.matmul(make(a, e, rng), make(b, e, rng), semiring=name), expected, e) + + +def test_matmul_semiring_dtypes(): + a, b = _rand((3, 4), np.inf, np.float32), _rand((4, 2), np.inf) + expected = np.min(a[:, :, None] + b[None], axis=1) + R = xp.matmul(make(a, np.inf), make(b, np.inf), semiring="min_plus") + assert_close(R, expected.astype(np.float64), np.inf) + # A typed semiring keeps its type + R = xp.matmul(make(a, np.inf), make(b, np.inf), semiring=gb.semiring.min_plus["FP32"]) + assert_close(R, expected.astype(np.float32), np.inf) + # Bools: lor_land (the default), and plus_times is the same thing + p, q = _rand((3, 4), False, bool), _rand((4, 5), False, bool) + for sr in [None, "lor_land", "plus_times", gb.semiring.lor_land]: + assert_close(xp.matmul(make(p, False), make(q, False), semiring=sr), p @ q, False) + # land_lor's identity is True + expected = np.all(p[:, :, None] | q[None], axis=1) + R = xp.matmul(make(p, True), make(q, True), semiring="land_lor") + assert_close(R, expected, True) + # Numbers through a bool semiring: nonzero values are True + f, g = _rand((3, 4), 0.0), _rand((4, 2), 0.0) + R = xp.matmul(make(f, 0.0), make(g, 0.0), semiring="lor_land") + assert_close(R, (f != 0).astype(int) @ (g != 0).astype(int) > 0, False) + # Unsigned max_times and plus_min: 0 annihilates times and min there + u, v = _rand((3, 4), 0, np.uint8), _rand((4, 5), 0, np.uint8) + expected = np.max(u[:, :, None] * v[None], axis=1) + assert_close(xp.matmul(make(u, 0), make(v, 0), semiring="max_times"), expected, 0) + expected = np.sum(np.minimum(u[:, :, None], v[None]), axis=1, dtype=np.uint8) + assert_close(xp.matmul(make(u, 0), make(v, 0), semiring="plus_min"), expected, 0) + # An explicit plus_times is the default, nonzero fills included + for fa, fb in [(0, 0), (1, 0), (2, 3)]: + a, b = _rand((3, 4), fa), _rand((4, 2), fb) + assert_close(xp.matmul(make(a, fa), make(b, fb), semiring="plus_times"), a @ b) + # Integer plus_times (any fill 0) + i, j = _rand((3, 4), 0, np.int16), _rand((4, 2), 0, np.int16) + assert_close(xp.matmul(make(i, 0), make(j, 0), semiring="plus_times"), i @ j, 0) + + +def test_contraction_skips_structurally(): + # As 0 * inf in plus_times (D6): an infinity meeting an unstored inf is no NaN + a = make(np.array([[-np.inf, np.inf]]), np.inf) + b = make(np.array([[np.inf], [-np.inf]]), np.inf) + assert a.nnz == b.nnz == 1 + assert_close(xp.matmul(a, b, semiring="min_plus"), [[np.inf]], np.inf) + # numpy: inf + -inf is NaN + with np.errstate(invalid="ignore"): + assert np.isnan(np.min(a.todense()[:, :, None] + b.todense()[None], axis=1)).all() + + +@pytest.mark.parametrize( + ("xshape", "yshape", "axes"), + [ + ((3, 4, 5), (5, 6), 1), + ((3, 4, 5), (4, 5, 2), 2), + ((3, 4, 5), (5, 6, 3), ([0, 2], [2, 0])), + ((3, 4), (2, 5), 0), + ((4,), (4,), 1), + ], + ids=str, +) +def test_tensordot_semirings(xshape, yshape, axes): + for name, (reduce, mult, e) in SEMIRINGS.items(): + a, b = _rand(xshape, e), _rand(yshape, e) + if isinstance(axes, int): + a_axes, b_axes = list(range(len(xshape) - axes, len(xshape))), list(range(axes)) + else: + a_axes, b_axes = axes + la = "abcde"[: len(xshape)] + lb = list("vwxyz"[: len(yshape)]) + for i, j in zip(a_axes, b_axes, strict=True): + lb[j] = la[i] + lb = "".join(lb) + out = "".join(c for c in la if c not in lb) + "".join(c for c in lb if c not in la) + subs = f"{la},{lb}->{out}" + expected = semiring_einsum(subs, a, b, reduce=reduce, mult=mult) + R = xp.tensordot(make(a, e), make(b, e, rng), axes=axes, semiring=name) + assert_close(R, expected, e) + assert_close(xp.linalg.tensordot(make(a, e), make(b, e), axes=axes, semiring=name), + expected, e) # fmt: skip + + +@pytest.mark.parametrize( + ("xshape", "yshape", "axis"), + [((3, 4), (3, 4), -1), ((3, 4), (4,), -1), ((2, 3, 4), (3, 4), -2), ((3, 1, 4), (5, 4), -1)], + ids=str, +) +def test_vecdot_semirings(xshape, yshape, axis): + for name, (reduce, mult, e) in SEMIRINGS.items(): + a, b = _rand(xshape, e), _rand(yshape, e) + a_b, b_b = np.broadcast_arrays(a, b) + expected = reduce(mult(a_b, b_b), axis=axis) + for func in [xp.vecdot, xp.linalg.vecdot]: + R = func(make(a, e, rng), make(b, e), axis=axis, semiring=name) + assert_close(R, expected, e) + # lor_land of bools is the default's any(x1 & x2) + p, q = _rand((3, 4), False, bool), _rand((3, 4), False, bool) + assert_close(xp.vecdot(make(p, False), make(q, False), semiring="lor_land"), + np.any(p & q, axis=-1), False) # fmt: skip + + +EINSUMS = [ + ("ij,jk->ik", [(3, 4), (4, 5)]), + ("ik,jk->ij", [(3, 4), (5, 4)]), + ("ij,jk,kl->il", [(3, 4), (4, 5), (5, 2)]), + ("ijk,jr,kr->ir", [(3, 4, 5), (4, 2), (5, 2)]), + ("ij,ij->i", [(3, 4), (3, 4)]), + ("i,j->ij", [(3,), (4,)]), + ("bij,bjk->bik", [(2, 3, 4), (2, 4, 5)]), + ("ij->ji", [(3, 4)]), + ("ij->i", [(3, 4)]), + ("ijk->", [(2, 3, 4)]), + ("ij,jk", [(3, 4), (4, 5)]), + ("ij,j->i", [(3, 4), (4,)]), +] + + +@pytest.mark.parametrize(("subscripts", "shapes"), EINSUMS, ids=[s for s, _ in EINSUMS]) +def test_einsum(subscripts, shapes): + # The default: numpy's einsum, zero fills structural, other fills densely through numpy + for dtype in [np.float64, np.int64, np.bool_, np.complex128]: + for fills in [(0,) * len(shapes), (1,) + (0,) * (len(shapes) - 1)]: + if np.dtype(dtype).kind == "b" and any(fills): + fills = (True,) * len(shapes) + arrays = [_rand(s, f, dtype) for s, f in zip(shapes, fills, strict=True)] + expected = np.einsum(subscripts, *arrays) + R = xp.einsum(subscripts, *[make(a, f) for a, f in zip(arrays, fills, strict=True)]) + assert_close(R, expected) + # Semirings + full = subscripts if "->" in subscripts else subscripts + "->ik" + for name, (reduce, mult, e) in SEMIRINGS.items(): + arrays = [_rand(s, e) for s in shapes] + expected = semiring_einsum(full, *arrays, reduce=reduce, mult=mult) + R = xp.einsum(subscripts, *[make(a, e, rng) for a in arrays], semiring=name) + assert_close(R, np.asarray(expected), e) + + +def test_einsum_operands_and_errors(): + a = _rand((3, 4), 0.0) + x = make(a, 0.0) + # numpy operands, empty operands and sizes + assert_close(xp.einsum("ij,kj->ik", a, x), a @ a.T) + assert_close(xp.einsum("ij,jk->ik", x, xp.zeros((4, 2))), np.zeros((3, 2))) + assert_close(xp.einsum("ij,jk->ik", xp.zeros((3, 0)), xp.zeros((0, 2))), np.zeros((3, 2))) + R = xp.einsum("ij,jk->ik", xp.full((3, 0), np.inf), xp.full((0, 2), np.inf), + semiring="min_plus") # fmt: skip + assert_close(R, np.full((3, 2), np.inf), np.inf) + # A 0-d operand + assert_close(xp.einsum(",ij->ij", xp.asarray(2.0), x), 2 * a) + assert_close(xp.einsum(",ij->ij", xp.asarray(0.0), x), 0 * a) + with pytest.raises(TypeError, match="subscripts"): + xp.einsum(x, x) + with pytest.raises(ValueError, match="at least one"): + xp.einsum("ij") + with pytest.raises(TypeError, match="Arrays"): + xp.einsum("ij->ij", [[1.0]]) + with pytest.raises(NotImplementedError, match="not supported"): + xp.einsum("...ij->...ji", x) + with pytest.raises(NotImplementedError, match="diagonal"): + xp.einsum("ii->i", make(np.eye(3), 0.0)) + with pytest.raises(ValueError, match="fill is inf"): + xp.einsum("ij,jk->ik", x, x.T, semiring="min_plus") + with pytest.raises(ValueError, match="operand 1"): + xp.einsum("ij,jk->ik", xp.full((3, 4), np.inf), x.T, semiring="min_plus") + # One operand only sums: the identity is all its fill needs + b = _rand((3, 4), -np.inf) + R = xp.einsum("ij->i", make(b, -np.inf), semiring="max_times") + assert_close(R, np.max(b, axis=1), -np.inf) + with gb.config.set(array_api_densify_limit=5), pytest.raises(xp.DensifyError): + xp.einsum("ij,jk->ik", make(a, 1.0), x.T) + + +def test_semiring_errors(): + a = make(_rand((3, 4), np.inf), np.inf) + b = make(_rand((4, 2), np.inf), np.inf) + z = make(_rand((4, 2), 0.0), 0.0) + # Fills that are not the identity + with pytest.raises(ValueError, match="x2 has the fill 0.0"): + xp.matmul(a, z, semiring="min_plus") + with pytest.raises(ValueError, match="x1 has the fill inf"): + xp.tensordot(a, z, axes=1, semiring="max_plus") + with pytest.raises(ValueError, match="fill is -inf"): + xp.vecdot(a, a, semiring="max_plus") + with pytest.raises(ValueError, match="fill is inf"): + xp.matmul(make(_rand((3, 4), np.nan), np.nan), b, semiring="min_plus") + # Semirings whose identity does not annihilate the multiply: no fill works + for sr, dtype in [ + ("max_times", np.float64), + ("plus_pair", np.float64), + ("any_pair", np.float64), + ("min_first", np.float64), + ("plus_first", np.float64), + ("min_plus", np.int64), + ("max_plus", np.uint8), + ("min_times", np.uint16), + ("lor_eq", np.float64), + ]: + e = np.asarray(0).astype(dtype) + x = make(_rand((3, 4), e, dtype), e) + y = make(_rand((4, 2), e, dtype), e) + with pytest.raises(ValueError, match="cannot skip"): + xp.matmul(x, y, semiring=sr) + with pytest.raises(ValueError, match="cannot skip"): + xp.einsum("ij,jk->ik", x, y, semiring=sr) + with pytest.raises(ValueError, match="cannot skip"): + xp.matmul(a, b, semiring=gb.semiring.ss.min_firsti) + with pytest.raises(ValueError, match="unknown semiring"): + xp.matmul(a, b, semiring="min.nope") + with pytest.raises(TypeError, match="Semiring or its name"): + xp.matmul(a, b, semiring=gb.binary.plus) + with pytest.raises(TypeError, match="does not work with complex128"): + xp.matmul(xp.astype(a, xp.complex128), b, semiring="min_plus") + with pytest.raises(TypeError, match="does not work with bool"): + xp.vecdot(xp.zeros(3, dtype=xp.bool), xp.zeros(3, dtype=xp.bool), semiring="min_plus") + + +def test_probes(): + from graphblas.array_api._linalg import _probes + + for dt in ["bool", "int8", "uint64", "float32", "complex64"]: + p = _probes(np.dtype(dt)) + assert p.dtype == np.dtype(dt) + if p.dtype.kind in "fc": + assert np.isfinite(p).all() + assert {-128, 127, 0} <= set(_probes(np.dtype("int8")).tolist()) diff --git a/graphblas/tests/test_tensor.py b/graphblas/tests/test_tensor.py index 0b7e4ded4..3b3a6cec5 100644 --- a/graphblas/tests/test_tensor.py +++ b/graphblas/tests/test_tensor.py @@ -2816,6 +2816,42 @@ def test_tensordot_semirings_masks_and_errors(): assert ("ttv", "entries") in {(r["point"], r["choice"]) for r in recs} +def test_masked_contraction_takes_a_transpose_descriptor(): + # Under a mask, an operand stored as the transpose of what the product needs is not + # copied (SuiteSparse's masked dot product, an SDDMM); the values are the masked product + from graphblas.core import tensor as tensor_mod + from graphblas.core import tensor_census as census + + rng = np.random.default_rng(32) + X, xv, xp = random_tensor((6, 4), rng, density=0.7) + Y, yv, yp = random_tensor((5, 4), rng, density=0.7) + Mk, _, mp = random_tensor((6, 5), rng, dtype=dtypes.BOOL, density=0.4) + ev, ep = np_matmul(xv, xp, yv.T, yp.T) + rv, rp = model_update(np.zeros((6, 5)), np.zeros((6, 5), bool), ev, ep, mask=mp) + XT = X.relayout(((1,), (0,))) # (X's transpose in storage) + for expr in [ + X.matmul(Y.T), + X.tensordot(Y, ([1], [1])), + tensor_mod.einsum("ik,jk->ij", X, Y), + XT.T.T.matmul(Y.T), + ]: + census.start() + try: + R = expr.new(mask=Mk.S) + recs = census.stop() + finally: + census.stop() + census.records.clear() + check(R, rv, rp, approx=True) + moved = [r["moved"] for r in recs if r["point"] == "contract" and "moved" in r] + assert moved == [[False, False]], moved + # Unmasked, the planner's relayout stays + check(X.matmul(Y.T).new(), ev, ep, approx=True) + C, cv, cp = random_tensor((6, 5), rng) + C(Mk.S, binary.plus) << X.matmul(Y.T) + check(C, *model_update(cv, cp, ev, ep, mask=mp, accum=np.add), approx=True) + + # -------------------------------------------------------------------------------------------- # Core hooks (Matrix, Vector, Scalar and operators meet a Tensor) From 7d09925c393145a43f0510200ea3c54d2d2bb947 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 03:08:29 -0600 Subject: [PATCH 30/77] array_api: keep each Array's shape and dtype on the Array shape, dtype, ndim and size no longer read the backing Tensor, so an Array whose Tensor is not computed yet can answer them. --- graphblas/array_api/_array.py | 18 ++++++++++++------ graphblas/tests/test_array_api_creation.py | 11 +++++++++++ 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py index ff72da66b..d8c872b92 100644 --- a/graphblas/array_api/_array.py +++ b/graphblas/array_api/_array.py @@ -26,7 +26,9 @@ class Array: fill value. """ - __slots__ = "_tensor", "_fill", "_facts", "__weakref__" + # The shape and dtype are kept here, not read from the Tensor, so that an Array whose + # Tensor is not computed yet can answer them + __slots__ = "_tensor", "_fill", "_facts", "_shape", "_dtype", "__weakref__" # numpy hands its operators over to ours (``ndarray + Array``), as it does for other # array types, rather than converting the Array to a dense ndarray __array_ufunc__ = None @@ -48,7 +50,9 @@ def _new(cls, tensor, fill): tensor = Tensor(tensor.dtype, ()) self = object.__new__(cls) self._tensor = tensor - self._fill = _utils.scalar(fill, np_dtype(tensor.dtype)) + self._shape = tensor.shape + self._dtype = np_dtype(tensor.dtype) + self._fill = _utils.scalar(fill, self._dtype) self._facts = None return self @@ -61,6 +65,8 @@ def _T(self): def _replace(self, other): """Make this Array hold ``other``'s values (for in-place operators and assignment).""" self._tensor = other._T() + self._shape = other._shape + self._dtype = other._dtype self._fill = other._fill self._facts = None @@ -81,7 +87,7 @@ def _copy(self): # Attributes @property def dtype(self): - return np_dtype(self._T().dtype) + return self._dtype @property def device(self): @@ -89,15 +95,15 @@ def device(self): @property def shape(self): - return self._T().shape + return self._shape @property def ndim(self): - return self._T().ndim + return len(self._shape) @property def size(self): - return math.prod(self._T().shape) + return math.prod(self._shape) @property def T(self): diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py index e7ff27ed7..6a2062667 100644 --- a/graphblas/tests/test_array_api_creation.py +++ b/graphblas/tests/test_array_api_creation.py @@ -388,6 +388,17 @@ def test_array_attributes(): hash(x) +def test_metadata_is_stored(): + # shape, dtype, ndim and size never read the Tensor (one may not be computed yet) + x = make(np.arange(24, dtype=np.int32).reshape(2, 3, 4), 0) + x._tensor = None # (any read of it would fail) + assert (x.shape, x.dtype, x.ndim, x.size) == ((2, 3, 4), xp.int32, 3, 24) + assert x.fill_value == 0 + y = xp.zeros((5, 0), dtype=xp.bool) + y += xp.ones((5, 0), dtype=xp.bool) + assert (y.shape, y.dtype) == ((5, 0), xp.bool) + + def test_python_scalars(): assert bool(xp.asarray(True)) is True assert int(xp.asarray(-3)) == -3 From 09de9b282f6d39e4c4e289f1ce019ee0c80b3258 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 03:13:14 -0600 Subject: [PATCH 31/77] array_api: defer contractions; multiply and slicing compute only what they need matmul, tensordot, vecdot and einsum, when they skip the operands' unstored positions (fills equal to the semiring's identity) and the result has a dimension, return a pending Array: it holds copy-on-write snapshots of its operands' Tensors (Tensor.dup, O(1); changing an operand afterwards copies the operand's data, not the snapshot's) and computes its Tensor on first use, once. shape, dtype, ndim, size, fill_value and repr do not compute it. Two consumers use the pending product instead of computing it whole: - multiply(S, P) or multiply(P, S), with S of P's shape and both fills 0, computes P under S's structure (a masked product), then the multiply: an SDDMM. It does so only when P's operands are finite (one parallel reduction each, or the facts the operand already cached), since outside S the answer is 0 * P. A finite product that overflows to inf outside S's stored positions then reads 0 where numpy gives NaN, as a contraction already reads 0 * inf at an unstored zero. - P[key] with ints and slices slices the operands along the indexed axes of the result, then contracts (pending again); matmul whose batch axes broadcast, vecdot and all-int keys compute the whole product first. Everything else computes the product, including __setitem__ (which then changes the computed Tensor) and iteration. graphblas.array_api._lazy.enabled = False turns deferral off for debugging. --- graphblas/array_api/_array.py | 30 +- graphblas/array_api/_elementwise.py | 4 +- graphblas/array_api/_indexing.py | 7 +- graphblas/array_api/_lazy.py | 168 ++++++++++ graphblas/array_api/_linalg.py | 117 ++++--- graphblas/tests/test_array_api_lazy.py | 364 +++++++++++++++++++++ graphblas/tests/test_array_api_semiring.py | 7 + 7 files changed, 652 insertions(+), 45 deletions(-) create mode 100644 graphblas/array_api/_lazy.py create mode 100644 graphblas/tests/test_array_api_lazy.py diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py index d8c872b92..f1033f53e 100644 --- a/graphblas/array_api/_array.py +++ b/graphblas/array_api/_array.py @@ -27,8 +27,8 @@ class Array: """ # The shape and dtype are kept here, not read from the Tensor, so that an Array whose - # Tensor is not computed yet can answer them - __slots__ = "_tensor", "_fill", "_facts", "_shape", "_dtype", "__weakref__" + # Tensor is not computed yet (``_pending``, a deferred contraction) can answer them + __slots__ = "_tensor", "_pending", "_fill", "_facts", "_shape", "_dtype", "__weakref__" # numpy hands its operators over to ours (``ndarray + Array``), as it does for other # array types, rather than converting the Array to a dense ndarray __array_ufunc__ = None @@ -50,17 +50,33 @@ def _new(cls, tensor, fill): tensor = Tensor(tensor.dtype, ()) self = object.__new__(cls) self._tensor = tensor + self._pending = None self._shape = tensor.shape self._dtype = np_dtype(tensor.dtype) self._fill = _utils.scalar(fill, self._dtype) self._facts = None return self + @classmethod + def _deferred(cls, pending, shape, dtype, fill): + """An Array whose Tensor ``pending.compute()`` gives on first use (see ``_lazy``).""" + self = object.__new__(cls) + self._tensor = None + self._pending = pending + self._shape = tuple(shape) + self._dtype = np.dtype(dtype) + self._fill = _utils.scalar(fill, self._dtype) + self._facts = None + return self + def _T(self): - """The Tensor of stored entries. Every read of it goes through here, so that an Array - may later hold an unevaluated expression and compute it on first use. + """The Tensor of stored entries. Every read of it goes through here: a pending Array + computes it here, once. """ - return self._tensor + if (T := self._tensor) is None: + T = self._tensor = self._pending.compute() + self._pending = None # (releases the operands' snapshots) + return T def _replace(self, other): """Make this Array hold ``other``'s values (for in-place operators and assignment).""" @@ -164,6 +180,8 @@ def __repr__(self): fill = self._fill.item() if self.ndim == 0: return f"Array({fill!r}, dtype={self.dtype})" + if self._tensor is None: # (its nnz would compute it) + return f"Array(shape={self.shape}, dtype={self.dtype}, fill_value={fill!r}, pending)" return f"Array(shape={self.shape}, dtype={self.dtype}, nnz={self.nnz}, fill_value={fill!r})" # Protocols @@ -232,6 +250,8 @@ def __len__(self): def __iter__(self): if self.ndim == 0: raise TypeError("iteration over a 0-d array") + # (computed first: each row of a pending product would be a product of its own) + self._T() return (self[i] for i in range(self.shape[0])) __hash__ = None diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index ae134630d..cc1f16ccc 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -17,7 +17,7 @@ from .. import binary, config, unary from ..core.tensor import Tensor -from . import _utils +from . import _lazy, _utils from ._array import Array from ._dtypes import check_supported from ._utils import gb_dtype, gb_scalar @@ -421,6 +421,8 @@ def __init__(self, name, ufunc, kind="arith", gb=None, annihilator=None, gb_valu def _binary(spec, x1, x2): a = _as_operand(x1, spec.name) b = _as_operand(x2, spec.name) + if spec.name == "multiply" and (out := _lazy.masked_multiply(a, b)) is not None: + return out # (an SDDMM) ufunc = spec.ufunc if not isinstance(a, Array) and not isinstance(b, Array): from ._creation import asarray diff --git a/graphblas/array_api/_indexing.py b/graphblas/array_api/_indexing.py index 7876379bd..4c619ecc9 100644 --- a/graphblas/array_api/_indexing.py +++ b/graphblas/array_api/_indexing.py @@ -20,7 +20,7 @@ import numpy as np from ..core.tensor import Tensor -from . import _utils +from . import _lazy, _utils from ._array import Array from ._utils import gb_dtype @@ -235,7 +235,10 @@ def getitem(x, key): if _has_adv(items): out, nadv = _advanced_extract(x, axes_key) else: - out, nadv = _basic_extract(x, tuple(axes_key)), 0 + # (a pending product slices its operands instead) + if (out := _lazy.sliced(x, axes_key)) is None: + out = _basic_extract(x, tuple(axes_key)) + nadv = 0 return _finish(out, items, nadv) diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py new file mode 100644 index 000000000..523280152 --- /dev/null +++ b/graphblas/array_api/_lazy.py @@ -0,0 +1,168 @@ +"""Deferred contractions (DESIGN D9, scoped to one level of deferral). + +matmul, tensordot, vecdot and einsum return an Array whose Tensor is not computed yet when +they skip the operands' unstored positions (their fills are the semiring's identity) and the +result has a dimension. Such a pending Array holds copy-on-write snapshots (``Tensor.dup``, +O(1)) of its operands' Tensors, so changing an operand afterwards does not change it (the +operand copies its data then, and the snapshots keep the old data alive until the result is +computed), and computes its Tensor on the first read (``Array._T``), once. + +Two consumers use a pending product P without computing it whole: + +- ``multiply(S, P)`` or ``multiply(P, S)``, with S of P's shape and both fills 0, computes P + only at S's stored positions (a masked product, then the multiply: an SDDMM). Elsewhere S + is 0, and 0 times P's value there is 0 when that value is finite, which holds when P's + operands are finite (checked, else P is computed whole). A product of finite values that + overflows to inf outside S's stored positions then reads 0 where numpy gives NaN: the same + structural reading of 0 as in D6's 0 * inf. +- ``P[key]`` with a basic index (ints and slices) slices the operands along the indexed axes + of the result, then contracts (pending again). + +Everything else reads the Tensor, which computes it. Values do not depend on whether a result +was deferred, but for the overflow above; ``enabled = False`` turns deferral off, for +debugging. +""" + +import numpy as np + +from .. import monoid +from ..core.tensor import Tensor +from ..dtypes import FC64, FP64 +from ._array import Array +from ._utils import gb_dtype, np_dtype + +# Private kill switch: False computes every contraction when it is called +enabled = True + + +class Contraction: + """A contraction not computed yet. + + ``build(*tensors)`` gives its TensorExpression; ``tensors`` are snapshots of the operands' + Tensors. ``spec``, ``(terms, output, semiring)`` of an einsum that computes it, lets a + slice of the result slice the operands instead (None: no such einsum). ``finite`` holds + each operand's finiteness: True, False, or None when not checked yet. + """ + + __slots__ = "tensors", "build", "gbdt", "spec", "finite" + + def __init__(self, tensors, build, gbdt, spec, finite): + self.tensors = tensors + self.build = build + self.gbdt = gbdt + self.spec = spec + self.finite = finite + + def compute(self, mask=None): + return self.build(*self.tensors).new(dtype=self.gbdt, mask=mask) + + def all_finite(self): + """Whether every operand value is finite: then so is every product, but for + overflow. + """ + for i, ok in enumerate(self.finite): + if ok is None: + ok = self.finite[i] = _all_finite(self.tensors[i]) + if not ok: + return False + return True + + +def _all_finite(T): + kind = np_dtype(T.dtype).kind + if kind not in "fc" or T.nvals == 0: + return True + # A NaN or an infinity makes the sum non-finite. Finite values whose sum overflows make + # this False too, which costs the masked product, not the answer. One parallel pass over + # the values, with no copy. + total = T.reduce_scalar(monoid.plus[FC64 if kind == "c" else FP64]).new().value + return bool(np.isfinite(total)) + + +def _known_finite(x): + """Whether ``x``'s values are finite, if ``x`` already knows (from its cached facts); + else None. + """ + if x.dtype.kind not in "fc": + return True + facts = x._facts + if facts is not None and "summary" in facts: + return facts["summary"]["finite"] + return None + + +def contraction(arrays, build, shape, dtype, fill, spec=None): + """The Array of ``build(*tensors)`` over the arrays' Tensors, with fill ``fill``: pending + when deferral is on and the result has a dimension, else computed now. + """ + tensors = [x._T() for x in arrays] + gbdt = gb_dtype(dtype) + if not enabled or not shape: + return Array._new(build(*tensors).new(dtype=gbdt), fill) + snapshots = [T.dup() for T in tensors] + node = Contraction(snapshots, build, gbdt, spec, [_known_finite(x) for x in arrays]) + return Array._deferred(node, shape, dtype, fill) + + +def masked_multiply(a, b): + """``multiply(a, b)`` when one of them is a pending product P and the other, S, lets P be + computed at S's stored positions only (see the module's docstring); else None. + """ + if not isinstance(a, Array) or not isinstance(b, Array): + return None + if b._pending is not None: + S, P = a, b + elif a._pending is not None: + S, P = b, a + else: + return None + if S.shape != P.shape or bool(S._fill) or bool(P._fill): + return None + node = P._pending + if not node.all_finite(): + return None + PS = Array._new(node.compute(mask=S._T().S), P._fill) + from ._elementwise import multiply + + return multiply(PS, S) if P is a else multiply(S, PS) + + +def sliced(x, axes_key): + """``x[axes_key]`` (an int or a slice per axis) of a pending ``x`` that an einsum + computes: its operands sliced along the indexed axes, contracted (pending again). None + when that does not apply. + """ + node = x._pending + if node is None or node.spec is None or all(k == slice(None) for k in axes_key): + return None + terms, output, typed = node.spec + dropped = {c for c, k in zip(output, axes_key, strict=True) if not isinstance(k, slice)} + new_terms = ["".join(c for c in term if c not in dropped) for term in terms] + out = "".join(c for c in output if c not in dropped) + if not out or not all(new_terms): + return None # (a 0-d result or operand: computed whole, then indexed) + shape = tuple( + len(range(*k.indices(n))) + for k, n in zip(axes_key, x.shape, strict=True) + if isinstance(k, slice) + ) + keys = dict(zip(output, axes_key, strict=True)) + tensors = [] + for term, T in zip(terms, node.tensors, strict=True): + key = tuple(keys.get(c, slice(None)) for c in term) + if any(k != slice(None) for k in key): + T = T[key].new() + tensors.append(T) + if 0 in shape or any(T.nvals == 0 for T in tensors): + return Array._new(Tensor(node.gbdt, shape), x._fill) + from ..core.tensor_einsum import einsum + + subscripts = ",".join(new_terms) + "->" + out + + def build(*tensors): + return einsum(subscripts, *tensors, op=typed) + + # (a slice of a finite operand is finite) + spec = (tuple(new_terms), out, typed) + sub = Contraction(tensors, build, node.gbdt, spec, list(node.finite)) + return Array._deferred(sub, shape, x.dtype, x._fill) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 0080be1f5..08a81b310 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -18,6 +18,10 @@ position makes no NaN (inf + -inf in min_plus). The annihilator is checked on a set of finite probe values of each operand type, once per typed semiring. The default (None) is plus_times, or lor_land for bool, under the rules of the first paragraph. + +A contraction that skips the unstored positions is deferred: its result computes on first use, +or only where a consumer needs it (``multiply`` by a zero-filled array, a slice; see +``_lazy``). """ import math @@ -29,7 +33,7 @@ from ..core.scalar import Scalar from ..core.tensor import Tensor from ..core.vector import Vector -from . import _utils +from . import _lazy, _utils from ._array import Array from ._dtypes import check_supported from ._utils import gb_dtype, np_dtype @@ -274,16 +278,59 @@ def matmul(x1, x2, /, *, semiring=None): def _product(a, b, typed, out_dt, fill): - """``a @ b`` with the typed semiring, skipping unstored positions; the result's fill is - ``fill``. + """``a @ b`` with the typed semiring, skipping unstored positions (deferred, see + ``_lazy``); the result's fill is ``fill``. """ Ta, Tb = a._T(), b._T() shape = np.broadcast_shapes(a.shape[:-2], b.shape[:-2]) shape = shape + (a.shape[-2:-1] if a.ndim > 1 else ()) + (b.shape[-1:] if b.ndim > 1 else ()) - gbdt = gb_dtype(out_dt) if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: - return Array._new(Tensor(gbdt, shape), fill) - return Array._new(Ta.matmul(Tb, typed).new(dtype=gbdt), fill) + return Array._new(Tensor(gb_dtype(out_dt), shape), fill) + + def build(Ta, Tb): + return Ta.matmul(Tb, typed) + + spec = _matmul_terms(a.shape, b.shape) + spec = None if spec is None else (*spec, typed) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec) + + +# Index letters for the einsum that computes a contraction (for slicing a pending result) +_LETTERS = "abcdefghlmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + + +def _matmul_terms(sa, sb): + """``(terms, output)`` of the einsum equal to matmul of shapes ``sa`` and ``sb``, or None + when a batch axis of size 1 broadcasts (einsum has no broadcasting). + """ + ba, bb = sa[:-2], sb[:-2] + nb = max(len(ba), len(bb)) + if nb > len(_LETTERS): + return None # pragma: no cover (safety) + batch = _LETTERS[:nb] + for b in (ba, bb): + for n, m in zip(b[::-1], np.broadcast_shapes(ba, bb)[::-1], strict=False): + if n != m: + return None + ta = batch[nb - len(ba) :] + ("ik" if len(sa) > 1 else "k") + tb = batch[nb - len(bb) :] + ("kj" if len(sb) > 1 else "k") + out = batch + ("i" if len(sa) > 1 else "") + ("j" if len(sb) > 1 else "") + return (ta, tb), out + + +def _tensordot_terms(na, nb, a_axes, b_axes): + """``(terms, output)`` of the einsum equal to tensordot of arrays of ``na`` and ``nb`` + axes contracting ``a_axes`` with ``b_axes``. + """ + if na + nb > len(_LETTERS): + return None # pragma: no cover (safety) + ta = _LETTERS[:na] + tb = list(_LETTERS[na : na + nb]) + for i, j in zip(a_axes, b_axes, strict=True): + tb[j] = ta[i] + tb = "".join(tb) + out = "".join(c for c in ta if c not in tb) + "".join(c for c in tb if c not in ta) + return (ta, tb), out def matrix_transpose(x, /): @@ -347,11 +394,16 @@ def tensordot(x1, x2, /, *, axes=2, semiring=None): def _tensordot(a, b, a_axes, b_axes, shape, typed, out_dt, fill): Ta, Tb = a._T(), b._T() - gbdt = gb_dtype(out_dt) if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: - return Array._new(Tensor(gbdt, shape), fill) - R = Ta.tensordot(Tb, (list(a_axes), list(b_axes)), typed).new(dtype=gbdt) - return Array._new(R, fill) + return Array._new(Tensor(gb_dtype(out_dt), shape), fill) + axes = (list(a_axes), list(b_axes)) + + def build(Ta, Tb): + return Ta.tensordot(Tb, axes, typed) + + spec = _tensordot_terms(a.ndim, b.ndim, a_axes, b_axes) + spec = None if spec is None else (*spec, typed) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec) def vecdot(x1, x2, /, *, axis=-1, semiring=None): @@ -388,13 +440,7 @@ def vecdot(x1, x2, /, *, axis=-1, semiring=None): x1 = conj(x1) a0, fa = _zero_fill(x1) b0, fb = _zero_fill(x2) - Ta, Tb = a0._T(), b0._T() - gbdt = gb_dtype(rdt) - if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: - out = Array._new(Tensor(gbdt, out_shape), 0) - else: - P = Array._new(_vecdot_product(Ta, Tb, shape, typed.binaryop, gbdt), 0) - out = sum_(P, axis=axis, dtype=rdt) if rdt.kind != "b" else _any(P, axis) + out = _vecdot(a0, b0, axis, shape, out_shape, typed, rdt, _utils.scalar(0, rdt)) if fa is not None: out = add(out, multiply(fa, sum_(b0, axis=axis, dtype=rdt))) if fb is not None: @@ -404,26 +450,20 @@ def vecdot(x1, x2, /, *, axis=-1, semiring=None): return out if out.dtype == rdt else cast(out, rdt) -def _vecdot_product(Ta, Tb, shape, mult, gbdt=None): - if _both_broadcast(Ta.shape, Tb.shape, shape): - # (an outer product: ewise_mult replicates one operand) - _utils.expansion_check(Ta, Tb, shape) - return Ta.ewise_mult(Tb, mult).new(dtype=gbdt) - - def _vecdot(a, b, axis, shape, out_shape, typed, out_dt, fill): + """The semiring's monoid over ``axis`` of ``mult(a, b)`` (broadcast to ``shape``).""" Ta, Tb = a._T(), b._T() - gbdt = gb_dtype(out_dt) if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: - return Array._new(Tensor(gbdt, out_shape), fill) - P = _vecdot_product(Ta, Tb, shape, typed.binaryop) - return Array._new(P.reduce(typed.monoid, axis=len(shape) + axis).new(dtype=gbdt), fill) - + return Array._new(Tensor(gb_dtype(out_dt), out_shape), fill) + if _both_broadcast(Ta.shape, Tb.shape, shape): + # (an outer product: ewise_mult replicates one operand) + _utils.expansion_check(Ta, Tb, shape) + axis += len(shape) -def _any(x, axis): - from ._statistical import any as any_ + def build(Ta, Tb): + return Ta.ewise_mult(Tb, typed.binaryop).new().reduce(typed.monoid, axis=axis) - return any_(x, axis=axis) + return _lazy.contraction((a, b), build, out_shape, out_dt, fill) def _both_broadcast(sa, sb, shape): @@ -464,8 +504,11 @@ def einsum(subscripts, /, *operands, semiring=None): out_dt, fill = rdt, _utils.scalar(0, rdt) else: out_dt, fill = _sound(typed, ops, "einsum", multiplies=len(ops) > 1) - tensors = [x._T() for x in ops] - gbdt = gb_dtype(out_dt) - if math.prod(shape) == 0 or any(T.nvals == 0 for T in tensors): - return Array._new(Tensor(gbdt, shape), fill) - return Array._new(tensor_einsum(subscripts, *tensors, op=typed).new(dtype=gbdt), fill) + if math.prod(shape) == 0 or any(x._T().nvals == 0 for x in ops): + return Array._new(Tensor(gb_dtype(out_dt), shape), fill) + + def build(*tensors): + return tensor_einsum(subscripts, *tensors, op=typed) + + terms, output = _einsum_plan(subscripts, shapes, "greedy")[:2] + return _lazy.contraction(ops, build, shape, out_dt, fill, (terms, output, typed)) diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py new file mode 100644 index 000000000..2050a1b70 --- /dev/null +++ b/graphblas/tests/test_array_api_lazy.py @@ -0,0 +1,364 @@ +"""graphblas.array_api: deferred contractions and their two push-down consumers (DESIGN D9).""" + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp +from graphblas.array_api import _lazy, _utils + +from .test_array_api_elementwise import assert_array_equal, check_invariants, make + +rng = np.random.default_rng(9) + + +def _rand(shape, fill=0.0, dtype=np.float64, density=0.5): + dtype = np.dtype(dtype) + if dtype.kind == "b": + a = rng.random(shape) < 0.5 + elif dtype.kind in "iu": + a = rng.integers(1, 9, shape) + else: + a = np.round(rng.normal(size=shape) * 4, 1) + return np.where(rng.random(shape) < density, a, fill).astype(dtype) + + +def assert_close(x, expected): + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape, (x.shape, expected.shape) + got = x.todense() + if got.dtype.kind in "fc": + tol = 1e-5 if got.dtype in (np.float32, np.complex64) else 1e-9 + np.testing.assert_allclose(got, expected, rtol=tol, atol=tol) + else: + np.testing.assert_array_equal(got, expected) + + +def pending(x): + return x._tensor is None + + +@pytest.fixture +def masks(monkeypatch): + """The masks the pending products were computed under (None: whole).""" + seen = [] + compute = _lazy.Contraction.compute + + def recording(self, mask=None): + seen.append(mask) + return compute(self, mask) + + monkeypatch.setattr(_lazy.Contraction, "compute", recording) + return seen + + +def test_contractions_are_pending(masks): + a, b = _rand((4, 5)), _rand((5, 3)) + x1, x2 = make(a, 0.0), make(b, 0.0) + for P, expected in [ + (x1 @ x2, a @ b), + (xp.matmul(x1, x2), a @ b), + (xp.tensordot(x1, x2, axes=1), a @ b), + (xp.einsum("ij,jk->ik", x1, x2), a @ b), + (xp.vecdot(x1, x1), np.vecdot(a, a)), + (xp.matmul(x1, x2, semiring="lor_land"), (a != 0).astype(int) @ (b != 0) > 0), + ]: + assert pending(P) + # Metadata and repr do not compute it + assert P.shape == expected.shape + assert P.ndim == expected.ndim + assert P.size == expected.size + assert P.dtype == expected.dtype + assert P.fill_value == 0 + assert P.device == "cpu" + assert repr(P) == ( + f"Array(shape={expected.shape}, dtype={expected.dtype}, " + f"fill_value={P.fill_value.item()!r}, pending)" + ) + assert pending(P) + assert_close(P, expected) # (computes it) + assert not pending(P) + assert P._pending is None # (the snapshots are released) + T = P._T() + assert P._T() is T # (computed once) + assert masks == [None] * 6 + # Not deferred: 0-d results, empty operands, the kill switch, nonzero fills + assert not pending(xp.vecdot(make(a[0], 0.0), make(a[1], 0.0))) + assert not pending(xp.einsum("ij,ij->", x1, x1)) + assert not pending(x1 @ xp.zeros((5, 3))) + assert not pending(x1 @ make(b, 1.0)) + with gb.config.set(), pytest.MonkeyPatch.context() as mp: + mp.setattr(_lazy, "enabled", False) + P = x1 @ x2 + assert not pending(P) + assert_close(P, a @ b) + + +def test_snapshots(masks): + a, b = _rand((4, 5)), _rand((5, 3)) + x1, x2 = make(a, 0.0), make(b, 0.0) + P = x1 @ x2 + # The snapshots share the operands' data (copy-on-write: O(1)) + T1, T2 = P._pending.tensors + assert T1._M is x1._T()._M + assert T2._M is x2._T()._M + # Changing an operand afterwards changes neither the pending result nor the operand's + # earlier readers + x1[0, :] = 7.0 + x1[1, 1] = 0.0 + x2 += 1.0 + assert T1._M is not x1._T()._M + assert_close(P, a @ b) + a[0, :] = 7.0 + a[1, 1] = 0.0 + assert_close(x1 @ x2, a @ (b + 1.0)) + # The same Array on both sides + P = x1 @ x1.T + x1[...] = 0.0 + assert_close(P, a @ a.T) + # __setitem__ on a pending Array computes it first + P = make(a, 0.0) @ make(b, 0.0) + P[0, 0] = -1.0 + expected = a @ b + expected[0, 0] = -1.0 + assert_close(P, expected) + P = make(a, 0.0) @ make(b, 0.0) + P += 1.0 + assert_close(P, a @ b + 1.0) + Q = make(a[:, :4], 0.0) + Q @= make(2 * np.eye(4), 0.0) + assert not pending(Q) + assert_close(Q, 2 * a[:, :4]) + + +def test_pending_operands_and_consumers(): + a, b, c = _rand((4, 5)), _rand((5, 3)), _rand((3, 2)) + x1, x2, x3 = make(a, 0.0), make(b, 0.0), make(c, 0.0) + # A pending operand is computed (one level of deferral) + P = x1 @ x2 + R = P @ x3 + assert not pending(P) + assert pending(R) + assert_close(R, a @ b @ c) + # Everything else computes it + for func, np_func in [ + (lambda p: p + 1, lambda p: p + 1), + (xp.sum, np.sum), + (lambda p: p.T, lambda p: p.T), + (lambda p: p[[0, 2]], lambda p: p[[0, 2]]), + (lambda p: xp.asarray(np.asarray(p)), lambda p: p), + (lambda p: xp.stack(list(p)), lambda p: p), + ]: + P = x1 @ x2 + assert_close(func(P), np_func(a @ b)) + assert not pending(P) + # (the stored entries: where some term meets) + structure = (a != 0).astype(int) @ (b != 0).astype(int) > 0 + P = x1 @ x2 + assert P.nnz == structure.sum() + assert not pending(P) + P = x1 @ x2 + assert_close(xp.asarray(P.tensor), a @ b) + assert not pending(P) + + +def _sddmm_inputs(n=12, k=3, dtype=np.float64): + S = _rand((n, n), 0.0, dtype, density=0.2) + A = _rand((n, k), 0.0, dtype, density=1.0) + B = _rand((n, k), 0.0, dtype, density=1.0) + return S, A, B + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32, np.int64, np.complex128]) +def test_sddmm(dtype, masks): + s, a, b = _sddmm_inputs(dtype=dtype) + S, A, B = make(s, 0), make(a, 0), make(b, 0) + expected = s * (a @ b.T) + for func in [ + lambda P: xp.multiply(S, P), + lambda P: xp.multiply(P, S), + lambda P: S * P, + lambda P: P * S, + ]: + P = xp.matmul(A, xp.matrix_transpose(B)) + masks.clear() + R = func(P) + assert len(masks) == 1 + assert masks[0] is not None # (computed under S's structure) + assert pending(P) # (the masked product is not kept) + assert_close(R, expected) + assert R.nnz <= S.nnz + # Other contractions, and S with explicit zeros + S = make(s, 0, rng, explicit=0.3) + for P in [ + xp.einsum("ik,jk->ij", A, B), + xp.tensordot(A, B, axes=([1], [1])), + xp.vecdot(A[:, None, :], B[None, :, :]), + ]: + masks.clear() + assert_close(xp.multiply(S, P), expected) + assert masks[0] is not None + # A pending S is computed, then serves as the mask + P = A @ xp.matrix_transpose(B) + Q = make(np.eye(a.shape[0], dtype=dtype), 0) @ S + assert_close(xp.multiply(Q, P), expected) + + +def test_sddmm_bool_and_semirings(masks): + s, a, b = _sddmm_inputs(dtype=np.bool_) + S, A, B = make(s, False), make(a, False), make(b, False) + masks.clear() + assert_close(xp.multiply(S, A @ xp.matrix_transpose(B)), s & (a @ b.T)) + assert masks[0] is not None + # A fill other than 0 is computed whole: min_plus's inf + s, a, b = _sddmm_inputs() + S = make(s, 0.0) + A, B = make(np.where(a == 0, np.inf, a), np.inf), make(np.where(b == 0, np.inf, b), np.inf) + masks.clear() + P = xp.matmul(A, xp.matrix_transpose(B), semiring="min_plus") + expected = np.min(A.todense()[:, None, :] + B.todense()[None, :, :], axis=2) + with np.errstate(invalid="ignore"): + assert_close(xp.multiply(S, P), s * expected) + assert masks == [None] + + +def test_sddmm_falls_back(masks): + s, a, b = _sddmm_inputs() + A, B = make(a, 0.0), make(b, 0.0) + expected = a @ b.T + # S's fill is not 0, S broadcasts, or S is a Python scalar: the product is computed whole + for S, s_np in [(make(s, 1.0), s), (make(s[:1], 0.0), s[:1]), (2.0, 2.0)]: + masks.clear() + P = A @ xp.matrix_transpose(B) + assert_close(xp.multiply(S, P), s_np * expected) + assert [m is None for m in masks] == [True] + # (an ndarray is an Array with fill 0) + masks.clear() + assert_close(xp.multiply(s, A @ xp.matrix_transpose(B)), s * expected) + assert [m is None for m in masks] == [False] + # An operand holding inf or NaN: exact, as numpy (0 * inf is NaN outside S) + for bad in [np.inf, np.nan]: + a2 = a.copy() + a2[0, 0] = bad + A2 = make(a2, 0.0) + masks.clear() + R = xp.multiply(make(s, 0.0), A2 @ xp.matrix_transpose(B)) + with np.errstate(invalid="ignore"): + assert_close(R, s * (a2 @ b.T)) + assert [m is None for m in masks] == [True] + # Overflow to inf outside S: masked, it reads 0 (numpy: NaN), as D6's 0 * inf + big = make(np.full((2, 1), 1e200), 0.0) + s = np.array([[1.0, 0.0], [0.0, 0.0]]) + R = xp.multiply(make(s, 0.0), big @ xp.matrix_transpose(big)) + assert_array_equal(R, [[np.inf, 0.0], [0.0, 0.0]]) + with pytest.MonkeyPatch.context() as mp, np.errstate(invalid="ignore", over="ignore"): + mp.setattr(_lazy, "enabled", False) + R = xp.multiply(make(s, 0.0), big @ xp.matrix_transpose(big)) + assert_array_equal(R, s * (np.full((2, 1), 1e200) @ np.full((1, 2), 1e200))) + + +def test_finite_facts(): + x = make(np.array([1.0, np.inf]), 0.0) + assert not _lazy._all_finite(x._T()) + assert not _lazy._all_finite(make(np.array([np.nan, 1.0]), 0.0)._T()) + assert not _lazy._all_finite(make(np.array([np.inf, -np.inf]), 0.0)._T()) + assert not _lazy._all_finite(make(np.array([1 + 0j, complex(0, np.inf)]), 0.0)._T()) + assert _lazy._all_finite(make(np.array([1.0, -2.0]), 0.0)._T()) + assert _lazy._all_finite(make(np.array([5, 7]), 0)._T()) + assert _lazy._all_finite(xp.zeros(3)._T()) + # Facts an Array already knows are taken as they are + y = make(np.array([1.0, 2.0]), 0.0) + assert _lazy._known_finite(y) is None + _utils.summary(y) + assert _lazy._known_finite(y) is True + assert _lazy._known_finite(make(np.array([1, 2]), 0)) is True + P = xp.matmul(y, make(np.array([[1.0], [2.0]]), 0.0)) + assert P._pending.finite == [True, None] + + +SLICES = [ + (1,), + (slice(1, 3),), + (slice(None), 2), + (slice(None, None, -2), slice(1, 3)), + (Ellipsis, 1), + (None, slice(0, 2)), + (2, None), +] + + +@pytest.mark.parametrize("key", SLICES, ids=str) +def test_getitem_slices_the_operands(key, masks): + a, b = _rand((4, 5)), _rand((5, 3)) + x1, x2 = make(a, 0.0), make(b, 0.0) + expected = (a @ b)[key] + for P in [x1 @ x2, xp.tensordot(x1, x2, axes=1), xp.einsum("ij,jk->ik", x1, x2)]: + masks.clear() + R = P[key] + assert pending(P) + if R.ndim and not any(k is None for k in key): + assert pending(R) + assert_close(R, expected) + assert masks == [None] # (only the slice was computed) + assert_close(P, a @ b) + + +def test_getitem_cases(masks): + a, b = _rand((2, 4, 5)), _rand((2, 5, 3)) + x1, x2 = make(a, 0.0), make(b, 0.0) + # Batched (equal batch axes), 1-d operands, three operands + for P, expected in [ + (x1 @ x2, a @ b), + (x1[0, 0] @ x2[0], a[0, 0] @ b[0]), + (x1[0] @ x2[0, :, 0], a[0] @ b[0, :, 0]), + ( + xp.einsum("bij,bjk,kl->bil", x1, x2, make(b[0].T, 0.0)), + np.einsum("bij,bjk,kl->bil", a, b, b[0].T), + ), + (xp.tensordot(x1, x2, axes=([0, 2], [0, 1])), np.tensordot(a, b, ([0, 2], [0, 1]))), + ]: + key = ( + (1,) + (slice(1, None),) * (expected.ndim - 1) + if expected.ndim > 1 + else (slice(1, None),) + ) + masks.clear() + R = P[key] + assert pending(P) + assert_close(R, expected[key]) + # Broadcast batch axes: no slicing of the operands, computed whole + y1 = make(a[:1], 0.0) + P = y1 @ x2 + assert_close(P[1, 1:], (a[:1] @ b)[1, 1:]) + assert not pending(P) + # vecdot has no slicing; all-int keys, full slices and index arrays compute whole + for P, expected, key in [ + (xp.vecdot(x1, x1), np.vecdot(a, a), (0, slice(1, 3))), + (x1 @ x2, a @ b, (0, 1, 2)), + (x1 @ x2, a @ b, (slice(None),)), + (x1 @ x2, a @ b, (Ellipsis,)), + (x1 @ x2, a @ b, (np.array([0, 1]),)), + (x1[0, 0] @ x2[0], a[0, 0] @ b[0], (1,)), + ]: + assert pending(P) + assert_close(P[key], expected[key]) + assert not pending(P) + # Slices of empty extent, and slicing then the SDDMM + P = make(a[0], 0.0) @ make(b[0], 0.0) + assert_close(P[2:2], (a[0] @ b[0])[2:2]) + s = _rand((2, 3), 0.0, density=0.5) + masks.clear() + R = xp.multiply(make(s, 0.0), P[1:3]) + assert_close(R, s * (a[0] @ b[0])[1:3]) + assert [m is not None for m in masks] == [True] + + +def test_iteration_computes_once(masks): + a, b = _rand((4, 5)), _rand((5, 3)) + P = make(a, 0.0) @ make(b, 0.0) + rows = list(P) + assert masks == [None] + for row, expected in zip(rows, a @ b, strict=True): + assert_close(row, expected) diff --git a/graphblas/tests/test_array_api_semiring.py b/graphblas/tests/test_array_api_semiring.py index c53179b6b..9e9de8f34 100644 --- a/graphblas/tests/test_array_api_semiring.py +++ b/graphblas/tests/test_array_api_semiring.py @@ -275,6 +275,13 @@ def test_semiring_errors(): xp.vecdot(a, a, semiring="max_plus") with pytest.raises(ValueError, match="fill is inf"): xp.matmul(make(_rand((3, 4), np.nan), np.nan), b, semiring="min_plus") + with pytest.raises(ValueError, match="fill_value=xp.inf"): + xp.matmul(z.T, z, semiring="min_plus") + with pytest.raises(ValueError, match="fill_value=-xp.inf"): + xp.matmul(z.T, z, semiring="max_plus") + t = xp.zeros((4, 2), dtype=xp.bool) + with pytest.raises(ValueError, match="fill_value=True"): + xp.matmul(t.T, t, semiring="land_lor") # Semirings whose identity does not annihilate the multiply: no fill works for sr, dtype in [ ("max_times", np.float64), From 1f3c1bb0c3ce4f0d2213c04ecc56bbd5a9340d7e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 03:28:56 -0600 Subject: [PATCH 32/77] array_api: copy.copy and copy.deepcopy of an Array copy its data Both shared the backing Tensor, so assigning into the copy changed the original. --- graphblas/array_api/_array.py | 7 +++++++ graphblas/tests/test_array_api_creation.py | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py index f1033f53e..9f46ed309 100644 --- a/graphblas/array_api/_array.py +++ b/graphblas/array_api/_array.py @@ -100,6 +100,13 @@ def _fact(self, name, compute): def _copy(self): return Array._new(self._T().dup(), self._fill) + # (copy.copy would share the Tensor, which __setitem__ changes in place) + def __copy__(self): + return self._copy() + + def __deepcopy__(self, memo): + return self._copy() + # Attributes @property def dtype(self): diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py index 6a2062667..928bb1359 100644 --- a/graphblas/tests/test_array_api_creation.py +++ b/graphblas/tests/test_array_api_creation.py @@ -2,6 +2,7 @@ functions, against numpy. """ +import copy import math import pickle @@ -431,6 +432,18 @@ def test_array_protocol_is_guarded(): assert xp.sum(big) == 0.0 +def test_copy_module(): + a = np.array([[1.0, 0.0], [2.0, 3.0]]) + x = make(a, 0.0) + for func in [copy.copy, copy.deepcopy]: + y = func(x) + y[0, 0] = 5.0 + assert_array_equal(x, a) + assert y[0, 0] == 5.0 + z = func(x @ x) + assert_array_equal(z, a @ a) + + def test_pickle_not_supported_yet_or_roundtrips(): x = make(np.array([1.0, 0.0, 2.0]), 0.0) try: From f9df0b2a6e7ec8aa6c54eb373de13e5710ec232e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 03:55:40 -0600 Subject: [PATCH 33/77] array_api: drop the array_api_signed_zeros option and its strict paths The sign of a zero at a position an array does not store is unspecified (the contract is the Array API's semantics with numpy as the oracle), so the option that kept every such sign exact goes, with the strict branches of multiply's intersection test and the one-sided product and quotient. The default behavior does not change: a product with an unstored zero is the fill, and inf and NaN stay exact. The tests' oracle compares exactly but for the sign of a zero at an unstored position. --- graphblas/array_api/_elementwise.py | 37 ++-- graphblas/graphblas.yaml | 4 - graphblas/tests/test_array_api_creation.py | 1 - graphblas/tests/test_array_api_elementwise.py | 161 ++++++++---------- graphblas/tests/test_array_api_errors.py | 1 - graphblas/tests/test_array_api_indexing.py | 1 - graphblas/tests/test_array_api_linalg.py | 1 - .../tests/test_array_api_manipulation.py | 1 - graphblas/tests/test_array_api_searching.py | 1 - graphblas/tests/test_array_api_set.py | 1 - graphblas/tests/test_array_api_sorting.py | 1 - graphblas/tests/test_array_api_statistical.py | 1 - 12 files changed, 85 insertions(+), 126 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index cc1f16ccc..c3fcc33d4 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -15,7 +15,7 @@ import numpy as np -from .. import binary, config, unary +from .. import binary, unary from ..core.tensor import Tensor from . import _lazy, _utils from ._array import Array @@ -557,20 +557,17 @@ def _both_fills_zero(a, b, fa, fb, ct): def _multiply_intersects(a, b, fa, fb, ct): """Whether ``a * b`` stores only the intersection of entries: both fills are 0 and, for - floats, ``value * 0.0`` is ``0.0`` for every stored value (not ``-0.0``, nor NaN for inf). + floats, no stored value is inf or NaN (whose product with 0.0 is NaN). - Unless ``graphblas.config["array_api_signed_zeros"]`` is True (strict), a ``-0.0`` - product counts as the fill: finite values suffice (the union of a row scaling - ``d[:, None] * A`` with signed ``d`` would store every position, to keep the signs of its - zeros; DESIGN D3a). + A ``-0.0`` product with an unstored zero reads as the fill, 0.0: the sign of a zero at an + unstored position is unspecified (DESIGN D3a, D3b), and keeping it would make a row + scaling ``d[:, None] * A`` with signed ``d`` store every position. """ if not _both_fills_zero(a, b, fa, fb, ct): return False if ct.kind != "f": return True - if not config.get("array_api_signed_zeros", False): - return _utils.summary(a)["finite"] and _utils.summary(b)["finite"] - return _finite_nonnegative(a) and _finite_nonnegative(b) + return _utils.summary(a)["finite"] and _utils.summary(b)["finite"] def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): @@ -580,7 +577,6 @@ def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): at the first operand's entries (a broadcast operand smaller than them is stored in full first, so that it is not replicated). None when that does not hold. """ - strict = config.get("array_api_signed_zeros", False) sides = [(a, b)] if spec.name == "divide" else [(a, b), (b, a)] for X, Y in sides: fx = fa if X is a else fb @@ -589,7 +585,7 @@ def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): continue if _stores_all(Y): fy = _utils.scalar(1, ct) # (never seen) - if not _zero_times(spec, Y, fy, strict) or (strict and np.signbit(fx)): + if not _zero_times(spec, Y, fy): continue Tx, Ty = X._T(), Y._T() with np.errstate(all="ignore"): @@ -610,24 +606,15 @@ def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): return None -def _zero_times(spec, y, fy, strict): +def _zero_times(spec, y, fy): """Whether 0 * v (or 0 / v) is a zero for every value v of ``y`` (its fill ``fy`` too): - finite values (for divide also no 0), and in the strict mode nonnegative ones (so that - the zero is +0.0, the fill). NaN never qualifies. + finite values (for divide, also no 0). NaN never qualifies; the sign of the zero does not + matter (an unstored position reads the fill). """ facts = _utils.summary(y) if spec.name == "divide": - ok = not facts["zero"] and not facts["nan"] and fy != 0 and not np.isnan(fy) - # (strict: 0.0 / -2.0 is -0.0, which an unstored position would read as 0.0) - return ok and (not strict or (facts["nonnegative"] and fy > 0)) - if not facts["finite"] or not np.isfinite(fy): - return False - return not strict or (facts["nonnegative"] and not np.signbit(fy)) - - -def _finite_nonnegative(x): - facts = _utils.summary(x) - return facts["finite"] and facts["nonnegative"] + return not facts["zero"] and not facts["nan"] and fy != 0 and not np.isnan(fy) + return facts["finite"] and bool(np.isfinite(fy)) def _max_min_gb(op): diff --git a/graphblas/graphblas.yaml b/graphblas/graphblas.yaml index 5bc7615fd..4711f66b2 100644 --- a/graphblas/graphblas.yaml +++ b/graphblas/graphblas.yaml @@ -10,7 +10,3 @@ jit_c_control: null tensor_max_expansion: 1000000000 # graphblas.array_api: most elements an array may hold densely (np.asarray, cumulative_sum, ...) array_api_densify_limit: 100000000 -# graphblas.array_api: keep the sign of every zero a product or quotient makes with an unstored -# zero, as numpy (True), even where that stores every position (a row scaling d[:, None] * A with -# negative d); by default (False) those positions read 0.0, and multiply stores the intersection -array_api_signed_zeros: false diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py index 928bb1359..4a0d76f0c 100644 --- a/graphblas/tests/test_array_api_creation.py +++ b/graphblas/tests/test_array_api_creation.py @@ -13,7 +13,6 @@ import graphblas.array_api as xp from graphblas import Matrix, Scalar, Tensor, Vector -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import ( DTYPES, assert_array_equal, diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index bbf73f2c7..cc09afd40 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -23,15 +23,6 @@ pytest.skip("graphblas.array_api needs SuiteSparse:GraphBLAS 10", allow_module_level=True) -@pytest.fixture(autouse=True) -def _strict_signed_zeros(): - """These tests compare with numpy bit for bit: keep every signed zero (DESIGN D3a's strict - mode). The default mode, where products with unstored zeros read 0.0, has its own tests. - """ - with gb.config.set(array_api_signed_zeros=True): - yield - - DTYPES = [ np.dtype(name) for name in ( @@ -130,24 +121,32 @@ def check_invariants(x): def assert_array_equal(x, expected): - """``x`` holds ``expected`` exactly: dtype, shape, NaN, signed zeros.""" + """``x`` holds ``expected`` exactly (dtype, shape, NaN, signed zeros), but for the sign of + a zero at a position ``x`` does not store, which the contract leaves unspecified (DESIGN + D3b). A 0-d array's value is exact. + """ expected = np.asarray(expected) check_invariants(x) assert x.dtype == expected.dtype, (x.dtype, expected.dtype) assert x.shape == expected.shape, (x.shape, expected.shape) dense = x.todense() assert dense.dtype == expected.dtype - _assert_identical(dense, expected) + unstored = np.zeros(x.shape, dtype=bool) + if x.ndim: + unstored[...] = True + unstored[x.to_coo()[0]] = False + _assert_identical(dense, expected, unstored) -def _assert_identical(a, b): +def _assert_identical(a, b, unstored): if a.dtype.kind == "c": - _assert_identical(a.real, b.real) - _assert_identical(a.imag, b.imag) + _assert_identical(a.real, b.real, unstored) + _assert_identical(a.imag, b.imag, unstored) return if a.dtype.kind == "f": nan_a, nan_b = np.isnan(a), np.isnan(b) - bad = (nan_a != nan_b) | (~nan_a & ((a != b) | (np.signbit(a) != np.signbit(b)))) + sign = (np.signbit(a) != np.signbit(b)) & ~(unstored & (a == 0) & (b == 0)) + bad = (nan_a != nan_b) | (~nan_a & ((a != b) | sign)) else: bad = a != b if bad.any(): @@ -528,8 +527,6 @@ def test_float_multiply_keeps_union(): def test_broadcast_expansion_is_guarded(): - import graphblas as gb - x = make(np.arange(1, 101, dtype=np.float64).reshape(100, 1), 0.0) y = make(np.arange(1, 101, dtype=np.float64).reshape(1, 100), 0.0) with gb.config.set(array_api_densify_limit=1000), pytest.raises(xp.DensifyError): @@ -540,8 +537,6 @@ def test_broadcast_expansion_is_guarded(): def test_values_mapped_on_every_storage_format(): # numpy maps the stored values through the Container in place: check bitmap, full, # hypersparse and iso storage - import graphblas as gb - rng = np.random.default_rng(5) for density in [0.0, 0.05, 0.5, 1.0]: for shape in [(40,), (30, 20), (6, 5, 4)]: @@ -563,7 +558,6 @@ def test_values_mapped_on_every_storage_format(): other = make(np.arange(8.0), 0.0) assert_array_equal(xp.atan2(iso, other), np.arctan2(iso.todense(), other.todense())) assert_array_equal(xp.atan2(other, iso), np.arctan2(other.todense(), iso.todense())) - del gb def test_unseen_fill_does_not_raise(): @@ -632,8 +626,6 @@ def test_unsupported_operands(): def test_intersection_of_broadcast_operands_is_guarded(): # an intersection of two broadcast operands is an outer product (review finding M4) - import graphblas as gb - x = make(np.arange(1, 101).reshape(100, 1), 0) y = make(np.arange(1, 101).reshape(1, 100), 0) b = x > 0 @@ -665,46 +657,37 @@ def test_clip_with_array_bounds_ties(): assert_array_equal(xp.clip(xp.asarray(2.0), min=make(np.array([1.0, 3.0]), 0.0)), [2.0, 3.0]) -def test_signed_zeros_give_way_by_default(): - # (review finding H1, DESIGN D3a) by default a product with an unstored zero is the fill +def test_signed_zeros_give_way(): + # (review finding H1, DESIGN D3a, D3b) a product with an unstored zero is the fill # whatever its sign: a row scaling stores only A's entries rng = np.random.default_rng(0) A = make(np.where(rng.random((50, 50)) < 0.05, rng.normal(size=(50, 50)), 0.0), 0.0) d = make(rng.normal(size=(50, 1)), 0.0) - exact = d * A - assert_array_equal(exact, d.todense() * A.todense()) - assert exact.nnz == 2500 - with gb.config.set(array_api_signed_zeros=False): - loose = d * A - assert loose.nnz == A.nnz - np.testing.assert_array_equal(loose.todense(), d.todense() * A.todense()) # (0.0 == -0.0) - # inf and NaN stay exact - inf = make(np.array([[np.inf, 0.0]]), 0.0) - assert_array_equal(inf * make(np.array([[0.0, 1.0]]), 0.0), np.array([[np.nan, 0.0]])) - # divide: the numerator's entries, when no value of the denominator is 0 or NaN - a = np.where(rng.random((4, 5)) < 0.5, 0.0, rng.normal(size=(4, 5))) - b = rng.normal(size=5) + 3.0 - for fb in [1.0, -2.0]: - q = make(a, 0.0) / make(b, fb) - assert q.nnz == np.count_nonzero(a) - np.testing.assert_array_equal(q.todense(), a / b) - q = make(a, 0.0) / make(np.r_[b[:4], 0.0], 1.0) - np.testing.assert_array_equal( - np_call(np.isnan, q.todense()), np.isnan(np_call(np.divide, a, np.r_[b[:4], 0.0])) - ) + out = d * A + assert out.nnz == A.nnz + assert_array_equal(out, d.todense() * A.todense()) + # inf and NaN stay exact + inf = make(np.array([[np.inf, 0.0]]), 0.0) + assert_array_equal(inf * make(np.array([[0.0, 1.0]]), 0.0), np.array([[np.nan, 0.0]])) + # divide: the numerator's entries, when no value of the denominator is 0 or NaN + a = np.where(rng.random((4, 5)) < 0.5, 0.0, rng.normal(size=(4, 5))) + b = rng.normal(size=5) + 3.0 + for fb in [1.0, -2.0]: + q = make(a, 0.0) / make(b, fb) + assert q.nnz == np.count_nonzero(a) + assert_array_equal(q, a / b) + q = make(a, 0.0) / make(np.r_[b[:4], 0.0], 1.0) + assert_array_equal(q, np_call(np.divide, a, np.r_[b[:4], 0.0])) @pytest.mark.parametrize("name", ["multiply", "divide"]) -def test_default_mode_matches_numpy_but_for_zero_signs(name): +def test_products_match_numpy_but_for_unstored_zero_signs(name): func = getattr(np, name) - with gb.config.set(array_api_signed_zeros=False): - for shapes in [((5,), (5,)), ((3, 4), (4,)), ((3, 1), (1, 4))]: - for dt1, dt2 in [(np.dtype("float64"),) * 2, (np.dtype("float32"), np.dtype("int8"))]: - for dense1, x1, dense2, x2 in _binary_cases(dt1, dt2, shapes, 5): - out = getattr(xp, name)(x1, x2).todense() - expected = np_call(func, dense1, dense2) - assert out.dtype == expected.dtype - np.testing.assert_array_equal(out, expected) # (NaN == NaN, 0.0 == -0.0) + for shapes in [((5,), (5,)), ((3, 4), (4,)), ((3, 1), (1, 4))]: + for dt1, dt2 in [(np.dtype("float64"),) * 2, (np.dtype("float32"), np.dtype("int8"))]: + for dense1, x1, dense2, x2 in _binary_cases(dt1, dt2, shapes, 5): + out = getattr(xp, name)(x1, x2) + assert_array_equal(out, np_call(func, dense1, dense2)) def test_operands_storing_every_position_keep_products_sparse(): @@ -714,28 +697,24 @@ def test_operands_storing_every_position_keep_products_sparse(): A = make(np.where(rng.random((40, 40)) < 0.05, rng.normal(size=(40, 40)), 0.0), 0.0) d = xp.asarray(rng.normal(size=40) + 5.0) deg = xp.asarray(rng.integers(1, 9, 40).astype(float)) - with gb.config.set(array_api_signed_zeros=False): - q = A / d - assert q.nnz == A.nnz - np.testing.assert_array_equal(q.todense(), A.todense() / d.todense()) - inv = 1 / deg - assert inv.fill_value == np.inf - p = A * xp.reshape(inv, (40, 1)) - assert p.nnz == A.nnz - assert p.fill_value == 0.0 - np.testing.assert_array_equal(p.todense(), A.todense() * inv.todense()[:, None]) - # integers too: a full operand's fill takes no part - i = make(np.where(rng.random((5, 6)) < 0.5, 0, rng.integers(1, 5, (5, 6))), 0) - full = xp.asarray(rng.integers(1, 5, 6), fill_value=7) - assert (i * full).nnz == i.nnz - assert_array_equal(i * full, i.todense() * full.todense()) - # strict mode: the intersection only for values whose product with 0.0 is 0.0 + q = A / d + assert q.nnz == A.nnz + assert_array_equal(q, A.todense() / d.todense()) + inv = 1 / deg + assert inv.fill_value == np.inf p = A * xp.reshape(inv, (40, 1)) assert p.nnz == A.nnz + assert p.fill_value == 0.0 assert_array_equal(p, A.todense() * inv.todense()[:, None]) + # integers too: a full operand's fill takes no part + i = make(np.where(rng.random((5, 6)) < 0.5, 0, rng.integers(1, 5, (5, 6))), 0) + full = xp.asarray(rng.integers(1, 5, 6), fill_value=7) + assert (i * full).nnz == i.nnz + assert_array_equal(i * full, i.todense() * full.todense()) assert (A / xp.exp(d)).nnz == A.nnz assert_array_equal(A / xp.exp(d), A.todense() / np.exp(d.todense())) signed = xp.asarray(rng.normal(size=40)) + assert (A / signed).nnz == A.nnz assert_array_equal(A / signed, A.todense() / signed.todense()) @@ -747,26 +726,32 @@ def test_one_sided_products(): v = make(np.where(rng.random(30) < 0.5, 0.0, rng.normal(size=30)), 0.0) e = xp.exp(v) # (fill 1, some positions unstored) w = make(np.where(rng.random(30) < 0.3, 1.0, rng.normal(size=30) + 4), 1.0) - for mode in [False, True]: - with gb.config.set(array_api_signed_zeros=mode): - for out, expected in [ - (A * xp.reshape(e, (1, 30)), A.todense() * e.todense()[None, :]), - (xp.reshape(e, (30, 1)) * A, e.todense()[:, None] * A.todense()), - (A / xp.reshape(w, (30, 1)), A.todense() / w.todense()[:, None]), - (A / e, A.todense() / e.todense()), - ]: - assert out.nnz <= A.nnz - if mode: - assert_array_equal(out, expected) - else: - np.testing.assert_array_equal(out.todense(), expected) - # a second operand as large as the result: the masked union, no full copy + for out, expected in [ + (A * xp.reshape(e, (1, 30)), A.todense() * e.todense()[None, :]), + (xp.reshape(e, (30, 1)) * A, e.todense()[:, None] * A.todense()), + (A / xp.reshape(w, (30, 1)), A.todense() / w.todense()[:, None]), + (A / e, A.todense() / e.todense()), + ]: + assert out.nnz <= A.nnz + assert_array_equal(out, expected) + # a second operand as large as the result: the masked union, no full copy B = make(np.where(rng.random((60, 60)) < 0.05, rng.normal(size=(60, 60)), 0.0), 0.0) big = make(np.where(rng.random((60, 60)) < 0.5, 2.0, rng.normal(size=(60, 60)) + 5), 2.0) - with gb.config.set(array_api_signed_zeros=False): - out = B / big - assert out.nnz == B.nnz - np.testing.assert_array_equal(out.todense(), B.todense() / big.todense()) + out = B / big + assert out.nnz == B.nnz + assert_array_equal(out, B.todense() / big.todense()) + + +def test_unstored_zero_signs_in_the_oracle(): + # The tests' oracle: a stored -0.0 must match, an unstored one may read 0.0 + x = make(np.array([-0.0, 1.0]), 0.0) + assert x.nnz == 2 + with pytest.raises(AssertionError, match="differ"): + assert_array_equal(x, np.array([0.0, 1.0])) + y = make(np.array([0.0, 1.0]), 0.0) + assert_array_equal(y, np.array([-0.0, 1.0])) + with pytest.raises(AssertionError, match="differ"): + assert_array_equal(xp.asarray(-0.0), np.array(0.0)) def test_clip_fill_uses_numpys_array_loop(): diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index d01a0d4c7..856c05f02 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -8,7 +8,6 @@ from graphblas.array_api import _utils from graphblas.core.tensor import Tensor -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, make diff --git a/graphblas/tests/test_array_api_indexing.py b/graphblas/tests/test_array_api_indexing.py index f55c0c028..243fa41d4 100644 --- a/graphblas/tests/test_array_api_indexing.py +++ b/graphblas/tests/test_array_api_indexing.py @@ -5,7 +5,6 @@ import graphblas.array_api as xp -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, inputs, make rng = np.random.default_rng(0) diff --git a/graphblas/tests/test_array_api_linalg.py b/graphblas/tests/test_array_api_linalg.py index 53e1792ef..f8d545bb3 100644 --- a/graphblas/tests/test_array_api_linalg.py +++ b/graphblas/tests/test_array_api_linalg.py @@ -8,7 +8,6 @@ import graphblas as gb import graphblas.array_api as xp -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, check_invariants, make, np_call rng = np.random.default_rng(21) diff --git a/graphblas/tests/test_array_api_manipulation.py b/graphblas/tests/test_array_api_manipulation.py index 4b002e55b..8fc56c488 100644 --- a/graphblas/tests/test_array_api_manipulation.py +++ b/graphblas/tests/test_array_api_manipulation.py @@ -5,7 +5,6 @@ import graphblas.array_api as xp -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, inputs, make diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py index 2fb53d07a..578b2ccbf 100644 --- a/graphblas/tests/test_array_api_searching.py +++ b/graphblas/tests/test_array_api_searching.py @@ -5,7 +5,6 @@ import graphblas.array_api as xp -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import DTYPES, assert_array_equal, inputs, make, np_call WHERE_SHAPES = [ diff --git a/graphblas/tests/test_array_api_set.py b/graphblas/tests/test_array_api_set.py index da41fb1ee..6bcd36084 100644 --- a/graphblas/tests/test_array_api_set.py +++ b/graphblas/tests/test_array_api_set.py @@ -8,7 +8,6 @@ import graphblas as gb import graphblas.array_api as xp -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, check_invariants, inputs, make DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint32", "int64", "float32", "float64")] diff --git a/graphblas/tests/test_array_api_sorting.py b/graphblas/tests/test_array_api_sorting.py index 3a3f8c85e..85cba8519 100644 --- a/graphblas/tests/test_array_api_sorting.py +++ b/graphblas/tests/test_array_api_sorting.py @@ -6,7 +6,6 @@ import graphblas as gb import graphblas.array_api as xp -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import assert_array_equal, inputs, make DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint16", "int64", "float32", "float64")] diff --git a/graphblas/tests/test_array_api_statistical.py b/graphblas/tests/test_array_api_statistical.py index 1182b16fb..68e26fd81 100644 --- a/graphblas/tests/test_array_api_statistical.py +++ b/graphblas/tests/test_array_api_statistical.py @@ -8,7 +8,6 @@ import graphblas as gb import graphblas.array_api as xp -from .test_array_api_elementwise import _strict_signed_zeros # noqa: F401 (autouse) from .test_array_api_elementwise import ( DTYPES, assert_array_equal, From 5d987aab972f6ba106754f6e461600a73cfecf47 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 04:09:14 -0600 Subject: [PATCH 34/77] array_api: facts about stored values from one GraphBLAS pass each Facts about an array's stored values (all finite, some NaN, some zero) were a numpy summary over a copy of the values. Each is now one parallel GraphBLAS pass with no copy (apply to bool, then a reduction that stops at the first deciding value), computed only when asked: integers and bools are finite and NaN-free without a look, NaN and finiteness are read off each other, and abs, negative, positive and conj keep all three facts of their input, the roundings two. The sign-bit facts go. The Array API leaves the order of -0.0 and 0.0 in max, min, maximum and minimum unspecified, so those take GraphBLAS's fmax and fmin whenever no NaN is involved, and the numpy reduction kept for signed zeros is gone. On 1e7 stored float64 values, the finiteness fact takes 3.9 ms at 1 thread and 0.9 ms at 8 (the summary took 11 and 10 ms), and x * y with fresh facts 11 and 4.4 ms (28 and 23 ms). An SDDMM through the namespace (n = 1e6, 1e7 stored in S, k = 64, 8 threads) went from 0.118 s to 0.110 s; the masked product and multiply alone take 0.100 s. --- graphblas/array_api/_elementwise.py | 56 ++++++----- graphblas/array_api/_lazy.py | 4 +- graphblas/array_api/_statistical.py | 36 +------ graphblas/array_api/_utils.py | 99 ++++++++++--------- graphblas/tests/test_array_api_elementwise.py | 35 ++++--- graphblas/tests/test_array_api_errors.py | 34 ++++++- graphblas/tests/test_array_api_lazy.py | 2 +- graphblas/tests/test_array_api_statistical.py | 10 +- 8 files changed, 145 insertions(+), 131 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index c3fcc33d4..65e22c032 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -166,12 +166,17 @@ def _zero_d(dtype, fill): # Unary functions -def _unary(name, ufunc, x, gb=None, *, const=None): +_ALL_FACTS = ("finite", "nan", "zero") + + +def _unary(name, ufunc, x, gb=None, *, const=None, keeps=()): """``ufunc`` of every position of ``x``. ``gb(dtype)`` gives a GraphBLAS UnaryOp equal to ``ufunc`` on that dtype, a callable taking the Tensor, or None; ``const(dtype)`` gives the value of every position when the result does not depend on the entries (``isnan`` of an integer array), else None. + ``keeps``: the facts about the stored values (``_utils.fact``) that hold for the result + when they hold for ``x``. """ x = _as_operand(x, name) if not isinstance(x, Array): @@ -198,7 +203,7 @@ def _unary(name, ufunc, x, gb=None, *, const=None): R = op(T) else: R = T.apply(op).new(dtype=gb_dtype(rdt)) - return _result(R, fill) + return _utils.keep_facts(x, _result(R, fill), keeps) def _typed(op, kinds): @@ -229,15 +234,20 @@ def get(dtype): def abs(x, /): # noqa: A001 (the Array API's name) # (GraphBLAS's integer abs of the most negative value is undefined in C; numpy wraps) - return _unary("abs", np.abs, x, _first_of(_typed(unary.abs, _FLOAT), _copy_if("ub"))) + gb = _first_of(_typed(unary.abs, _FLOAT), _copy_if("ub")) + return _unary("abs", np.abs, x, gb, keeps=_ALL_FACTS) def negative(x, /): - return _unary("negative", np.negative, x, _typed(unary.ainv, "iufc")) + return _unary("negative", np.negative, x, _typed(unary.ainv, "iufc"), keeps=_ALL_FACTS) def positive(x, /): - return _unary("positive", np.positive, x, _copy_if("iufc")) + return _unary("positive", np.positive, x, _copy_if("iufc"), keeps=_ALL_FACTS) + + +# (a rounding keeps a value finite, or NaN; it may make a zero) +_ROUNDING_KEEPS = ("finite", "nan") def sqrt(x, /): @@ -246,20 +256,23 @@ def sqrt(x, /): def floor(x, /): - return _unary("floor", np.floor, x, _first_of(_typed(unary.floor, _FLOAT), _copy_if("biu"))) + gb = _first_of(_typed(unary.floor, _FLOAT), _copy_if("biu")) + return _unary("floor", np.floor, x, gb, keeps=_ROUNDING_KEEPS) def ceil(x, /): - return _unary("ceil", np.ceil, x, _first_of(_typed(unary.ceil, _FLOAT), _copy_if("biu"))) + gb = _first_of(_typed(unary.ceil, _FLOAT), _copy_if("biu")) + return _unary("ceil", np.ceil, x, gb, keeps=_ROUNDING_KEEPS) def trunc(x, /): - return _unary("trunc", np.trunc, x, _first_of(_typed(unary.trunc, _FLOAT), _copy_if("biu"))) + gb = _first_of(_typed(unary.trunc, _FLOAT), _copy_if("biu")) + return _unary("trunc", np.trunc, x, gb, keeps=_ROUNDING_KEEPS) def round(x, /): # noqa: A001 (the Array API's name) # numpy rounds half to even; C's round (GraphBLAS's) rounds half away from zero - return _unary("round", np.round, x, _copy_if("biu")) + return _unary("round", np.round, x, _copy_if("biu"), keeps=_ROUNDING_KEEPS) def _const_unless(kinds, value): @@ -315,7 +328,8 @@ def reciprocal(x, /): def conj(x, /): - return _unary("conj", np.conj, x, _first_of(_typed(unary.conj, "c"), _copy_if(_REAL))) + gb = _first_of(_typed(unary.conj, "c"), _copy_if(_REAL)) + return _unary("conj", np.conj, x, gb, keeps=_ALL_FACTS) def real(x, /): @@ -567,7 +581,7 @@ def _multiply_intersects(a, b, fa, fb, ct): return False if ct.kind != "f": return True - return _utils.summary(a)["finite"] and _utils.summary(b)["finite"] + return _utils.fact(a, "finite") and _utils.fact(b, "finite") def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): @@ -611,16 +625,16 @@ def _zero_times(spec, y, fy): finite values (for divide, also no 0). NaN never qualifies; the sign of the zero does not matter (an unstored position reads the fill). """ - facts = _utils.summary(y) if spec.name == "divide": - return not facts["zero"] and not facts["nan"] and fy != 0 and not np.isnan(fy) - return facts["finite"] and bool(np.isfinite(fy)) + ok = fy != 0 and not np.isnan(fy) + return ok and not _utils.fact(y, "zero") and not _utils.fact(y, "nan") + return bool(np.isfinite(fy)) and _utils.fact(y, "finite") def _max_min_gb(op): - """GraphBLAS's max and min of floats are C's fmax and fmin, which skip NaN and may take - either zero of 0.0 and -0.0; numpy's propagate NaN and order -0.0 below 0.0. They agree - when neither operand holds NaN or -0.0 (cached per Array), fills included. + """GraphBLAS's max and min of floats are C's fmax and fmin, which skip NaN; numpy's + propagate it. They agree when neither operand holds NaN (a cached fact), fills included. + Either may take either zero of 0.0 and -0.0: the Array API leaves that order unspecified. """ def get(a, b, ct): @@ -630,13 +644,9 @@ def get(a, b, ct): return None for x in (a, b): if isinstance(x, Array): - fill = np.asarray(x._fill) - if np.isnan(fill) or (fill == 0 and np.signbit(fill)): - return None - facts = _utils.summary(x) - if facts["nan"] or facts["negzero"]: + if np.isnan(x._fill) or _utils.fact(x, "nan"): return None - elif isinstance(x, float) and (np.isnan(x) or (x == 0 and np.signbit(x))): + elif isinstance(x, float) and np.isnan(x): return None return op[gb_dtype(ct)] diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py index 523280152..3c7d2432c 100644 --- a/graphblas/array_api/_lazy.py +++ b/graphblas/array_api/_lazy.py @@ -86,9 +86,7 @@ def _known_finite(x): if x.dtype.kind not in "fc": return True facts = x._facts - if facts is not None and "summary" in facts: - return facts["summary"]["finite"] - return None + return None if facts is None else facts.get("finite") def contraction(arrays, build, shape, dtype, fill, spec=None): diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index b78ec5737..3776c7957 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -245,10 +245,6 @@ def _reduce_main(name, x, axis, keepdims, dtype): fill_out = spec.repeat(fill, L, ct) if T.nvals == 0 or math.prod(shape) == 0: return _keepdims(Array._new(Tensor(gb_dtype(ct), shape), fill_out), x, axes, keepdims) - if name in {"min", "max"} and ct.kind == "f" and _signed_zeros(x, fill): - # GraphBLAS's min and max are C's fmin and fmax, which may take either of 0.0 and -0.0 - out = _reduce_numpy(x, axes, np.minimum if name == "min" else np.maximum, fill, ct) - return _keepdims(out, x, axes, keepdims) nan_fix = name in {"min", "max"} and ct.kind == "f" and _has_nan(x) full = T.nvals == x.size if len(axes) == x.ndim: @@ -303,38 +299,8 @@ def _gb_monoid(spec, ct): return spec.gb[gb_dtype(ct)] -def _signed_zeros(x, fill): - """Whether a min or max might meet both zeros: a stored -0.0, or a fill of -0.0.""" - return _utils.summary(x)["negzero"] or bool(fill == 0 and np.signbit(fill)) - - -def _reduce_numpy(x, axes, ufunc, fill, ct): - """``ufunc.reduce`` over ``axes`` by numpy on the stored entries of each fiber, then with - the fill once where a fiber has an unstored position: O(entries log entries). - """ - from ._manipulation import permute_dims - - others = tuple(a for a in range(x.ndim) if a not in axes) - shape = tuple(x.shape[a] for a in others) - L = math.prod(x.shape[a] for a in axes) - coords, values = permute_dims(x, others + tuple(axes))._T().to_coo() - if shape: - fiber = np.ravel_multi_index(coords[: len(others)], shape) - else: - fiber = np.zeros(len(values), np.int64) - starts = np.flatnonzero(np.r_[True, fiber[1:] != fiber[:-1]]) - counts = np.diff(np.r_[starts, len(fiber)]) - with np.errstate(all="ignore"): - r = ufunc.reduceat(values, starts) - r = np.where(counts < L, ufunc(r, fill), r).astype(ct) - if not shape: - return Array._new(Tensor(gb_dtype(ct), ()), r[0]) - idx = np.unravel_index(fiber[starts], shape) - return Array._new(Tensor.from_coo(idx, r, gb_dtype(ct), shape=shape), fill) - - def _has_nan(x): - return _utils.summary(x)["nan"] + return _utils.fact(x, "nan") def sum(x, /, *, axis=None, dtype=None, keepdims=False): # noqa: A001 (the Array API's name) diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index c7c2c898e..5e0dd56c2 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -5,7 +5,7 @@ import numpy as np -from .. import binary, config +from .. import binary, config, monoid, unary from ..core import tensor_layout as tl from ..core.scalar import Scalar from ..core.tensor import Tensor @@ -177,55 +177,60 @@ def with_fill(x, fill): return Array._new(T, fill) -def tensor_values(T): - """The stored values of ``T`` in storage order (a copy).""" - return T._M.to_coo(rows=False, columns=False, sort=False)[2] +# Facts about an Array's stored values (not its fill), each found by one parallel GraphBLAS +# pass and cached on the Array; ``fact`` documents them. For each: the unary op that says +# whether a value has the property, the monoid that combines those answers, and the answer when +# the Array stores nothing. +_FACT_SCANS = { + "finite": (unary.isfinite, monoid.land, True), + "nan": (unary.isnan, monoid.lor, False), + "zero": (None, monoid.lor, False), +} -def summary(x): - """Facts about the stored values of the Array ``x``, computed once (cached on ``x``): - ``finite`` (all finite), ``nonnegative`` (no sign bit set: no negative value, no -0.0), - ``nan`` (some NaN), ``negzero`` (some -0.0) and ``zero`` (some 0.0 or -0.0). +def fact(x, name): + """A fact about the stored values of the Array ``x``, computed once (cached on ``x``): + "finite" (every value finite), "nan" (some value is NaN) or "zero" (some value is 0 or + -0.0). Integers and bools are finite and not NaN without a look; each other fact is one + parallel GraphBLAS pass with no copy of the values (an apply to bool, then a reduction + that stops at the first deciding value). """ - - def compute(x): - T = x._T() - kind = x.dtype.kind - if T.nvals == 0: - return { - "finite": True, - "nonnegative": True, - "nan": False, - "negzero": False, - "zero": False, - } - values = tensor_values(T) - if kind in "biu": - nonneg = kind in "bu" or bool((values >= 0).all()) - zero = not bool(values.all()) - return { - "finite": True, - "nonnegative": nonneg, - "nan": False, - "negzero": False, - "zero": zero, - } - if kind == "c": - parts = [values.real, values.imag] - else: - parts = [values] - signbit = [np.signbit(p) for p in parts] - return { - "finite": all(bool(np.isfinite(p).all()) for p in parts), - "nonnegative": not any(bool(sb.any()) for sb in signbit), - "nan": any(bool(np.isnan(p).any()) for p in parts), - "negzero": any( - bool((sb & (p == 0)).any()) for p, sb in zip(parts, signbit, strict=True) - ), - "zero": not bool(values.all()), - } - - return x._fact("summary", compute) + if name not in _FACT_SCANS: # pragma: no cover (safety) + raise KeyError(name) + return x._fact(name, lambda x: _compute_fact(x, name)) + + +def _compute_fact(x, name): + facts = x._facts + if name != "zero" and x.dtype.kind not in "fc": + return name == "finite" + if name == "nan" and facts.get("finite"): + return False + if name == "finite" and facts.get("nan"): + return False + op, combine, empty = _FACT_SCANS[name] + T = x._T() + if T.nvals == 0: + return empty + dt = gb_dtype(x.dtype) + if op is None: + B = T.apply(binary.eq[dt], right=gb_scalar(0, x.dtype)).new() + else: + B = T.apply(op[dt]).new() + return bool(B.reduce_scalar(combine[lookup_dtype(bool)]).new().value) + + +def keep_facts(src, dst, names=("finite", "nan", "zero")): + """Give ``dst`` the facts ``src`` already knows among ``names`` (for results whose stored + values keep them: ``abs`` and ``negative`` keep all three, a rounding keeps finiteness). + """ + if facts := src._facts: + kept = {k: v for k, v in facts.items() if k in names} + if kept: + if dst._facts is None: + dst._facts = {} + dst._facts.update(kept) + return dst def _x_array(cont, dtype, stack): diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index cc09afd40..eed6a3c3f 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -120,10 +120,11 @@ def check_invariants(x): T._check_invariants() -def assert_array_equal(x, expected): +def assert_array_equal(x, expected, *, zero_signs=True): """``x`` holds ``expected`` exactly (dtype, shape, NaN, signed zeros), but for the sign of a zero at a position ``x`` does not store, which the contract leaves unspecified (DESIGN - D3b). A 0-d array's value is exact. + D3b). A 0-d array's value is exact. ``zero_signs=False`` ignores every zero's sign (as + for max and min, where the Array API leaves the order of -0.0 and 0.0 unspecified). """ expected = np.asarray(expected) check_invariants(x) @@ -131,8 +132,8 @@ def assert_array_equal(x, expected): assert x.shape == expected.shape, (x.shape, expected.shape) dense = x.todense() assert dense.dtype == expected.dtype - unstored = np.zeros(x.shape, dtype=bool) - if x.ndim: + unstored = np.full(x.shape, not zero_signs) + if x.ndim and zero_signs: unstored[...] = True unstored[x.to_coo()[0]] = False _assert_identical(dense, expected, unstored) @@ -338,7 +339,7 @@ def test_binary(name, func, shapes): continue out = getattr(xp, name)(x1, x2) try: - assert_array_equal(out, expected) + assert_array_equal(out, expected, zero_signs=name not in MAX_MIN) except AssertionError as exc: raise AssertionError( f"{name} of {dt1} fill {x1.fill_value!r} and {dt2} fill {x2.fill_value!r}: " @@ -346,6 +347,7 @@ def test_binary(name, func, shapes): ) from None +MAX_MIN = {"maximum", "minimum"} PY_SCALARS = [0, 1, 2, -3, 200, 2**40, True, False, 0.0, -0.0, 0.5, -1.5, np.nan, np.inf, 1j] @@ -367,7 +369,7 @@ def test_binary_python_scalar(name, func, shape): continue out = np_call(getattr(xp, name), *args) try: - assert_array_equal(out, expected) + assert_array_equal(out, expected, zero_signs=name not in MAX_MIN) except AssertionError as exc: raise AssertionError(f"{name}{tuple(map(repr, dargs))}: {exc}") from None @@ -497,27 +499,32 @@ def test_float_multiply_of_nonnegative_values_intersects(): @pytest.mark.parametrize("name", ["maximum", "minimum"]) def test_float_maximum_minimum(name): - # Without NaN or -0.0 GraphBLAS's max and min are numpy's; with them, numpy computes + # Without NaN GraphBLAS's max and min are numpy's; with it, numpy computes. The order of + # -0.0 and 0.0 is unspecified. rng = np.random.default_rng(8) func = getattr(np, name) + + def check(x, expected): + assert_array_equal(x, expected, zero_signs=False) + for fill1, fill2 in [(0.0, 0.0), (2.5, -1.0), (np.inf, 0.0), (0.0, -0.0), (np.nan, 1.0)]: a = np.where(rng.random((4, 5)) < 0.5, fill1, np.round(rng.normal(size=(4, 5)), 1)) b = np.where(rng.random(5) < 0.5, fill2, np.round(rng.normal(size=5), 1)) x, y = make(a, fill1), make(b, fill2) - assert_array_equal(getattr(xp, name)(x, y), func(a, b)) + check(getattr(xp, name)(x, y), func(a, b)) for s in [0.0, -0.0, 1.5, np.nan, -np.inf, 3]: - assert_array_equal(getattr(xp, name)(x, s), func(a, s)) - assert_array_equal(getattr(xp, name)(s, y), func(s, b)) + check(getattr(xp, name)(x, s), func(a, s)) + check(getattr(xp, name)(s, y), func(s, b)) a[0, 0] = -0.0 a[1, 1] = np.nan x = make(a, fill1) - assert_array_equal(getattr(xp, name)(x, y), func(a, b)) - assert_array_equal(getattr(xp, name)(y, x), func(b, a)) + check(getattr(xp, name)(x, y), func(a, b)) + check(getattr(xp, name)(y, x), func(b, a)) x = make(np.array([1.0, 0.0, 2.0]), 0.0) y = make(np.array([0.5, 0.5, 0.5]), 0.0) - assert_array_equal(getattr(xp, name)(x, y), func(x.todense(), y.todense())) + check(getattr(xp, name)(x, y), func(x.todense(), y.todense())) x *= -0.0 # (the cached facts about x's values must not survive) - assert_array_equal(getattr(xp, name)(x, y), func(x.todense(), y.todense())) + check(getattr(xp, name)(x, y), func(x.todense(), y.todense())) def test_float_multiply_keeps_union(): diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 856c05f02..141a9b6bb 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -199,12 +199,36 @@ def test_scalar_cache_is_bounded(): assert _utils.gb_scalar(-0.0, np.dtype(float)).value == 0.0 -def test_summary_of_complex_values(): +def test_facts_of_stored_values(): x = make(np.array([1 + 0j, complex(-0.0, 1), complex(np.nan, 0)]), 0j) - facts = _utils.summary(x) - assert facts == { - "finite": False, "nonnegative": False, "nan": True, "negzero": True, "zero": False - } # fmt: skip + assert [_utils.fact(x, k) for k in ["finite", "nan", "zero"]] == [False, True, False] + assert x._facts == {"finite": False, "nan": True, "zero": False} + y = make(np.array([-0.0, np.inf, 2.0]), 1.0) + assert [_utils.fact(y, k) for k in ["finite", "nan", "zero"]] == [False, False, True] + # NaN and finiteness exclude each other: one is known from the other + z = make(np.array([np.nan, 2.0]), 0.0) + assert _utils.fact(z, "nan") + assert not _utils.fact(z, "finite") + w = make(np.array([3.0, 2.0]), 0.0) + assert _utils.fact(w, "finite") + assert not _utils.fact(w, "nan") + # integers are finite without a look; empty arrays hold nothing + i = make(np.array([0, 5]), 1) + assert _utils.fact(i, "finite") + assert not _utils.fact(i, "nan") + assert _utils.fact(i, "zero") + e = xp.zeros(4) + assert [_utils.fact(e, k) for k in ["finite", "nan", "zero"]] == [True, False, False] + # abs, negative, positive and conj keep all three; roundings keep two + v = make(np.array([-1.5, np.nan, 0.0]), 0.0) + for k in ["finite", "nan", "zero"]: + _utils.fact(v, k) + for func in [xp.abs, xp.negative, xp.positive, xp.conj]: + assert func(v)._facts == v._facts + for func in [xp.floor, xp.ceil, xp.trunc, xp.round]: + assert func(v)._facts == {"finite": False, "nan": True} + assert xp.exp(v)._facts is None + assert xp.abs(xp.zeros(3))._facts is None def test_thread3_argument_checks(): diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 2050a1b70..13bafa84b 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -271,7 +271,7 @@ def test_finite_facts(): # Facts an Array already knows are taken as they are y = make(np.array([1.0, 2.0]), 0.0) assert _lazy._known_finite(y) is None - _utils.summary(y) + _utils.fact(y, "finite") assert _lazy._known_finite(y) is True assert _lazy._known_finite(make(np.array([1, 2]), 0)) is True P = xp.matmul(y, make(np.array([[1.0], [2.0]]), 0.0)) diff --git a/graphblas/tests/test_array_api_statistical.py b/graphblas/tests/test_array_api_statistical.py index 68e26fd81..b915d1f3b 100644 --- a/graphblas/tests/test_array_api_statistical.py +++ b/graphblas/tests/test_array_api_statistical.py @@ -327,8 +327,10 @@ def test_sum_starts_from_positive_zero(): assert_array_equal(xp.sum(xp.full((3,), -0.0)), np.sum(np.full(3, -0.0))) -def test_min_max_order_signed_zeros(): - # GraphBLAS's fmin/fmax take either zero; numpy orders -0.0 below 0.0 (review finding M2) +def test_min_max_of_signed_zeros(): + # GraphBLAS's fmin and fmax take either zero, and numpy orders -0.0 below 0.0: the Array + # API leaves that order unspecified (review finding M2 closed by the contract); NaN + # propagates cases = [ (np.array([-0.0, 0.0]), 1.0), (np.array([0.0, -0.0]), 1.0), @@ -343,7 +345,9 @@ def test_min_max_order_signed_zeros(): for name in ["min", "max"]: for axis in [None, 0, -1]: assert_array_equal( - getattr(xp, name)(x, axis=axis), getattr(np, name)(dense, axis=axis) + getattr(xp, name)(x, axis=axis), + getattr(np, name)(dense, axis=axis), + zero_signs=False, ) From f7bab1b72f44c792f09295e6cc6d2ec139a22898 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 04:14:28 -0600 Subject: [PATCH 35/77] array_api: an SDDMM is computed under S only when no product can overflow multiply(S, P) with a pending product P computed P at S's stored positions whenever P's operands were finite, so a sum of finite products that overflowed to inf outside S read 0 where numpy has 0 * inf, NaN. The push-down now also needs the number of terms of each sum times the largest magnitude of each operand (one GraphBLAS pass per operand) to stay below half the dtype's largest value; otherwise P is computed whole, as numpy would, under the densify limit, since the whole of a product meant for an SDDMM may be dense. --- graphblas/array_api/_lazy.py | 66 ++++++++++++++++---------- graphblas/array_api/_linalg.py | 13 +++-- graphblas/array_api/_utils.py | 10 ++-- graphblas/tests/test_array_api_lazy.py | 65 +++++++++++++++++++++---- 4 files changed, 115 insertions(+), 39 deletions(-) diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py index 3c7d2432c..34aea2059 100644 --- a/graphblas/array_api/_lazy.py +++ b/graphblas/array_api/_lazy.py @@ -11,23 +11,24 @@ - ``multiply(S, P)`` or ``multiply(P, S)``, with S of P's shape and both fills 0, computes P only at S's stored positions (a masked product, then the multiply: an SDDMM). Elsewhere S - is 0, and 0 times P's value there is 0 when that value is finite, which holds when P's - operands are finite (checked, else P is computed whole). A product of finite values that - overflows to inf outside S's stored positions then reads 0 where numpy gives NaN: the same - structural reading of 0 as in D6's 0 * inf. + is 0, and 0 times P's value there is 0 when that value is finite. That holds when the + operands are finite and no sum of products can overflow: the number of terms of each sum + times the largest magnitude of each operand stays below the dtype's largest value (one + GraphBLAS pass per operand for each). Otherwise P is computed whole, as numpy would, under + the densify limit: the whole of a product meant for an SDDMM may be dense. - ``P[key]`` with a basic index (ints and slices) slices the operands along the indexed axes of the result, then contracts (pending again). Everything else reads the Tensor, which computes it. Values do not depend on whether a result -was deferred, but for the overflow above; ``enabled = False`` turns deferral off, for -debugging. +was deferred; ``enabled = False`` turns deferral off, for debugging. """ import numpy as np -from .. import monoid +from .. import monoid, unary from ..core.tensor import Tensor from ..dtypes import FC64, FP64 +from . import _utils from ._array import Array from ._utils import gb_dtype, np_dtype @@ -41,42 +42,53 @@ class Contraction: ``build(*tensors)`` gives its TensorExpression; ``tensors`` are snapshots of the operands' Tensors. ``spec``, ``(terms, output, semiring)`` of an einsum that computes it, lets a slice of the result slice the operands instead (None: no such einsum). ``finite`` holds - each operand's finiteness: True, False, or None when not checked yet. + each operand's finiteness: True, False, or None when not checked yet. ``nterms``: the + number of terms each result position sums (the product of the contracted sizes). """ - __slots__ = "tensors", "build", "gbdt", "spec", "finite" + __slots__ = "tensors", "build", "gbdt", "spec", "finite", "nterms" - def __init__(self, tensors, build, gbdt, spec, finite): + def __init__(self, tensors, build, gbdt, spec, finite, nterms): self.tensors = tensors self.build = build self.gbdt = gbdt self.spec = spec self.finite = finite + self.nterms = nterms def compute(self, mask=None): return self.build(*self.tensors).new(dtype=self.gbdt, mask=mask) def all_finite(self): - """Whether every operand value is finite: then so is every product, but for - overflow. + """Whether every value of the product is finite: every operand value is finite, and + ``nterms`` times the largest magnitude of each operand is below half the largest + value of the result's dtype (the half covers the rounding of the sums). """ for i, ok in enumerate(self.finite): if ok is None: ok = self.finite[i] = _all_finite(self.tensors[i]) if not ok: return False - return True + dt = np_dtype(self.gbdt) + if dt.kind not in "fc": + return True # (integers wrap, as numpy's do: 0 times a wrapped sum is 0) + bound = float(self.nterms) + for T in self.tensors: + bound *= _max_abs(T) # (an overflow here is inf, and fails the test) + return bound < float(np.finfo(dt).max) / 2 def _all_finite(T): + return np_dtype(T.dtype).kind not in "fc" or _utils.tensor_fact(T, "finite") + + +def _max_abs(T): + """The largest magnitude of the stored values of ``T``, as a float.""" + if T.nvals == 0: + return 0.0 kind = np_dtype(T.dtype).kind - if kind not in "fc" or T.nvals == 0: - return True - # A NaN or an infinity makes the sum non-finite. Finite values whose sum overflows make - # this False too, which costs the masked product, not the answer. One parallel pass over - # the values, with no copy. - total = T.reduce_scalar(monoid.plus[FC64 if kind == "c" else FP64]).new().value - return bool(np.isfinite(total)) + A = T.apply(unary.abs[FC64 if kind == "c" else FP64]).new() + return float(A.reduce_scalar(monoid.max[FP64]).new().value) def _known_finite(x): @@ -89,16 +101,18 @@ def _known_finite(x): return None if facts is None else facts.get("finite") -def contraction(arrays, build, shape, dtype, fill, spec=None): +def contraction(arrays, build, shape, dtype, fill, spec=None, nterms=1): """The Array of ``build(*tensors)`` over the arrays' Tensors, with fill ``fill``: pending - when deferral is on and the result has a dimension, else computed now. + when deferral is on and the result has a dimension, else computed now. ``nterms``: the + number of terms each result position sums. """ tensors = [x._T() for x in arrays] gbdt = gb_dtype(dtype) if not enabled or not shape: return Array._new(build(*tensors).new(dtype=gbdt), fill) snapshots = [T.dup() for T in tensors] - node = Contraction(snapshots, build, gbdt, spec, [_known_finite(x) for x in arrays]) + finite = [_known_finite(x) for x in arrays] + node = Contraction(snapshots, build, gbdt, spec, finite, nterms) return Array._deferred(node, shape, dtype, fill) @@ -118,6 +132,10 @@ def masked_multiply(a, b): return None node = P._pending if not node.all_finite(): + _utils.check_dense( + P.size, + "multiply(S, P) where P may hold inf or NaN outside S's entries (P computed whole)", + ) return None PS = Array._new(node.compute(mask=S._T().S), P._fill) from ._elementwise import multiply @@ -162,5 +180,5 @@ def build(*tensors): # (a slice of a finite operand is finite) spec = (tuple(new_terms), out, typed) - sub = Contraction(tensors, build, node.gbdt, spec, list(node.finite)) + sub = Contraction(tensors, build, node.gbdt, spec, list(node.finite), node.nterms) return Array._deferred(sub, shape, x.dtype, x._fill) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 08a81b310..25589c10e 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -292,7 +292,7 @@ def build(Ta, Tb): spec = _matmul_terms(a.shape, b.shape) spec = None if spec is None else (*spec, typed) - return _lazy.contraction((a, b), build, shape, out_dt, fill, spec) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, a.shape[-1]) # Index letters for the einsum that computes a contraction (for slicing a pending result) @@ -403,7 +403,8 @@ def build(Ta, Tb): spec = _tensordot_terms(a.ndim, b.ndim, a_axes, b_axes) spec = None if spec is None else (*spec, typed) - return _lazy.contraction((a, b), build, shape, out_dt, fill, spec) + nterms = math.prod(a.shape[i] for i in a_axes) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, nterms) def vecdot(x1, x2, /, *, axis=-1, semiring=None): @@ -463,7 +464,7 @@ def _vecdot(a, b, axis, shape, out_shape, typed, out_dt, fill): def build(Ta, Tb): return Ta.ewise_mult(Tb, typed.binaryop).new().reduce(typed.monoid, axis=axis) - return _lazy.contraction((a, b), build, out_shape, out_dt, fill) + return _lazy.contraction((a, b), build, out_shape, out_dt, fill, nterms=shape[axis]) def _both_broadcast(sa, sb, shape): @@ -511,4 +512,8 @@ def build(*tensors): return tensor_einsum(subscripts, *tensors, op=typed) terms, output = _einsum_plan(subscripts, shapes, "greedy")[:2] - return _lazy.contraction(ops, build, shape, out_dt, fill, (terms, output, typed)) + sizes = { + c: n for term, s in zip(terms, shapes, strict=True) for c, n in zip(term, s, strict=True) + } + nterms = math.prod(n for c, n in sizes.items() if c not in output) + return _lazy.contraction(ops, build, shape, out_dt, fill, (terms, output, typed), nterms) diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index 5e0dd56c2..0f1986af6 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -208,13 +208,17 @@ def _compute_fact(x, name): return False if name == "finite" and facts.get("nan"): return False + return tensor_fact(x._T(), name) + + +def tensor_fact(T, name): + """``fact`` of the stored values of the Tensor ``T``, by one pass, not cached.""" op, combine, empty = _FACT_SCANS[name] - T = x._T() if T.nvals == 0: return empty - dt = gb_dtype(x.dtype) + dt = T.dtype if op is None: - B = T.apply(binary.eq[dt], right=gb_scalar(0, x.dtype)).new() + B = T.apply(binary.eq[dt], right=gb_scalar(0, np_dtype(dt))).new() else: B = T.apply(op[dt]).new() return bool(B.reduce_scalar(combine[lookup_dtype(bool)]).new().value) diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 13bafa84b..65c2959f3 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -248,15 +248,64 @@ def test_sddmm_falls_back(masks): with np.errstate(invalid="ignore"): assert_close(R, s * (a2 @ b.T)) assert [m is None for m in masks] == [True] - # Overflow to inf outside S: masked, it reads 0 (numpy: NaN), as D6's 0 * inf - big = make(np.full((2, 1), 1e200), 0.0) + + +def test_sddmm_overflow_outside_s(masks): + # A product that may overflow to inf outside S's entries is computed whole: 0 * inf there + # is NaN, as numpy has it s = np.array([[1.0, 0.0], [0.0, 0.0]]) - R = xp.multiply(make(s, 0.0), big @ xp.matrix_transpose(big)) - assert_array_equal(R, [[np.inf, 0.0], [0.0, 0.0]]) - with pytest.MonkeyPatch.context() as mp, np.errstate(invalid="ignore", over="ignore"): - mp.setattr(_lazy, "enabled", False) - R = xp.multiply(make(s, 0.0), big @ xp.matrix_transpose(big)) - assert_array_equal(R, s * (np.full((2, 1), 1e200) @ np.full((1, 2), 1e200))) + a = np.full((2, 1), 1e200) + with np.errstate(invalid="ignore", over="ignore"): + expected = s * (a @ a.T) + assert np.isnan(expected[0, 1]) + for func in [ + lambda A, B: A @ xp.matrix_transpose(B), + lambda A, B: xp.einsum("ik,jk->ij", A, B), + lambda A, B: xp.vecdot(A[:, None, :], B[None, :, :]), + ]: + masks.clear() + R = xp.multiply(make(s, 0.0), func(make(a, 0.0), make(a, 0.0))) + assert_array_equal(R, expected) + assert [m is None for m in masks] == [True] + # The bound counts the terms of each sum (and keeps half the range for rounding): one + # term of 6e153 * 6e153 is 3.6e307, three may reach 1.1e308 + a = np.full((3, 3), 6e153) + for k, pushed in [(1, True), (3, False)]: + masks.clear() + A = make(a[:, :k], 0.0) + with np.errstate(over="ignore", invalid="ignore"): + expected = s[:1, :1] * (a[:1, :k] @ a[:1, :k].T) + R = xp.multiply(make(s[:1, :1], 0.0), A[:1] @ xp.matrix_transpose(A[:1])) + assert_array_equal(R, expected) + assert [m is not None for m in masks] == [pushed] + # Large values well within the bound, float32 too, and complex + for dtype, value in [(np.float64, 1e100), (np.float32, 1e15), (np.complex128, 1e100j)]: + a = np.full((3, 2), value, dtype=dtype) + s3 = np.eye(3, dtype=dtype) + masks.clear() + R = xp.multiply(make(s3, 0), make(a, 0) @ xp.matrix_transpose(make(a, 0))) + assert_array_equal(R, s3 * (a @ a.T)) + assert [m is not None for m in masks] == [True] + masks.clear() + a = np.full((3, 2), 1e20, dtype=np.float32) + with np.errstate(over="ignore", invalid="ignore"): + expected = np.eye(3, dtype=np.float32) * (a @ a.T) + R = xp.multiply(make(np.eye(3, dtype=np.float32), 0), make(a, 0) @ make(a.T, 0)) + assert_array_equal(R, expected) + assert [m is None for m in masks] == [True] + # Computing it whole is guarded: that whole may be dense + big = make(np.full((4, 1), 1e200), 0.0) + eye = make(np.eye(4), 0.0) + with gb.config.set(array_api_densify_limit=10), pytest.raises(xp.DensifyError, match="inf"): + xp.multiply(eye, big @ xp.matrix_transpose(big)) + + +def test_max_abs(): + assert _lazy._max_abs(xp.zeros(3)._T()) == 0.0 + assert _lazy._max_abs(make(np.array([-5, 3], np.int64), 0)._T()) == 5.0 + assert _lazy._max_abs(make(np.array([True, False]), False)._T()) == 1.0 + assert _lazy._max_abs(make(np.array([3 + 4j, 1j]), 0)._T()) == 5.0 + assert _lazy._max_abs(make(np.array([-2.5, np.nan, 1.0]), 0.0)._T()) == 2.5 def test_finite_facts(): From 6b8ccfc64860c1edf538ef35134210b1caa0f6d3 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 04:41:21 -0600 Subject: [PATCH 36/77] array_api: transcendental functions by SuiteSparse's unary ops acos, acosh, asin, asinh, atan, atanh, cos, cosh, exp, log, log2, log10, sin, sinh, tan and tanh of float32, float64, complex64 and complex128 arrays, and expm1 and log1p of the real ones, map their stored values by SuiteSparse's unary ops (C's libm, in parallel) instead of numpy over a copy of the values. They agree with numpy on every special value (NaN, inf, signed zeros, all pairs of them for complex) and within one ulp elsewhere on macOS arm64. SuiteSparse's complex expm1 and log1p are its own formulas, inexact near 0 and off numpy's special values, so they stay on numpy, as do integer inputs (numpy's result dtype there, such as float16 for int8, is its own). A unary function computed by a GraphBLAS op now computes the fill by the same op, so that a stored entry equal to the fill maps exactly to the new fill. Tests compare these functions with numpy within 4 ulp (each of libm and numpy's loops is within about 2 ulp of the exact value), special values exactly. --- graphblas/array_api/_elementwise.py | 62 +++++++++------ graphblas/tests/test_array_api_elementwise.py | 79 ++++++++++++++++++- 2 files changed, 116 insertions(+), 25 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 65e22c032..e3cf7ae9e 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -1,4 +1,6 @@ -"""Element-wise functions: numpy's answer, bit for bit, on every position. +"""Element-wise functions: numpy's answer on every position (DESIGN D3b: NaN, inf and signed +zeros exact, transcendental functions within libm's rounding, the sign of a zero at an unstored +position unspecified). A function of one array maps the stored entries and the fill alike. A function of two arrays takes the union of their entries, an unstored position of one operand taking its fill there, @@ -16,6 +18,7 @@ import numpy as np from .. import binary, unary +from ..core.scalar import Scalar from ..core.tensor import Tensor from . import _lazy, _utils from ._array import Array @@ -183,9 +186,16 @@ def _unary(name, ufunc, x, gb=None, *, const=None, keeps=()): from ._creation import asarray x = asarray(x) + op = gb(x.dtype) if gb is not None else None with np.errstate(all="ignore"): rdt = ufunc(np.empty(0, x.dtype)).dtype - fill = _fill_of(ufunc, rdt, (x,)) + if getattr(op, "opclass", None) == "UnaryOp": + # The fill by the op that maps the entries, so that an entry equal to the fill maps + # to the new fill (libm's rounding may differ from numpy's by an ulp) + s = Scalar.from_value(x._fill, gb_dtype(x.dtype)) + fill = s.apply(op).new(dtype=gb_dtype(rdt)).value + else: + fill = _fill_of(ufunc, rdt, (x,)) check_supported(rdt, f"{name} of {x.dtype}") T = x._T() if T.ndim == 0: @@ -194,7 +204,6 @@ def _unary(name, ufunc, x, gb=None, *, const=None, keeps=()): return _result(_utils.empty_tensor(rdt, T.shape), value) if T.nvals == 0: return _result(_utils.empty_tensor(rdt, T.shape), fill) - op = gb(x.dtype) if gb is not None else None if op is None: R = _utils.map_values(ufunc, [T], rdt) elif op == "copy": @@ -351,33 +360,40 @@ def signbit(x, /): return _unary("signbit", np.signbit, x) -def _numpy_unary(name, ufunc): +def _libm_unary(name, ufunc, op, kinds=_INEXACT): + """``ufunc`` computed by the SuiteSparse unary op ``op`` (C's libm, in parallel) for + floating dtypes of ``kinds``, by numpy otherwise: within libm's rounding of numpy's + answer, special values (NaN, inf, signed zeros) exact (DESIGN D3b). + """ + def func(x, /): - return _unary(name, ufunc, x) + return _unary(name, ufunc, x, _typed(op, kinds)) func.__name__ = name func.__qualname__ = name return func -acos = _numpy_unary("acos", np.arccos) -acosh = _numpy_unary("acosh", np.arccosh) -asin = _numpy_unary("asin", np.arcsin) -asinh = _numpy_unary("asinh", np.arcsinh) -atan = _numpy_unary("atan", np.arctan) -atanh = _numpy_unary("atanh", np.arctanh) -cos = _numpy_unary("cos", np.cos) -cosh = _numpy_unary("cosh", np.cosh) -exp = _numpy_unary("exp", np.exp) -expm1 = _numpy_unary("expm1", np.expm1) -log = _numpy_unary("log", np.log) -log1p = _numpy_unary("log1p", np.log1p) -log2 = _numpy_unary("log2", np.log2) -log10 = _numpy_unary("log10", np.log10) -sin = _numpy_unary("sin", np.sin) -sinh = _numpy_unary("sinh", np.sinh) -tan = _numpy_unary("tan", np.tan) -tanh = _numpy_unary("tanh", np.tanh) +acos = _libm_unary("acos", np.arccos, unary.acos) +acosh = _libm_unary("acosh", np.arccosh, unary.acosh) +asin = _libm_unary("asin", np.arcsin, unary.asin) +asinh = _libm_unary("asinh", np.arcsinh, unary.asinh) +atan = _libm_unary("atan", np.arctan, unary.atan) +atanh = _libm_unary("atanh", np.arctanh, unary.atanh) +cos = _libm_unary("cos", np.cos, unary.cos) +cosh = _libm_unary("cosh", np.cosh, unary.cosh) +exp = _libm_unary("exp", np.exp, unary.exp) +log = _libm_unary("log", np.log, unary.log) +log2 = _libm_unary("log2", np.log2, unary.log2) +log10 = _libm_unary("log10", np.log10, unary.log10) +sin = _libm_unary("sin", np.sin, unary.sin) +sinh = _libm_unary("sinh", np.sinh, unary.sinh) +tan = _libm_unary("tan", np.tan, unary.tan) +tanh = _libm_unary("tanh", np.tanh, unary.tanh) +# SuiteSparse's complex expm1 and log1p are formulas of its own, not libm's: inexact near 0 +# and off numpy's special values (expm1(-inf + nan j) is -1 there, not nan + nan j) +expm1 = _libm_unary("expm1", np.expm1, unary.expm1, _FLOAT) +log1p = _libm_unary("log1p", np.log1p, unary.log1p, _FLOAT) # Binary functions diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index eed6a3c3f..d98bc9a8e 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -156,6 +156,49 @@ def _assert_identical(a, b, unstored): raise AssertionError(f"{bad.sum()} positions differ (position, got, expected): {detail}") +# Functions computed by libm (SuiteSparse's unary ops) for floating dtypes: special values +# (NaN, inf, zeros and their signs) exactly as numpy has them, other values within MAXULP units +# in the last place. libm and numpy's own loops are each within 1 to 2 ulp of the exact value +# for these (glibc's libm-test-ulps lists at most 2 for the real ones in double; numpy's float32 +# SIMD loops document up to 1.5), so two of them differ by up to 4; on macOS arm64 the largest +# difference seen was 1 (tanh, float32 sin and cos, complex log2 and log10). +LIBM = {"acos", "acosh", "asin", "asinh", "atan", "atanh", "cos", "cosh", "exp", "expm1", "log", + "log1p", "log2", "log10", "sin", "sinh", "tan", "tanh"} # fmt: skip +MAXULP = 4 + + +def assert_libm_close(x, expected): + """``x`` holds ``expected`` but for libm's rounding: special values exact, the rest within + MAXULP (complex: each part's special values exact, the rest within MAXULP * eps of the + magnitude). + """ + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape, (x.shape, expected.shape) + got = x.todense() + if got.dtype.kind not in "fc": + _assert_identical(got, expected, np.zeros(got.shape, bool)) + return + if got.dtype.kind == "f": + parts = [(got, expected)] + else: + parts = [(got.real, expected.real), (got.imag, expected.imag)] + with np.errstate(all="ignore"): + scale = np.abs(expected) + for g, e in parts: + special = ~np.isfinite(g) | ~np.isfinite(e) | (g == 0) | (e == 0) + _assert_identical(g[special], e[special], np.zeros(special.sum(), bool)) + g, e = g[~special], e[~special] + if got.dtype.kind == "f": + np.testing.assert_array_max_ulp(g, e, maxulp=MAXULP) + else: + # (a part much smaller than the whole is only as exact as the whole) + s = scale[~special] + s = np.where(np.isfinite(s), s, np.abs(e)) + assert (np.abs(g - e) <= MAXULP * np.finfo(e.dtype).eps * s).all(), (g, e) + + def np_call(func, *args, **kwargs): with np.errstate(all="ignore"), warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -287,7 +330,7 @@ def test_unary(name, func, shape): continue out = getattr(xp, name)(x) try: - assert_array_equal(out, expected) + (assert_libm_close if name in LIBM else assert_array_equal)(out, expected) except AssertionError as exc: raise AssertionError(f"{name} of {dtype} fill {x.fill_value!r}: {exc}") from None @@ -595,8 +638,40 @@ def test_results_do_not_share_writes(): assert_array_equal(y, np.array([2.0, 1.0, 3.0])) +@pytest.mark.parametrize("dtype", ["float32", "float64", "complex64", "complex128"]) +def test_libm_functions_on_wide_values(dtype): + # Values over the whole range, the edges of each domain and the special values + rng = np.random.default_rng(11) + dt = np.dtype(dtype) + info = np.finfo(dt) + edges = [0.0, -0.0, 1.0, -1.0, 0.5, -0.5, 2.0, info.tiny, -info.tiny, info.tiny / 8, + info.max, -info.max, 88.0, 89.0, 709.0, 711.0, -745.0, np.pi, np.nan, np.inf, + -np.inf] # fmt: skip + n = 3000 + v = np.concatenate([ + np.array(edges), + rng.normal(size=n) * 3, + np.exp(rng.uniform(-60, 60, n)) * rng.choice([-1.0, 1.0], n), + rng.uniform(-1.2, 1.2, n), + ]) # fmt: skip + with np.errstate(all="ignore"): + if dt.kind == "c": + w = np.empty(v.size, dt) + w.real, w.imag = v, rng.permutation(v) + v = w + v = v.astype(dt) + x = make(v, 0) + for name in sorted(LIBM): + func = dict(UNARY)[name] + try: + assert_libm_close(getattr(xp, name)(x), np_call(func, v)) + except AssertionError as exc: + raise AssertionError(f"{name} of {dtype}: {exc}") from None + + def test_math_consistency_of_fill(): - # An entry equal to the fill maps to the new fill, bit for bit + # An entry equal to the fill maps to the new fill, bit for bit (DESIGN D10: no value + # depends on which positions are stored), also where libm computes the entries for name, _ in UNARY: if name in {"logical_not", "bitwise_invert"}: continue From 476da9322b41c8b72f02852ffb3bf7f423b21003 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 04:42:57 -0600 Subject: [PATCH 37/77] array_api: a one-sided product stores a broadcast operand in full X * y or X / y, where X's fill is 0 and y broadcasts and does not store every position, took the union of X with y masked by X's entries, which replicated y over the axes it broadcasts along: X of shape (1000, 1e6) with 4 entries times exp(v)[None, :] raised DensifyError at the default limit (and at 100 rows took 500 MB). y is small along those axes, so it is now stored in full (within the densify limit) and the product is an intersection with X's entries: 4 entries, 20 ms. --- graphblas/array_api/_elementwise.py | 10 ++++-- graphblas/tests/test_array_api_elementwise.py | 33 +++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index e3cf7ae9e..b80ebaaf7 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -620,8 +620,14 @@ def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): Tx, Ty = X._T(), Y._T() with np.errstate(all="ignore"): fill = spec.ufunc(fx, fy) if X is a else spec.ufunc(fy, fx) - if not _stores_all(Y) and Y.size <= Tx.nvals + 1024: - # (O(size of Y): every position of Y stored, so the product is an intersection) + broadcast = Ty.shape != shape + if not _stores_all(Y) and ( + Y.size <= Tx.nvals + 1024 or (broadcast and Y.size <= _utils.densify_limit()) + ): + # (O(size of Y): every position of Y stored, so the product is an intersection. + # A broadcast Y is small along its broadcast axes; the union with it would + # replicate it over them: exp(v)[None, :] times X of shape (1000, 1e6) with 4 + # entries replicated v 1000 times) other = 0 if Y._fill != 0 else 1 Ty = _utils.with_fill(Y, other)._T() Ta, Tb = (Tx, Ty) if X is a else (Ty, Tx) diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index d98bc9a8e..12f456b33 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -836,6 +836,39 @@ def test_unstored_zero_signs_in_the_oracle(): assert_array_equal(xp.asarray(-0.0), np.array(0.0)) +def test_one_sided_broadcast_operand_is_stored_in_full(): + # (review finding M9) a broadcast operand that does not store every position is stored in + # full, not replicated over the axes it broadcasts along: X of shape (1000, 1e6) with 4 + # entries times exp(v)[None, :] (fill 1) stores 4 entries, under a densify limit of 2e6 + rng = np.random.default_rng(12) + n = 10**6 + rows, cols = [0, 10, 500, 999], [5, 77, n - 1, 123_456] + X = xp.from_coo((rows, cols), [1.0, 2.0, 3.0, 4.0], (1000, n)) + v = xp.from_coo((np.sort(rng.choice(n, n // 2, replace=False)),), rng.normal(size=n // 2), (n,)) + w = make(np.where(rng.random(n) < 0.5, 2.0, rng.random(n) + 1), 2.0) + vals = np.array([1.0, 2.0, 3.0, 4.0]) + with gb.config.set(array_api_densify_limit=2 * n): + e = xp.exp(v) + ed, wd = e.todense()[cols], w.todense()[cols] + for out, expected in [ + (X * e[None, :], vals * ed), + (e[None, :] * X, ed * vals), + (X / w, vals / wd), + ]: + assert out.nnz == 4 + assert out.fill_value == 0.0 + coords, values = out.to_coo() + assert coords[0].tolist() == rows + np.testing.assert_array_equal(values, expected) + # numpy's answer at every position, on a small one + A = make(np.where(rng.random((6, 50)) < 0.05, rng.normal(size=(6, 50)), 0.0), 0.0) + u = make(np.where(rng.random(50) < 0.5, 1.0, rng.normal(size=50)), 1.0) + with gb.config.set(array_api_densify_limit=60): + out = A * u[None, :] + assert out.nnz == A.nnz + assert_array_equal(out, A.todense() * u.todense()[None, :]) + + def test_clip_fill_uses_numpys_array_loop(): # numpy's 0-d clip breaks a tie differently from its array loop (review finding N3) x = xp.full((3,), 0.0) From 3482f8278e6424d833893d9ecc520e06c86d2922 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 04:45:51 -0600 Subject: [PATCH 38/77] array_api: pow follows the Array API for an exponent of 0.5 numpy computes x ** 0.5 by a square root when the exponent repeats (a Python scalar, a 0-d array, or a stride of 0), so pow(-inf, 0.5) was NaN and pow(-0.0, 0.5) was -0.0 there, where the Array API (and numpy's own general loop) give +inf and +0.0. pow now corrects those two cases. A test checks every special case of pow in the standard with array, 0-d and Python scalar exponents and a repeated base, for float32 and float64. --- graphblas/array_api/_elementwise.py | 23 ++++++-- graphblas/tests/test_array_api_elementwise.py | 56 ++++++++++++++++--- 2 files changed, 66 insertions(+), 13 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index b80ebaaf7..ac6e12b93 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -140,10 +140,9 @@ def _fill_of(ufunc, rdt, operands): (integers to a negative power) and one operand is such an Array, the result's fill is never seen either, and 0 stands in for it. """ - # numpy picks its loop by its operands' shapes (an array and a 0-d exponent of 0.5 is - # sqrt, so pow(-inf, 0.5) is NaN there and inf between two arrays; clip breaks ties by - # its 0-d path otherwise): a fill stands in for an array as a 1-element array, a 0-d - # Array stays 0-d, and a Python scalar stays weak + # numpy picks its loop by its operands' shapes (clip breaks ties by its 0-d path + # otherwise): a fill stands in for an array as a 1-element array, a 0-d Array stays 0-d, + # and a Python scalar stays weak args = [ ( (np.reshape(x._fill, (1,)) if x.ndim else np.asarray(x._fill)) @@ -561,6 +560,20 @@ def _binary_arrays(spec, a, b, op, ct, rdt, fill): return _result(_utils.map_values(spec.ufunc, [A, B], rdt), fill) +def _power(x1, x2): + """The power by numpy, but where numpy takes a square root for an exponent of 0.5 that + repeats (a scalar, or an exponent of stride 0): the Array API's pow(-inf, 0.5) is +inf and + pow(-0.0, 0.5) is +0.0, where sqrt gives NaN and -0.0 (review finding L11). + """ + out = np.power(x1, x2) + if out.dtype.kind == "f": + x1 = np.asarray(x1) + fix = (np.asarray(x2) == 0.5) & ((x1 == -np.inf) | (x1 == 0)) + if fix.any(): + out = np.where(fix, np.where(x1 == 0, 0.0, np.inf), out).astype(out.dtype) + return out + + def _typed_for(op_by_kind): """``gb`` for _Binary: ``op_by_kind[dtype.kind]`` typed for the dtype, else None.""" @@ -742,7 +755,7 @@ def get(a, b, ct): # by zero), integer powers, shifts past the width, and libm-free special values "floor_divide": _Binary("floor_divide", np.floor_divide), "remainder": _Binary("remainder", np.remainder), - "pow": _Binary("pow", np.power), + "pow": _Binary("pow", _power), "atan2": _Binary("atan2", np.arctan2), "hypot": _Binary("hypot", np.hypot), "logaddexp": _Binary("logaddexp", np.logaddexp), diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 12f456b33..b5864b43f 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -878,11 +878,51 @@ def test_clip_fill_uses_numpys_array_loop(): assert_array_equal(xp.clip(x, min=-0.0, max=0.5), np.clip(a, -0.0, 0.5)) -def test_fill_follows_numpys_loop_choice(): - # numpy's power of an array by a 0-d 0.5 is sqrt (pow(-inf, 0.5) is NaN), and between two - # arrays it is pow (inf): an array's fill takes the loop the dense array would - x = make(np.array([[-np.inf, 2.0], [-np.inf, -np.inf]], np.float32), np.float32(-np.inf)) - half = xp.asarray(np.float32(0.5)) - halves = make(np.full((2, 2), 0.5, np.float32), np.float32(0.5)) - assert_array_equal(np_call(xp.pow, x, half), np_call(np.power, x.todense(), np.float32(0.5))) - assert_array_equal(np_call(xp.pow, x, halves), np_call(np.power, x.todense(), halves.todense())) +def _spec_pow(x1, x2): + """The Array API's pow of two Python floats, case by case (its special cases).""" + odd = math.isfinite(x2) and x2 % 2 == 1 + if math.isnan(x2): + return 1.0 if x1 == 1 else math.nan + if x2 == 0: + return 1.0 + if math.isnan(x1): + return math.nan + if math.isinf(x2): + if abs(x1) == 1: + return 1.0 + return math.inf if (abs(x1) > 1) == (x2 > 0) else 0.0 + if x1 == math.inf: + return math.inf if x2 > 0 else 0.0 + if x1 == -math.inf: + if x2 > 0: + return -math.inf if odd else math.inf + return -0.0 if odd else 0.0 + if x1 == 0: + negative = math.copysign(1.0, x1) < 0 and odd + if x2 > 0: + return -0.0 if negative else 0.0 + return -math.inf if negative else math.inf + if x1 < 0 and x2 != int(x2): + return math.nan + return x1**x2 + + +def test_pow_special_cases(): + # The Array API's special cases of pow, with arrays, 0-d arrays and Python scalars as + # exponents (numpy takes sqrt for an exponent of 0.5 that repeats: pow(-inf, 0.5) was NaN + # there, review finding L11) + bases = [np.nan, np.inf, -np.inf, 0.0, -0.0, 2.0, 0.5, 1.0, -1.0, -2.0, -0.5] + exponents = [np.nan, 0.0, -0.0, np.inf, -np.inf, 0.5, 2.0, 3.0, -1.0, -2.0, -3.0, 1.5, -0.5] + for dt in [np.float64, np.float32]: + b = np.array(bases, dt) + for e in exponents: + expected = np.array([_spec_pow(float(v), e) for v in b], dt) + x = make(b, dt(-np.inf)) + for exponent in [e, xp.asarray(dt(e)), xp.full((len(b),), dt(e))]: + assert_array_equal(np_call(xp.pow, x, exponent), expected) + assert_array_equal(np_call(xp.pow, make(b, dt(0.0)), e), expected) + # a base that repeats: a scalar, or a 0-d array + for v in bases: + e = make(np.array(exponents, dt), dt(0.5)) + expected = np.array([_spec_pow(v, float(k)) for k in e.todense()], dt) + assert_array_equal(np_call(xp.pow, xp.asarray(dt(v)), e), expected) From 39e38aa67223c5883a069af55c7f32fda4e5e320 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 04:48:24 -0600 Subject: [PATCH 39/77] array_api: xp.prune drops stored entries equal to the fill An extension: the same values, fewer stored entries. "Equal" is "identical", so with fill 0.0 a stored -0.0 stays, and with fill NaN the stored NaNs go. Integers, bools and a nonzero float fill take one GraphBLAS select; a zero fill with no stored zero (one parallel check) is a copy; other cases compare the values by numpy in place. nnz keeps counting stored entries, explicit fills included, as its docstring now says. --- graphblas/array_api/_array.py | 4 ++- graphblas/array_api/_manipulation.py | 27 ++++++++++++++++- .../tests/test_array_api_manipulation.py | 29 ++++++++++++++++++- 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py index 9f46ed309..e0ef9d871 100644 --- a/graphblas/array_api/_array.py +++ b/graphblas/array_api/_array.py @@ -153,7 +153,9 @@ def fill_value(self): @property def nnz(self): - """The number of stored entries (some may equal the fill value).""" + """The number of stored entries, entries equal to the fill value included (as in + pydata/sparse); ``prune`` drops those. No other function's value depends on it. + """ return self._T().nvals @property diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py index 86c5007dc..c48663cff 100644 --- a/graphblas/array_api/_manipulation.py +++ b/graphblas/array_api/_manipulation.py @@ -4,7 +4,7 @@ import numpy as np -from .. import binary +from .. import binary, select from ..core.tensor import Tensor from . import _utils from ._array import Array @@ -19,6 +19,7 @@ "flip", "moveaxis", "permute_dims", + "prune", "repeat", "reshape", "roll", @@ -354,3 +355,27 @@ def tile(x, repetitions, /): for c, t, n in zip(coords, tiles, x.shape, strict=True): out.append((c[None, :] + (t * n)[:, None]).reshape(-1)) return _from_entries(out, np.tile(values, ntiles), shape, x) + + +def prune(x, /): + """``x`` without the stored entries that equal its fill value (an extension): the same + values, fewer stored entries. "Equal" is "identical": with fill 0.0 a stored -0.0 stays, + and with fill NaN the stored NaNs go. ``nnz`` counts stored entries, these included, + until they are pruned. + """ + if not isinstance(x, Array): + raise TypeError(f"prune needs an Array, not {type(x).__name__}") + T = x._T() + fill = x._fill + if T.nvals == 0: + return x._copy() + kind = x.dtype.kind + if kind in "biu" or (kind == "f" and fill != 0 and not np.isnan(fill)): + # (a value equal to a nonzero number is that number) + return Array._new(T.select(select.valuene, fill).new(), fill) + if fill == 0 and not _utils.fact(x, "zero"): + return x._copy() # (no stored zero of either sign) + keep = _utils.map_values(lambda v: _utils.differs(v, fill), [T], np.dtype(bool)) + R = Tensor(T.dtype, T.shape) + R(keep.V) << T + return Array._new(R, fill) diff --git a/graphblas/tests/test_array_api_manipulation.py b/graphblas/tests/test_array_api_manipulation.py index 8fc56c488..8970be933 100644 --- a/graphblas/tests/test_array_api_manipulation.py +++ b/graphblas/tests/test_array_api_manipulation.py @@ -5,7 +5,7 @@ import graphblas.array_api as xp -from .test_array_api_elementwise import assert_array_equal, inputs, make +from .test_array_api_elementwise import assert_array_equal, inputs, make, same_value @pytest.mark.parametrize( @@ -148,3 +148,30 @@ def test_tile(reps): assert_array_equal(xp.tile(xp.asarray(2.0), (3,)), np.tile(np.array(2.0), 3)) with pytest.raises(ValueError, match="nonnegative"): xp.tile(xp.zeros(2), (-1,)) + + +def test_prune(): + cases = [ + (np.array([3, 1, 3, 0]), 3), + (np.array([True, False, True]), True), + (np.array([2.5, 1.0, 2.5, np.nan, -0.0]), 2.5), + (np.array([1.0, np.nan, 2.0]), 0.0), + (np.array([0.0, -0.0, 1.0, np.nan, 0.0]), 0.0), + (np.array([-0.0, 0.0, 2.0]), -0.0), + (np.array([np.nan, 1.0, np.nan, 0.0]), np.nan), + (np.array([0j, complex(-0.0, 0.0), 1j, complex(0.0, -0.0)]), 0j), + (np.zeros(0), 0.0), + ] + for dense, fill in cases: + # every position stored, explicit fills included + idx = np.nonzero(np.ones(dense.shape, bool)) + x = xp.from_coo(idx, dense[idx], dense.shape, fill_value=fill, dtype=dense.dtype) + out = xp.prune(x) + assert_array_equal(out, dense) + assert out.fill_value == x.fill_value or np.isnan(fill) + expected = sum(not same_value(v, out.fill_value) for v in dense.reshape(-1)) + assert out.nnz == expected + assert x.nnz == dense.size + assert_array_equal(xp.prune(xp.asarray(5.0)), np.array(5.0)) + with pytest.raises(TypeError, match="Array"): + xp.prune(np.zeros(3)) From 526fdeaeea008794a59dd2bde2aab4a7603b7dd6 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 04:52:05 -0600 Subject: [PATCH 40/77] Tensor: a private _map_values; the namespace stops reaching into Tensor internals graphblas.array_api mapped numpy over a Tensor's stored values through the Tensor's Matrix, layout and Container helpers directly. That code is now Tensor._map_values(func, others=(), dtype=None), private, with its own tests, and the namespace calls it. No public Tensor API is added. --- graphblas/array_api/_utils.py | 67 ++---------------------- graphblas/core/tensor.py | 66 +++++++++++++++++++++++ graphblas/tests/test_array_api_errors.py | 6 --- graphblas/tests/test_tensor.py | 30 +++++++++++ 4 files changed, 100 insertions(+), 69 deletions(-) diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index 0f1986af6..522d550d7 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -1,12 +1,10 @@ """Private helpers shared by the modules of ``graphblas.array_api``.""" -import contextlib import math import numpy as np from .. import binary, config, monoid, unary -from ..core import tensor_layout as tl from ..core.scalar import Scalar from ..core.tensor import Tensor from ..core.tensor_ewise import _broadcast_factor @@ -237,66 +235,9 @@ def keep_facts(src, dst, names=("finite", "nan", "zero")): return dst -def _x_array(cont, dtype, stack): - """The value array of an unloaded Container as a numpy view (None when it is empty).""" - return stack.enter_context(tl._VectorArray(cont.x, np.dtype(dtype))).array - - def map_values(func, tensors, dtype): - """A Tensor with the entries of ``tensors`` (all of one shape, layout and pattern) whose - values are ``func(*value_arrays)``, computed by numpy over the stored values. - - The pattern is copied in GraphBLAS and the values written into the copy's own array - through a GxB_Container, so the entries are never rebuilt from coordinates. + """A Tensor with the entries of ``tensors`` (all of one shape and pattern) whose values + are ``func(*value_arrays)`` in ``dtype``, computed by numpy over the stored values + (``Tensor._map_values``: the values in place, the pattern copied in GraphBLAS). """ - T0 = tensors[0] - shape, layout = T0._shape, T0._layout - gbdt = gb_dtype(dtype) - Ms = [T._M for T in tensors] - if T0._M._nvals == 0: - return Tensor._from_M(tl.new_matrix(gbdt, Ms[0]._nrows, Ms[0]._ncols), shape, layout) - if any(T._layout != layout for T in tensors[1:]): - Ms = [Ms[0]] + [tl.relayout(T._M, shape, T._layout, layout)[0] for T in tensors[1:]] - # An iso Matrix has one value; the template for the result is one that holds them all, - # if any does, so that its value array has the length of func's output. - isos = [M.ss.is_iso for M in Ms] - template = Ms[isos.index(False)] if False in isos else Ms[0] - W = tl.new_matrix(gbdt, template._nrows, template._ncols) - W << template - with contextlib.ExitStack() as stack: - # (tensors sharing one Matrix, copy-on-write, unload it once) - unique = {id(M): M for M in Ms} - conts = {key: stack.enter_context(tl._Unloaded(M)) for key, M in unique.items()} - wcont = stack.enter_context(tl._Unloaded(W)) - formats = {c.format for c in conts.values()} | {wcont.format} - nvals = {c.nvals for c in conts.values()} | {wcont.nvals} - if len(formats) == 1 and len(nvals) == 1: - values = { - key: _x_array(conts[key], np_dtype(M.dtype), stack) for key, M in unique.items() - } - arrays = [values[id(M)] for M in Ms] - out = _x_array(wcont, dtype, stack) - with np.errstate(all="ignore"): - result = func(*arrays) - result = np.asarray(result) - if out is not None and result.shape == out.shape: - out[...] = result - return Tensor._from_M(W, shape, layout) - # (formats differ: a fallback through coordinates, which aligns the values by position) - return _map_values_coo(func, tensors, dtype) # pragma: no cover (safety) - - -def _map_values_coo(func, tensors, dtype): - coords = None - arrays = [] - for T in tensors: - idx, vals = T.to_coo() - if coords is None: - coords = idx - arrays.append(vals) - with np.errstate(all="ignore"): - result = np.asarray(func(*arrays), dtype=dtype) - if result.ndim == 0: - result = np.broadcast_to(result, coords[0].shape if coords else (1,)) - T0 = tensors[0] - return Tensor.from_coo(coords, result, gb_dtype(dtype), shape=T0._shape) + return tensors[0]._map_values(func, tensors[1:], gb_dtype(dtype)) diff --git a/graphblas/core/tensor.py b/graphblas/core/tensor.py index 817e12dae..59a351254 100644 --- a/graphblas/core/tensor.py +++ b/graphblas/core/tensor.py @@ -28,6 +28,7 @@ (77M entries), Uber (3.3M), NIPS (3.1M) and Enron (54M) are FROSTT tensors (frostt.io). """ +import contextlib import itertools import math import operator @@ -1119,6 +1120,51 @@ def apply(self, op, right=None, *, left=None): expr_repr="{0.name}.apply({op})", ) + def _map_values(self, func, others=(), dtype=None): + """A Tensor with this Tensor's entries whose values are ``func(*value_arrays)``: the + stored values of this Tensor and of ``others`` (Tensors of the same shape and pattern), + as numpy arrays in one order, in the result's ``dtype`` (default: this Tensor's). + + Private, for ``graphblas.array_api`` (numpy over the values of entries GraphBLAS has + no operator for). The pattern is copied in GraphBLAS and the values written into the + copy's own array through a GxB_Container, so the entries are never rebuilt from + coordinates; the inputs' arrays are views, not copies. + """ + tensors = [self, *others] + gbdt = self.dtype if dtype is None else lookup_dtype(dtype) + shape, layout = self._shape, self._layout + Ms = [T._M for T in tensors] + if self._M._nvals == 0: + return Tensor._from_M(tl.new_matrix(gbdt, Ms[0]._nrows, Ms[0]._ncols), shape, layout) + if any(T._layout != layout for T in others): + Ms = [Ms[0]] + [tl.relayout(T._M, shape, T._layout, layout)[0] for T in others] + # An iso Matrix has one value; the template for the result is one that holds them all, + # if any does, so that its value array has the length of func's output. + isos = [M.ss.is_iso for M in Ms] + template = Ms[isos.index(False)] if False in isos else Ms[0] + W = tl.new_matrix(gbdt, template._nrows, template._ncols) + W << template + with contextlib.ExitStack() as stack: + # (tensors sharing one Matrix, copy-on-write, unload it once) + unique = {id(M): M for M in Ms} + conts = {key: stack.enter_context(tl._Unloaded(M)) for key, M in unique.items()} + wcont = stack.enter_context(tl._Unloaded(W)) + formats = {c.format for c in conts.values()} | {wcont.format} + nvals = {c.nvals for c in conts.values()} | {wcont.nvals} + if len(formats) == 1 and len(nvals) == 1: + values = { + key: _value_array(conts[key], M.dtype, stack) for key, M in unique.items() + } + arrays = [values[id(M)] for M in Ms] + out = _value_array(wcont, gbdt, stack) + with np.errstate(all="ignore"): + result = np.asarray(func(*arrays)) + if out is not None and result.shape == out.shape: + out[...] = result + return Tensor._from_M(W, shape, layout) + # (formats differ: a fallback through coordinates, which aligns the values by position) + return _map_values_coo(func, tensors, gbdt) # pragma: no cover (safety) + # Reductions def reduce(self, op=monoid.plus, axis=None, *, keepdims=False): """Reduce over ``axis`` (None: all axes, giving a 0-d Tensor) with the monoid ``op``, @@ -1254,6 +1300,26 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * ) +def _value_array(cont, dtype, stack): + """The value array of an unloaded Container as a numpy view (None when it is empty).""" + return stack.enter_context(tl._VectorArray(cont.x, np.dtype(dtype.np_type))).array + + +def _map_values_coo(func, tensors, dtype): + coords = None + arrays = [] + for T in tensors: + idx, vals = T.to_coo() + if coords is None: + coords = idx + arrays.append(vals) + with np.errstate(all="ignore"): + result = np.asarray(func(*arrays), dtype=dtype.np_type) + if result.ndim == 0: + result = np.broadcast_to(result, coords[0].shape if coords else (1,)) + return Tensor.from_coo(coords, result, dtype, shape=tensors[0]._shape) + + def _ewise_infix(left, right, op, kind): ltensor = _as_tensor(left, "infix operator") rtensor = _as_tensor(right, "infix operator") diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 141a9b6bb..0117a7a90 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -183,12 +183,6 @@ def test_map_values_storage_cases(): np.testing.assert_array_equal(out.to_coo()[1], [2.0, 4.0]) E = Tensor(T.dtype, (2, 3)) assert _utils.map_values(np.sin, [E], np.dtype(float)).nvals == 0 - out = _utils._map_values_coo(np.add, [T, T], np.dtype(float)) - np.testing.assert_array_equal(out.to_coo()[1], [2.0, 4.0]) - out = _utils._map_values_coo(np.negative, [T], np.dtype(float)) - np.testing.assert_array_equal(out.to_coo()[1], [-1.0, -2.0]) - out = _utils._map_values_coo(lambda v: np.float64(7.0), [T], np.dtype(float)) - np.testing.assert_array_equal(out.to_coo()[1], [7.0, 7.0]) assert _utils.full_tensor(np.dtype(float), (0, 3), 1.0).nvals == 0 diff --git a/graphblas/tests/test_tensor.py b/graphblas/tests/test_tensor.py index 3b3a6cec5..bad951547 100644 --- a/graphblas/tests/test_tensor.py +++ b/graphblas/tests/test_tensor.py @@ -2816,6 +2816,36 @@ def test_tensordot_semirings_masks_and_errors(): assert ("ttv", "entries") in {(r["point"], r["choice"]) for r in recs} +def test_private_map_values(): + # Tensor._map_values (for graphblas.array_api): numpy over the stored values, in place + rng = np.random.default_rng(33) + X, xv, xp = random_tensor((4, 5, 3), rng, density=0.4) + Y = Tensor(dtypes.FP64, (4, 5, 3)) + Y(X.S) << X.apply(binary.plus, right=1.0) # (X's pattern) + Y = Y.relayout(((2,), (0, 1))) # (another layout) + R = X._map_values(np.hypot, [Y]) + check(R, np.where(xp, np.hypot(xv, xv + 1), 0), xp, approx=True) + assert R.layout == X.layout + # one Matrix shared by two Tensors, a dtype change, an iso Tensor, and no entries + R = X._map_values(lambda a, b: a * b > 1, [X.dup()], dtype=dtypes.BOOL) + check(R, np.where(xp, xv * xv > 1, False), xp) + iso = Tensor.from_coo(([0, 2], [1, 3]), 2.5, shape=(3, 4)) + assert iso.ss.is_iso + check(iso._map_values(np.sqrt), np.where(iso.to_dense(0) != 0, np.sqrt(2.5), 0), + iso.to_dense(0) != 0, approx=True) # fmt: skip + E = Tensor(dtypes.INT64, (2, 3)) + assert E._map_values(np.negative, dtype=dtypes.FP32).dtype == dtypes.FP32 + assert E._map_values(np.negative).nvals == 0 + # The fallback through coordinates + from graphblas.core.tensor import _map_values_coo + + T = Tensor.from_coo([[0, 1], [1, 2]], [1.0, 2.0], shape=(2, 3)) + assert _map_values_coo(np.add, [T, T], dtypes.FP64).to_coo()[1].tolist() == [2.0, 4.0] + assert _map_values_coo(np.negative, [T], dtypes.FP64).to_coo()[1].tolist() == [-1.0, -2.0] + out = _map_values_coo(lambda v: np.float64(7.0), [T], dtypes.FP64) + assert out.to_coo()[1].tolist() == [7.0, 7.0] + + def test_masked_contraction_takes_a_transpose_descriptor(): # Under a mask, an operand stored as the transpose of what the product needs is not # copied (SuiteSparse's masked dot product, an SDDMM); the values are the masked product From 742ec69204d7f1ebc2fd07c7731604bf360377ef Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 05:21:35 -0600 Subject: [PATCH 41/77] array_api: concat by one GxB_Matrix_concat of the stored Matrices concat (and stack, which concatenates) exported every part's coordinates, joined them in numpy and rebuilt the result. A private Tensor._concat now joins the parts' stored Matrices in one GxB_Matrix_concat: as row blocks when the axis leads the row group of the first part's layout, as column blocks when it leads the column group, after moving parts stored in another layout (or every part, to a layout with the axis leading the rows, when it leads neither). Two arrays of 1e7 stored values each: 5e6 x 5e6 along axis 0, 0.239 s to 0.085 s at 1 thread and 0.111 s to 0.016 s at 8; along axis 1, 0.818 s to 0.077 s and 0.178 s to 0.014 s; 1000^3 along its middle axis (moved first), 0.618 s to 0.068 s and 0.224 s to 0.024 s. --- graphblas/array_api/_manipulation.py | 20 ++----------- graphblas/core/tensor.py | 30 +++++++++++++++++++ .../tests/test_array_api_manipulation.py | 24 +++++++++++++++ 3 files changed, 57 insertions(+), 17 deletions(-) diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py index c48663cff..dc39b62e3 100644 --- a/graphblas/array_api/_manipulation.py +++ b/graphblas/array_api/_manipulation.py @@ -152,7 +152,7 @@ def _broadcast_to(x, shape): def concat(arrays, /, *, axis=0): - """Join arrays along ``axis`` (None: flattened first), O(entries). + """Join arrays along ``axis`` (None: flattened first), O(entries), in parallel. The result's fill is the one held by the most positions; a part with another fill stores its fill explicitly (O(that part's size), under the densify limit). @@ -181,22 +181,8 @@ def concat(arrays, /, *, axis=0): dtype = result_type(*arrays) arrays = [a if a.dtype == dtype else cast(a, dtype) for a in arrays] fill = _common_fill(arrays) - coords, values = [], [] - offset = 0 - for a in arrays: - T = _utils.with_fill(a, fill)._T() - if T.nvals: - # (stored order: no relayout to C order; the build sorts) - idx, vals = T.to_coo(sort=False) - idx = list(idx) - idx[axis] = idx[axis] + offset - coords.append(idx) - values.append(vals) - offset += a.shape[axis] - if not values: - return Array._new(Tensor(gb_dtype(dtype), shape), fill) - coords = tuple(np.concatenate([c[i] for c in coords]) for i in range(ndim)) - T = Tensor.from_coo(coords, np.concatenate(values), gb_dtype(dtype), shape=shape) + # (one GxB_Matrix_concat of the stored Matrices; a part stored in another layout moves) + T = Tensor._concat([_utils.with_fill(a, fill)._T() for a in arrays], axis) return Array._new(T, fill) diff --git a/graphblas/core/tensor.py b/graphblas/core/tensor.py index 59a351254..ef43e6862 100644 --- a/graphblas/core/tensor.py +++ b/graphblas/core/tensor.py @@ -1165,6 +1165,36 @@ def _map_values(self, func, others=(), dtype=None): # (formats differ: a fallback through coordinates, which aligns the values by position) return _map_values_coo(func, tensors, gbdt) # pragma: no cover (safety) + @classmethod + def _concat(cls, tensors, axis): + """``tensors`` (of one dtype, and of one shape but along ``axis``) joined along + ``axis``, by one GxB_Matrix_concat of their stored Matrices: row blocks when ``axis`` + leads the row group of the first one's layout, column blocks when it leads the column + group. Otherwise every part moves to a layout with ``axis`` leading the rows first. + The others take the first one's layout. Private, for ``graphblas.array_api``. + """ + T0 = tensors[0] + shape = list(T0._shape) + shape[axis] = sum(T._shape[axis] for T in tensors) + shape = tuple(shape) + rows, cols = T0._layout + if rows[:1] == (axis,) or cols[:1] == (axis,): + layout = T0._layout + else: + layout = Layout( + (axis, *(a for a in rows if a != axis)), tuple(a for a in cols if a != axis) + ) + by_rows = layout.rows[:1] == (axis,) + Ms = [ + T._M if T._layout == layout else tl.relayout(T._M, T._shape, T._layout, layout)[0] + for T in tensors + if T._shape[axis] # (a part of size 0 along the axis adds nothing) + ] + R = tl.new_matrix(T0.dtype, tl.prod(shape, layout.rows), tl.prod(shape, layout.cols)) + if Ms: + R.ss.concat([[M] for M in Ms] if by_rows else [Ms]) + return cls._from_M(R, shape, layout) + # Reductions def reduce(self, op=monoid.plus, axis=None, *, keepdims=False): """Reduce over ``axis`` (None: all axes, giving a 0-d Tensor) with the monoid ``op``, diff --git a/graphblas/tests/test_array_api_manipulation.py b/graphblas/tests/test_array_api_manipulation.py index 8970be933..fecdce538 100644 --- a/graphblas/tests/test_array_api_manipulation.py +++ b/graphblas/tests/test_array_api_manipulation.py @@ -175,3 +175,27 @@ def test_prune(): assert_array_equal(xp.prune(xp.asarray(5.0)), np.array(5.0)) with pytest.raises(TypeError, match="Array"): xp.prune(np.zeros(3)) + + +@pytest.mark.parametrize("axis", [0, 1, 2]) +def test_concat_any_layout(axis): + # One GxB concat of the stored Matrices, whatever layout each part is stored in (a part in + # another layout moves), parts of size 0 along the axis included + from graphblas import Tensor + + rng = np.random.default_rng(7) + shapes = [[4, 3, 5], [4, 3, 5], [4, 3, 5]] + shapes[1][axis] = 2 + shapes[2][axis] = 0 + dense = [np.where(rng.random(s) < 0.4, rng.normal(size=s), 0.0) for s in shapes] + layouts = [((0,), (1, 2)), ((0, 1), (2,)), ((2, 0), (1,)), ((1,), (2, 0))] + for la in layouts: + for lb in layouts: + parts = [] + for d, lay in zip(dense, [la, lb, la], strict=True): + idx = np.nonzero(d) + T = Tensor.from_coo(idx, d[idx], shape=d.shape).relayout(lay) + parts.append(xp.asarray(T)) + assert_array_equal(xp.concat(parts, axis=axis), np.concatenate(dense, axis=axis)) + empty = [xp.zeros((0, 3)), xp.zeros((0, 3))] + assert_array_equal(xp.concat(empty), np.zeros((0, 3))) From e16d7cd449777fa44fd0796cf67d0c2dd4839908 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 05:27:19 -0600 Subject: [PATCH 42/77] array_api: complex transcendental functions stay on numpy The complex functions of C libraries handle the many special cases of the Array API (pairs of infinities, NaNs and signed zeros in both parts) differently from platform to platform, and SuiteSparse's complex expm1 and log1p are formulas of its own, inexact near 0. On macOS they matched numpy on every special value probed, but numpy is the oracle the tests can check on every platform, so only real floating arrays take SuiteSparse's unary ops for these functions; complex ones keep numpy's answers, as before. --- graphblas/array_api/_elementwise.py | 21 +++++++++++-------- graphblas/tests/test_array_api_elementwise.py | 2 +- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index ac6e12b93..b1b7bdcec 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -359,14 +359,19 @@ def signbit(x, /): return _unary("signbit", np.signbit, x) -def _libm_unary(name, ufunc, op, kinds=_INEXACT): - """``ufunc`` computed by the SuiteSparse unary op ``op`` (C's libm, in parallel) for - floating dtypes of ``kinds``, by numpy otherwise: within libm's rounding of numpy's - answer, special values (NaN, inf, signed zeros) exact (DESIGN D3b). +def _libm_unary(name, ufunc, op): + """``ufunc`` computed by the SuiteSparse unary op ``op`` (C's libm, in parallel) for real + floating dtypes, by numpy otherwise: within libm's rounding of numpy's answer, special + values (NaN, inf, signed zeros) exact (DESIGN D3b). + + Complex values stay on numpy: C libraries treat the many special cases of complex + arguments differently from platform to platform, and SuiteSparse's complex expm1 and + log1p are formulas of its own, inexact near 0. Integer inputs stay too (numpy's result + dtype there, such as float16 for int8, is its own). """ def func(x, /): - return _unary(name, ufunc, x, _typed(op, kinds)) + return _unary(name, ufunc, x, _typed(op, _FLOAT)) func.__name__ = name func.__qualname__ = name @@ -389,10 +394,8 @@ def func(x, /): sinh = _libm_unary("sinh", np.sinh, unary.sinh) tan = _libm_unary("tan", np.tan, unary.tan) tanh = _libm_unary("tanh", np.tanh, unary.tanh) -# SuiteSparse's complex expm1 and log1p are formulas of its own, not libm's: inexact near 0 -# and off numpy's special values (expm1(-inf + nan j) is -1 there, not nan + nan j) -expm1 = _libm_unary("expm1", np.expm1, unary.expm1, _FLOAT) -log1p = _libm_unary("log1p", np.log1p, unary.log1p, _FLOAT) +expm1 = _libm_unary("expm1", np.expm1, unary.expm1) +log1p = _libm_unary("log1p", np.log1p, unary.log1p) # Binary functions diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index b5864b43f..edad45721 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -161,7 +161,7 @@ def _assert_identical(a, b, unstored): # in the last place. libm and numpy's own loops are each within 1 to 2 ulp of the exact value # for these (glibc's libm-test-ulps lists at most 2 for the real ones in double; numpy's float32 # SIMD loops document up to 1.5), so two of them differ by up to 4; on macOS arm64 the largest -# difference seen was 1 (tanh, float32 sin and cos, complex log2 and log10). +# difference seen was 1 (tanh, float32 sin and cos). Complex values stay on numpy. LIBM = {"acos", "acosh", "asin", "asinh", "atan", "atanh", "cos", "cosh", "exp", "expm1", "log", "log1p", "log2", "log10", "sin", "sinh", "tan", "tanh"} # fmt: skip MAXULP = 4 From 8f152c462c96b1a4ca5c039f437847e2049a0c08 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 05:32:02 -0600 Subject: [PATCH 43/77] array_api: a "..." that covers no axis separates advanced indices, as in numpy x[:, 1, ..., [0, 2]] of a 3-d x put the broadcast advanced axis in place, as if the int and the index array were next to each other; numpy puts it first, since the "..." between them separates them even when it covers no axis. getitem and setitem now keep that separation. A test compares both with numpy over every key of up to 4 parts mixing ints, slices, integer arrays, "..." and None that numpy accepts for a 3-d and a 4-d array (setitem without None). --- graphblas/array_api/_indexing.py | 31 +++++++++++---- graphblas/tests/test_array_api_indexing.py | 44 ++++++++++++++++++++++ 2 files changed, 68 insertions(+), 7 deletions(-) diff --git a/graphblas/array_api/_indexing.py b/graphblas/array_api/_indexing.py index 4c619ecc9..4f8432125 100644 --- a/graphblas/array_api/_indexing.py +++ b/graphblas/array_api/_indexing.py @@ -46,6 +46,18 @@ def __init__(self, length=1): self.length = length +class _Gap: + """A ``...`` that covers no axis. It takes no axis, but separates the advanced indices on + either side of it, as in numpy: ``x[:, 1, ..., [0, 2]]`` of a 3-d ``x`` puts the + broadcast advanced axis first. + """ + + __slots__ = () + + +_GAP = _Gap() + + def _is_int(k): return isinstance(k, (int, np.integer)) and not isinstance(k, (bool, np.bool_)) @@ -138,11 +150,11 @@ def width(k): if i is None: items.extend([slice(None)] * (x.ndim - consumed)) else: - items[i : i + 1] = [slice(None)] * (x.ndim - consumed) + items[i : i + 1] = [slice(None)] * (x.ndim - consumed) or [_GAP] out = [] axis = 0 for k in items: - if isinstance(k, _NewAxis): + if isinstance(k, (_NewAxis, _Gap)): out.append(k) continue if isinstance(k, (np.ndarray, Array)) and k.dtype.kind == "b": @@ -196,6 +208,8 @@ def _finish(out, items, nadv): for k in items: if isinstance(k, (_NewAxis, slice)): comps.append(k) + elif isinstance(k, _Gap): + continue elif not placed: comps.append("adv") placed = True @@ -231,7 +245,7 @@ def _basic_shape(shape, axes_key): # getitem def getitem(x, key): items = _normalize(x, key) - axes_key = [k for k in items if not isinstance(k, _NewAxis)] + axes_key = [k for k in items if not isinstance(k, (_NewAxis, _Gap))] if _has_adv(items): out, nadv = _advanced_extract(x, axes_key) else: @@ -291,7 +305,7 @@ def setitem(x, key, value): raise TypeError(type(x)) value = _setitem_value(x, value) items = _normalize(x, key) - axes_key = [k for k in items if not isinstance(k, _NewAxis)] + axes_key = [k for k in items if not isinstance(k, (_NewAxis, _Gap))] if any(isinstance(k, _NewAxis) and k.length == 0 for k in items): return # (a False index selects nothing) if x.ndim == 0: @@ -310,7 +324,7 @@ def setitem(x, key, value): x, tuple(k.index if isinstance(k, _Adv) else k for k in axes_key), value ) elif adv: - _scatter_setitem(x, axes_key, value) + _scatter_setitem(x, [k for k in items if not isinstance(k, _NewAxis)], value) else: _region_setitem(x, tuple(axes_key), value) x._facts = None @@ -371,8 +385,11 @@ def _region_setitem(x, key, value): def _scatter_setitem(x, axes_key, value): """Assign by several index arrays: write every selected position.""" - # numpy lays out the selection; index broadcast views of the coordinates with the key - key = tuple(k.index if isinstance(k, _Adv) else k for k in axes_key) + # numpy lays out the selection; index broadcast views of the coordinates with the key (a + # ``...`` that covers no axis stays: it decides where the advanced axes go) + key = tuple( + k.index if isinstance(k, _Adv) else Ellipsis if isinstance(k, _Gap) else k for k in axes_key + ) coords = [] for a, n in enumerate(x.shape): grid = np.arange(n, dtype=np.int64).reshape([n if b == a else 1 for b in range(x.ndim)]) diff --git a/graphblas/tests/test_array_api_indexing.py b/graphblas/tests/test_array_api_indexing.py index 243fa41d4..09bac62d4 100644 --- a/graphblas/tests/test_array_api_indexing.py +++ b/graphblas/tests/test_array_api_indexing.py @@ -1,5 +1,7 @@ """graphblas.array_api indexing (get and set), take and take_along_axis against numpy.""" +import itertools + import numpy as np import pytest @@ -220,3 +222,45 @@ def test_take_along_axis(axis): ) with pytest.raises(ValueError, match="dimensions"): xp.take_along_axis(xp.zeros((2, 2)), xp.asarray([0])) + + +def _index_patterns(): + """Keys mixing ints, slices, integer arrays, ``...`` and ``None`` with at least one + integer array, as numpy accepts them for a 4-d array. + """ + parts = [slice(None), 1, np.array([0, 2]), np.array([[1], [0]]), Ellipsis, None, slice(0, 2)] + for n in range(1, 5): + for key in itertools.product(parts, repeat=n): + if not any(isinstance(k, np.ndarray) for k in key): + continue + if sum(k is Ellipsis for k in key) > 1: + continue + yield key + + +def test_mixed_keys_order_axes_as_numpy(): + # (shim finding B1) the advanced axes go first when the advanced indices are apart, a + # "..." that covers no axis included; in place when they are next to each other + count = 0 + for a, key in itertools.chain( + ((np.arange(36.0).reshape(3, 4, 3) + 1, k) for k in _index_patterns()), + ((np.arange(108.0).reshape(3, 4, 3, 3) + 1, k) for k in _index_patterns()), + ): + x = make(a, 0.0) + try: + expected = a[key] + except IndexError: + continue + count += 1 + assert_array_equal(x[key], expected) + if any(k is None for k in key): + continue # (a value for a key with new axes is not read with them yet) + b = a.copy() + value = np.arange(expected.size, dtype=np.float64).reshape(expected.shape) - 50 + b[key] = value + y = make(a, 0.0) + y[key] = make(value, 0.0) if value.ndim else xp.asarray(value) + assert_array_equal(y, b) + assert count > 2000 + x = make(np.arange(24.0).reshape(2, 3, 4), 0.0) + assert_array_equal(x[:, 1, ..., [0, 2]], np.array([[4.0, 16.0], [6.0, 18.0]])) From 97405747b0d692f9da5e1381115170b66a424a45 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 05:34:27 -0600 Subject: [PATCH 44/77] array_api: an SDDMM asks for dot products when they are much cheaper multiply(S, A @ B) with B held by row (B of shape (k, n), not the transpose of an (n, k) array) ran SuiteSparse's masked saxpy, all of A @ B's work: full A (4000 x 32) and B (32 x 4000) under a mask of 4e4 entries took 0.6 s at 1 thread. When the masked product is one product of its two operands (matmul, tensordot, an einsum with no permutation or sum after it) and dot products at the mask's entries cost under a quarter of a saxpy over the whole product (by the operands' and the mask's entry counts), the product now asks for SuiteSparse's dot method, which transposes B itself: 0.6 ms. The SAPS run found it. --- graphblas/array_api/_lazy.py | 48 +++++++++++++++++++++----- graphblas/array_api/_linalg.py | 9 +++-- graphblas/tests/test_array_api_lazy.py | 46 ++++++++++++++++++++++-- 3 files changed, 88 insertions(+), 15 deletions(-) diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py index 34aea2059..86e201831 100644 --- a/graphblas/array_api/_lazy.py +++ b/graphblas/array_api/_lazy.py @@ -44,20 +44,22 @@ class Contraction: slice of the result slice the operands instead (None: no such einsum). ``finite`` holds each operand's finiteness: True, False, or None when not checked yet. ``nterms``: the number of terms each result position sums (the product of the contracted sizes). + ``product``: whether it is one product of its two operands, which takes the mask. """ - __slots__ = "tensors", "build", "gbdt", "spec", "finite", "nterms" + __slots__ = "tensors", "build", "gbdt", "spec", "finite", "nterms", "product" - def __init__(self, tensors, build, gbdt, spec, finite, nterms): + def __init__(self, tensors, build, gbdt, spec, finite, nterms, product=False): self.tensors = tensors self.build = build self.gbdt = gbdt self.spec = spec self.finite = finite self.nterms = nterms + self.product = product - def compute(self, mask=None): - return self.build(*self.tensors).new(dtype=self.gbdt, mask=mask) + def compute(self, mask=None, **opts): + return self.build(*self.tensors).new(dtype=self.gbdt, mask=mask, **opts) def all_finite(self): """Whether every value of the product is finite: every operand value is finite, and @@ -101,10 +103,10 @@ def _known_finite(x): return None if facts is None else facts.get("finite") -def contraction(arrays, build, shape, dtype, fill, spec=None, nterms=1): +def contraction(arrays, build, shape, dtype, fill, spec=None, nterms=1, product=False): """The Array of ``build(*tensors)`` over the arrays' Tensors, with fill ``fill``: pending when deferral is on and the result has a dimension, else computed now. ``nterms``: the - number of terms each result position sums. + number of terms each result position sums; ``product``: see ``Contraction``. """ tensors = [x._T() for x in arrays] gbdt = gb_dtype(dtype) @@ -112,7 +114,7 @@ def contraction(arrays, build, shape, dtype, fill, spec=None, nterms=1): return Array._new(build(*tensors).new(dtype=gbdt), fill) snapshots = [T.dup() for T in tensors] finite = [_known_finite(x) for x in arrays] - node = Contraction(snapshots, build, gbdt, spec, finite, nterms) + node = Contraction(snapshots, build, gbdt, spec, finite, nterms, product) return Array._deferred(node, shape, dtype, fill) @@ -137,12 +139,30 @@ def masked_multiply(a, b): "multiply(S, P) where P may hold inf or NaN outside S's entries (P computed whole)", ) return None - PS = Array._new(node.compute(mask=S._T().S), P._fill) + TS = S._T() + PS = Array._new(node.compute(mask=TS.S, **_dot_hint(node, TS.nvals)), P._fill) from ._elementwise import multiply return multiply(PS, S) if P is a else multiply(S, PS) +def _dot_hint(node, nmask): + """``axb_method="dot"`` for a product under a mask of ``nmask`` entries when dot products + at the mask's entries cost well below a saxpy over the whole product, else nothing. + + SuiteSparse takes a masked saxpy for A @ B with B held by row, all of A @ B's work: an + SDDMM with full A (n x 32) and B (32 x n), n = 4000 and 4e4 mask entries, took 0.6 s, + against 0.6 ms as dot products (SuiteSparse transposes B for them). + """ + if not node.product: + return {} + Ta, Tb = node.tensors + k = max(node.nterms, 1) + saxpy = Ta.nvals * Tb.nvals / k + dot = nmask * k * (Ta.nvals / max(Ta.size, 1) + Tb.nvals / max(Tb.size, 1)) + return {"axb_method": "dot"} if 4 * dot < saxpy else {} + + def sliced(x, axes_key): """``x[axes_key]`` (an int or a slice per axis) of a pending ``x`` that an einsum computes: its operands sliced along the indexed axes, contracted (pending again). None @@ -180,5 +200,15 @@ def build(*tensors): # (a slice of a finite operand is finite) spec = (tuple(new_terms), out, typed) - sub = Contraction(tensors, build, node.gbdt, spec, list(node.finite), node.nterms) + product = _one_product(new_terms, out) + sub = Contraction(tensors, build, node.gbdt, spec, list(node.finite), node.nterms, product) return Array._deferred(sub, shape, x.dtype, x._fill) + + +def _one_product(terms, output): + """Whether an einsum of ``terms`` into ``output`` is one product of two operands, whose + result needs no permutation or sum after it (so a mask goes into the product). + """ + from ..core.tensor_einsum import _direct_axes + + return len(terms) == 2 and _direct_axes(*terms, output) is not None diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 25589c10e..1f378924f 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -292,7 +292,7 @@ def build(Ta, Tb): spec = _matmul_terms(a.shape, b.shape) spec = None if spec is None else (*spec, typed) - return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, a.shape[-1]) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, a.shape[-1], True) # Index letters for the einsum that computes a contraction (for slicing a pending result) @@ -404,7 +404,7 @@ def build(Ta, Tb): spec = _tensordot_terms(a.ndim, b.ndim, a_axes, b_axes) spec = None if spec is None else (*spec, typed) nterms = math.prod(a.shape[i] for i in a_axes) - return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, nterms) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, nterms, True) def vecdot(x1, x2, /, *, axis=-1, semiring=None): @@ -516,4 +516,7 @@ def build(*tensors): c: n for term, s in zip(terms, shapes, strict=True) for c, n in zip(term, s, strict=True) } nterms = math.prod(n for c, n in sizes.items() if c not in output) - return _lazy.contraction(ops, build, shape, out_dt, fill, (terms, output, typed), nterms) + product = _lazy._one_product(terms, output) + return _lazy.contraction( + ops, build, shape, out_dt, fill, (terms, output, typed), nterms, product + ) diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 65c2959f3..1da0b8fc6 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -40,15 +40,23 @@ def pending(x): return x._tensor is None +class _Seen(list): + hints = None + + @pytest.fixture def masks(monkeypatch): """The masks the pending products were computed under (None: whole).""" - seen = [] + seen = _Seen() compute = _lazy.Contraction.compute - def recording(self, mask=None): + def recording(self, mask=None, **opts): seen.append(mask) - return compute(self, mask) + hints.append(opts) + return compute(self, mask, **opts) + + hints = [] + seen.hints = hints # (the options each computation got) monkeypatch.setattr(_lazy.Contraction, "compute", recording) return seen @@ -411,3 +419,35 @@ def test_iteration_computes_once(masks): assert masks == [None] for row, expected in zip(rows, a @ b, strict=True): assert_close(row, expected) + + +def test_sddmm_dot_products(masks): + # A @ B with B held by row (k x n) under a sparse mask: dot products at the mask's + # entries, not SuiteSparse's masked saxpy over the whole product (found by the SAPS run) + rng = np.random.default_rng(13) + n, k = 60, 4 + s = np.where(rng.random((n, n)) < 0.02, rng.random((n, n)), 0.0) + a, b = rng.random((n, k)) + 0.5, rng.random((k, n)) + 0.5 + S, A, B = make(s, 0.0), xp.asarray(a), xp.asarray(b) + for P, expected in [ + (A @ B, s * (a @ b)), + (xp.tensordot(A, B, axes=1), s * (a @ b)), + (xp.einsum("ik,kj->ij", A, B), s * (a @ b)), + (xp.einsum("ik,jk->ij", A, xp.asarray(b.T)), s * (a @ b)), + ((A @ B)[: n // 2], s[: n // 2] * (a @ b)[: n // 2]), + ]: + masks.clear() + masks.hints.clear() + S2 = S if expected.shape == s.shape else S[: n // 2] + assert_close(xp.multiply(S2, P), expected) + assert masks.hints == [{"axb_method": "dot"}] + # Not for a dense mask, a vecdot, or an einsum whose product is permuted after + for S2, P, expected in [ + (make(np.ones((n, n)), 0.0), A @ B, a @ b), + (S, xp.vecdot(A[:, None, :], xp.asarray(b.T)[None, :, :]), s * (a @ b)), + (S, xp.einsum("ik,kj->ji", A, B), s * (a @ b).T), + ]: + masks.clear() + masks.hints.clear() + assert_close(xp.multiply(S2, P), expected) + assert masks.hints == [{}] From 4701b656db259a553ed3051f7f5ef2e58bab0fdb Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 05:35:47 -0600 Subject: [PATCH 45/77] array_api: test the empty-ellipsis key with list, ndarray and Array indices x[:, 1, ..., index] of a 3-d x, the key of the report behind "array_api: a "..." that covers no axis separates advanced indices, as in numpy", with each kind of index, for getitem and setitem. --- graphblas/tests/test_array_api_indexing.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/graphblas/tests/test_array_api_indexing.py b/graphblas/tests/test_array_api_indexing.py index 09bac62d4..f4d92b13a 100644 --- a/graphblas/tests/test_array_api_indexing.py +++ b/graphblas/tests/test_array_api_indexing.py @@ -263,4 +263,11 @@ def test_mixed_keys_order_axes_as_numpy(): assert_array_equal(y, b) assert count > 2000 x = make(np.arange(24.0).reshape(2, 3, 4), 0.0) - assert_array_equal(x[:, 1, ..., [0, 2]], np.array([[4.0, 16.0], [6.0, 18.0]])) + expected = np.array([[4.0, 16.0], [6.0, 18.0]]) + for index in [[0, 2], np.array([0, 2]), xp.asarray([0, 2])]: + assert_array_equal(x[:, 1, ..., index], expected) + y = make(np.arange(24.0).reshape(2, 3, 4), 0.0) + y[:, 1, ..., index] = make(-expected, 0.0) + b = np.arange(24.0).reshape(2, 3, 4) + b[:, 1, ..., [0, 2]] = -expected + assert_array_equal(y, b) From 28d02ec87501bdab9a87d1f3228cc2fb1844b535 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 05:39:46 -0600 Subject: [PATCH 46/77] array_api: a fully stored bool result takes its common value as the fill An array that stores every position never shows its fill. For a bool result of an element-wise function (a comparison, a logical function, isnan, astype) that stores every position, the value most of its entries hold becomes the fill and those entries are dropped, when at most a quarter of them stay: one count and one select. equal(x, 1) of a fully stored x now stores only the positions equal to 1, with fill False, instead of every position with fill True, so a later product with it stays sparse. Values do not change. Results with unstored positions are left as they are (their fill is seen), and so are results where more than a quarter of the entries would stay, where dropping costs about as much as the comparison and saves a consumer little. Numeric results are not re-filled: a fully stored numeric result is mostly dense data, and finding the common value would cost a pass on every result for a rare gain. --- graphblas/array_api/_elementwise.py | 20 ++++++++++- graphblas/tests/test_array_api_elementwise.py | 35 +++++++++++++++---- 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index b1b7bdcec..fbb1d87ce 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -17,9 +17,10 @@ import numpy as np -from .. import binary, unary +from .. import binary, monoid, select, unary from ..core.scalar import Scalar from ..core.tensor import Tensor +from ..dtypes import BOOL, INT64 from . import _lazy, _utils from ._array import Array from ._dtypes import check_supported @@ -160,9 +161,26 @@ def _fill_of(ufunc, rdt, operands): def _result(T, fill): + if T.dtype == BOOL and T.ndim and T.nvals == T.size and T.nvals: + T, fill = _normalize_bool(T, fill) return Array._new(T, fill) +def _normalize_bool(T, fill): + """A bool Tensor that stores every position never shows its fill: make the fill the value + most entries hold and drop those entries, when at most a quarter of them stay (so + ``equal(deg, 0)`` of a fully stored ``deg`` stores only its True entries, with fill + False). One count, then one select; else ``T`` and ``fill`` as they were (dropping more + than a quarter costs about as much as the comparison and saves a later consumer little). + """ + n = T.nvals + ntrue = T.reduce_scalar(monoid.plus[INT64]).new().value + major = 2 * ntrue > n + if 4 * (n - ntrue if major else ntrue) > n: + return T, fill + return T.select(select.valueeq, not major).new(), major + + def _zero_d(dtype, fill): return Array._new(Tensor(gb_dtype(dtype), ()), fill) diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index edad45721..5864f8670 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -710,15 +710,38 @@ def test_intersection_of_broadcast_operands_is_guarded(): # an intersection of two broadcast operands is an outer product (review finding M4) x = make(np.arange(1, 101).reshape(100, 1), 0) y = make(np.arange(1, 101).reshape(1, 100), 0) - b = x > 0 - for func, args in [ - (xp.multiply, (x, y)), - (xp.bitwise_and, (x, y)), - (xp.logical_and, (b, b.mT)), + b = make(np.ones((100, 1), bool), False) + for func, args, nnz in [ + (xp.multiply, (x, y), 10000), + (xp.bitwise_and, (x, y), 10000), + (xp.logical_and, (b, b.mT), 0), # (all True, every position stored: fill True) ]: with gb.config.set(array_api_densify_limit=1000), pytest.raises(xp.DensifyError): func(*args) - assert func(*args).nnz == 10000 + assert func(*args).nnz == nnz + + +def test_fully_stored_bool_results_take_the_common_value_as_fill(): + # A bool result that stores every position never shows its fill: the value most of its + # entries hold becomes the fill and those entries go, when at most a quarter stay + # (equal(deg, 0) of a fully stored deg stores the zero-degree positions only) + deg = xp.asarray(np.array([3, 1, 0, 2, 5, 0, 4, 1, 2, 7])) + assert deg.nnz == 8 + assert deg.fill_value == 0 + full = xp.asarray(np.array([3, 1, 9, 2, 5, 6, 4, 1, 2, 7]), fill_value=-1) + assert full.nnz == 10 + for out, expected, nnz, fill in [ + (xp.equal(full, 1), full.todense() == 1, 2, False), + (xp.greater(full, 1), full.todense() > 1, 2, True), + (xp.less(full, 5), full.todense() < 5, 10, (-1 < 5)), # (6 of 10 True: kept) + (xp.isfinite(xp.astype(full, xp.float64)), np.ones(10, bool), 0, True), + (xp.astype(full == 9, xp.bool), full.todense() == 9, 1, False), + ]: + assert_array_equal(out, expected) + assert (out.nnz, out.fill_value) == (nnz, fill) + # a 0-d result and a result with unstored positions are left as they are + assert xp.equal(deg, 0).nnz == 8 + assert_array_equal(xp.equal(xp.asarray(3), 3), np.array(True)) def test_clip_with_array_bounds_ties(): From 43c43032b3e99be2efa9d0edb6747dd16e373960 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 05:45:51 -0600 Subject: [PATCH 47/77] array_api: the SDDMM overflow bound reads magnitudes by two reductions The largest magnitude of a real operand is now the larger of |max| and |min|, two reductions in place, instead of a reduction over a copy of the absolute values. An SDDMM through the namespace with full A and B of 1e6 x 64 values and 1e7 entries in S, at 8 threads: 0.154 s to 0.136 s (the checks of A and B, finiteness and magnitude, 38 ms to 30 ms of it). --- graphblas/array_api/_lazy.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py index 86e201831..f202a2664 100644 --- a/graphblas/array_api/_lazy.py +++ b/graphblas/array_api/_lazy.py @@ -85,12 +85,16 @@ def _all_finite(T): def _max_abs(T): - """The largest magnitude of the stored values of ``T``, as a float.""" + """The largest magnitude of the stored values of ``T``, as a float (NaN skipped).""" if T.nvals == 0: return 0.0 - kind = np_dtype(T.dtype).kind - A = T.apply(unary.abs[FC64 if kind == "c" else FP64]).new() - return float(A.reduce_scalar(monoid.max[FP64]).new().value) + if np_dtype(T.dtype).kind == "c": + A = T.apply(unary.abs[FC64]).new() + return float(A.reduce_scalar(monoid.max[FP64]).new().value) + # (two reductions in place: a copy of |values| cost 40 ms of 155 at n = 1e6, k = 64) + hi = T.reduce_scalar(monoid.max[FP64]).new().value + lo = T.reduce_scalar(monoid.min[FP64]).new().value + return max(abs(hi), abs(lo)) def _known_finite(x): From 51d8df21969078516bb7eba6122292ea54c78210 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:09:46 -0600 Subject: [PATCH 48/77] array_api: an in-place operator on a pending product forgets the product P *= S with P a pending product went through the SDDMM push-down, which does not compute P, and then replaced P's Tensor while P kept its pending product: a later S2 * P or P[1:] read the old A @ B. Replacing an Array's value now clears the pending product. A test runs every in-place operator on a pending Array and reads it back through both consumers. Found by review. --- graphblas/array_api/_array.py | 1 + graphblas/tests/test_array_api_lazy.py | 29 ++++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py index e0ef9d871..f9ae5ef81 100644 --- a/graphblas/array_api/_array.py +++ b/graphblas/array_api/_array.py @@ -81,6 +81,7 @@ def _T(self): def _replace(self, other): """Make this Array hold ``other``'s values (for in-place operators and assignment).""" self._tensor = other._T() + self._pending = None # (a pending product it was is not its value any more) self._shape = other._shape self._dtype = other._dtype self._fill = other._fill diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 1da0b8fc6..a600371ca 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -451,3 +451,32 @@ def test_sddmm_dot_products(masks): masks.hints.clear() assert_close(xp.multiply(S2, P), expected) assert masks.hints == [{}] + + +@pytest.mark.parametrize( + "op", + ["iadd", "isub", "imul", "itruediv", "ifloordiv", "imod", "ipow", "imatmul", "iand", "ior", + "ixor", "ilshift", "irshift"], +) # fmt: skip +def test_inplace_operators_on_a_pending_array(op): + # (review finding H1) an in-place operator replaces a pending product: later consumers + # (the SDDMM, slicing) must read the new value, not the old product + import operator + + func = getattr(operator, op) + rng = np.random.default_rng(14) + dtype = np.int64 if op in {"iand", "ior", "ixor", "ilshift", "irshift"} else np.float64 + a = rng.integers(1, 5, (4, 4)).astype(dtype) + s = np.zeros((4, 4), dtype) + s[0, 0], s[1, 2], s[3, 3] = 2, 3, 1 + s2 = np.eye(4, dtype=dtype) + P = make(a, 0) @ make(a, 0) + assert pending(P) + with np.errstate(all="ignore"): + expected = func((a @ a).copy(), s if op != "imatmul" else s2) + P = func(P, make(s, 0) if op != "imatmul" else make(s2, 0)) + assert not pending(P) + assert P._pending is None + assert_close(P, expected) + assert_close(xp.multiply(make(s2, 0), P), s2 * expected) + assert_close(P[1:], expected[1:]) From 6bdea2853530b60fcec1325d0f2e9a5b605e2d34 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:11:10 -0600 Subject: [PATCH 49/77] array_api: a fully stored bool result takes fill False, never True A mostly-True bool result that stores every position took fill True, and the contractions cannot skip a True fill: a mostly-True mask times a sparse bool matrix (m @ A, einsum, lxor_land, max_times) raised DensifyError. A fully stored bool result now always takes fill False (an O(1) relabel: no position shows the fill), and drops its False entries only when at most a quarter of the entries are True. Found by review. --- graphblas/array_api/_elementwise.py | 19 +++++---- graphblas/tests/test_array_api_elementwise.py | 39 +++++++++++++++---- 2 files changed, 40 insertions(+), 18 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index fbb1d87ce..928d57d6e 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -167,18 +167,17 @@ def _result(T, fill): def _normalize_bool(T, fill): - """A bool Tensor that stores every position never shows its fill: make the fill the value - most entries hold and drop those entries, when at most a quarter of them stay (so - ``equal(deg, 0)`` of a fully stored ``deg`` stores only its True entries, with fill - False). One count, then one select; else ``T`` and ``fill`` as they were (dropping more - than a quarter costs about as much as the comparison and saves a later consumer little). + """A bool Tensor that stores every position never shows its fill, so it takes fill False + (the fill contractions and masks work with; a fill True made a mostly-True mask times a + sparse matrix dense). Its False entries are dropped when at most a quarter of the + entries are True (``equal(deg, 0)`` of a fully stored ``deg`` then stores only its True + entries): one count, then one select. Otherwise they stay; dropping more costs about as + much as the comparison and saves a later consumer little. """ n = T.nvals - ntrue = T.reduce_scalar(monoid.plus[INT64]).new().value - major = 2 * ntrue > n - if 4 * (n - ntrue if major else ntrue) > n: - return T, fill - return T.select(select.valueeq, not major).new(), major + if 4 * T.reduce_scalar(monoid.plus[INT64]).new().value > n: + return T, False + return T.select(select.valueeq, True).new(), False def _zero_d(dtype, fill): diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 5864f8670..97dad6201 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -714,17 +714,17 @@ def test_intersection_of_broadcast_operands_is_guarded(): for func, args, nnz in [ (xp.multiply, (x, y), 10000), (xp.bitwise_and, (x, y), 10000), - (xp.logical_and, (b, b.mT), 0), # (all True, every position stored: fill True) + (xp.logical_and, (b, b.mT), 10000), ]: with gb.config.set(array_api_densify_limit=1000), pytest.raises(xp.DensifyError): func(*args) assert func(*args).nnz == nnz -def test_fully_stored_bool_results_take_the_common_value_as_fill(): - # A bool result that stores every position never shows its fill: the value most of its - # entries hold becomes the fill and those entries go, when at most a quarter stay - # (equal(deg, 0) of a fully stored deg stores the zero-degree positions only) +def test_fully_stored_bool_results_take_fill_false(): + # A bool result that stores every position never shows its fill: it takes fill False, + # and its False entries go when at most a quarter of the entries are True (equal(deg, 0) + # of a fully stored deg stores the zero-degree positions only) deg = xp.asarray(np.array([3, 1, 0, 2, 5, 0, 4, 1, 2, 7])) assert deg.nnz == 8 assert deg.fill_value == 0 @@ -732,9 +732,9 @@ def test_fully_stored_bool_results_take_the_common_value_as_fill(): assert full.nnz == 10 for out, expected, nnz, fill in [ (xp.equal(full, 1), full.todense() == 1, 2, False), - (xp.greater(full, 1), full.todense() > 1, 2, True), - (xp.less(full, 5), full.todense() < 5, 10, (-1 < 5)), # (6 of 10 True: kept) - (xp.isfinite(xp.astype(full, xp.float64)), np.ones(10, bool), 0, True), + (xp.greater(full, 1), full.todense() > 1, 10, False), # (8 of 10 True: kept) + (xp.less(full, 5), full.todense() < 5, 10, False), + (xp.isfinite(xp.astype(full, xp.float64)), np.ones(10, bool), 10, False), (xp.astype(full == 9, xp.bool), full.todense() == 9, 1, False), ]: assert_array_equal(out, expected) @@ -744,6 +744,29 @@ def test_fully_stored_bool_results_take_the_common_value_as_fill(): assert_array_equal(xp.equal(xp.asarray(3), 3), np.array(True)) +def test_mostly_true_masks_stay_sparse_in_contractions(): + # (review finding H3) a mostly-True fully stored mask keeps fill False, so a contraction + # with a sparse bool matrix stays sparse (it raised DensifyError with fill True) + rng = np.random.default_rng(15) + n = 300 + a = np.zeros((n, n), bool) + a[rng.integers(0, n, 900), rng.integers(0, n, 900)] = True + A = make(a, False) + d = rng.integers(1, 10, n).astype(float) + m = xp.greater(xp.asarray(d), 1.5) + assert m.fill_value == False # noqa: E712 + mv = d > 1.5 + with gb.config.set(array_api_densify_limit=n * 10): + assert_array_equal(m @ A, mv @ a) + assert_array_equal(xp.einsum("i,ij->j", m, A), np.einsum("i,ij->j", mv, a)) + assert_array_equal( + xp.matmul(m, A, semiring="lxor_land"), + np.logical_xor.reduce(mv[:, None] & a, axis=0), + ) + out = xp.matmul(xp.astype(m, xp.uint8), xp.astype(A, xp.uint8), semiring="max_times") + assert_array_equal(out, np.max(mv[:, None].astype(np.uint8) * a, axis=0)) + + def test_clip_with_array_bounds_ties(): # numpy's clip gives the bound on a tie, so the sign of a zero is the bound's (review M6) x = make(np.array([0.0, -0.0, 0.0, 1.0, np.nan]), 0.0) From 2e47d32ba6391cebabf8886ab86efe130545fa63 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:13:50 -0600 Subject: [PATCH 50/77] array_api: an SDDMM with inf, or that may overflow, computes the sparse product When an operand of a pending product held inf or NaN, or its products could overflow, multiply(S, P) checked P's logical size against the densify limit before computing P whole, so any large sparse product raised DensifyError where computing without deferral succeeds. It now computes the product (a sparse product, as without deferral) and multiplies, with no densify check. Found by review. --- graphblas/array_api/_lazy.py | 10 +++------ graphblas/tests/test_array_api_lazy.py | 29 +++++++++++++++++++++----- 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py index f202a2664..b4ffe8c46 100644 --- a/graphblas/array_api/_lazy.py +++ b/graphblas/array_api/_lazy.py @@ -14,8 +14,8 @@ is 0, and 0 times P's value there is 0 when that value is finite. That holds when the operands are finite and no sum of products can overflow: the number of terms of each sum times the largest magnitude of each operand stays below the dtype's largest value (one - GraphBLAS pass per operand for each). Otherwise P is computed whole, as numpy would, under - the densify limit: the whole of a product meant for an SDDMM may be dense. + GraphBLAS pass per operand for each). Otherwise P is computed whole (a sparse product, as + without deferral), then multiplied. - ``P[key]`` with a basic index (ints and slices) slices the operands along the indexed axes of the result, then contracts (pending again). @@ -138,11 +138,7 @@ def masked_multiply(a, b): return None node = P._pending if not node.all_finite(): - _utils.check_dense( - P.size, - "multiply(S, P) where P may hold inf or NaN outside S's entries (P computed whole)", - ) - return None + return None # (P is computed whole, a sparse product as without deferral) TS = S._T() PS = Array._new(node.compute(mask=TS.S, **_dot_hint(node, TS.nvals)), P._fill) from ._elementwise import multiply diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index a600371ca..5e0cfbf08 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -301,11 +301,30 @@ def test_sddmm_overflow_outside_s(masks): R = xp.multiply(make(np.eye(3, dtype=np.float32), 0), make(a, 0) @ make(a.T, 0)) assert_array_equal(R, expected) assert [m is None for m in masks] == [True] - # Computing it whole is guarded: that whole may be dense - big = make(np.full((4, 1), 1e200), 0.0) - eye = make(np.eye(4), 0.0) - with gb.config.set(array_api_densify_limit=10), pytest.raises(xp.DensifyError, match="inf"): - xp.multiply(eye, big @ xp.matrix_transpose(big)) + + +def test_sddmm_with_inf_stays_sparse(masks): + # (review finding H2) an operand holding inf, or values whose products may overflow, + # computes the product whole: a sparse product, as without deferral, with no densify + # check on its logical size + rng = np.random.default_rng(16) + n = 400 + r, c = rng.integers(0, n, 400), rng.integers(0, n, 400) + a = np.zeros((n, n)) + a[r, c] = rng.random(400) + 0.5 + a[r[0], c[0]] = np.inf + s = np.zeros((n, n)) + s[rng.integers(0, n, 400), rng.integers(0, n, 400)] = 1.0 + A = xp.from_coo(np.nonzero(a), a[np.nonzero(a)], (n, n)) + S = xp.from_coo(np.nonzero(s), s[np.nonzero(s)], (n, n)) + with gb.config.set(array_api_densify_limit=n * n // 10): + R = xp.multiply(S, A @ A) + # (the product as computed without deferral: D6 skips an unstored zero times inf) + P = A @ A + with np.errstate(invalid="ignore"): + assert_close(R, s * P.todense()) + assert [m is None for m in masks] == [True, True] + assert np.isinf(P.todense()).any() def test_max_abs(): From b919461c10adab634e26a6a3257ce164fa92d7e2 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:16:16 -0600 Subject: [PATCH 51/77] array_api: a one-sided product gathers a broadcast operand at the entries X * y or X / y, where X's fill is 0 and y broadcasts and does not store every position, stored y in full (O(size of y), within the densify limit) and past the limit replicated it over its broadcast axes. It now reads y's values at X's stored coordinates (projected onto y's axes, with y's fill where y stores nothing): O(entries of X). X of shape (1e4, 1e4, 2) with 4 entries times exp(W)[:, :, None] (W 1e4 x 1e4 with 1e5 entries): peak RSS 1634 MB to 150 MB, for a result of 4 entries. Found by review. --- graphblas/array_api/_elementwise.py | 32 ++++++++++++---- graphblas/tests/test_array_api_elementwise.py | 38 ++++++++++++------- 2 files changed, 48 insertions(+), 22 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 928d57d6e..5c7ff4037 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -653,14 +653,14 @@ def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): Tx, Ty = X._T(), Y._T() with np.errstate(all="ignore"): fill = spec.ufunc(fx, fy) if X is a else spec.ufunc(fy, fx) - broadcast = Ty.shape != shape - if not _stores_all(Y) and ( - Y.size <= Tx.nvals + 1024 or (broadcast and Y.size <= _utils.densify_limit()) - ): - # (O(size of Y): every position of Y stored, so the product is an intersection. - # A broadcast Y is small along its broadcast axes; the union with it would - # replicate it over them: exp(v)[None, :] times X of shape (1000, 1e6) with 4 - # entries replicated v 1000 times) + if Ty.shape != shape and not _stores_all(Y): + # Y broadcasts and shows its fill: its value at each of X's entries, gathered + # from Y itself at the projected coordinates, O(entries of X). (Storing Y in full + # cost O(size of Y), and the union replicated Y over its broadcast axes: + # exp(v)[None, :] times X of shape (1000, 1e6) with 4 entries made v 1000 times.) + Ty = _gather(Tx, Y, shape) + elif not _stores_all(Y) and Y.size <= Tx.nvals + 1024: + # (O(size of Y): every position of Y stored, so the product is an intersection) other = 0 if Y._fill != 0 else 1 Ty = _utils.with_fill(Y, other)._T() Ta, Tb = (Tx, Ty) if X is a else (Ty, Tx) @@ -675,6 +675,22 @@ def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): return None +def _gather(Tx, Y, shape): + """A Tensor with ``Tx``'s entries (of ``shape``) holding the values of ``Y`` (an Array + that broadcasts to ``shape``) there: its stored value, or its fill. + """ + coords = Tx.to_coo(values=False)[0] + Ty = Y._T() + lead = len(shape) - Y.ndim + # (a broadcast axis of Y has size 1: index 0 there) + index = tuple( + coords[lead + i] if n != 1 else np.zeros(len(coords[0]), np.int64) + for i, n in enumerate(Ty.shape) + ) + values = Ty.vindex[index].to_dense(fill_value=Y._fill) + return Tensor.from_coo(coords, values, Ty.dtype, shape=shape) + + def _zero_times(spec, y, fy): """Whether 0 * v (or 0 / v) is a zero for every value v of ``y`` (its fill ``fy`` too): finite values (for divide, also no 0). NaN never qualifies; the sign of the zero does not diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 97dad6201..c8e66a4f2 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -882,10 +882,10 @@ def test_unstored_zero_signs_in_the_oracle(): assert_array_equal(xp.asarray(-0.0), np.array(0.0)) -def test_one_sided_broadcast_operand_is_stored_in_full(): - # (review finding M9) a broadcast operand that does not store every position is stored in - # full, not replicated over the axes it broadcasts along: X of shape (1000, 1e6) with 4 - # entries times exp(v)[None, :] (fill 1) stores 4 entries, under a densify limit of 2e6 +def test_one_sided_broadcast_operand_is_gathered(): + # (review findings rev1 M9, rev2 M5) a broadcast operand that shows its fill is read at the + # other operand's entries only, O(those entries): X of shape (1000, 1e6) with 4 entries + # times exp(v)[None, :] (fill 1) stores 4 entries, under a densify limit of 1000 rng = np.random.default_rng(12) n = 10**6 rows, cols = [0, 10, 500, 999], [5, 77, n - 1, 123_456] @@ -893,26 +893,36 @@ def test_one_sided_broadcast_operand_is_stored_in_full(): v = xp.from_coo((np.sort(rng.choice(n, n // 2, replace=False)),), rng.normal(size=n // 2), (n,)) w = make(np.where(rng.random(n) < 0.5, 2.0, rng.random(n) + 1), 2.0) vals = np.array([1.0, 2.0, 3.0, 4.0]) - with gb.config.set(array_api_densify_limit=2 * n): - e = xp.exp(v) - ed, wd = e.todense()[cols], w.todense()[cols] + e = xp.exp(v) + ed, wd = e.todense()[cols], w.todense()[cols] + with gb.config.set(array_api_densify_limit=1000): for out, expected in [ (X * e[None, :], vals * ed), (e[None, :] * X, ed * vals), (X / w, vals / wd), + (X * xp.asarray(-2.0), vals * -2.0), ]: assert out.nnz == 4 assert out.fill_value == 0.0 coords, values = out.to_coo() assert coords[0].tolist() == rows np.testing.assert_array_equal(values, expected) - # numpy's answer at every position, on a small one - A = make(np.where(rng.random((6, 50)) < 0.05, rng.normal(size=(6, 50)), 0.0), 0.0) - u = make(np.where(rng.random(50) < 0.5, 1.0, rng.normal(size=50)), 1.0) - with gb.config.set(array_api_densify_limit=60): - out = A * u[None, :] - assert out.nnz == A.nnz - assert_array_equal(out, A.todense() * u.todense()[None, :]) + # numpy's answer at every position, on small ones: broadcast along either axis, 3-d, a + # 0-d operand, explicit fills, an X with no entries + A = make(np.where(rng.random((6, 50)) < 0.05, rng.normal(size=(6, 50)), 0.0), 0.0, rng) + u = make(np.where(rng.random(50) < 0.5, 1.0, rng.normal(size=50)), 1.0, rng) + c = make(np.where(rng.random((6, 1)) < 0.5, 3.0, rng.normal(size=(6, 1))), 3.0) + B = make(np.where(rng.random((4, 6, 5)) < 0.2, rng.normal(size=(4, 6, 5)), 0.0), 0.0) + g = make(np.where(rng.random((6, 1)) < 0.5, 2.0, rng.normal(size=(6, 1)) + 3), 2.0) + for out, expected in [ + (A * u[None, :], A.todense() * u.todense()[None, :]), + (A * c, A.todense() * c.todense()), + (B * c, B.todense() * c.todense()), + (B / g, B.todense() / g.todense()), + (xp.zeros((6, 50)) * u, np.zeros((6, 50)) * u.todense()), + ]: + assert out.nnz <= max(A.nnz, B.nnz) + assert_array_equal(out, expected) def test_clip_fill_uses_numpys_array_loop(): From db32ee99f29af73c857d002d7d1bd98f1a24d681 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:21:41 -0600 Subject: [PATCH 52/77] array_api: contractions ignore explicit fills, propagate NaN terms, and drop the fill split Three changes to the contractions, and two small ones, all found by review: - Explicit fills no longer change values: when an operand holds inf or NaN, the operands' stored entries equal to their fill are pruned first, so such a value meets an explicit fill as it meets an unstored one, with no term ([explicit 0] @ [inf] was NaN, the pruned form 0). - A stored NaN term propagates through a float min or max monoid, as numpy's min and max do; SuiteSparse's fmin and fmax skipped it, so min_plus of [[nan, 5]] and [[1], [3]] was 8. The result is computed at once and NaN set where a NaN term is: one operand's NaN entries times the others' patterns. - A nonzero fill no longer splits A into A0 + fill (which turned inf into NaN and lost every digit when the fill dwarfs the values): an operand that stores every position takes fill 0, any other nonzero fill goes dense. - The dense path checks the result's size against the densify limit, not only the operands'. - tensordot refuses an axis out of range instead of wrapping it. --- graphblas/array_api/_lazy.py | 56 ++++++- graphblas/array_api/_linalg.py | 176 ++++++++------------- graphblas/tests/test_array_api_lazy.py | 5 +- graphblas/tests/test_array_api_linalg.py | 35 ++++ graphblas/tests/test_array_api_semiring.py | 43 +++++ 5 files changed, 199 insertions(+), 116 deletions(-) diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py index b4ffe8c46..a94c79a49 100644 --- a/graphblas/array_api/_lazy.py +++ b/graphblas/array_api/_lazy.py @@ -25,9 +25,9 @@ import numpy as np -from .. import monoid, unary +from .. import binary, monoid, select, semiring, unary from ..core.tensor import Tensor -from ..dtypes import FC64, FP64 +from ..dtypes import FC64, FP64, lookup_dtype from . import _utils from ._array import Array from ._utils import gb_dtype, np_dtype @@ -107,13 +107,31 @@ def _known_finite(x): return None if facts is None else facts.get("finite") -def contraction(arrays, build, shape, dtype, fill, spec=None, nterms=1, product=False): - """The Array of ``build(*tensors)`` over the arrays' Tensors, with fill ``fill``: pending - when deferral is on and the result has a dimension, else computed now. ``nterms``: the - number of terms each result position sums; ``product``: see ``Contraction``. +def contraction(arrays, build, shape, dtype, fill, spec=None, nterms=1, product=False, typed=None): + """The Array of ``build(*tensors)`` over the arrays' Tensors (``build(*tensors, op=)`` + with another semiring), with fill ``fill``: pending when deferral is on and the result + has a dimension, else computed now. ``nterms``: the number of terms each result position + sums; ``product``: see ``Contraction``; ``typed``: the semiring. + + When an operand holds inf or NaN, the operands' stored entries equal to their fill go + first, so that such a value meets an explicit fill as it meets an unstored one, with no + term (DESIGN D12a). When a float min or max monoid would drop a stored NaN term, the + result is computed now and NaN set where such a term is (D12b). """ + if any(x.dtype.kind in "fc" and not _utils.fact(x, "finite") for x in arrays): + from ._manipulation import prune + + arrays = [prune(x) for x in arrays] tensors = [x._T() for x in arrays] gbdt = gb_dtype(dtype) + nan = ( + typed is not None + and typed.monoid.name in {"min", "max"} + and np_dtype(typed.return_type).kind == "f" + and any(x.dtype.kind == "f" and _utils.fact(x, "nan") for x in arrays) + ) + if nan: + return Array._new(_with_nan_terms(build, tensors, gbdt), fill) if not enabled or not shape: return Array._new(build(*tensors).new(dtype=gbdt), fill) snapshots = [T.dup() for T in tensors] @@ -122,6 +140,28 @@ def contraction(arrays, build, shape, dtype, fill, spec=None, nterms=1, product= return Array._deferred(node, shape, dtype, fill) +def _with_nan_terms(build, tensors, gbdt): + """``build(*tensors)`` with NaN wherever a term holds a stored NaN: SuiteSparse's float + min and max are fmin and fmax, which skip NaN, where numpy's propagate it. Those positions + are the product of one operand's NaN entries with the others' patterns (any_pair). + """ + R = build(*tensors).new(dtype=gbdt) + bool_ = lookup_dtype(bool) + patterns = [T.apply(unary.one[bool_]).new() for T in tensors] + N = None + for i, T in enumerate(tensors): + # (every operand is float here: its fill is the identity of min or max, an infinity) + nan = T.apply(unary.isnan).new().select(select.valueeq, True).new() + if nan.nvals == 0: + continue + ops = [nan if j == i else P for j, P in enumerate(patterns)] + Ni = build(*ops, op=semiring.any_pair[bool_]).new() + N = Ni if N is None else N.ewise_add(Ni, binary.lor).new() + if N is not None and N.nvals: + R(N.S) << _utils.gb_scalar(np.nan, np_dtype(gbdt)) + return R + + def masked_multiply(a, b): """``multiply(a, b)`` when one of them is a pending product P and the other, S, lets P be computed at S's stored positions only (see the module's docstring); else None. @@ -195,8 +235,8 @@ def sliced(x, axes_key): subscripts = ",".join(new_terms) + "->" + out - def build(*tensors): - return einsum(subscripts, *tensors, op=typed) + def build(*tensors, op=typed): + return einsum(subscripts, *tensors, op=op) # (a slice of a finite operand is finite) spec = (tuple(new_terms), out, typed) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 1f378924f..9fff78639 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -3,10 +3,10 @@ Contractions treat a fill of zero as a structural zero (DESIGN D6), as scipy.sparse and pydata/sparse do: a product with an unstored position is skipped, so 0 * inf there is not NaN, -unlike numpy. Sums may round differently from numpy's (the order of the terms differs). A -nonzero fill is the fill times a matrix of ones plus an array with fill zero, which turns into -sums of the operands broadcast over the result: the result is then dense in general (under the -densify limit). A fill that is not finite, or True, makes the result go through numpy densely. +unlike numpy. Sums may round differently from numpy's (the order of the terms differs). An +operand that stores every position never shows its fill, so a nonzero fill there is taken as +0; another nonzero fill makes the contraction go through numpy densely, under the densify +limit on its operands and its result (DESIGN D12c). Semirings (an extension, DESIGN D11): ``semiring=`` takes a ``graphblas`` Semiring or its name ("min_plus", "max_min", ...), whose monoid takes the place of the sum and whose binary @@ -187,34 +187,36 @@ def _is_zero(fill): return not bool(fill) -def _dense_ok(*arrays): - """Whether the decomposition A = A0 + fill works: finite fills, not bool.""" +def _zero_fills(*arrays): + """The arrays with fill 0, or None when one cannot have it. An array that stores every + position never shows its fill, so its fill is relabeled 0 (O(1)); any other nonzero fill + gives None (the contraction is then dense). + """ + out = [] for x in arrays: - if _is_zero(x._fill): - continue - if x.dtype.kind == "b" or (x.dtype.kind in "fc" and not np.isfinite(x._fill)): - return False - return True + if not _is_zero(x._fill): + T = x._T() + if T.nvals != x.size: + return None + x = Array._new(T, _utils.scalar(0, x.dtype)) + out.append(x) + return out -def _dense(func, *arrays): +def _dense(func, shape, *arrays): + """``func`` of the arrays densely by numpy, under the densify limit on the operands and on + the result of ``shape``. + """ from ._creation import _from_numpy + what = "a contraction with a nonzero fill that does not store every position" for x in arrays: - _utils.check_dense(x.size, "a contraction with a fill that is not finite or is True") + _utils.check_dense(x.size, what) + _utils.check_dense(math.prod(shape), what) out = np.asarray(func(*[x.todense() for x in arrays])) return _from_numpy(out) -def _zero_fill(x): - """``x`` minus its fill (exact where the fill is 0), and the fill.""" - from ._elementwise import subtract - - if _is_zero(x._fill): - return x, None - return subtract(x, x._fill), x._fill - - def _rdt(name, func, x1, x2): probe = func( np.ones((1,) * max(x1.ndim, 1), x1.dtype), np.ones((1,) * max(x2.ndim, 1), x2.dtype) @@ -226,10 +228,6 @@ def matmul(x1, x2, /, *, semiring=None): """The matrix product (numpy's matmul rules). ``semiring`` (an extension): a graphblas Semiring or its name, under the module's rules for fills. """ - from ._elementwise import add, multiply - from ._manipulation import expand_dims - from ._statistical import sum as sum_ - x1, x2 = _operand(x1), _operand(x2) _check(x1, "matmul") _check(x2, "matmul") @@ -249,32 +247,14 @@ def matmul(x1, x2, /, *, semiring=None): if not default: out_dt, fill = _sound(typed, (x1, x2), "matmul") return _product(x1, x2, typed, out_dt, fill) - if not _dense_ok(x1, x2): - return _dense(np.matmul, x1, x2) - a0, fa = _zero_fill(x1) - b0, fb = _zero_fill(x2) - out = _product(a0, b0, typed, rdt, _utils.scalar(0, rdt)) - if fa is None and fb is None: - return out - terms = [] - if fa is not None: - # fa * (ones @ b0): b0 summed over its rows, one row of the result - s = sum_(b0, axis=-2 if x2.ndim > 1 else 0, dtype=rdt) - if x1.ndim > 1 and x2.ndim > 1: - s = expand_dims(s, axis=-2) - terms.append(multiply(fa, s)) - if fb is not None: - r = sum_(a0, axis=-1, dtype=rdt) - if x2.ndim > 1 and x1.ndim > 1: - r = expand_dims(r, axis=-1) - terms.append(multiply(fb, r)) - if fa is not None and fb is not None: - terms.append(multiply(multiply(fa, fb), K)) - for t in terms: - out = add(out, t) - from ._elementwise import cast - - return out if out.dtype == rdt else cast(out, rdt) + if (ops := _zero_fills(x1, x2)) is None: + return _dense(np.matmul, _matmul_shape(x1, x2), x1, x2) + return _product(*ops, typed, rdt, _utils.scalar(0, rdt)) + + +def _matmul_shape(a, b): + shape = np.broadcast_shapes(a.shape[:-2], b.shape[:-2]) + return shape + (a.shape[-2:-1] if a.ndim > 1 else ()) + (b.shape[-1:] if b.ndim > 1 else ()) def _product(a, b, typed, out_dt, fill): @@ -282,17 +262,16 @@ def _product(a, b, typed, out_dt, fill): ``_lazy``); the result's fill is ``fill``. """ Ta, Tb = a._T(), b._T() - shape = np.broadcast_shapes(a.shape[:-2], b.shape[:-2]) - shape = shape + (a.shape[-2:-1] if a.ndim > 1 else ()) + (b.shape[-1:] if b.ndim > 1 else ()) + shape = _matmul_shape(a, b) if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: return Array._new(Tensor(gb_dtype(out_dt), shape), fill) - def build(Ta, Tb): - return Ta.matmul(Tb, typed) + def build(Ta, Tb, op=typed): + return Ta.matmul(Tb, op) spec = _matmul_terms(a.shape, b.shape) spec = None if spec is None else (*spec, typed) - return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, a.shape[-1], True) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, a.shape[-1], True, typed) # Index letters for the einsum that computes a contraction (for slicing a pending result) @@ -346,21 +325,24 @@ def _tensordot_axes(x1, x2, axes): raise ValueError("tensordot: more axes than the arrays have") return tuple(range(x1.ndim - axes, x1.ndim)), tuple(range(axes)) a, b = axes - a = tuple(int(i) % x1.ndim for i in (a if isinstance(a, (list, tuple)) else (a,))) - b = tuple(int(i) % x2.ndim for i in (b if isinstance(b, (list, tuple)) else (b,))) + a = tuple(_axis(i, x1.ndim) for i in (a if isinstance(a, (list, tuple)) else (a,))) + b = tuple(_axis(i, x2.ndim) for i in (b if isinstance(b, (list, tuple)) else (b,))) if len(a) != len(b): raise ValueError("tensordot needs as many axes of each array") return a, b +def _axis(i, ndim): + i = int(i) + if not -ndim <= i < ndim: + raise IndexError(f"tensordot: axis {i} is out of bounds for an array of {ndim} dimensions") + return i % ndim + + def tensordot(x1, x2, /, *, axes=2, semiring=None): """The tensor dot product (numpy's tensordot). ``semiring`` (an extension): a graphblas Semiring or its name, under the module's rules for fills. """ - from ._elementwise import add, cast, multiply - from ._manipulation import reshape - from ._statistical import sum as sum_ - _check(x1, "tensordot") _check(x2, "tensordot") a_axes, b_axes = _tensordot_axes(x1, x2, axes) @@ -375,21 +357,10 @@ def tensordot(x1, x2, /, *, axes=2, semiring=None): if not default: out_dt, fill = _sound(typed, (x1, x2), "tensordot") return _tensordot(x1, x2, a_axes, b_axes, shape, typed, out_dt, fill) - if not _dense_ok(x1, x2): - return _dense(lambda a, b: np.tensordot(a, b, axes=(a_axes, b_axes)), x1, x2) - a0, fa = _zero_fill(x1) - b0, fb = _zero_fill(x2) - out = _tensordot(a0, b0, a_axes, b_axes, shape, typed, rdt, _utils.scalar(0, rdt)) - if fa is not None: - s = sum_(b0, axis=b_axes, dtype=rdt) if b_axes else cast(b0, rdt) - out = add(out, multiply(fa, reshape(s, (1,) * len(free1) + free2))) - if fb is not None: - r = sum_(a0, axis=a_axes, dtype=rdt) if a_axes else cast(a0, rdt) - out = add(out, multiply(fb, reshape(r, free1 + (1,) * len(free2)))) - if fa is not None and fb is not None: - K = math.prod(x1.shape[a] for a in a_axes) - out = add(out, multiply(multiply(fa, fb), K)) - return out if out.dtype == rdt else cast(out, rdt) + if (ops := _zero_fills(x1, x2)) is None: + func = lambda a, b: np.tensordot(a, b, axes=(a_axes, b_axes)) # noqa: E731 + return _dense(func, shape, x1, x2) + return _tensordot(*ops, a_axes, b_axes, shape, typed, rdt, _utils.scalar(0, rdt)) def _tensordot(a, b, a_axes, b_axes, shape, typed, out_dt, fill): @@ -398,13 +369,13 @@ def _tensordot(a, b, a_axes, b_axes, shape, typed, out_dt, fill): return Array._new(Tensor(gb_dtype(out_dt), shape), fill) axes = (list(a_axes), list(b_axes)) - def build(Ta, Tb): - return Ta.tensordot(Tb, axes, typed) + def build(Ta, Tb, op=typed): + return Ta.tensordot(Tb, axes, op) spec = _tensordot_terms(a.ndim, b.ndim, a_axes, b_axes) spec = None if spec is None else (*spec, typed) nterms = math.prod(a.shape[i] for i in a_axes) - return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, nterms, True) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, nterms, True, typed) def vecdot(x1, x2, /, *, axis=-1, semiring=None): @@ -413,8 +384,7 @@ def vecdot(x1, x2, /, *, axis=-1, semiring=None): ``semiring`` (an extension): a graphblas Semiring or its name, whose monoid reduces ``mult(x1, x2)`` over ``axis`` (no conjugate), under the module's rules for fills. """ - from ._elementwise import add, cast, conj, multiply - from ._statistical import sum as sum_ + from ._elementwise import conj _check(x1, "vecdot") _check(x2, "vecdot") @@ -435,20 +405,12 @@ def vecdot(x1, x2, /, *, axis=-1, semiring=None): if not default: out_dt, fill = _sound(typed, (x1, x2), "vecdot") return _vecdot(x1, x2, axis, shape, out_shape, typed, out_dt, fill) - if not _dense_ok(x1, x2): - return _dense(lambda a, b: np.vecdot(a, b, axis=axis), x1, x2) - if x1.dtype.kind == "c": - x1 = conj(x1) - a0, fa = _zero_fill(x1) - b0, fb = _zero_fill(x2) - out = _vecdot(a0, b0, axis, shape, out_shape, typed, rdt, _utils.scalar(0, rdt)) - if fa is not None: - out = add(out, multiply(fa, sum_(b0, axis=axis, dtype=rdt))) - if fb is not None: - out = add(out, multiply(fb, sum_(a0, axis=axis, dtype=rdt))) - if fa is not None and fb is not None: - out = add(out, multiply(multiply(fa, fb), shape[axis])) - return out if out.dtype == rdt else cast(out, rdt) + if (ops := _zero_fills(x1, x2)) is None: + return _dense(lambda a, b: np.vecdot(a, b, axis=axis), out_shape, x1, x2) + a0, b0 = ops + if a0.dtype.kind == "c": + a0 = conj(a0) + return _vecdot(a0, b0, axis, shape, out_shape, typed, rdt, _utils.scalar(0, rdt)) def _vecdot(a, b, axis, shape, out_shape, typed, out_dt, fill): @@ -461,10 +423,12 @@ def _vecdot(a, b, axis, shape, out_shape, typed, out_dt, fill): _utils.expansion_check(Ta, Tb, shape) axis += len(shape) - def build(Ta, Tb): - return Ta.ewise_mult(Tb, typed.binaryop).new().reduce(typed.monoid, axis=axis) + def build(Ta, Tb, op=typed): + return Ta.ewise_mult(Tb, op.binaryop).new().reduce(op.monoid, axis=axis) - return _lazy.contraction((a, b), build, out_shape, out_dt, fill, nterms=shape[axis]) + return _lazy.contraction( + (a, b), build, out_shape, out_dt, fill, nterms=shape[axis], typed=typed + ) def _both_broadcast(sa, sb, shape): @@ -499,17 +463,17 @@ def einsum(subscripts, /, *operands, semiring=None): shapes = tuple(x.shape for x in ops) shape = _einsum_plan(subscripts, shapes, "greedy")[2] if default: - if not all(_is_zero(x._fill) for x in ops): - _utils.check_dense(math.prod(shape), "einsum with a nonzero fill") - return _dense(lambda *a: np.einsum(subscripts, *a), *ops) + if (zeroed := _zero_fills(*ops)) is None: + return _dense(lambda *a: np.einsum(subscripts, *a), shape, *ops) + ops = zeroed out_dt, fill = rdt, _utils.scalar(0, rdt) else: out_dt, fill = _sound(typed, ops, "einsum", multiplies=len(ops) > 1) if math.prod(shape) == 0 or any(x._T().nvals == 0 for x in ops): return Array._new(Tensor(gb_dtype(out_dt), shape), fill) - def build(*tensors): - return tensor_einsum(subscripts, *tensors, op=typed) + def build(*tensors, op=typed): + return tensor_einsum(subscripts, *tensors, op=op) terms, output = _einsum_plan(subscripts, shapes, "greedy")[:2] sizes = { @@ -518,5 +482,5 @@ def build(*tensors): nterms = math.prod(n for c, n in sizes.items() if c not in output) product = _lazy._one_product(terms, output) return _lazy.contraction( - ops, build, shape, out_dt, fill, (terms, output, typed), nterms, product + ops, build, shape, out_dt, fill, (terms, output, typed), nterms, product, typed ) diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 5e0cfbf08..5fa24da74 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -96,7 +96,8 @@ def test_contractions_are_pending(masks): assert not pending(xp.vecdot(make(a[0], 0.0), make(a[1], 0.0))) assert not pending(xp.einsum("ij,ij->", x1, x1)) assert not pending(x1 @ xp.zeros((5, 3))) - assert not pending(x1 @ make(b, 1.0)) + assert not pending(x1 @ make(np.where(b == 0, 1.0, b), 1.0)) # (dense: a fill shows) + assert pending(x1 @ make(b, 1.0)) # (every position stored: its fill is taken as 0) with gb.config.set(), pytest.MonkeyPatch.context() as mp: mp.setattr(_lazy, "enabled", False) P = x1 @ x2 @@ -351,7 +352,7 @@ def test_finite_facts(): assert _lazy._known_finite(y) is True assert _lazy._known_finite(make(np.array([1, 2]), 0)) is True P = xp.matmul(y, make(np.array([[1.0], [2.0]]), 0.0)) - assert P._pending.finite == [True, None] + assert P._pending.finite == [True, True] # (every contraction asks: DESIGN D12a) SLICES = [ diff --git a/graphblas/tests/test_array_api_linalg.py b/graphblas/tests/test_array_api_linalg.py index f8d545bb3..0db53f38a 100644 --- a/graphblas/tests/test_array_api_linalg.py +++ b/graphblas/tests/test_array_api_linalg.py @@ -243,3 +243,38 @@ def test_dense_algebra(): xp.linalg.det(xp.zeros(3)) with pytest.raises(TypeError, match="Array"): xp.linalg.det(a) + + +def test_nonzero_fills_follow_d6(): + # (review finding M3, DESIGN D12c) a nonzero fill of an operand that stores every + # position is taken as 0; any other nonzero fill goes dense (no rank-1 split, which turned + # inf into NaN and lost every digit when the fill dwarfs the values) + inf = make(np.array([np.inf]), 0.0) + five = xp.from_coo(([0],), [1.0], (1,), fill_value=5.0) + assert_close(xp.vecdot(inf, five), np.array(np.inf)) + m = make(np.array([[np.inf, 0.0]]), 0.0) + k = xp.from_coo(([0, 1], [0, 0]), [1.0, 5.0], (2, 1), fill_value=5.0) + assert_close(m @ k, [[np.inf]]) + big = xp.from_coo(([0, 1],), [1.0, 2.0], (2,), fill_value=1e20) + assert_close(xp.vecdot(make(np.array([1.0, -1.0]), 0.0), big), np.array(-1.0)) + big32 = xp.from_coo(([0, 1],), np.array([1.0, 2.0], np.float32), (2,), fill_value=1e8) + out = xp.vecdot(make(np.array([1.0, -1.0], np.float32), 0.0), big32) + assert_close(out, np.array(-1.0, np.float32)) + # dense under the limit on the result, not only on the operands (review finding L8) + with gb.config.set(array_api_densify_limit=1000): + c = make(np.full((100, 1), np.inf), np.inf) + d = make(np.full((1, 100), 2.0), np.inf) + with pytest.raises(xp.DensifyError): + c @ d + with pytest.raises(xp.DensifyError): + xp.tensordot(c, d, axes=0) + + +def test_tensordot_axes_out_of_bounds(): + # (review finding L2) as numpy: no wrapping modulo the number of axes + x = make(np.ones((2, 3)), 0.0) + with pytest.raises(IndexError, match="out of bounds"): + xp.tensordot(x, x.T, axes=([5], [0])) + with pytest.raises(IndexError, match="out of bounds"): + xp.tensordot(x, x.T, axes=([1], [-3])) + assert_close(xp.tensordot(x, x.T, axes=([-1], [0])), np.ones((2, 3)) @ np.ones((3, 2))) diff --git a/graphblas/tests/test_array_api_semiring.py b/graphblas/tests/test_array_api_semiring.py index 9e9de8f34..a4d5cf11c 100644 --- a/graphblas/tests/test_array_api_semiring.py +++ b/graphblas/tests/test_array_api_semiring.py @@ -322,3 +322,46 @@ def test_probes(): if p.dtype.kind in "fc": assert np.isfinite(p).all() assert {-128, 127, 0} <= set(_probes(np.dtype("int8")).tolist()) + + +def test_explicit_fills_do_not_change_values(): + # (review finding M2, DESIGN D12a) a stored entry equal to the fill meets an inf as an + # unstored position does: no term + a = xp.from_coo(([0], [0]), [0.0], (1, 1)) + b = make(np.array([[np.inf]]), 0.0) + assert_close(a @ b, [[0.0]]) + assert_close(xp.prune(a) @ b, [[0.0]]) + a = xp.from_coo(([0], [0]), [np.inf], (1, 1), fill_value=np.inf) + b = make(np.array([[-np.inf]]), np.inf) + for x in [a, xp.prune(a)]: + assert_close(xp.matmul(x, b, semiring="min_plus"), [[np.inf]], np.inf) + # with no inf or NaN anywhere, explicit fills stay (nothing to prune) + a = xp.from_coo(([0, 0], [0, 1]), [0.0, 2.0], (1, 2)) + b = make(np.array([[1.0], [3.0]]), 0.0) + assert_close(a @ b, [[6.0]]) + + +def test_nan_propagates_through_min_and_max(): + # (review finding M1, DESIGN D12b) a stored NaN term makes the min or max NaN, as numpy's + # do, although SuiteSparse's fmin and fmax skip it; positions it does not reach keep theirs + inf = np.inf + a = np.array([[np.nan, 5.0], [inf, 2.0]]) + b = np.array([[1.0, inf], [3.0, 4.0]]) + A, B = make(a, inf), make(b, inf) + for name, (reduce, mult, e) in [("min_plus", SEMIRINGS["min_plus"]), + ("min_max", SEMIRINGS["min_max"])]: # fmt: skip + with np.errstate(invalid="ignore"): + expected = reduce(mult(a[:, :, None], b[None]), axis=1) + # (D6: a NaN meeting an unstored inf is no term, so [0, 1] keeps its value) + expected[0, 1] = reduce(mult(a[0, 1], b[1, 1])) + assert_close(xp.matmul(A, B, semiring=name), expected, e) + assert_close(xp.tensordot(A, B, axes=1, semiring=name), expected, e) + assert_close(xp.einsum("ij,jk->ik", A, B, semiring=name), expected, e) + m = make(np.array([[np.nan, 5.0]]), -inf) + n = make(np.array([[1.0], [3.0]]), -inf) + assert_close(xp.matmul(m, n, semiring="max_plus"), [[np.nan]], -inf) + assert_close(xp.vecdot(m, xp.matrix_transpose(n), semiring="max_plus"), [np.nan], -inf) + x = make(np.array([2.0, np.nan, 1.0]), inf) + assert_close(xp.einsum("b->", x, semiring="min_max"), np.array(np.nan)) + assert_close(xp.einsum("b->", make(np.array([2.0, 1.0]), inf), semiring="min_max"), + np.array(1.0)) # fmt: skip From 88a39a8c6a3d9a92f667e6c14bef426b4b7538af Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:22:47 -0600 Subject: [PATCH 53/77] array_api: abs of a complex array keeps only the "zero" fact |z| of a finite complex z may overflow to inf (|1.5e308 + 1.5e308j|), so abs of a complex array no longer passes a cached "finite" (or "nan") fact on to its result: a multiply after it took the intersection and read 0 where numpy has inf * 0, NaN. Found by review. --- graphblas/array_api/_elementwise.py | 4 +++- graphblas/tests/test_array_api_errors.py | 8 ++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 5c7ff4037..cdfcecfab 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -260,7 +260,9 @@ def get(dtype): def abs(x, /): # noqa: A001 (the Array API's name) # (GraphBLAS's integer abs of the most negative value is undefined in C; numpy wraps) gb = _first_of(_typed(unary.abs, _FLOAT), _copy_if("ub")) - return _unary("abs", np.abs, x, gb, keeps=_ALL_FACTS) + # (|z| of a finite complex z may overflow, and |inf + nan j| is inf: only "zero" holds) + keeps = ("zero",) if isinstance(x, Array) and x.dtype.kind == "c" else _ALL_FACTS + return _unary("abs", np.abs, x, gb, keeps=keeps) def negative(x, /): diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 0117a7a90..11a097961 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -223,6 +223,14 @@ def test_facts_of_stored_values(): assert func(v)._facts == {"finite": False, "nan": True} assert xp.exp(v)._facts is None assert xp.abs(xp.zeros(3))._facts is None + # (review finding M4) |z| of a finite complex z may be inf: abs keeps only "zero" there + z = make(np.array([1.5e308 + 1.5e308j, 1, 0, 0]), 0j) + assert _utils.fact(z, "finite") + _utils.fact(z, "zero") + assert set(xp.abs(z)._facts) == {"zero"} + zz = make(np.array([0.0, 0.0, 2.0, 3.0]), 0.0) + with np.errstate(over="ignore", invalid="ignore"): + assert_array_equal(xp.abs(z) * zz, np.abs(z.todense()) * zz.todense()) def test_thread3_argument_checks(): From 64ffd58bab4ec1f1aa641afc547a46ca0c6faba7 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:25:36 -0600 Subject: [PATCH 54/77] array_api: semiring checks per object, a batched dot hint, and P * P computed once Four small fixes, found by review: - The soundness check of a semiring is cached per typed semiring object, not by name: two user semirings of one name shared the first one's verdict. - The SDDMM's dot-product hint divides its saxpy estimate by the number of batches (a batched product pairs entries within a batch only). 2000 batches of 64 x 64 at density 0.01 under a mask of density 0.5, 8 threads: the deferred SDDMM went from 20.4 ms to 7.5 ms (7.8 ms without deferral). - multiply(P, P) of a pending product computes it once. - vecdot's docstring says the default semiring, given or not, conjugates x1. --- graphblas/array_api/_lazy.py | 19 ++++++++++---- graphblas/array_api/_linalg.py | 11 +++++--- graphblas/tests/test_array_api_lazy.py | 29 ++++++++++++++++++++++ graphblas/tests/test_array_api_semiring.py | 21 ++++++++++++++++ 4 files changed, 71 insertions(+), 9 deletions(-) diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py index a94c79a49..cab174c41 100644 --- a/graphblas/array_api/_lazy.py +++ b/graphblas/array_api/_lazy.py @@ -23,6 +23,8 @@ was deferred; ``enabled = False`` turns deferral off, for debugging. """ +import math + import numpy as np from .. import binary, monoid, select, semiring, unary @@ -174,19 +176,20 @@ def masked_multiply(a, b): S, P = b, a else: return None - if S.shape != P.shape or bool(S._fill) or bool(P._fill): - return None + if S is P or S.shape != P.shape or bool(S._fill) or bool(P._fill): + return None # (P * P: computed once, as a whole) node = P._pending if not node.all_finite(): return None # (P is computed whole, a sparse product as without deferral) TS = S._T() - PS = Array._new(node.compute(mask=TS.S, **_dot_hint(node, TS.nvals)), P._fill) + hint = _dot_hint(node, TS.nvals, P.shape) + PS = Array._new(node.compute(mask=TS.S, **hint), P._fill) from ._elementwise import multiply return multiply(PS, S) if P is a else multiply(S, PS) -def _dot_hint(node, nmask): +def _dot_hint(node, nmask, shape): """``axb_method="dot"`` for a product under a mask of ``nmask`` entries when dot products at the mask's entries cost well below a saxpy over the whole product, else nothing. @@ -198,7 +201,13 @@ def _dot_hint(node, nmask): return {} Ta, Tb = node.tensors k = max(node.nterms, 1) - saxpy = Ta.nvals * Tb.nvals / k + # (a batched product pairs entries within each batch only) + if node.spec is not None: + (ta, tb), output = node.spec[:2] + nbatch = math.prod(shape[output.index(c)] for c in output if c in ta and c in tb) + else: + nbatch = math.prod(shape[:-2]) # (a matmul whose batch axes broadcast) + saxpy = Ta.nvals * Tb.nvals / (k * max(nbatch, 1)) dot = nmask * k * (Ta.nvals / max(Ta.size, 1) + Tb.nvals / max(Tb.size, 1)) return {"axb_method": "dot"} if 4 * dot < saxpy else {} diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 9fff78639..773f482ce 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -25,6 +25,7 @@ """ import math +import weakref import numpy as np @@ -110,14 +111,15 @@ def _probes(dt): return np.array(values, dtype=dt) -_ANNIHILATES = {} +_ANNIHILATES = weakref.WeakKeyDictionary() def _annihilates(typed): """Whether the identity of the typed semiring's monoid annihilates its multiply from both - sides on the probe values (cached). + sides on the probe values (cached per typed semiring object: two user semirings may share + a name). """ - key = (typed.name, typed.type.name, typed.type2.name) + key = typed if (ok := _ANNIHILATES.get(key)) is None: ok = _ANNIHILATES[key] = _check_annihilates(typed) return ok @@ -382,7 +384,8 @@ def vecdot(x1, x2, /, *, axis=-1, semiring=None): """Sum over ``axis`` of ``conj(x1) * x2`` (broadcast), with zero fills structural. ``semiring`` (an extension): a graphblas Semiring or its name, whose monoid reduces - ``mult(x1, x2)`` over ``axis`` (no conjugate), under the module's rules for fills. + ``mult(x1, x2)`` over ``axis``, under the module's rules for fills. The default semiring + (plus_times, or lor_land for bool), given or not, conjugates ``x1``; others do not. """ from ._elementwise import conj diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 5fa24da74..993483e64 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -500,3 +500,32 @@ def test_inplace_operators_on_a_pending_array(op): assert_close(P, expected) assert_close(xp.multiply(make(s2, 0), P), s2 * expected) assert_close(P[1:], expected[1:]) + + +def test_sddmm_dot_hint_counts_batches(masks): + # (review finding L9) a batched product pairs entries within a batch only: the saxpy + # estimate divides by the number of batches, so sparse batches keep SuiteSparse's choice + rng = np.random.default_rng(17) + for dens_ab, dens_s, hint in [(1.0, 0.02, {"axb_method": "dot"}), (0.02, 0.5, {})]: + a = np.where(rng.random((8, 16, 4)) < dens_ab, rng.random((8, 16, 4)) + 0.5, 0.0) + b = np.where(rng.random((8, 4, 16)) < dens_ab, rng.random((8, 4, 16)) + 0.5, 0.0) + s = np.where(rng.random((8, 16, 16)) < dens_s, 1.0, 0.0) + masks.clear() + masks.hints.clear() + assert_close(xp.multiply(make(s, 0.0), make(a, 0.0) @ make(b, 0.0)), s * (a @ b)) + assert masks.hints == [hint] + # a matmul whose batch axes broadcast (no einsum terms): the batches of the result + a = rng.random((1, 16, 4)) + 0.5 + b = rng.random((8, 4, 16)) + 0.5 + s = np.where(rng.random((8, 16, 16)) < 0.005, 1.0, 0.0) + masks.hints.clear() + assert_close(xp.multiply(make(s, 0.0), xp.asarray(a) @ xp.asarray(b)), s * (a @ b)) + assert masks.hints == [{"axb_method": "dot"}] + + +def test_product_times_itself_computes_once(masks): + # (review finding L6) P * P: no push-down, one computation + a = _rand((5, 5)) + P = make(a, 0.0) @ make(a, 0.0) + assert_close(P * P, (a @ a) * (a @ a)) + assert masks == [None] diff --git a/graphblas/tests/test_array_api_semiring.py b/graphblas/tests/test_array_api_semiring.py index a4d5cf11c..fa1ecd885 100644 --- a/graphblas/tests/test_array_api_semiring.py +++ b/graphblas/tests/test_array_api_semiring.py @@ -365,3 +365,24 @@ def test_nan_propagates_through_min_and_max(): assert_close(xp.einsum("b->", x, semiring="min_max"), np.array(np.nan)) assert_close(xp.einsum("b->", make(np.array([2.0, 1.0]), inf), semiring="min_max"), np.array(1.0)) # fmt: skip + + +def test_soundness_is_cached_per_semiring_object(): + # (review finding L1) two semirings of one name: each is checked + from graphblas.core.operator import Semiring + + s1 = Semiring.register_anonymous(gb.monoid.min, gb.binary.plus, name="same_name") + s2 = Semiring.register_anonymous(gb.monoid.min, gb.binary.minus, name="same_name") + a = make(np.array([[1.0, 2.0]]), np.inf) + b = make(np.array([[1.0], [3.0]]), np.inf) + assert_close(xp.matmul(a, b, semiring=s1), [[2.0]], np.inf) + with pytest.raises(ValueError, match="cannot skip"): + xp.matmul(a, b, semiring=s2) + + +def test_vecdot_plus_times_conjugates(): + # (review finding L4) the default semiring, given or not, conjugates x1 + x = make(np.array([1j, 2.0]), 0j) + y = make(np.array([1j, 1.0]), 0j) + for semiring in [None, "plus_times"]: + assert_close(xp.vecdot(x, y, semiring=semiring), np.array(3 + 0j)) From 86b1b80d039d8eb9ccddf831048324a2628c3092 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:27:52 -0600 Subject: [PATCH 55/77] array_api: setitem with None in a basic key takes numpy's value shape x[0, None] = v with v shaped like x[0, None] raised: the value was sized without the new axes. For a key of ints, slices, "..." and None, the value's axes that line up with the new axes (of length 1) are now dropped first. Keys that mix None with index arrays keep their former behavior. Found by review. --- graphblas/array_api/_indexing.py | 24 ++++++++++++++++ graphblas/tests/test_array_api_indexing.py | 32 ++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/graphblas/array_api/_indexing.py b/graphblas/array_api/_indexing.py index 4f8432125..7d7b51bac 100644 --- a/graphblas/array_api/_indexing.py +++ b/graphblas/array_api/_indexing.py @@ -308,6 +308,10 @@ def setitem(x, key, value): axes_key = [k for k in items if not isinstance(k, (_NewAxis, _Gap))] if any(isinstance(k, _NewAxis) and k.length == 0 for k in items): return # (a False index selects nothing) + if isinstance(value, Array) and not _has_adv(items): + value = _without_new_axes(value, items) + if value.ndim == 0: + value = _scalar_value(value._fill, x.dtype) if x.ndim == 0: if isinstance(value, Array): raise ValueError(f"cannot assign an array of shape {value.shape} to a 0-d array") @@ -330,6 +334,26 @@ def setitem(x, key, value): x._facts = None +def _without_new_axes(value, items): + """``value`` for a basic key with new axes (``None``): the selection numpy broadcasts it + against has those axes, of length 1, which the region assigned to does not; drop the + value's axes that line up with them. + """ + kinds = [k for k in items if isinstance(k, (_NewAxis, slice))] + offset = len(kinds) - value.ndim + drop = [p - offset for p, k in enumerate(kinds) if isinstance(k, _NewAxis) and p - offset >= 0] + if not drop: + return value + if any(value.shape[q] != 1 for q in drop): + raise ValueError( + f"cannot assign an array of shape {value.shape} to a selection of " + f"{len(kinds)} dimensions whose new axes have length 1" + ) + from ._manipulation import squeeze + + return squeeze(value, axis=tuple(drop)) + + def _setitem_value(x, value): """The value cast to ``x``'s dtype: a numpy scalar, or an Array of at least 1 dimension.""" from ._creation import asarray diff --git a/graphblas/tests/test_array_api_indexing.py b/graphblas/tests/test_array_api_indexing.py index f4d92b13a..0a3454c3d 100644 --- a/graphblas/tests/test_array_api_indexing.py +++ b/graphblas/tests/test_array_api_indexing.py @@ -1,6 +1,7 @@ """graphblas.array_api indexing (get and set), take and take_along_axis against numpy.""" import itertools +import math import numpy as np import pytest @@ -271,3 +272,34 @@ def test_mixed_keys_order_axes_as_numpy(): b = np.arange(24.0).reshape(2, 3, 4) b[:, 1, ..., [0, 2]] = -expected assert_array_equal(y, b) + + +def test_setitem_with_new_axes(): + # (review finding L3) the value of a basic key with None is shaped like numpy's selection, + # new axes included + parts = [slice(None), 1, slice(0, 2), Ellipsis, None] + count = 0 + for a in [np.arange(12.0).reshape(3, 4) + 1, np.arange(24.0).reshape(2, 3, 4) + 1]: + for n in range(1, 5): + for key in itertools.product(parts, repeat=n): + if None not in key or sum(k is Ellipsis for k in key) > 1: + continue + try: + shape = a[key].shape + except IndexError: + continue + for vshape in [shape, shape[1:], (1,) * len(shape)]: + value = np.arange(math.prod(vshape), dtype=np.float64).reshape(vshape) - 50 + b = a.copy() + b[key] = value + y = make(a, 0.0) + y[key] = make(value, 0.0) if value.ndim else xp.asarray(value) + assert_array_equal(y, b) + count += 1 + assert count > 300 + z = xp.asarray(1.0) + z[None] = xp.asarray([5.0]) + assert_array_equal(z, np.array(5.0)) + x = make(np.ones((3, 4)), 0.0) + with pytest.raises(ValueError, match="new axes"): + x[None] = xp.ones((2, 3, 4)) From 767d33b1dc83afb9866b69fad03feda142f366b7 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:32:40 -0600 Subject: [PATCH 56/77] array_api: row sums and scalar factors keep contractions sparse A contraction with an operand whose nonzero fill shows went dense, so A @ xp.ones(n), einsum("ij,j->i", A, ones) and einsum("ij,->ij", A, 2.0) of a large sparse A raised DensifyError. An operand with a finite nonzero fill that is no larger than the largest operand's entries is now stored in full first (O(its size); an unstored zero of another operand times its finite values is 0 either way), and einsum takes a 0-d operand with a nonzero fill out as a scalar factor. A non-finite fill still goes dense, as numpy computes it. Found by review. --- graphblas/array_api/_linalg.py | 40 ++++++++++++++++++++-- graphblas/tests/test_array_api_lazy.py | 3 +- graphblas/tests/test_array_api_semiring.py | 31 +++++++++++++++++ 3 files changed, 70 insertions(+), 4 deletions(-) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 773f482ce..dfa689180 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -191,15 +191,23 @@ def _is_zero(fill): def _zero_fills(*arrays): """The arrays with fill 0, or None when one cannot have it. An array that stores every - position never shows its fill, so its fill is relabeled 0 (O(1)); any other nonzero fill - gives None (the contraction is then dense). + position never shows its fill, so its fill is relabeled 0 (O(1)). One with a finite fill + that is no larger than the entries of the largest operand (a vector of ones by a matrix: + row sums) is stored in full first, O(its size): an unstored zero of another operand times + its finite values is 0 either way. Any other nonzero fill gives None: the contraction is + dense (a non-finite fill times another's unstored zero is NaN there, as numpy has it). """ + most = max(x._T().nvals for x in arrays) out = [] for x in arrays: if not _is_zero(x._fill): T = x._T() if T.nvals != x.size: - return None + small = x.size <= min(most + 1024, _utils.densify_limit()) + if x.ndim == 0 or not small or not np.isfinite(x._fill): + return None + # (its fill stored at every position it does not store, the new fill 0) + T = _utils.with_fill(x, _utils.scalar(0, x.dtype))._T() x = Array._new(T, _utils.scalar(0, x.dtype)) out.append(x) return out @@ -465,6 +473,9 @@ def einsum(subscripts, /, *operands, semiring=None): typed, default = _resolve(semiring, rdt, "einsum") shapes = tuple(x.shape for x in ops) shape = _einsum_plan(subscripts, shapes, "greedy")[2] + if default and any(x.ndim == 0 and not _is_zero(x._fill) for x in ops): + # A 0-d operand is a scalar factor of every term (and of the sum, for plus_times) + return _einsum_scalars(subscripts, ops) if default: if (zeroed := _zero_fills(*ops)) is None: return _dense(lambda *a: np.einsum(subscripts, *a), shape, *ops) @@ -487,3 +498,26 @@ def build(*tensors, op=typed): return _lazy.contraction( ops, build, shape, out_dt, fill, (terms, output, typed), nterms, product, typed ) + + +def _einsum_scalars(subscripts, ops): + """Einsum with the 0-d operands taken out as a product of scalars: the einsum of the + others, times that product. + """ + from ._elementwise import multiply + + inputs, arrow, output = subscripts.replace(" ", "").partition("->") + terms = inputs.split(",") + keep = [i for i, x in enumerate(ops) if x.ndim > 0] + factor = None + for x in ops: + if x.ndim == 0: + factor = x if factor is None else multiply(factor, x) + if not keep: + return _dense(lambda *a: np.einsum(subscripts, *a), (), *ops) + rest = ",".join(terms[i] for i in keep) + arrow + output + if not arrow: + # (implicit output: the letters that appear once, in order, as numpy picks them) + counts = {c: "".join(terms).count(c) for c in "".join(terms)} + rest += "->" + "".join(sorted(c for c, k in counts.items() if k == 1)) + return multiply(einsum(rest, *(ops[i] for i in keep)), factor) diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 993483e64..8d05b1ee3 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -96,8 +96,9 @@ def test_contractions_are_pending(masks): assert not pending(xp.vecdot(make(a[0], 0.0), make(a[1], 0.0))) assert not pending(xp.einsum("ij,ij->", x1, x1)) assert not pending(x1 @ xp.zeros((5, 3))) - assert not pending(x1 @ make(np.where(b == 0, 1.0, b), 1.0)) # (dense: a fill shows) + assert not pending(x1 @ make(np.where(b == 0, np.inf, b), np.inf)) # (dense: inf shows) assert pending(x1 @ make(b, 1.0)) # (every position stored: its fill is taken as 0) + assert pending(x1 @ make(np.where(b == 0, 1.0, b), 1.0)) # (small: stored in full) with gb.config.set(), pytest.MonkeyPatch.context() as mp: mp.setattr(_lazy, "enabled", False) P = x1 @ x2 diff --git a/graphblas/tests/test_array_api_semiring.py b/graphblas/tests/test_array_api_semiring.py index fa1ecd885..b35e3ae4f 100644 --- a/graphblas/tests/test_array_api_semiring.py +++ b/graphblas/tests/test_array_api_semiring.py @@ -386,3 +386,34 @@ def test_vecdot_plus_times_conjugates(): y = make(np.array([1j, 1.0]), 0j) for semiring in [None, "plus_times"]: assert_close(xp.vecdot(x, y, semiring=semiring), np.array(3 + 0j)) + + +def test_einsum_with_scalars_and_ones_stays_sparse(): + # (review finding L5) a 0-d operand is a scalar factor, and a small operand with a finite + # nonzero fill (a vector of ones: row sums) is stored in full, not densified with the rest + rng = np.random.default_rng(18) + n = 400 + a = np.zeros((n, n)) + a[rng.integers(0, n, 600), rng.integers(0, n, 600)] = rng.random(600) + 0.5 + A = xp.from_coo(np.nonzero(a), a[np.nonzero(a)], (n, n)) + two, ones = xp.asarray(2.0), xp.ones(n) + with gb.config.set(array_api_densify_limit=10 * n): + outs = [ + xp.einsum("ij,j->i", A, ones), + A @ ones, + xp.tensordot(A, ones, axes=1), + xp.einsum("ij,->ij", A, two), + xp.einsum(",ij,->ji", two, A, two), + xp.einsum("ij,", A, two), + ] + for out, expected in zip( + outs, [a.sum(1), a.sum(1), a.sum(1), 2 * a, 4 * a.T, 2 * a], strict=True + ): + assert_close(out, expected) + assert_close(xp.einsum(",->", two, xp.asarray(3.0)), np.array(6.0)) + # a non-finite fill that shows goes dense, as numpy computes it (inf * 0 is NaN there) + c = np.array([[1.0, np.inf], [np.inf, 2.0]]) + d = np.array([[0.0, 1.0], [1.0, 0.0]]) + with np.errstate(invalid="ignore"): + expected = np.einsum("ij,jk->ik", c, d) + assert_close(xp.einsum("ij,jk->ik", make(c, np.inf), make(d, 0.0)), expected) From 3a14c188b9ba6e4557c42a44d1eb3e5490ebb23a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:33:48 -0600 Subject: [PATCH 57/77] array_api: Arrays pickle through their coordinates An Array could not be pickled: its Tensor holds weak references. Array now pickles as its stored coordinates and values, shape, dtype and fill, and is rebuilt by from_coo (explicit fills stay stored; a pending product is computed first). Found by review. --- graphblas/array_api/_array.py | 13 +++++++++++ graphblas/tests/test_array_api_creation.py | 25 ++++++++++++++++------ 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py index f9ae5ef81..72ec43a49 100644 --- a/graphblas/array_api/_array.py +++ b/graphblas/array_api/_array.py @@ -101,6 +101,11 @@ def _fact(self, name, compute): def _copy(self): return Array._new(self._T().dup(), self._fill) + def __reduce__(self): + # (a Tensor holds weak references, which do not pickle: rebuild from coordinates) + coords, values = self._T().to_coo() + return _unpickle, (coords, values, self._shape, self._dtype.str, self._fill) + # (copy.copy would share the Tensor, which __setitem__ changes in place) def __copy__(self): return self._copy() @@ -331,6 +336,14 @@ def __invert__(self): return bitwise_invert(self) +def _unpickle(coords, values, shape, dtype, fill): + if not shape: + return Array._new(Tensor(_utils.gb_dtype(dtype), ()), fill) + from ._creation import from_coo + + return from_coo(coords, values, shape, fill_value=fill, dtype=np.dtype(dtype)) + + def _is_operand(x): return isinstance(x, (Array, bool, int, float, complex, np.generic, np.ndarray)) diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py index 4a0d76f0c..c1708bf4b 100644 --- a/graphblas/tests/test_array_api_creation.py +++ b/graphblas/tests/test_array_api_creation.py @@ -443,13 +443,26 @@ def test_copy_module(): assert_array_equal(z, a @ a) -def test_pickle_not_supported_yet_or_roundtrips(): - x = make(np.array([1.0, 0.0, 2.0]), 0.0) - try: +def test_pickle_roundtrips(): + # (review finding L7) through coordinates, values, shape, dtype and fill; explicit fills + # stay stored, and a pending product is computed first + rng = np.random.default_rng(19) + for dense, fill in [ + (np.array([1.0, 0.0, 2.0]), 0.0), + (np.array([[1, 7], [7, 3]], np.int16), 7), + (np.array([True, False, True]), True), + (np.array([1j, 0, np.nan]), 0j), + (np.array(2.5), 0.0), + (np.zeros((0, 3)), 1.0), + ]: + x = make(dense, fill, rng) y = pickle.loads(pickle.dumps(x)) - except (TypeError, AttributeError, pickle.PicklingError): - return - assert_array_equal(y, x.todense()) + assert_array_equal(y, dense) + assert y.nnz == x.nnz + assert y.fill_value == x.fill_value or np.isnan(fill) + a = make(np.array([[1.0, 2.0], [0.0, 3.0]]), 0.0) + P = a @ a + assert_array_equal(pickle.loads(pickle.dumps(P)), a.todense() @ a.todense()) # Indexing and shape functions From 06bc3f8765eedb06ccf4aa383c0b8a1bddc3324c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:37:30 -0600 Subject: [PATCH 58/77] array_api: einsum factors out only finite scalars einsum took a 0-d operand out as a scalar factor of the whole sum; for an infinite scalar that differs from numpy, which multiplies each term (the unstored zeros times inf are NaN). A non-finite 0-d operand now takes the dense path, as before the factoring. The SDDMM's dot hint takes the result's shape as optional (one batch without it). --- graphblas/array_api/_lazy.py | 6 ++++-- graphblas/array_api/_linalg.py | 7 +++++-- graphblas/tests/test_array_api_lazy.py | 3 +++ graphblas/tests/test_array_api_semiring.py | 3 +++ 4 files changed, 15 insertions(+), 4 deletions(-) diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py index cab174c41..a9c5afd41 100644 --- a/graphblas/array_api/_lazy.py +++ b/graphblas/array_api/_lazy.py @@ -189,7 +189,7 @@ def masked_multiply(a, b): return multiply(PS, S) if P is a else multiply(S, PS) -def _dot_hint(node, nmask, shape): +def _dot_hint(node, nmask, shape=()): """``axb_method="dot"`` for a product under a mask of ``nmask`` entries when dot products at the mask's entries cost well below a saxpy over the whole product, else nothing. @@ -202,7 +202,9 @@ def _dot_hint(node, nmask, shape): Ta, Tb = node.tensors k = max(node.nterms, 1) # (a batched product pairs entries within each batch only) - if node.spec is not None: + if not shape: + nbatch = 1 + elif node.spec is not None: (ta, tb), output = node.spec[:2] nbatch = math.prod(shape[output.index(c)] for c in output if c in ta and c in tb) else: diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index dfa689180..7ca70c411 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -473,8 +473,11 @@ def einsum(subscripts, /, *operands, semiring=None): typed, default = _resolve(semiring, rdt, "einsum") shapes = tuple(x.shape for x in ops) shape = _einsum_plan(subscripts, shapes, "greedy")[2] - if default and any(x.ndim == 0 and not _is_zero(x._fill) for x in ops): - # A 0-d operand is a scalar factor of every term (and of the sum, for plus_times) + scalars = [x._fill for x in ops if x.ndim == 0] + if default and any(map(bool, scalars)) and all(np.isfinite(scalars)): + # A finite 0-d operand is a scalar factor of every term (and of the sum, for + # plus_times); a non-finite one goes dense, as numpy has it (inf times the unstored + # zeros is NaN) return _einsum_scalars(subscripts, ops) if default: if (zeroed := _zero_fills(*ops)) is None: diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 8d05b1ee3..cdbc9ab92 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -530,3 +530,6 @@ def test_product_times_itself_computes_once(masks): P = make(a, 0.0) @ make(a, 0.0) assert_close(P * P, (a @ a) * (a @ a)) assert masks == [None] + # (without the result's shape, one batch is assumed) + P = make(a, 0.0) @ make(a, 0.0) + assert _lazy._dot_hint(P._pending, 1) == {"axb_method": "dot"} diff --git a/graphblas/tests/test_array_api_semiring.py b/graphblas/tests/test_array_api_semiring.py index b35e3ae4f..872e00dba 100644 --- a/graphblas/tests/test_array_api_semiring.py +++ b/graphblas/tests/test_array_api_semiring.py @@ -414,6 +414,9 @@ def test_einsum_with_scalars_and_ones_stays_sparse(): # a non-finite fill that shows goes dense, as numpy computes it (inf * 0 is NaN there) c = np.array([[1.0, np.inf], [np.inf, 2.0]]) d = np.array([[0.0, 1.0], [1.0, 0.0]]) + with np.errstate(invalid="ignore"): + assert_close(xp.einsum("ij,->", make(d, 0.0), xp.asarray(np.inf)), + np.einsum("ij,->", d, np.inf)) # fmt: skip with np.errstate(invalid="ignore"): expected = np.einsum("ij,jk->ik", c, d) assert_close(xp.einsum("ij,jk->ik", make(c, np.inf), make(d, 0.0)), expected) From 237b54dec04c00a03a1d6ed0ed1ce9417c6423c7 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:43:56 -0600 Subject: [PATCH 59/77] array_api: a fully stored bool result decides on dropping False by a sample Counting the True values of a fully stored bool result, to decide whether to drop its False entries, cost as much as the comparison that made it (2e7 values at 1 thread: greater 3.4 ms, with the count 6.4 ms), on every result. A sample of the leading 2^16 entries now decides: over a third True, the result is kept as it is (3.3 ms); otherwise one select of the True values decides, and is the result when at most a quarter of the entries are True. Dropping still costs the select (2e7 values, 1% True: 22 ms at 1 thread, 4.2 ms at 8, against 3.2 and 1.4 for the comparison alone). Values never change; a sample unlike the rest only changes what is stored. Found by review. --- graphblas/array_api/_elementwise.py | 21 ++++++++++++++----- graphblas/tests/test_array_api_elementwise.py | 19 +++++++++++++++++ 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index cdfcecfab..df81c6f86 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -171,13 +171,24 @@ def _normalize_bool(T, fill): (the fill contractions and masks work with; a fill True made a mostly-True mask times a sparse matrix dense). Its False entries are dropped when at most a quarter of the entries are True (``equal(deg, 0)`` of a fully stored ``deg`` then stores only its True - entries): one count, then one select. Otherwise they stay; dropping more costs about as - much as the comparison and saves a later consumer little. + entries). Otherwise they stay; dropping more costs about as much as the comparison and + saves a later consumer little. + + A sample of the leading entries says whether dropping is likely: a mostly-True result pays + only for the sample (a count over all 2e7 values cost as much as the comparison at 1 + thread). Then one select of the True values decides, and is the result when few. """ n = T.nvals - if 4 * T.reduce_scalar(monoid.plus[INT64]).new().value > n: - return T, False - return T.select(select.valueeq, True).new(), False + rows = -(-_SAMPLE * T.shape[0] // n) # (leading rows holding about _SAMPLE entries) + if rows < T.shape[0]: + sample = T[(slice(0, rows),) + (slice(None),) * (T.ndim - 1)].new() + if 3 * sample.reduce_scalar(monoid.plus[INT64]).new().value > sample.nvals: + return T, False # (over a third True in the sample: dropping is unlikely) + S = T.select(select.valueeq, True).new() + return (S if 4 * S.nvals <= n else T), False + + +_SAMPLE = 2**16 def _zero_d(dtype, fill): diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index c8e66a4f2..eab19bca4 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -744,6 +744,25 @@ def test_fully_stored_bool_results_take_fill_false(): assert_array_equal(xp.equal(xp.asarray(3), 3), np.array(True)) +def test_bool_normalization_decides_by_a_sample(monkeypatch): + # (review finding L10) a sample of the leading entries says whether dropping the False + # entries is likely; values never change, only what is stored + from graphblas.array_api import _elementwise + + monkeypatch.setattr(_elementwise, "_SAMPLE", 4) + v = np.arange(40.0).reshape(10, 4) + full = xp.asarray(v + 1) + for out, expected, nnz in [ + (xp.greater(full, 4), v + 1 > 4, 40), # (sample mostly True: kept, fill False) + (xp.greater(full, 36), v + 1 > 36, 4), # (few True: dropped) + (xp.less(full, 37), v + 1 < 37, 40), # (sample all True, most True: kept) + (xp.less(full, 5), v + 1 < 5, 40), # (the sample is all True: kept, though few are) + (xp.greater(full, 3), v + 1 > 3, 40), # (a quarter of the sample: the select decides) + ]: + assert_array_equal(out, expected) + assert (out.nnz, out.fill_value) == (nnz, False) + + def test_mostly_true_masks_stay_sparse_in_contractions(): # (review finding H3) a mostly-True fully stored mask keeps fill False, so a contraction # with a sparse bool matrix stays sparse (it raised DensifyError with fill True) From 55c5c515ef7961e1ca2f07d02145eb974c630250 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:04:26 -0600 Subject: [PATCH 60/77] array_api: finite fills store in full under the limit; gather only past y's size Two review findings in the paths that keep products sparse: - A contraction stored an operand with a finite nonzero fill in full only when it was no larger than the other operands' stored entries, so A @ xp.ones((n, 8)) of a large sparse A raised DensifyError, and the answer could depend on the other operand's explicit zeros (stored in full: no term where an unstored zero meets inf; dense: NaN). Such an operand is now stored in full whenever its size fits under the densify limit, whatever the others store; a non-finite fill still goes dense. - The one-sided product gathered a broadcast operand's values at every entry of the other even when that operand is smaller: A (1e6 x 1e6, 1e7 entries) times y[:, None] with y 10% stored took 147 ms and 480 MB more RSS at 8 threads. It now gathers only when the other has fewer entries than y has positions, and otherwise stores y in full for the broadcast product: 6 ms, 36 MB. --- graphblas/array_api/_elementwise.py | 14 +++++----- graphblas/array_api/_linalg.py | 12 ++++----- graphblas/tests/test_array_api_elementwise.py | 14 ++++++++++ graphblas/tests/test_array_api_linalg.py | 26 +++++++++++++++++++ 4 files changed, 54 insertions(+), 12 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index df81c6f86..c2a6424e0 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -666,13 +666,15 @@ def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): Tx, Ty = X._T(), Y._T() with np.errstate(all="ignore"): fill = spec.ufunc(fx, fy) if X is a else spec.ufunc(fy, fx) - if Ty.shape != shape and not _stores_all(Y): - # Y broadcasts and shows its fill: its value at each of X's entries, gathered - # from Y itself at the projected coordinates, O(entries of X). (Storing Y in full - # cost O(size of Y), and the union replicated Y over its broadcast axes: - # exp(v)[None, :] times X of shape (1000, 1e6) with 4 entries made v 1000 times.) + if Ty.shape != shape and not _stores_all(Y) and Tx.nvals < Y.size: + # Y broadcasts, shows its fill, and is larger than X's entries: its value at each + # of X's entries, gathered from Y itself at the projected coordinates, O(entries + # of X). (Storing Y in full cost O(size of Y), and the union replicated Y over its + # broadcast axes: exp(v)[None, :] times X of shape (1000, 1e6) with 4 entries made + # v 1000 times.) A Y no larger is stored in full below: the broadcast ewise_mult + # then needs no coordinates (a row scaling of 1e7 entries: 5 ms, against 143) Ty = _gather(Tx, Y, shape) - elif not _stores_all(Y) and Y.size <= Tx.nvals + 1024: + elif not _stores_all(Y) and (Y.size <= Tx.nvals + 1024 or Ty.shape != shape): # (O(size of Y): every position of Y stored, so the product is an intersection) other = 0 if Y._fill != 0 else 1 Ty = _utils.with_fill(Y, other)._T() diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 7ca70c411..e6adf4644 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -192,18 +192,18 @@ def _is_zero(fill): def _zero_fills(*arrays): """The arrays with fill 0, or None when one cannot have it. An array that stores every position never shows its fill, so its fill is relabeled 0 (O(1)). One with a finite fill - that is no larger than the entries of the largest operand (a vector of ones by a matrix: - row sums) is stored in full first, O(its size): an unstored zero of another operand times - its finite values is 0 either way. Any other nonzero fill gives None: the contraction is - dense (a non-finite fill times another's unstored zero is NaN there, as numpy has it). + whose size fits under the densify limit (a vector or matrix of ones by a sparse matrix) is + stored in full first, O(its size): an unstored zero of another operand times its finite + values is 0 either way, so the answer does not depend on what the others store. Any other + nonzero fill gives None: the contraction is dense (a non-finite fill times another's + unstored zero is NaN there, as numpy has it). """ - most = max(x._T().nvals for x in arrays) out = [] for x in arrays: if not _is_zero(x._fill): T = x._T() if T.nvals != x.size: - small = x.size <= min(most + 1024, _utils.densify_limit()) + small = x.size <= _utils.densify_limit() if x.ndim == 0 or not small or not np.isfinite(x._fill): return None # (its fill stored at every position it does not store, the new fill 0) diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index eab19bca4..93ac21117 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -763,6 +763,20 @@ def test_bool_normalization_decides_by_a_sample(monkeypatch): assert (out.nnz, out.fill_value) == (nnz, False) +def test_one_sided_stores_a_small_broadcast_operand(): + # (review finding R2) a broadcast operand no larger than the other's entries is stored in + # full and taken by the broadcast product, with no coordinates exported + rng = np.random.default_rng(21) + a = np.where(rng.random((50, 40)) < 0.6, rng.random((50, 40)) + 0.5, 0.0) + y = np.where(rng.random(50) < 0.1, rng.random(50), 0.0) + A, Y = make(a, 0.0), xp.exp(make(y, 0.0)) + assert Y.nnz < 50 + assert Y.fill_value == 1.0 + out = A * Y[:, None] + assert out.nnz == A.nnz + assert_array_equal(out, a * np.exp(y)[:, None]) + + def test_mostly_true_masks_stay_sparse_in_contractions(): # (review finding H3) a mostly-True fully stored mask keeps fill False, so a contraction # with a sparse bool matrix stays sparse (it raised DensifyError with fill True) diff --git a/graphblas/tests/test_array_api_linalg.py b/graphblas/tests/test_array_api_linalg.py index 0db53f38a..fc2f96a9d 100644 --- a/graphblas/tests/test_array_api_linalg.py +++ b/graphblas/tests/test_array_api_linalg.py @@ -278,3 +278,29 @@ def test_tensordot_axes_out_of_bounds(): with pytest.raises(IndexError, match="out of bounds"): xp.tensordot(x, x.T, axes=([1], [-3])) assert_close(xp.tensordot(x, x.T, axes=([-1], [0])), np.ones((2, 3)) @ np.ones((3, 2))) + + +def test_finite_fills_are_stored_in_full(): + # (review findings R1, R3) an operand with a finite nonzero fill that fits under the + # densify limit is stored in full, whatever the other operands store: matrices of ones + # by a large sparse matrix stay sparse, and explicit zeros in the other operand do not + # change the answer (D10) + rng = np.random.default_rng(20) + n = 400 + a = np.zeros((n, n)) + a[rng.integers(0, n, 600), rng.integers(0, n, 600)] = rng.random(600) + 0.5 + A = xp.from_coo(np.nonzero(a), a[np.nonzero(a)], (n, n)) + with gb.config.set(array_api_densify_limit=10 * n): + r1 = A @ xp.ones((n, 8)) + r2 = A @ xp.full((n, 8), 2.5) + assert_close(r1, a @ np.ones((n, 8))) + assert_close(r2, a @ np.full((n, 8), 2.5)) + b = np.ones((2, 1500)) + b[0, 0] = np.inf + B = xp.asarray(b, fill_value=1.0) + A1 = xp.from_coo(([0], [0]), np.array([1.0]), (3000, 2)) + r = np.arange(1, 2001) + A2 = xp.from_coo((np.r_[0, r], np.r_[0, np.zeros(2000, int)]), np.r_[1.0, np.zeros(2000)], + (3000, 2)) # fmt: skip + for x in [A1, A2, xp.prune(A2)]: + assert (x @ B).todense()[1, 0] == 0.0 # (D6: an unstored zero times inf is no term) From 7f3080e136576c69d05a4778f22bbb318f68cb9f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:50:03 -0600 Subject: [PATCH 61/77] array_api: a three-operand einsum computes the pair under the third's pattern einsum("ij,jk,ik->", U, U, U) (triangle counting) built all of U @ U before the third operand met it. When one operand's indices are exactly what the other two leave after contracting their shared ones (batch indices included), the pair is now computed at that operand's stored positions only (a masked product, with the SDDMM dot hint), times the operand, then summed into the output. Only for semirings whose multiply distributes over the monoid (plus_times, lor_land, min_plus, max_plus, min_max, max_min). Operands holding inf or NaN prune their explicit fills first, as the two-operand path does; a NaN under a float min or max monoid, a complex operand holding inf, and a non-finite masking operand under plus_times (inf * (x + y) is not inf * x + inf * y) take the full einsum. Values do not depend on the path. --- graphblas/array_api/_linalg.py | 59 +++++++++++++ graphblas/tests/test_array_api_lazy.py | 115 +++++++++++++++++++++++++ 2 files changed, 174 insertions(+) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index e6adf4644..5a59401d0 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -493,6 +493,10 @@ def build(*tensors, op=typed): return tensor_einsum(subscripts, *tensors, op=op) terms, output = _einsum_plan(subscripts, shapes, "greedy")[:2] + if len(ops) == 3: + masked = _masked_three(terms, output, ops, typed, out_dt, fill) + if masked is not None: + return masked sizes = { c: n for term, s in zip(terms, shapes, strict=True) for c, n in zip(term, s, strict=True) } @@ -503,6 +507,61 @@ def build(*tensors, op=typed): ) +# (monoid, multiply) pairs where the multiply distributes over the monoid, so that a term's +# third factor can multiply the sum of the other two's products: s * (a + b) = s*a + s*b +_DISTRIBUTES = { + ("plus", "times"), + ("lor", "land"), + ("min", "plus"), + ("max", "plus"), + ("min", "max"), + ("max", "min"), +} + + +def _masked_three(terms, output, ops, typed, out_dt, fill): + """An einsum of three operands in which one operand's indices are exactly those the + other two leave after contracting their shared ones (triangle counting's + "ij,jk,ik->", an SDDMM's "ik,ij,jk->ik"), computed as that pair's product at the third + operand's stored positions only, times the third, then summed into the output; else + None. Every term of the full einsum is a term of this one: a position the third operand + does not store gives no term either way (DESIGN D6). Sums differ only in their order + (rounding, and where a sum overflows). An operand holding inf or NaN prunes the explicit + fills first, as every contraction does (D12a); a NaN under a float min or max monoid + takes the full einsum, which propagates it (D12b). With plus_times the third operand + must be finite: inf * (x + y) is not inf * x + inf * y when x and y differ in sign. + """ + if (typed.monoid.name, typed.binaryop.name) not in _DISTRIBUTES: + return None + finite = [x.dtype.kind not in "fc" or _utils.fact(x, "finite") for x in ops] + if not all(finite): + if typed.monoid.name in {"min", "max"} and any(_utils.fact(x, "nan") for x in ops): + return None + if any(x.dtype.kind == "c" for x in ops): + return None # (complex infinities do not distribute) + from ._manipulation import prune + + ops = [prune(x) for x in ops] + for m, tm in enumerate(terms): + a, b = (t for i, t in enumerate(terms) if i != m) + contracted = (set(a) & set(b)) - set(tm) + # (the pair's other indices, batch ones included, are exactly the third's) + if not contracted or set(a) | set(b) != set(tm) | contracted or not set(output) <= set(tm): + continue + if typed.monoid.name == "plus" and not finite[m]: + continue + A, B = (x for i, x in enumerate(ops) if i != m) + P = einsum(f"{a},{b}->{tm}", A, B, semiring=typed) + TS = ops[m]._T() + if (node := P._pending) is not None: + TP = node.compute(mask=TS.S, **_lazy._dot_hint(node, TS.nvals, P.shape)) + else: + TP = P._T() # (computed already: a NaN or deferral off) + R = Array._new(TS.ewise_mult(TP, typed.binaryop).new(dtype=gb_dtype(out_dt)), fill) + return R if output == tm else einsum(f"{tm}->{output}", R, semiring=typed) + return None + + def _einsum_scalars(subscripts, ops): """Einsum with the 0-d operands taken out as a product of scalars: the einsum of the others, times that product. diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index cdbc9ab92..a994d6516 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -533,3 +533,118 @@ def test_product_times_itself_computes_once(masks): # (without the result's shape, one batch is assumed) P = make(a, 0.0) @ make(a, 0.0) assert _lazy._dot_hint(P._pending, 1) == {"axb_method": "dot"} + + +THREE = [ + ("ij,jk,ik->", [(5, 6), (6, 7), (5, 7)]), # (triangle counting) + ("ij,jk,ik->i", [(5, 6), (6, 7), (5, 7)]), + ("ij,jk,ik->ki", [(5, 6), (6, 7), (5, 7)]), + ("ik,ij,jk->ik", [(5, 7), (5, 6), (6, 7)]), # (an SDDMM) + ("qn,qf,nf->qn", [(5, 7), (5, 3), (7, 3)]), + ("ij,j,i->", [(5, 6), (6,), (5,)]), + ("bij,bjk,bik->b", [(2, 3, 4), (2, 4, 5), (2, 3, 5)]), +] + + +@pytest.mark.parametrize(("subscripts", "shapes"), THREE, ids=[s for s, _ in THREE]) +def test_einsum_three_operands_masked(subscripts, shapes, masks): + # One operand's indices are what the other two leave: their product is computed at that + # operand's stored positions only + for dtype in [np.float64, np.int64, np.bool_, np.complex128]: + arrays = [_rand(s, 0, dtype, density=0.4) for s in shapes] + fill = np.zeros((), dtype)[()] + Xs = [make(a, fill, rng) for a in arrays] + masks.clear() + R = xp.einsum(subscripts, *Xs) + assert_close(R, np.einsum(subscripts, *arrays)) + if all(xp.prune(x).nnz for x in Xs): # (an empty operand: an empty pair, no product) + assert masks[0] is not None + # Semirings whose multiply distributes over the monoid + from .test_array_api_semiring import SEMIRINGS, semiring_einsum + + for name in ["min_plus", "max_plus", "min_max", "max_min"]: + reduce, mult, e = SEMIRINGS[name] + arrays = [_rand(s, e, density=0.4) for s in shapes] + masks.clear() + Xs = [make(a, e, rng) for a in arrays] + R = xp.einsum(subscripts, *Xs, semiring=name) + assert_close(R, semiring_einsum(subscripts, *arrays, reduce=reduce, mult=mult)) + assert R.fill_value == e or not R.ndim + if all(xp.prune(x).nnz for x in Xs): + assert masks[0] is not None + + +def _structural_three(a, b, s, mult=np.multiply): + """``ij,jk,ik->ik`` with a term only where all three operands are nonzero (DESIGN D6).""" + stored = (a[:, :, None] != 0) & (b[None] != 0) & (s[:, None, :] != 0) + with np.errstate(invalid="ignore", over="ignore"): + terms = mult(mult(a[:, :, None], b[None]), s[:, None, :]) + return np.where(stored, terms, 0).sum(axis=1) + + +def test_einsum_three_operands_nonfinite(masks, monkeypatch): + # An operand holding inf or NaN: explicit fills go first (D12a), and the terms are the + # full einsum's (D6): no NaN where the masking operand stores nothing. The values do not + # depend on which path ran. + from graphblas.array_api import _linalg + + a, b, s = _rand((4, 5)), _rand((5, 6)), _rand((4, 6)) + for bad in [np.inf, -np.inf, np.nan]: + for which in range(3): + arrays = [a.copy(), b.copy(), s.copy()] + arrays[which][0, 1] = bad + Xs = [make(x, 0.0, rng) for x in arrays] + masks.clear() + R = xp.einsum("ij,jk,ik->ik", *Xs) + # (only the third operand covers the output "ik": when it is the non-finite one, + # the full einsum runs) + assert any(m is not None for m in masks) == (which != 2) + if which != 2: + assert_close(R, _structural_three(*arrays)) + with monkeypatch.context() as m: + m.setattr(_linalg, "_masked_three", lambda *args: None) + assert_close(R, xp.einsum("ij,jk,ik->ik", *Xs).todense()) + # With plus_times, a non-finite operand is never the one the pair is computed under + # (inf * (x + y) differs from inf * x + inf * y), and complex goes whole + c = [x.astype(np.complex128) for x in (a, b, s)] + c[0][0, 1] = np.inf + masks.clear() + R = xp.einsum("ij,jk,ik->ik", *[make(x, 0j) for x in c]) + assert all(m is None for m in masks) + assert_close(R, _structural_three(*c)) + # A NaN under a float min or max monoid: the full einsum, which propagates it (D12b) + inf = np.inf + arrays = [np.where(x == 0, inf, x) for x in (a, b, s)] + arrays[0][0, 0] = np.nan + masks.clear() + R = xp.einsum("ij,jk,ik->ik", *[make(x, inf) for x in arrays], semiring="min_plus") + assert all(m is None for m in masks) + x, y, z = arrays[0][:, :, None], arrays[1][None], arrays[2][:, None, :] + stored = (x != inf) & (y != inf) & (z != inf) # (a NaN is stored: the term is NaN) + assert_close(R, np.min(np.where(stored, x + y + z, inf), axis=1)) + + +def test_einsum_three_operands_not_masked(masks): + a, b, s = _rand((4, 5)), _rand((5, 6)), _rand((4, 6)) + A, B, S = make(a, 0.0), make(b, 0.0), make(s, 0.0) + # No operand is what the other two leave, or no index is shared + for subscripts, arrays in [ + ("ij,jk,kl->il", [a, b, _rand((6, 3))]), + ("ij,kl,ik->", [a, _rand((3, 2)), _rand((4, 3))]), + ]: + masks.clear() + R = xp.einsum(subscripts, *[make(x, 0.0) for x in arrays]) + assert_close(R, np.einsum(subscripts, *arrays)) + assert all(m is None for m in masks) + # A semiring whose multiply does not distribute over its monoid: plus_min (unsigned) + p, q, r = (np.abs(x).astype(np.uint8) for x in (a, b, s)) + masks.clear() + R = xp.einsum("ij,jk,ik->ik", make(p, 0), make(q, 0), make(r, 0), semiring="plus_min") + assert all(m is None for m in masks) + assert_close(R, _structural_three(p, q, r, np.minimum).astype(np.uint8)) + # Deferral off: the pair is computed whole, then intersected + _lazy.enabled = False + try: + assert_close(xp.einsum("ij,jk,ik->", A, B, S), np.einsum("ij,jk,ik->", a, b, s)) + finally: + _lazy.enabled = True From ad859dd71cf0f9700e501052943e7dcf9a75f76a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 06:55:06 -0600 Subject: [PATCH 62/77] array_api: where stores the condition's entries over the operand its fill selects Where the condition stores nothing it reads its fill, so the result there is the operand that fill selects. When that operand needs no broadcast (or is a scalar, or stores nothing), where now starts from it and writes the condition's stored positions over it, reading the other operand only at those positions. SAPS's connected components, min(where(graph, labels, int_max), axis=1), replicated labels over every row of graph (n^2 entries; a DensifyError past about 10k nodes); it now stores graph's entries. Otherwise the union as before. --- graphblas/array_api/_searching.py | 52 ++++++++++++++++++++- graphblas/tests/test_array_api_searching.py | 36 ++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index e0847cb78..8d7f69f56 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -37,8 +37,13 @@ def _value(x): def where(condition, x1, x2, /): """``x1`` where ``condition`` is true, else ``x2`` (broadcast); numpy's ``where``. - The entries are the union of the three arrays' entries, O(entries) (with the broadcast - operands' entries replicated, as for any union). + Where the condition stores nothing it reads its fill, so the result there is the operand + that fill selects. When that operand needs no broadcast (or is a scalar, or stores + nothing), the result is that operand with the condition's stored positions written over: + O(entries of the condition and of that operand), the other operand read only at the + condition's entries (``xp.where(graph, labels, big)``, with ``labels`` broadcast along an + axis of ``graph``, stores ``graph``'s entries). Otherwise the entries are the union of + the three arrays' entries, with the broadcast operands' entries replicated. """ cond = _operand(condition, "where") if not isinstance(cond, Array): @@ -54,6 +59,8 @@ def where(condition, x1, x2, /): arrays = [x for x in (cond, a, b) if isinstance(x, Array) and x.ndim] if not shape or all(x._T().nvals == 0 for x in arrays): return Array._new(Tensor(gb_dtype(rdt), shape), fill) + if (R := _where_at_condition(cond, a, b, shape, rdt)) is not None: + return Array._new(R, fill) tensors = iter(_utils.aligned(arrays, shape)) slots = [next(tensors) if isinstance(x, Array) and x.ndim else None for x in (cond, a, b)] constants = [_value(x) for x in (cond, a, b)] @@ -69,6 +76,47 @@ def func(*values): return Array._new(R, fill) +def _where_at_condition(cond, a, b, shape, rdt): + """``where(cond, a, b)``'s Tensor built from the operand ``cond``'s fill selects, with + ``cond``'s stored positions written over (see ``where``), or None when that operand + would have to be broadcast over stored entries, or ``cond`` itself broadcasts. + """ + from .. import binary + + Tc = cond._T() + if Tc.shape != shape: + return None + base = a if bool(cond._fill) else b + if isinstance(base, Array) and base.ndim and base.shape != shape and base._T().nvals: + return None + coords, c = Tc.to_coo() + values = [] + for x in (a, b): + if isinstance(x, Array) and x.ndim: + values.append(_values_at(coords, x, shape) if len(c) else c) + else: + values.append(_value(x)) + V = np.where(np.asarray(c, dtype=bool), *values).astype(rdt, copy=False) + gdt = gb_dtype(rdt) + R = Tensor.from_coo(coords, V, gdt, shape=shape) if len(c) else Tensor(gdt, shape) + if isinstance(base, Array) and base.ndim and base._T().nvals: + R = base._T().ewise_add(R, binary.second[gdt]).new(dtype=gdt) + return R + + +def _values_at(coords, x, shape): + """The values of ``x`` (an Array that broadcasts to ``shape``) at ``coords`` of + ``shape``: its stored value, or its fill. O(len(coords)), read from x's own entries. + """ + lead = len(shape) - x.ndim + n = len(coords[0]) + # (a broadcast axis of x has size 1: index 0 there) + index = tuple( + coords[lead + i] if size != 1 else np.zeros(n, np.int64) for i, size in enumerate(x.shape) + ) + return x._T().vindex[index].to_dense(fill_value=x._fill) + + def _check_array(x, name): if not isinstance(x, Array): raise TypeError(f"{name} needs an Array, not {type(x).__name__}") diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py index 578b2ccbf..2c5be1196 100644 --- a/graphblas/tests/test_array_api_searching.py +++ b/graphblas/tests/test_array_api_searching.py @@ -46,6 +46,42 @@ def test_where(shapes): xp.where(make(np.ones(3, bool), False), x, x) +def test_where_stores_the_condition_and_the_selected_operand(): + # Where the condition stores nothing, its fill selects one operand: the result is that + # operand with the condition's stored positions written over, and the other operand is + # read only there (no replication of a broadcast operand) + rng = np.random.default_rng(7) + n = 60 + c = rng.random((n, n)) < 0.05 + labels = np.arange(n) + big = np.iinfo(np.int64).max + G = make(c, False, rng) # (with explicit False entries) + W = xp.where(G, xp.arange(n), big) + assert_array_equal(W, np.where(c, labels, big)) + assert W.nnz == G.nnz + assert W.fill_value == big + # other layouts of the condition, and a fill of True (the base is then x1) + for cond, cn in [(xp.matrix_transpose(G), c.T), (xp.logical_not(G), ~c)]: + a = rng.normal(size=(n, n)) + a[rng.random((n, n)) < 0.5] = 0 + a[0, 0] = np.nan + for x1, x2, e in [ + (make(a, 0.0), xp.arange(n) * 1.0, np.where(cn, a, labels * 1.0)), + (xp.arange(n) * 1.0, make(a, 0.0, rng), np.where(cn, labels * 1.0, a)), + (xp.arange(n)[:, None] * 1.0, -1.0, np.where(cn, labels[:, None] * 1.0, -1.0)), + ]: + assert_array_equal(xp.where(cond, x1, x2), e) + # The selected operand broadcasts over stored entries, or the condition broadcasts: the + # union (the result really holds the operand's entries along the broadcast axis) + v = np.where(rng.random(n) < 0.5, rng.normal(size=n), 0.0) + assert_array_equal(xp.where(G, -1.0, make(v, 0.0)), np.where(c, -1.0, v)) + assert_array_equal(xp.where(make(c[0], False), make(a, 0.0), 2.0), np.where(c[0], a, 2.0)) + # An empty condition: the selected operand + E = xp.zeros((n, n), dtype=xp.bool) + assert_array_equal(xp.where(E, 1.0, make(a, 0.0)), a) + assert xp.where(E, 1.0, make(a, 0.0)).nnz == make(a, 0.0).nnz + + def test_where_dtypes(): for dt in DTYPES: for dense, x in inputs(dt, (5,), seed=6): From 9c81ea41d83413761c1325f8f66e35d0b1bcec36 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:00:16 -0600 Subject: [PATCH 63/77] array_api: fills that absorb the op make add, maximum and minimum intersect inf + y is inf for every y but -inf and NaN, so when both operands' fills are inf the result can differ from inf only where both store an entry. Add (inf or -inf), maximum (inf, or an integer's maximum, or True), minimum (the opposites) and logical_or (True) now take the intersection when both fills absorb the op, which never replicates a broadcast operand. An operand that stores every position counts as holding that fill. Floats need no NaN in either operand and, for add, no opposite infinity: new cached facts "posinf" and "neginf" beside "finite", "nan" and "zero". SAPS's Bellman-Ford step, expand_dims(D, 1) + G with inf fills, wrote D across every column of G. --- graphblas/array_api/_elementwise.py | 46 ++++++++++++++ graphblas/array_api/_utils.py | 16 +++-- graphblas/tests/test_array_api_elementwise.py | 60 +++++++++++++++++++ graphblas/tests/test_array_api_errors.py | 6 ++ 4 files changed, 122 insertions(+), 6 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index c2a6424e0..8cabaef07 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -572,6 +572,12 @@ def _binary_arrays(spec, a, b, op, ct, rdt, fill): R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) with np.errstate(all="ignore"): return _result(R, spec.ufunc(fa_seen, fb_seen)) + if (e := _absorbing(spec, a, b, fa, fb, ct)) is not None: + # Each operand's fill absorbs the op: where either holds its fill the result is + # that fill, so the intersection, which never replicates a broadcast operand + _utils.mult_expansion_check(Ta, Tb, shape) + R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) + return _result(R, _utils.scalar(e, rdt)) if spec.name in {"multiply", "divide"} and ct.kind == "f": out = _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape) if out is not None: @@ -593,6 +599,46 @@ def _binary_arrays(spec, a, b, op, ct, rdt, fill): return _result(_utils.map_values(spec.ufunc, [A, B], rdt), fill) +def _absorbing_values(name, ct): + """The values ``e`` with ``op(e, y) == op(y, e) == e`` for the binary ``name`` working in + ``ct``, for every y but those ``_absorbing`` excludes (NaN, the opposite infinity). + """ + if name in {"add", "maximum", "logical_or"} and ct.kind == "b": + return (True,) + if name == "minimum" and ct.kind == "b": + return (False,) + if ct.kind == "f": + return {"add": (np.inf, -np.inf), "maximum": (np.inf,), "minimum": (-np.inf,)}.get(name, ()) + if ct.kind in "iu": + info = np.iinfo(ct) + return {"maximum": (info.max,), "minimum": (info.min,)}.get(name, ()) + return () + + +def _absorbing(spec, a, b, fa, fb, ct): + """The value both operands' fills take when it absorbs the op (inf for add and maximum, + -inf for add and minimum, True for logical_or, the integer extremes), so that the result + is the intersection of their entries with that fill; else None. An operand that stores + every position never shows its fill, so it counts as having that value, but not both. + Floats: no operand may hold NaN (numpy propagates it) or, for add, the opposite infinity + (inf + -inf is NaN): cached facts. + """ + full_a, full_b = _stores_all(a), _stores_all(b) + if full_a and full_b: + return None # (the union is the intersection already) + for e in _absorbing_values(spec.name, ct): + e = _utils.scalar(e, ct) + if not (full_a or _utils.same_value(fa, e)) or not (full_b or _utils.same_value(fb, e)): + continue + if ct.kind == "f": + opposite = "neginf" if e > 0 else "posinf" + for x in (a, b): + if _utils.fact(x, "nan") or (spec.name == "add" and _utils.fact(x, opposite)): + return None + return e + return None + + def _power(x1, x2): """The power by numpy, but where numpy takes a square root for an exponent of 0.5 that repeats (a scalar, or an exponent of stride 0): the Array API's pow(-inf, 0.5) is +inf and diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index 522d550d7..ff585ef05 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -179,17 +179,21 @@ def with_fill(x, fill): # pass and cached on the Array; ``fact`` documents them. For each: the unary op that says # whether a value has the property, the monoid that combines those answers, and the answer when # the Array stores nothing. +# (A number in place of the unary op: whether a value equals it.) _FACT_SCANS = { "finite": (unary.isfinite, monoid.land, True), "nan": (unary.isnan, monoid.lor, False), - "zero": (None, monoid.lor, False), + "zero": (0, monoid.lor, False), + "posinf": (np.inf, monoid.lor, False), + "neginf": (-np.inf, monoid.lor, False), } def fact(x, name): """A fact about the stored values of the Array ``x``, computed once (cached on ``x``): - "finite" (every value finite), "nan" (some value is NaN) or "zero" (some value is 0 or - -0.0). Integers and bools are finite and not NaN without a look; each other fact is one + "finite" (every value finite), "nan" (some value is NaN), "zero" (some value is 0 or + -0.0), "posinf" or "neginf" (some value is inf, or -inf). Integers and bools are finite, + with no NaN or infinity, without a look; each other fact is one parallel GraphBLAS pass with no copy of the values (an apply to bool, then a reduction that stops at the first deciding value). """ @@ -202,7 +206,7 @@ def _compute_fact(x, name): facts = x._facts if name != "zero" and x.dtype.kind not in "fc": return name == "finite" - if name == "nan" and facts.get("finite"): + if name in {"nan", "posinf", "neginf"} and facts.get("finite"): return False if name == "finite" and facts.get("nan"): return False @@ -215,8 +219,8 @@ def tensor_fact(T, name): if T.nvals == 0: return empty dt = T.dtype - if op is None: - B = T.apply(binary.eq[dt], right=gb_scalar(0, np_dtype(dt))).new() + if not callable(op): + B = T.apply(binary.eq[dt], right=gb_scalar(op, np_dtype(dt))).new() else: B = T.apply(op[dt]).new() return bool(B.reduce_scalar(combine[lookup_dtype(bool)]).new().value) diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 93ac21117..19d08080b 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -721,6 +721,66 @@ def test_intersection_of_broadcast_operands_is_guarded(): assert func(*args).nnz == nnz +def test_absorbing_fills_intersect(): + # A fill that absorbs the op (inf + y is inf) on both operands: the result is the + # intersection of their entries, and a broadcast operand is not replicated + rng = np.random.default_rng(31) + n = 40 + g = np.where(rng.random((n, n)) < 0.1, rng.random((n, n)) + 1, np.inf) + d = np.where(rng.random(n) < 0.5, rng.random(n), np.inf) + G, D = make(g, np.inf, rng), make(d, np.inf, rng) + for func, nf in [(xp.add, np.add), (xp.maximum, np.maximum)]: + R = func(xp.expand_dims(D, 1), G) + assert_array_equal(R, nf(d[:, None], g)) + assert R.nnz <= G.nnz + assert R.fill_value == np.inf + R = xp.minimum(make(-g, -np.inf), make(-d[:, None], -np.inf)) + assert_array_equal(R, np.minimum(-g, -d[:, None])) + assert R.nnz <= G.nnz + R = xp.add(make(-g, -np.inf), make(-d, -np.inf)) # (-inf absorbs add too) + assert_array_equal(R, -g + -d) + assert R.nnz <= G.nnz + # integers' extremes, and bools + big = np.iinfo(np.int64).max + gi = np.where(np.isinf(g), big, (np.where(np.isinf(g), 0, g) * 10).astype(np.int64)) + di = np.where(np.isinf(d), big, (np.where(np.isinf(d), 0, d) * 10).astype(np.int64)) + R = xp.maximum(make(gi, big), make(di, big)) + assert_array_equal(R, np.maximum(gi, di)) + assert R.nnz <= G.nnz + small = np.iinfo(np.int8).min + R = xp.minimum(make(np.int8(-3) * (gi % 2), 0), make(np.full((n, n), small, np.int8), small)) + assert_array_equal(R, np.minimum(np.int8(-3) * (gi % 2), np.int8(small))) + gb_, db_ = np.isinf(g), np.isinf(d) + for func, nf in [(xp.logical_or, np.logical_or), (xp.add, np.add), (xp.maximum, np.maximum)]: + R = func(make(gb_, True), make(db_, True)) + assert_array_equal(R, nf(gb_, db_)) + assert R.nnz <= (~gb_).sum() + R = xp.minimum(make(~gb_, False), make(~db_, False)) + assert_array_equal(R, np.minimum(~gb_, ~db_)) + # An operand that stores every position counts as holding the absorbing fill + full = rng.random(n) + 1 + R = xp.expand_dims(make(full, 0.0), 1) + G + assert_array_equal(R, full[:, None] + g) + assert R.nnz <= G.nnz + assert R.fill_value == np.inf + # NaN, or for add the opposite infinity, anywhere: the union, exact + for bad in [np.nan, -np.inf]: + d2 = d.copy() + d2[0] = bad + R = xp.add(xp.expand_dims(make(d2, np.inf), 1), G) + with np.errstate(invalid="ignore"): + assert_array_equal(R, d2[:, None] + g) + g2 = g.copy() + g2[0, 0] = -np.inf + assert_array_equal(xp.add(make(d, np.inf), make(g2, np.inf)), d + g2) + # Fills that do not both absorb, and two fully stored operands: the union as before + assert_array_equal(xp.add(make(d, np.inf), make(d, 0.0)), d + d) + R = xp.add(make(full, 0.0), make(full, 1.0)) + assert_array_equal(R, full + full) + with np.errstate(invalid="ignore"): + assert_array_equal(xp.subtract(make(d, np.inf), make(d, np.inf)), d - d) + + def test_fully_stored_bool_results_take_fill_false(): # A bool result that stores every position never shows its fill: it takes fill False, # and its False entries go when at most a quarter of the entries are True (equal(deg, 0) diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 11a097961..8fbfd63a0 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -213,6 +213,12 @@ def test_facts_of_stored_values(): assert _utils.fact(i, "zero") e = xp.zeros(4) assert [_utils.fact(e, k) for k in ["finite", "nan", "zero"]] == [True, False, False] + # infinities by sign; known absent from finiteness, and for integers + u = make(np.array([np.inf, 2.0]), 0.0) + assert [_utils.fact(u, k) for k in ["posinf", "neginf"]] == [True, False] + assert _utils.fact(make(np.array([-np.inf]), 0.0), "neginf") + assert not _utils.fact(w, "posinf") + assert not _utils.fact(i, "neginf") # abs, negative, positive and conj keep all three; roundings keep two v = make(np.array([-1.5, np.nan, 0.0]), 0.0) for k in ["finite", "nan", "zero"]: From ca4f1c82f5365d127036ea627890cf842d4257c9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:00:19 -0600 Subject: [PATCH 64/77] array_api: a semiring accepts an operand that stores every position Such an operand never shows its fill, so matmul, tensordot, vecdot and einsum with a semiring now relabel its fill as the monoid's identity (O(1), and no value can tell the fills apart) instead of refusing it, when its dtype holds the identity (an integer cannot hold inf). A dense distance matrix with no zero, given as an ndarray, now runs min_plus. --- graphblas/array_api/_linalg.py | 34 +++++++++++++++++----- graphblas/tests/test_array_api_semiring.py | 29 ++++++++++++++++++ 2 files changed, 55 insertions(+), 8 deletions(-) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 5a59401d0..a06fbcb0f 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -155,9 +155,10 @@ def _spell(value): def _sound(typed, arrays, name, *, multiplies=True): """The result dtype and fill of a contraction of ``arrays`` with a typed semiring other - than the default, after checking that it may skip their unstored positions (ValueError - if not). ``multiplies``: whether values meet the multiply (einsum of one operand only - sums). + than the default, and the arrays, after checking that it may skip their unstored + positions (ValueError if not). An array that stores every position never shows its fill, + so it is relabeled with the identity when its dtype holds it (O(1)). ``multiplies``: + whether values meet the multiply (einsum of one operand only sums). """ e = typed.monoid.identity out_dt = check_supported(np_dtype(typed.return_type), name) @@ -169,9 +170,13 @@ def _sound(typed, arrays, name, *, multiplies=True): f"annihilate {mult} ({mult}(e, y) and {mult}(y, e) must be e for every y), so no " "fill works. Use it on Tensors (x.tensor), which skip missing entries." ) + out = [] for i, x in enumerate(arrays): in_dt = np_dtype(typed.type if i == 0 else typed.type2) fill = _utils.scalar(x._fill, in_dt) + if fill != e and (relabeled := _relabel_full(x, e)) is not None: + x, fill = relabeled, e + out.append(x) if fill != e: # (a NaN fill too) spelled = _spell(e) raise ValueError( @@ -182,7 +187,20 @@ def _sound(typed, arrays, name, *, multiplies=True): f"with fill_value={spelled} (xp.asarray(x.tensor, fill_value={spelled}) reads " f"x's unstored positions as {e!r})." ) - return out_dt, _utils.scalar(e, out_dt) + return out_dt, _utils.scalar(e, out_dt), out + + +def _relabel_full(x, e): + """``x`` with fill ``e`` when it stores every position (its fill is never seen) and its + dtype holds ``e`` exactly; else None. + """ + if x._T().nvals != x.size: + return None + with np.errstate(all="ignore"): + fill = np.asarray(e).astype(x.dtype)[()] # (inf as an integer: no longer inf) + if not _utils.same_value(np.asarray(fill).astype(np.asarray(e).dtype)[()], e): + return None + return Array._new(x._T(), fill) def _is_zero(fill): @@ -255,7 +273,7 @@ def matmul(x1, x2, /, *, semiring=None): rdt = _rdt("matmul", np.matmul, x1, x2) typed, default = _resolve(semiring, rdt, "matmul") if not default: - out_dt, fill = _sound(typed, (x1, x2), "matmul") + out_dt, fill, (x1, x2) = _sound(typed, (x1, x2), "matmul") return _product(x1, x2, typed, out_dt, fill) if (ops := _zero_fills(x1, x2)) is None: return _dense(np.matmul, _matmul_shape(x1, x2), x1, x2) @@ -365,7 +383,7 @@ def tensordot(x1, x2, /, *, axes=2, semiring=None): free2 = tuple(n for i, n in enumerate(x2.shape) if i not in b_axes) shape = free1 + free2 if not default: - out_dt, fill = _sound(typed, (x1, x2), "tensordot") + out_dt, fill, (x1, x2) = _sound(typed, (x1, x2), "tensordot") return _tensordot(x1, x2, a_axes, b_axes, shape, typed, out_dt, fill) if (ops := _zero_fills(x1, x2)) is None: func = lambda a, b: np.tensordot(a, b, axes=(a_axes, b_axes)) # noqa: E731 @@ -414,7 +432,7 @@ def vecdot(x1, x2, /, *, axis=-1, semiring=None): typed, default = _resolve(semiring, rdt, "vecdot") out_shape = tuple(n for i, n in enumerate(shape) if i != len(shape) + axis) if not default: - out_dt, fill = _sound(typed, (x1, x2), "vecdot") + out_dt, fill, (x1, x2) = _sound(typed, (x1, x2), "vecdot") return _vecdot(x1, x2, axis, shape, out_shape, typed, out_dt, fill) if (ops := _zero_fills(x1, x2)) is None: return _dense(lambda a, b: np.vecdot(a, b, axis=axis), out_shape, x1, x2) @@ -485,7 +503,7 @@ def einsum(subscripts, /, *operands, semiring=None): ops = zeroed out_dt, fill = rdt, _utils.scalar(0, rdt) else: - out_dt, fill = _sound(typed, ops, "einsum", multiplies=len(ops) > 1) + out_dt, fill, ops = _sound(typed, ops, "einsum", multiplies=len(ops) > 1) if math.prod(shape) == 0 or any(x._T().nvals == 0 for x in ops): return Array._new(Tensor(gb_dtype(out_dt), shape), fill) diff --git a/graphblas/tests/test_array_api_semiring.py b/graphblas/tests/test_array_api_semiring.py index 872e00dba..42dbe7976 100644 --- a/graphblas/tests/test_array_api_semiring.py +++ b/graphblas/tests/test_array_api_semiring.py @@ -367,6 +367,35 @@ def test_nan_propagates_through_min_and_max(): np.array(1.0)) # fmt: skip +def test_fully_stored_operands_take_the_identity(): + # An operand that stores every position never shows its fill: a semiring accepts it + # whatever its fill, when its dtype holds the identity + a = np.array([[1.0, 2.0], [np.inf, 3.0]]) + b = np.array([[0.5, 1.0], [2.0, 4.0]]) + A, B = make(a, 0.0), make(b, 0.0) + assert A.nnz == A.size + expected = np.min(a[:, :, None] + b[None], axis=1) + for R in [ + xp.matmul(A, B, semiring="min_plus"), + xp.tensordot(A, B, axes=1, semiring="min_plus"), + xp.einsum("ij,jk->ik", A, B, semiring="min_plus"), + xp.vecdot(A[:, :, None], B[None], axis=1, semiring="min_plus"), + ]: + assert_close(R, expected, np.inf) + # (the inf entry is an explicit fill once relabeled; a NaN elsewhere still propagates) + a2 = a.copy() + a2[0, 0] = np.nan + R = xp.matmul(make(a2, 0.0), B, semiring="min_plus") + assert np.isnan(R.todense()[0]).all() + # An integer operand cannot hold inf: still refused + i = make(np.array([[1, 2], [3, 4]]), 0) + with pytest.raises(ValueError, match="needs operands whose fill"): + xp.matmul(i, make(b, np.inf), semiring="min_plus") + # Partly stored: refused as before + with pytest.raises(ValueError, match="needs operands whose fill"): + xp.matmul(make(np.array([[0.0, 1.0]]), 0.0), B, semiring="min_plus") + + def test_soundness_is_cached_per_semiring_object(): # (review finding L1) two semirings of one name: each is checked from graphblas.core.operator import Semiring From 73aa5af2aa260c8d3eba7b76765cea3e14bf087a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:03:28 -0600 Subject: [PATCH 65/77] array_api: xp.refill changes an Array's fill and keeps its values refill(x, fill_value) (an extension) stores the positions that held the old fill (O(size), under the densify limit; O(1) when x stores every position) and drops the entries equal to the new fill. A dense distance matrix with inf for "no edge" and zeros on the diagonal becomes fill inf with its zeros stored, which a min_plus contraction accepts. The semiring error's hint, which suggested asarray(x.tensor, fill_value=...) (that turns the positions holding the old fill into the new one), and asarray's refusal to change an Array's fill now point at refill. --- graphblas/array_api/_creation.py | 5 ++- graphblas/array_api/_linalg.py | 4 +- graphblas/array_api/_manipulation.py | 27 ++++++++++++ .../tests/test_array_api_manipulation.py | 42 +++++++++++++++++++ 4 files changed, 75 insertions(+), 3 deletions(-) diff --git a/graphblas/array_api/_creation.py b/graphblas/array_api/_creation.py index 25144966e..7ed5e2e77 100644 --- a/graphblas/array_api/_creation.py +++ b/graphblas/array_api/_creation.py @@ -137,7 +137,10 @@ def asarray(obj, /, *, dtype=None, device=None, copy=None, fill_value=None): check_device(device) if isinstance(obj, Array): if fill_value is not None and not _utils.same_value(fill_value, obj._fill): - raise ValueError("asarray of an Array keeps its fill value; use full_like to refill") + raise ValueError( + "asarray of an Array keeps its fill value; xp.refill(x, fill_value) gives one " + "with another (the same values)" + ) if dtype is not None and as_dtype(dtype) != obj.dtype: if copy is False: raise ValueError("changing the data type needs a copy, and copy=False") diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index a06fbcb0f..f9bfe597d 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -184,8 +184,8 @@ def _sound(typed, arrays, name, *, multiplies=True): f"the identity of {typed.monoid.name}" + (f", which annihilates {typed.binaryop.name}" if multiplies else "") + f"; {_operand_name(name, i)} has the fill {x._fill.item()!r}. Make the array " - f"with fill_value={spelled} (xp.asarray(x.tensor, fill_value={spelled}) reads " - f"x's unstored positions as {e!r})." + f"with fill_value={spelled}, or use xp.refill(x, {spelled}), which keeps its " + "values (it stores the positions that held the old fill)." ) return out_dt, _utils.scalar(e, out_dt), out diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py index dc39b62e3..880ba409b 100644 --- a/graphblas/array_api/_manipulation.py +++ b/graphblas/array_api/_manipulation.py @@ -20,6 +20,7 @@ "moveaxis", "permute_dims", "prune", + "refill", "repeat", "reshape", "roll", @@ -343,6 +344,32 @@ def tile(x, repetitions, /): return _from_entries(out, np.tile(values, ntiles), shape, x) +def refill(x, /, fill_value): + """``x`` with the fill value ``fill_value`` (an extension): the same value at every + position. Positions ``x`` does not store held its old fill, so they are stored now + (O(size), under the densify limit; O(1) when ``x`` stores every position), and stored + entries equal to ``fill_value`` go (``prune``). A 0-d array's value is its fill, so it + is returned as it is. ``fill_value`` must be a value of ``x``'s dtype. + """ + if not isinstance(x, Array): + raise TypeError(f"refill needs an Array, not {type(x).__name__}") + with np.errstate(all="ignore"): + fill = _utils.scalar(fill_value, x.dtype) + back = np.asarray(fill).astype(np.result_type(np.asarray(fill_value), x.dtype))[()] + if not _utils.same_value(back, np.asarray(fill_value)[()]): + raise ValueError(f"refill: {fill_value!r} is not a value of {x.dtype}") + if x.ndim == 0: + return x._copy() + T = x._T() + if _utils.same_value(fill, x._fill): + y = x + elif T.nvals == x.size: + y = Array._new(T, fill) + else: + y = _utils.with_fill(x, fill) + return prune(y) + + def prune(x, /): """``x`` without the stored entries that equal its fill value (an extension): the same values, fewer stored entries. "Equal" is "identical": with fill 0.0 a stored -0.0 stays, diff --git a/graphblas/tests/test_array_api_manipulation.py b/graphblas/tests/test_array_api_manipulation.py index fecdce538..61b3dc9d3 100644 --- a/graphblas/tests/test_array_api_manipulation.py +++ b/graphblas/tests/test_array_api_manipulation.py @@ -3,6 +3,7 @@ import numpy as np import pytest +import graphblas as gb import graphblas.array_api as xp from .test_array_api_elementwise import assert_array_equal, inputs, make, same_value @@ -150,6 +151,47 @@ def test_tile(reps): xp.tile(xp.zeros(2), (-1,)) +def test_refill(): + rng = np.random.default_rng(12) + a = np.where(rng.random((6, 7)) < 0.3, rng.normal(size=(6, 7)), 0.0) + a[0, 0] = np.inf + x = make(a, 0.0, rng) + for fill in [np.inf, 1.5, 0.0, np.nan, -0.0]: + y = xp.refill(x, fill) + assert_array_equal(y, a) + assert same_value(y.fill_value, np.float64(fill)) + expected = sum(not same_value(v, y.fill_value) for v in a.reshape(-1)) + assert y.nnz == expected + # a dense distance matrix: inf becomes the fill, the zeros stay (the hint of the + # semiring error) + g = np.full((4, 4), np.inf) + np.fill_diagonal(g, 0.0) + g[0, 1] = 1.0 + G = xp.refill(xp.asarray(g), np.inf) + assert (G.nnz, G.fill_value) == (5, np.inf) + assert_array_equal(xp.matmul(G, G, semiring="min_plus"), np.min(g[:, :, None] + g[None], 1)) + # stores every position: relabeled, then pruned; other dtypes; a 0-d array keeps its value + full = xp.asarray(np.array([1, 2, 2, 3])) + out = xp.refill(full, 2) + assert_array_equal(out, np.array([1, 2, 2, 3])) + assert (out.nnz, out.fill_value) == (2, 2) + b = xp.refill(make(np.array([True, False, False]), False), True) + assert_array_equal(b, np.array([True, False, False])) + assert (b.nnz, b.fill_value) == (2, True) + assert_array_equal(xp.refill(xp.asarray(5.0), 1.0), np.array(5.0)) + # errors: not an Array, a value the dtype does not hold, past the densify limit + with pytest.raises(TypeError, match="needs an Array"): + xp.refill(a, 1.0) + with pytest.raises(ValueError, match="not a value of int64"): + xp.refill(full, np.inf) + with pytest.raises(ValueError, match="not a value of int64"): + xp.refill(full, 1.5) + with gb.config.set(array_api_densify_limit=10), pytest.raises(xp.DensifyError): + xp.refill(x, 1.0) + with pytest.raises(ValueError, match="xp.refill"): + xp.asarray(x, fill_value=1.0) + + def test_prune(): cases = [ (np.array([3, 1, 3, 0]), 3), From 3a978224226fce78a2b3bc09810322f1d90c22c3 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:03:31 -0600 Subject: [PATCH 66/77] array_api: a bool result with fill True that stores mostly False flips to fill False equal(deg, 0) of a deg that leaves its few zero degrees unstored has fill True and stores a False at every other position. When those False entries are more than half of all positions, the result now takes fill False and stores only its True positions (its unstored positions and stored True entries): O(size), under twice its entries. In SAPS's PageRank on a graph with few dangling nodes, (1/N) * zero_deg then has fill 0 and M stays sparse (n=2000, 2 dangling: 19,965 entries instead of 3,996,000). A graph with mostly dangling nodes keeps fill True, which is already the smaller. --- graphblas/array_api/_elementwise.py | 23 ++++++++++++-- graphblas/tests/test_array_api_elementwise.py | 30 +++++++++++++++++-- 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 8cabaef07..205a7b5ce 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -161,11 +161,30 @@ def _fill_of(ufunc, rdt, operands): def _result(T, fill): - if T.dtype == BOOL and T.ndim and T.nvals == T.size and T.nvals: - T, fill = _normalize_bool(T, fill) + if T.dtype == BOOL and T.ndim and T.nvals: + if T.nvals == T.size: + T, fill = _normalize_bool(T, fill) + elif fill and 2 * T.nvals > T.size: + T, fill = _majority_false(T) return Array._new(T, fill) +def _majority_false(T): + """A bool Tensor with fill True that stores mostly False entries, more than half its + positions (``equal(deg, 0)`` of a ``deg`` that leaves a few zero degrees unstored): + the same values with fill False, storing only the True positions (its unstored ones and + its stored True entries). O(size), which is under twice its entries. Else as it is. + """ + ntrue = T.reduce_scalar(monoid.plus[INT64]).new().value + if 2 * (T.nvals - ntrue) <= T.size: + return T, True + R = Tensor(BOOL, T.shape) + R(~T.S)[...] << True + if ntrue: + R(T.select(select.valueeq, True).new().S)[...] << True + return R, False + + def _normalize_bool(T, fill): """A bool Tensor that stores every position never shows its fill, so it takes fill False (the fill contractions and masks work with; a fill True made a mostly-True mask times a diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py index 19d08080b..12cf698ce 100644 --- a/graphblas/tests/test_array_api_elementwise.py +++ b/graphblas/tests/test_array_api_elementwise.py @@ -799,8 +799,7 @@ def test_fully_stored_bool_results_take_fill_false(): ]: assert_array_equal(out, expected) assert (out.nnz, out.fill_value) == (nnz, fill) - # a 0-d result and a result with unstored positions are left as they are - assert xp.equal(deg, 0).nnz == 8 + # a 0-d result is left as it is assert_array_equal(xp.equal(xp.asarray(3), 3), np.array(True)) @@ -837,6 +836,33 @@ def test_one_sided_stores_a_small_broadcast_operand(): assert_array_equal(out, a * np.exp(y)[:, None]) +def test_partly_stored_bool_results_with_mostly_false_take_fill_false(): + # equal(deg, 0) of a deg that leaves its zero degrees unstored has fill True and stores + # a False at every other position; when those are more than half of all positions, it + # takes fill False and stores the True positions only + deg = xp.asarray(np.array([3, 1, 0, 2, 5, 0, 4, 1, 2, 7])) + out = xp.equal(deg, 0) + assert_array_equal(out, deg.todense() == 0) + assert (out.nnz, out.fill_value) == (2, False) + # with stored True entries too (explicit fills of a fill-True operand) + x = make(np.array([True, False, False, False, True, False]), True, np.random.default_rng(3)) + out = xp.logical_and(x, x) + assert_array_equal(out, x.todense()) + assert out.fill_value == np.False_ or out.nnz * 2 <= out.size + for dense in [np.eye(6, dtype=bool), np.array([[True, False], [False, False]])]: + out = xp.logical_not(xp.logical_not(make(dense, False))) + assert_array_equal(out, dense) + assert out.fill_value == np.False_ + assert out.nnz == dense.sum() + # mostly True, or not more than half stored: left with fill True + sparse_deg = xp.asarray(np.array([0, 0, 0, 2, 0, 0])) + assert (xp.equal(sparse_deg, 0).nnz, xp.equal(sparse_deg, 0).fill_value) == (1, True) + y = make( + np.array([True, True, True, False, True]), True, np.random.default_rng(4), explicit=0.9 + ) + assert xp.logical_and(y, y).fill_value == np.True_ or y.nnz * 2 <= y.size + + def test_mostly_true_masks_stay_sparse_in_contractions(): # (review finding H3) a mostly-True fully stored mask keeps fill False, so a contraction # with a sparse bool matrix stays sparse (it raised DensifyError with fill True) From e0be977a70cdb277c7915d2938190cdaea93083f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:05:59 -0600 Subject: [PATCH 67/77] array_api: xp.compute computes a pending Array now matmul, tensordot, vecdot and einsum return an Array that computes on first use. compute(x) (an extension) computes it at once and returns x, for timing and for frameworks with a compute step (SAPS's Framework.compute, which had only the private x._T()). Any other Array is returned as it is. --- graphblas/array_api/_linalg.py | 13 ++++++++++++- graphblas/tests/test_array_api_lazy.py | 13 +++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index f9bfe597d..490e5d014 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -39,7 +39,7 @@ from ._dtypes import check_supported from ._utils import gb_dtype, np_dtype -__all__ = ["einsum", "matmul", "matrix_transpose", "tensordot", "vecdot"] +__all__ = ["compute", "einsum", "matmul", "matrix_transpose", "tensordot", "vecdot"] def _check(x, name): @@ -56,6 +56,17 @@ def _operand(x): return x +def compute(x, /): + """``x``, computed now (an extension). matmul, tensordot, vecdot and einsum return an Array + that computes its value on first use, or only where a consumer needs it (DESIGN D9); + ``compute`` computes it at once (to time it, or before handing it on). Any other Array + is returned as it is. + """ + _check(x, "compute") + x._T() + return x + + def _semiring(dt): return (semirings.lor_land if dt.kind == "b" else semirings.plus_times)[gb_dtype(dt)] diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index a994d6516..0326c0222 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -648,3 +648,16 @@ def test_einsum_three_operands_not_masked(masks): assert_close(xp.einsum("ij,jk,ik->", A, B, S), np.einsum("ij,jk,ik->", a, b, s)) finally: _lazy.enabled = True + + +def test_compute(): + a, b = _rand((4, 5)), _rand((5, 3)) + A, B = make(a, 0.0), make(b, 0.0) + P = A @ B + assert pending(P) + assert xp.compute(P) is P + assert not pending(P) + assert_close(P, a @ b) + assert xp.compute(A) is A + with pytest.raises(TypeError, match="needs Arrays"): + xp.compute(a) From aa1a776a54d925e21f8da72ff3c092e4aa839c8d Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:06:02 -0600 Subject: [PATCH 68/77] array_api: from_coo sums the values at a repeated coordinate from_coo raised "Duplicate indices found, must provide dup_op" for a repeated coordinate; pydata/sparse and scipy.sparse sum such values. It now sums them (logical or for bool; integers wrap as numpy's sum does). One value for every coordinate stays stored once (iso) and is built again from the repeated value only when coordinates repeat. --- graphblas/array_api/_creation.py | 17 +++++++++++++++-- graphblas/tests/test_array_api_creation.py | 12 ++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/graphblas/array_api/_creation.py b/graphblas/array_api/_creation.py index 7ed5e2e77..2c467ad51 100644 --- a/graphblas/array_api/_creation.py +++ b/graphblas/array_api/_creation.py @@ -11,6 +11,7 @@ import numpy as np +from .. import binary from ..core.tensor import Tensor from . import _utils from ._array import Array, check_device @@ -208,7 +209,9 @@ def _unwrap_nested(obj): def from_coo(coords, data, shape=None, *, fill_value=None, dtype=None): """An Array storing ``data`` at ``coords`` (one index array per axis, or an (ndim, nnz) - array), and ``fill_value`` (0 by default) elsewhere. ``data`` may be a scalar. + array), and ``fill_value`` (0 by default) elsewhere. ``data`` may be a scalar. Values at a + repeated coordinate are summed (logical or for bool), as pydata/sparse and scipy.sparse + do. """ data = np.asarray(data) dt = as_dtype(dtype) if dtype is not None else check_supported(data.dtype, "from_coo data") @@ -225,7 +228,17 @@ def from_coo(coords, data, shape=None, *, fill_value=None, dtype=None): value = data.reshape(-1)[0] if data.size else fill return Array._new(Tensor(gb_dtype(dt), ()), value) if coords and coords[0].size: - T = Tensor.from_coo(coords, data, gb_dtype(dt), shape=shape) + gdt = gb_dtype(dt) + plus = binary.plus[gdt] # (lor for bool) + if data.ndim: + T = Tensor.from_coo(coords, data, gdt, shape=shape, dup_op=plus) + else: + # One value everywhere (iso, which keeps no duplicates); when coordinates + # repeat, each sums its copies: built again from the value repeated + T = Tensor.from_coo(coords, data, gdt, shape=shape) + if T.nvals < coords[0].size and dt.kind != "b": + data = np.broadcast_to(data, coords[0].shape) + T = Tensor.from_coo(coords, data, gdt, shape=shape, dup_op=plus) else: T = Tensor(gb_dtype(dt), shape) return Array._new(T, fill) diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py index c1708bf4b..48778f4eb 100644 --- a/graphblas/tests/test_array_api_creation.py +++ b/graphblas/tests/test_array_api_creation.py @@ -147,6 +147,18 @@ def test_asarray_scipy_and_pydata_sparse(): def test_from_coo(): x = xp.from_coo(([0, 2], [1, 0]), [1.0, 2.0], (3, 2), fill_value=5.0) assert_array_equal(x, np.array([[5.0, 1.0], [5.0, 5.0], [2.0, 5.0]])) + # a repeated coordinate sums its values (or for bool), as pydata/sparse and scipy do + x = xp.from_coo(([0, 0, 1], [1, 1, 0]), [1.0, 2.0, 5.0], (2, 2)) + assert_array_equal(x, np.array([[0.0, 3.0], [5.0, 0.0]])) + x = xp.from_coo(([0, 0],), [True, False], (3,)) + assert_array_equal(x, np.array([True, False, False])) + x = xp.from_coo(([1, 1],), np.array([100, 100], np.int8), (2,)) + assert_array_equal(x, np.array([0, -56], np.int8)) # (wraps, as numpy's int8 sum) + # one value for every coordinate: stored once (iso), summed where coordinates repeat + x = xp.from_coo(([0, 2, 0],), 7, (3,)) + assert_array_equal(x, np.array([14, 0, 7])) + assert_array_equal(xp.from_coo(([0, 0],), True, (2,)), np.array([True, False])) + assert_array_equal(xp.from_coo(([1, 0],), 2.5, (2,)), np.array([2.5, 2.5])) x = xp.from_coo(np.array([[0, 2], [1, 0]]), 7, (3, 2)) assert x.dtype == np.dtype(int) assert_array_equal(x, np.array([[0, 7], [0, 0], [7, 0]])) From e244ac012f679f54b06a3eba992e3674971b91d3 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:11:24 -0600 Subject: [PATCH 69/77] array_api: 0-d Arrays share one empty Tensor per dtype, copy-on-write A 0-d Array stores nothing (its value is its fill), but each one made a new empty Tensor: a 1 x 1 Matrix plus a per-Matrix format setting, about 6 us, half of a 0-d multiply. They now take an O(1) copy-on-write copy of one empty Tensor per dtype (_utils.empty_0d), which a change to any copy does not reach. A 0-d multiply went from 11.9 to 6.3 us; SAPS's time-stepping loops make one per operation. --- graphblas/array_api/_array.py | 5 ++--- graphblas/array_api/_creation.py | 6 +++--- graphblas/array_api/_elementwise.py | 2 +- graphblas/array_api/_manipulation.py | 2 +- graphblas/array_api/_searching.py | 2 +- graphblas/array_api/_set.py | 4 +--- graphblas/array_api/_statistical.py | 8 ++++---- graphblas/array_api/_utils.py | 20 ++++++++++++++++++-- graphblas/tests/test_array_api_errors.py | 16 ++++++++++++++++ 9 files changed, 47 insertions(+), 18 deletions(-) diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py index 72ec43a49..8f165ac69 100644 --- a/graphblas/array_api/_array.py +++ b/graphblas/array_api/_array.py @@ -4,7 +4,6 @@ import numpy as np -from ..core.tensor import Tensor from . import _utils from ._utils import np_dtype @@ -47,7 +46,7 @@ def _new(cls, tensor, fill): """ if tensor.ndim == 0 and tensor.nvals: fill = tensor.to_scalar().value - tensor = Tensor(tensor.dtype, ()) + tensor = _utils.empty_0d(tensor.dtype) self = object.__new__(cls) self._tensor = tensor self._pending = None @@ -338,7 +337,7 @@ def __invert__(self): def _unpickle(coords, values, shape, dtype, fill): if not shape: - return Array._new(Tensor(_utils.gb_dtype(dtype), ()), fill) + return Array._new(_utils.empty_0d(dtype), fill) from ._creation import from_coo return from_coo(coords, values, shape, fill_value=fill, dtype=np.dtype(dtype)) diff --git a/graphblas/array_api/_creation.py b/graphblas/array_api/_creation.py index 2c467ad51..a4a32a9e3 100644 --- a/graphblas/array_api/_creation.py +++ b/graphblas/array_api/_creation.py @@ -59,7 +59,7 @@ def _from_numpy(a, fill_value=None): dt = check_supported(a.dtype, "an array") fill = _default_fill(dt) if fill_value is None else _utils.scalar(fill_value, dt) if a.ndim == 0: - return Array._new(Tensor(gb_dtype(dt), ()), a[()]) + return Array._new(_utils.empty_0d(dt), a[()]) if a.size == 0: return Array._new(Tensor(gb_dtype(dt), a.shape), fill) keep = _utils.differs(a, fill) @@ -79,7 +79,7 @@ def _from_graphblas(obj, dtype, fill_value): if isinstance(obj, Scalar): dt = _utils.np_dtype(obj.dtype) value = obj.value if obj.nvals else fill_value if fill_value is not None else 0 - out = Array._new(Tensor(gb_dtype(dt), ()), value) + out = Array._new(_utils.empty_0d(dt), value) else: if isinstance(obj, Tensor): T = obj.dup() @@ -226,7 +226,7 @@ def from_coo(coords, data, shape=None, *, fill_value=None, dtype=None): fill = _default_fill(dt) if fill_value is None else _utils.scalar(fill_value, dt) if len(shape) == 0: value = data.reshape(-1)[0] if data.size else fill - return Array._new(Tensor(gb_dtype(dt), ()), value) + return Array._new(_utils.empty_0d(dt), value) if coords and coords[0].size: gdt = gb_dtype(dt) plus = binary.plus[gdt] # (lor for bool) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py index 205a7b5ce..b52fd44ef 100644 --- a/graphblas/array_api/_elementwise.py +++ b/graphblas/array_api/_elementwise.py @@ -211,7 +211,7 @@ def _normalize_bool(T, fill): def _zero_d(dtype, fill): - return Array._new(Tensor(gb_dtype(dtype), ()), fill) + return Array._new(_utils.empty_0d(dtype), fill) # Unary functions diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py index 880ba409b..db163c21c 100644 --- a/graphblas/array_api/_manipulation.py +++ b/graphblas/array_api/_manipulation.py @@ -71,7 +71,7 @@ def reshape(x, /, shape, *, copy=None): return Array._new(Tensor(T.dtype, shape), x._fill) if shape == (): # (one position: its entry, as a 0-d Array's value) - return Array._new(Tensor(T.dtype, ()), T.to_coo()[1][0]) + return Array._new(_utils.empty_0d(T.dtype), T.to_coo()[1][0]) if shape == T.shape: return x._copy() if copy else Array._new(T.dup(), x._fill) return Array._new(T.reshape(shape, copy=copy), x._fill) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index 8d7f69f56..656ad5065 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -201,7 +201,7 @@ def _arg(name, x, axis, keepdims): ) # fmt: skip out = Array._new(R, 0) else: - out = Array._new(Tensor(gb_dtype(np.int64), ()), result[0]) + out = Array._new(_utils.empty_0d(np.int64), result[0]) return expand_dims(out, axis) if keepdims else out diff --git a/graphblas/array_api/_set.py b/graphblas/array_api/_set.py index 9e2d27faa..ef95a37b2 100644 --- a/graphblas/array_api/_set.py +++ b/graphblas/array_api/_set.py @@ -84,9 +84,7 @@ def _unique(x, *, index=False, inverse=False, counts=False): if k: uinverse = np.delete(uinverse, at) if x.ndim == 0: - inv = Array._new( - Tensor(gb_dtype(np.int64), ()), uinverse[0] if len(uinverse) else inv_fill - ) + inv = Array._new(_utils.empty_0d(np.int64), uinverse[0] if len(uinverse) else inv_fill) elif T.nvals: R = Tensor.from_coo( coords, uinverse.astype(np.int64), gb_dtype(np.int64), shape=x.shape diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py index 3776c7957..111fe4593 100644 --- a/graphblas/array_api/_statistical.py +++ b/graphblas/array_api/_statistical.py @@ -229,7 +229,7 @@ def _reduce_main(name, x, axis, keepdims, dtype): kwargs = {"dtype": ct} if name in {"sum", "prod"} else {} with np.errstate(all="ignore"): value = spec.npfunc(np.asarray(x._fill), **kwargs) - return Array._new(Tensor(gb_dtype(ct), ()), value) + return Array._new(_utils.empty_0d(ct), value) if not axes: return _reduce_no_axes(name, x, ct) T = x._T() @@ -258,7 +258,7 @@ def _reduce_main(name, x, axis, keepdims, dtype): with np.errstate(all="ignore"): count = np.array([T.nvals], dtype=np.int64) value = spec.combine(L, fill, ct)(np.array([r], dtype=ct), count)[0] - out = Array._new(Tensor(gb_dtype(ct), ()), value) + out = Array._new(_utils.empty_0d(ct), value) return _keepdims(out, x, axes, keepdims) R = _reduce_tensor(T, _gb_monoid(spec, ct), axes) if not nan_fix and (full or _utils.same_value(fill, spec.identity(ct))): @@ -286,7 +286,7 @@ def _reduce_no_axes(name, x, ct): from ._elementwise import add, cast, multiply if name in {"sum", "prod"}: - identity = Array._new(Tensor(gb_dtype(ct), ()), 0 if name == "sum" else 1) + identity = Array._new(_utils.empty_0d(ct), 0 if name == "sum" else 1) out = (add if name == "sum" else multiply)(identity, x if x.dtype == ct else cast(x, ct)) return out if out.dtype == ct else cast(out, ct) return cast(x, ct) if x.dtype != ct else x._copy() @@ -379,7 +379,7 @@ def var(x, /, *, axis=None, correction=0.0, keepdims=False): S2 = Tensor(gb_dtype(real_dt), m.shape) C = Tensor(gb_dtype(np.int64), m.shape) k = subtract(L, Array._new(C, 0)) - dev = subtract(Array._new(Tensor(gb_dtype(rdt), ()), x._fill), m) + dev = subtract(Array._new(_utils.empty_0d(rdt), x._fill), m) # k unstored positions, each (fill - mean) ** 2 away; none for k == 0 (the fill may be inf) missing = binary_numpy("var", _times_count, k, _squared_magnitude(dev)) out = divide(add(Array._new(S2, 0), missing), denominator) diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py index ff585ef05..39a4827cb 100644 --- a/graphblas/array_api/_utils.py +++ b/graphblas/array_api/_utils.py @@ -8,7 +8,7 @@ from ..core.scalar import Scalar from ..core.tensor import Tensor from ..core.tensor_ewise import _broadcast_factor -from ..dtypes import lookup_dtype +from ..dtypes import DataType, lookup_dtype class DensifyError(ValueError): @@ -113,7 +113,23 @@ def mult_expansion_check(Ta, Tb, shape): def empty_tensor(dtype, shape): - return Tensor(gb_dtype(dtype), tuple(shape)) + shape = tuple(shape) + return empty_0d(dtype) if not shape else Tensor(gb_dtype(dtype), shape) + + +_EMPTY_0D = {} + + +def empty_0d(dtype): + """An empty 0-d Tensor of ``dtype`` (a numpy dtype or a graphblas one): a 0-d Array + stores nothing, its value is its fill. An O(1) copy-on-write copy of one made once per + dtype: making one costs a Matrix and a format setting, about 6 us, half of a 0-d + operation (SAPS's time-stepping loops make one per operation). + """ + gdt = dtype if isinstance(dtype, DataType) else gb_dtype(dtype) + if (T := _EMPTY_0D.get(gdt)) is None: + T = _EMPTY_0D[gdt] = Tensor(gdt, ()) + return T.dup() def full_tensor(dtype, shape, value): diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py index 8fbfd63a0..6041aa04a 100644 --- a/graphblas/tests/test_array_api_errors.py +++ b/graphblas/tests/test_array_api_errors.py @@ -330,3 +330,19 @@ def test_thread3_edge_cases(): assert_array_equal(xp.where(c, np.array([1, 2, 3]), 0), np.where(c, [1, 2, 3], 0)) assert_array_equal(xp.where(True, x, 0.0), x.todense()) assert_array_equal(np.eye(3) @ x.mT, np.eye(3) @ x.todense().T) + + +def test_zero_d_tensors_are_copies_of_one_empty_tensor(): + # (a 0-d Array stores nothing; its Tensor is an O(1) copy-on-write copy of one per dtype, + # which no change to a copy reaches) + x = xp.asarray(1.5) + y = xp.multiply(x, 2.0) + assert float(y) == 3.0 + t = x.tensor + t[()] = 7.0 + assert t.nvals == 1 + template = _utils._EMPTY_0D[_utils.gb_dtype(np.float64)] + assert template.nvals == 0 + assert xp.asarray(2.5)._T().nvals == 0 + assert float(xp.asarray(2.5)) == 2.5 + assert _utils.empty_tensor(np.int8, ()).nvals == 0 From c1594596316a8c281d8fa84c427306539f9a9bba Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 07:21:26 -0600 Subject: [PATCH 70/77] array_api: the masked three-operand einsum leaves out 0-d operands A 0-d operand stores nothing (its value is its fill), so it cannot mask the pair's product; einsum with one takes the full path. Only an identity-valued 0-d operand reached the masked path before, where the answer was the same. --- graphblas/array_api/_linalg.py | 2 +- graphblas/tests/test_array_api_lazy.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py index 490e5d014..8e5adfa00 100644 --- a/graphblas/array_api/_linalg.py +++ b/graphblas/array_api/_linalg.py @@ -560,7 +560,7 @@ def _masked_three(terms, output, ops, typed, out_dt, fill): takes the full einsum, which propagates it (D12b). With plus_times the third operand must be finite: inf * (x + y) is not inf * x + inf * y when x and y differ in sign. """ - if (typed.monoid.name, typed.binaryop.name) not in _DISTRIBUTES: + if (typed.monoid.name, typed.binaryop.name) not in _DISTRIBUTES or not all(x.ndim for x in ops): return None finite = [x.dtype.kind not in "fc" or _utils.fact(x, "finite") for x in ops] if not all(finite): diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py index 0326c0222..c7564c117 100644 --- a/graphblas/tests/test_array_api_lazy.py +++ b/graphblas/tests/test_array_api_lazy.py @@ -642,6 +642,15 @@ def test_einsum_three_operands_not_masked(masks): R = xp.einsum("ij,jk,ik->ik", make(p, 0), make(q, 0), make(r, 0), semiring="plus_min") assert all(m is None for m in masks) assert_close(R, _structural_three(p, q, r, np.minimum).astype(np.uint8)) + # A 0-d operand (its value is the identity here): the full einsum + inf = np.inf + m = [np.where(x == 0, inf, x) for x in (a, a)] + masks.clear() + R = xp.einsum( + "ij,ij,->", make(m[0], inf), make(m[1], inf), xp.asarray(inf), semiring="min_plus" + ) + assert all(x is None for x in masks) + assert_close(R, np.array(inf)) # (every term holds the identity inf) # Deferral off: the pair is computed whole, then intersected _lazy.enabled = False try: From be267a3757e512f30fcd83248a79704637fdd020 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 08:55:02 -0600 Subject: [PATCH 71/77] array_api: where stores only the condition's entries that select the other operand When the operand the condition's fill selects holds only the result's fill (a scalar, or an array that stores nothing), a condition entry equal to the fill would store that fill: it is left out, and the other operand is read only at the remaining entries. where(x > t, x, 0.0) with x of fill 0 stored every entry of x > t, explicit False ones included (an explicit 0.0 at each x <= t); it now stores x's entries above t only, as a select would. --- graphblas/array_api/_searching.py | 34 +++++++++++++++------ graphblas/tests/test_array_api_searching.py | 11 ++++++- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index 656ad5065..22f5eabc4 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -90,18 +90,32 @@ def _where_at_condition(cond, a, b, shape, rdt): if isinstance(base, Array) and base.ndim and base.shape != shape and base._T().nvals: return None coords, c = Tc.to_coo() - values = [] - for x in (a, b): - if isinstance(x, Array) and x.ndim: - values.append(_values_at(coords, x, shape) if len(c) else c) - else: - values.append(_value(x)) - V = np.where(np.asarray(c, dtype=bool), *values).astype(rdt, copy=False) + c = np.asarray(c, dtype=bool) gdt = gb_dtype(rdt) + if not (isinstance(base, Array) and base.ndim and base._T().nvals): + # The base holds only the result's fill: the condition's entries that select it + # (those equal to its fill) would store the fill, so only the others are kept, and + # the other operand is read there only (where(x > t, x, 0) stores x's entries > t) + keep = c != bool(cond._fill) + coords = tuple(i[keep] for i in coords) + other = b if bool(cond._fill) else a + if isinstance(other, Array) and other.ndim: + V = _values_at(coords, other, shape) if len(coords[0]) else c[keep] + else: + V = np.full(len(coords[0]), _value(other)) + V = np.asarray(V).astype(rdt, copy=False) + return Tensor.from_coo(coords, V, gdt, shape=shape) if len(V) else Tensor(gdt, shape) + values = [ + ( + (_values_at(coords, x, shape) if len(c) else c) + if isinstance(x, Array) and x.ndim + else _value(x) + ) + for x in (a, b) + ] + V = np.where(c, *values).astype(rdt, copy=False) R = Tensor.from_coo(coords, V, gdt, shape=shape) if len(c) else Tensor(gdt, shape) - if isinstance(base, Array) and base.ndim and base._T().nvals: - R = base._T().ewise_add(R, binary.second[gdt]).new(dtype=gdt) - return R + return base._T().ewise_add(R, binary.second[gdt]).new(dtype=gdt) def _values_at(coords, x, shape): diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py index 2c5be1196..e06ad026d 100644 --- a/graphblas/tests/test_array_api_searching.py +++ b/graphblas/tests/test_array_api_searching.py @@ -58,8 +58,17 @@ def test_where_stores_the_condition_and_the_selected_operand(): G = make(c, False, rng) # (with explicit False entries) W = xp.where(G, xp.arange(n), big) assert_array_equal(W, np.where(c, labels, big)) - assert W.nnz == G.nnz + assert W.nnz == c.sum() # (G's explicit False entries select big, the fill: not stored) assert W.fill_value == big + # where(x > t, x, 0): x's entries above t only, read there only + a = np.where(rng.random((n, n)) < 0.3, rng.normal(size=(n, n)), 0.0) + X = make(a, 0.0, rng) + R = xp.where(xp.greater(X, 0.5), X, 0.0) + assert_array_equal(R, np.where(a > 0.5, a, 0.0)) + assert R.nnz == (a > 0.5).sum() + R = xp.where(xp.less_equal(X, 0.5), 0.0, X) # (a condition with fill True) + assert_array_equal(R, np.where(a <= 0.5, 0.0, a)) + assert R.nnz == (a > 0.5).sum() # other layouts of the condition, and a fill of True (the base is then x1) for cond, cn in [(xp.matrix_transpose(G), c.T), (xp.logical_not(G), ~c)]: a = rng.normal(size=(n, n)) From e45ac15212583f359a17bd8cfe7b013132dbd213 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 09:31:36 -0600 Subject: [PATCH 72/77] array_api: where with a bool condition and an unbroadcast branch stays in GraphBLAS The previous commit kept only the condition's entries that select the other operand, but did it in numpy (export the coordinates, gather, rebuild), single-threaded. When the condition is bool and that operand has the result's shape, it is now a select of the condition's entries, holding the operand's fill, overwritten by the operand's entries there (an intersection). where(x > t, x, 0.0) on 1e7 entries, including the comparison: 0.32 s at 1 thread and 0.058 s at 8 (the numpy path: 0.19 and 0.13 s; before either: 0.33 and 0.068 s, storing twice the entries); the Tensor's select alone takes 0.081 and 0.012 s. --- graphblas/array_api/_searching.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index 22f5eabc4..67f387d14 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -81,7 +81,8 @@ def _where_at_condition(cond, a, b, shape, rdt): ``cond``'s stored positions written over (see ``where``), or None when that operand would have to be broadcast over stored entries, or ``cond`` itself broadcasts. """ - from .. import binary + from .. import binary, select + from ..dtypes import BOOL Tc = cond._T() if Tc.shape != shape: @@ -96,9 +97,17 @@ def _where_at_condition(cond, a, b, shape, rdt): # The base holds only the result's fill: the condition's entries that select it # (those equal to its fill) would store the fill, so only the others are kept, and # the other operand is read there only (where(x > t, x, 0) stores x's entries > t) + other = b if bool(cond._fill) else a + if Tc.dtype == BOOL and isinstance(other, Array) and other.shape == shape: + # (all in GraphBLAS, in parallel: the entries that select ``other``, holding + # its value, or its fill where it stores nothing) + sel = Tc.select(select.valueeq, not bool(cond._fill)).new() + second = binary.second[gdt] + R = sel.apply(second, right=_utils.gb_scalar(other._fill, rdt)).new(dtype=gdt) + R(accum=second) << sel.ewise_mult(other._T(), second) + return R keep = c != bool(cond._fill) coords = tuple(i[keep] for i in coords) - other = b if bool(cond._fill) else a if isinstance(other, Array) and other.ndim: V = _values_at(coords, other, shape) if len(coords[0]) else c[keep] else: From ef2030e018e5f9c4ba5712d970d8ea7da67fdd30 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 10:05:46 -0600 Subject: [PATCH 73/77] array_api: argmax and argmin along an axis run in GraphBLAS They exported every coordinate and ranked the values with np.lexsort: O(entries log entries) on one thread, 357x a max along the same axis at 1e7 entries. Now each step is a GraphBLAS pass over the stored entries, with the fibers as the rows (or columns) of the stored Matrix: the extreme of each fiber, the entries equal to it (a row or column scaling by the extremes under eq, with a value at every diagonal position so SuiteSparse runs it as a scaling), and the smallest position among those (a positional min_firstj or min_secondi product). The first stored NaN wins over it, as in numpy. Where the fill beats or ties a fiber's extreme (or is NaN), the fiber's first unstored position takes part. It is 0 unless the fiber stores position 0, which one positional product shows; only the fibers that do are read out, to find the first position missing from their entries. The result stays a Vector until it becomes the output Tensor, and stores only the positions that are not 0 (its fill). --- graphblas/array_api/_searching.py | 179 +++++++++++++------- graphblas/tests/test_array_api_searching.py | 33 ++++ 2 files changed, 147 insertions(+), 65 deletions(-) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index 67f387d14..991af0e76 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -6,6 +6,7 @@ import numpy as np from ..core.tensor import Tensor +from ..core.vector import Vector from . import _utils from ._array import Array from ._dtypes import check_supported @@ -147,10 +148,17 @@ def _check_array(x, name): def _arg(name, x, axis, keepdims): """The first position of the extreme value of each fiber (argmax, argmin), NaN counting - as the extreme (numpy's rule). O(entries log entries): the stored entries of each fiber - are ranked, and the fill competes once if the fiber has an unstored position. + as the extreme (numpy's rule), computed in GraphBLAS on the stored entries. + + With the fibers as the rows (or columns) of the stored Matrix: the extreme of each fiber + (a reduction), the entries equal to it (the Matrix scaled by the extremes as a diagonal, + under ``eq``), and the smallest position among those (a positional reduction); likewise + the first stored NaN. A fiber with an unstored position holds the fill there: when the + fill beats or ties the stored extreme, or is NaN, the fiber's first unstored position + takes part. That is 0 unless the fiber stores position 0, and otherwise is found from + that fiber's entries alone. Each step is a pass over the entries; nothing is sorted. """ - from ._manipulation import _normalize_axis, expand_dims, permute_dims, reshape + from ._manipulation import _normalize_axis, expand_dims, reshape _check_array(x, name) if x.dtype.kind == "c": @@ -165,75 +173,116 @@ def _arg(name, x, axis, keepdims): if L == 0: raise ValueError(f"{name} of an empty axis") out_shape = tuple(n for a, n in enumerate(x.shape) if a != axis) - others = tuple(a for a in range(x.ndim) if a != axis) - T = permute_dims(x, (*others, axis))._T() - nfibers = math.prod(out_shape) - if T.nvals == 0 or nfibers == 0: + T = x._T() + if T.nvals == 0 or math.prod(out_shape) == 0: out = Array._new(Tensor(gb_dtype(np.int64), out_shape), 0) return expand_dims(out, axis) if keepdims else out - coords, values = T.to_coo() # (C order: by fiber, then by position along the axis) - fiber = ( - np.ravel_multi_index(coords[:-1], out_shape) - if out_shape - else np.zeros(len(values), np.int64) - ) - pos = coords[-1] - starts = np.flatnonzero(np.r_[True, fiber[1:] != fiber[:-1]]) - counts = np.diff(np.r_[starts, len(fiber)]) - # Rank the values, NaN first: then the lowest key, earliest position, wins - isnan = np.isnan(values) if values.dtype.kind == "f" else np.zeros(len(values), bool) - _, rank = np.unique(values, return_inverse=True) - rank = rank.reshape(-1) - if name == "argmax": - rank = rank.max(initial=0) - rank - key = np.where(isnan, -1, rank) - order = np.lexsort((pos, key, fiber)) - best = order[starts] - best_pos, best_val = pos[best], values[best] - # The first position each fiber does not store (L if it stores them all) - within = np.arange(len(pos)) - np.repeat(starts, counts) - gap = pos != within - first_gap = np.where( - np.add.reduceat(gap, starts) > 0, - _first_true(gap, starts, counts, within), - counts, - ) - fill = x._fill - with np.errstate(invalid="ignore"): - if fill != fill: # (NaN) - better = ~np.isnan(best_val) if best_val.dtype.kind == "f" else np.ones(len(best), bool) - tie = ~better - else: - best_is_nan = ( - np.isnan(best_val) if best_val.dtype.kind == "f" else np.zeros(len(best), bool) - ) - beats = (fill > best_val) if name == "argmax" else (fill < best_val) - better = beats & ~best_is_nan - tie = (fill == best_val) & ~best_is_nan - unstored = first_gap < L - result = np.where( - unstored & better, - first_gap, - np.where(unstored & tie, np.minimum(first_gap, best_pos), best_pos), - ) - fibers = fiber[starts] + R = _arg_fibers(name, x, T, axis) if out_shape: - R = Tensor.from_coo( - np.unravel_index(fibers, out_shape), result.astype(np.int64), gb_dtype(np.int64), - shape=out_shape, - ) # fmt: skip - out = Array._new(R, 0) + out = Array._new(Tensor.from_vector(R).reshape(out_shape), 0) else: - out = Array._new(_utils.empty_0d(np.int64), result[0]) + out = Array._new(_utils.empty_0d(np.int64), np.int64(R.get(0, 0))) return expand_dims(out, axis) if keepdims else out -def _first_true(flags, starts, counts, within): - """Per group (``starts``, ``counts``): ``within`` at the group's first True flag.""" - idx = np.where(flags, np.arange(len(flags)), len(flags)) - first = np.minimum.reduceat(idx, starts) - first = np.minimum(first, len(flags) - 1) - return within[first] +def _arg_fibers(name, x, T, axis): + """A Vector over the fibers of ``T`` along ``axis`` (in C order) holding the first + position of each one's extreme (see ``_arg``), storing only the positions that are not 0. + """ + from .. import binary, monoid, select, semiring, unary + from ..dtypes import BOOL, INT64 + + L = T.shape[axis] + others = tuple(a for a in range(T.ndim) if a != axis) + # The fibers as the rows of M (byrow) or its columns: as stored when the layout allows + if T._layout == (others, (axis,)): + M, byrow = T._M, True + elif T._layout == ((axis,), others): + M, byrow = T._M, False + else: + M, byrow = T.to_matrix(rows=others), True + nfibers = M.nrows if byrow else M.ncols + if M.dtype == BOOL: + top = monoid.lor if name == "argmax" else monoid.land + else: + top = monoid.max if name == "argmax" else monoid.min + every = Vector(BOOL, L) + every << True # (iso: no values stored) + + def first(P, mask=None): + # (the smallest position along each fiber of P's entries) + v = Vector(INT64, nfibers) + if byrow: + v(mask=mask) << P.mxv(every, semiring.ss.min_firstj) + else: + v(mask=mask) << every.vxm(P, semiring.ss.min_secondi) + return v + + ext = (M.reduce_rowwise(top) if byrow else M.reduce_columnwise(top)).new() # (skips NaN) + # The entries equal to their fiber's extreme. With a value at every diagonal position, + # SuiteSparse runs the product as a row (column) scaling, one pass over M. + if nfibers <= 4 * M.nvals: + diag = Vector(ext.dtype, nfibers) + diag << 0 # (the value at a fiber with no entries scales nothing) + diag(accum=binary.second) << ext + else: + diag = ext + D = diag.diag() + eq = semiring.any_eq[M.dtype] + hits = (D.mxm(M, eq) if byrow else M.mxm(D, eq)).new() + R = first(hits.select(select.valueeq, True).new()) + del hits + # The first stored NaN, which beats every value + nan = None + if x.dtype.kind == "f" and _utils.fact(x, "nan"): + nan = first(M.apply(unary.isnan).new().select(select.valueeq, True).new()) + R(accum=binary.second) << nan + # The fibers where the fill takes part, if they have an unstored position: those whose + # stored extreme the fill beats or ties (and that store no NaN), or all, for a NaN fill + fill = x._fill + fill_nan = x.dtype.kind == "f" and np.isnan(fill) + if fill_nan: + need = ext + else: + beats = select.valuelt if name == "argmax" else select.valuegt + need = ext.select(beats, fill).new() + need(accum=binary.first) << ext.select(select.valueeq, fill) + if nan is not None: + keep = Vector(need.dtype, nfibers) + keep(~nan.S, replace=True) << need + need = keep + if need.nvals: + # A fiber whose first entry is past position 0 has its first unstored position at 0 + firsts = first(M, need.S) + R(~firsts.select(select.valuegt, 0).new().S, replace=True) << R # (there: 0) + at_0 = firsts.select(select.valueeq, 0).new() + if at_0.nvals: + _fill_at_0(name, M, byrow, R, ext, nan, fill, fill_nan, at_0, L) + return R.select(select.valuene, 0).new() + + +def _fill_at_0(name, M, byrow, R, ext, nan, fill, fill_nan, at_0, L): + """Update R at the fibers ``at_0`` (those that take the fill and store position 0) with + the fill's first position, read from those fibers' entries: the first position not equal + to its rank among the fiber's entries, or the number of entries when they fill a prefix + (none when that number is L: the fiber stores every position). + """ + fibers = at_0.to_coo(values=False)[0] + sub = M[fibers, :].new() if byrow else M[:, fibers].new().T.new() + r, c, _ = sub.to_coo() # (by fiber, positions sorted) + starts = np.flatnonzero(np.r_[True, r[1:] != r[:-1]]) + counts = np.diff(np.r_[starts, len(r)]) + rank = np.arange(len(r)) - np.repeat(starts, counts) + gap = np.minimum.reduceat(np.where(c != rank, rank, np.repeat(counts, counts)), starts) + best = R[fibers].new().to_dense(fill_value=L) + if fill_nan: + stored_nan = nan[fibers].new().to_dense(fill_value=L) if nan is not None else L + new = np.where(gap < L, np.minimum(gap, stored_nan), best) + else: + e = ext[fibers].new().to_dense() + wins = (fill > e) if name == "argmax" else (fill < e) + new = np.where(gap < L, np.where(wins, gap, np.minimum(gap, best)), best) + R[fibers] << Vector.from_coo(np.arange(len(fibers)), new.astype(np.int64), size=len(fibers)) def argmax(x, /, *, axis=None, keepdims=False): diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py index e06ad026d..22305129a 100644 --- a/graphblas/tests/test_array_api_searching.py +++ b/graphblas/tests/test_array_api_searching.py @@ -137,6 +137,39 @@ def test_argmax_ties_with_the_fill(): assert int(xp.argmin(x)) == np.argmin(a) +def test_argmax_argmin_random(): + # numpy's rules (first occurrence, NaN the extreme, the fill at unstored positions) + # over fills, densities, shapes, axes and stored layouts + rng = np.random.default_rng(12) + pools = { + "float64": [0.0, -0.0, 1.0, -1.0, 2.5, np.nan, np.inf, -np.inf], + "float32": [0.0, -0.0, 1.0, -1.0, np.nan, np.inf], + "int8": [-3, -1, 0, 1, 2], + "uint16": [0, 1, 2, 3], + "bool": [False, True], + } + shapes = [(1,), (6,), (1, 1), (3, 4), (6, 1), (1, 7), (2, 3, 4), (3, 1, 2, 2)] + for _ in range(250): + dt = np.dtype(list(pools)[rng.integers(len(pools))]) + pool = np.array(pools[dt.name], dt) + shape = shapes[rng.integers(len(shapes))] + fill = pool[rng.integers(len(pool))] + values = rng.choice(pool[: rng.integers(1, len(pool) + 1)], shape) + dense = np.where(rng.random(shape) < rng.choice([0.0, 0.2, 0.6, 1.0]), values, fill) + x = make(dense, fill, rng, explicit=0.2) + if dense.ndim > 1 and rng.random() < 0.5: + # (a stored layout other than the default) + perm = tuple(int(a) for a in rng.permutation(dense.ndim)) + inverse = tuple(int(a) for a in np.argsort(perm)) + x = xp.permute_dims(xp.asarray(xp.permute_dims(x, perm), copy=True), inverse) + axis = int(rng.integers(-dense.ndim, dense.ndim)) if rng.random() < 0.8 else None + keepdims = rng.random() < 0.3 + for name in ["argmax", "argmin"]: + expected = getattr(np, name)(dense, axis=axis, keepdims=keepdims) + out = getattr(xp, name)(x, axis=axis, keepdims=keepdims) + assert_array_equal(out, np.asarray(expected, dtype=np.int64)) + + @pytest.mark.parametrize("axis", [None, 0, 1, (0, 1), ()]) def test_count_nonzero(axis): for dtype in [*SEARCH_DTYPES, np.dtype("complex64")]: From 8bf0c0b2eae6206acbd6c15ec09c84079f7b9618 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 15:50:14 -0600 Subject: [PATCH 74/77] array_api: where reads a broadcast branch in GraphBLAS Where the condition stores nothing its fill selects one operand; when that operand holds only the fill, the result is the condition's selecting entries holding the other operand. A broadcast other operand (connected components' where(graph, labels, big)) was read by exporting the condition's coordinates and gathering in numpy, on one thread. Now the selecting entries are a select on the condition (truthy ones, or falsy ones under a truthy fill, for any real dtype), given the other's fill, with its stored values written over by a broadcast ewise_mult, which the Tensor runs as a diagonal scaling without replicating the operand. A complex condition still goes through numpy. --- graphblas/array_api/_searching.py | 22 +++++++++++++-------- graphblas/tests/test_array_api_searching.py | 8 ++++++++ 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index 991af0e76..dac586bae 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -91,22 +91,26 @@ def _where_at_condition(cond, a, b, shape, rdt): base = a if bool(cond._fill) else b if isinstance(base, Array) and base.ndim and base.shape != shape and base._T().nvals: return None - coords, c = Tc.to_coo() - c = np.asarray(c, dtype=bool) gdt = gb_dtype(rdt) if not (isinstance(base, Array) and base.ndim and base._T().nvals): # The base holds only the result's fill: the condition's entries that select it # (those equal to its fill) would store the fill, so only the others are kept, and # the other operand is read there only (where(x > t, x, 0) stores x's entries > t) other = b if bool(cond._fill) else a - if Tc.dtype == BOOL and isinstance(other, Array) and other.shape == shape: - # (all in GraphBLAS, in parallel: the entries that select ``other``, holding - # its value, or its fill where it stores nothing) - sel = Tc.select(select.valueeq, not bool(cond._fill)).new() + if cond.dtype.kind != "c": + # All in GraphBLAS, in parallel: the entries that select ``other`` (the truthy + # ones, or the falsy ones under a truthy fill; NaN is truthy), holding its fill, + # then its stored values written over. A broadcast ``other`` is read through a + # diagonal scaling, without replicating it (xp.where(graph, labels, big)). + pick = select.valueeq if bool(cond._fill) else select.valuene + sel = Tc.select(pick, False if Tc.dtype == BOOL else 0).new() second = binary.second[gdt] - R = sel.apply(second, right=_utils.gb_scalar(other._fill, rdt)).new(dtype=gdt) - R(accum=second) << sel.ewise_mult(other._T(), second) + R = sel.apply(second, right=_utils.gb_scalar(_value(other), rdt)).new(dtype=gdt) + if isinstance(other, Array) and other.ndim and other._T().nvals: + R(accum=second) << sel.ewise_mult(other._T(), second) return R + coords, c = Tc.to_coo() + c = np.asarray(c, dtype=bool) keep = c != bool(cond._fill) coords = tuple(i[keep] for i in coords) if isinstance(other, Array) and other.ndim: @@ -115,6 +119,8 @@ def _where_at_condition(cond, a, b, shape, rdt): V = np.full(len(coords[0]), _value(other)) V = np.asarray(V).astype(rdt, copy=False) return Tensor.from_coo(coords, V, gdt, shape=shape) if len(V) else Tensor(gdt, shape) + coords, c = Tc.to_coo() + c = np.asarray(c, dtype=bool) values = [ ( (_values_at(coords, x, shape) if len(c) else c) diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py index 22305129a..e0d86efdb 100644 --- a/graphblas/tests/test_array_api_searching.py +++ b/graphblas/tests/test_array_api_searching.py @@ -80,6 +80,14 @@ def test_where_stores_the_condition_and_the_selected_operand(): (xp.arange(n)[:, None] * 1.0, -1.0, np.where(cn, labels[:, None] * 1.0, -1.0)), ]: assert_array_equal(xp.where(cond, x1, x2), e) + # A condition that is not bool selects by truth value (NaN is true), with a broadcast + # operand read at its entries; a complex condition goes through numpy + f = np.where(c, np.where(rng.random((n, n)) < 0.3, np.nan, 2.5), 0.0) + f[0, 0] = -0.0 + for cond, fill in [(f, 0.0), (np.where(c, 0, 3).astype(np.int8), 3), (f + 1j * f, 0j)]: + for x1, x2 in [(xp.arange(n), big), (big, xp.arange(n)[:, None]), (labels[0], -1)]: + e = np.where(cond, np.asarray(x1), np.asarray(x2)) + assert_array_equal(xp.where(make(cond, fill), x1, x2), e) # The selected operand broadcasts over stored entries, or the condition broadcasts: the # union (the result really holds the operand's entries along the broadcast axis) v = np.where(rng.random(n) < 0.5, rng.normal(size=n), 0.0) From 35c602bdb6a657c15e0e2654afb9f9fd606272ef Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 15:57:46 -0600 Subject: [PATCH 75/77] array_api: where stores a broadcast branch in full before scaling by it A broadcast ewise_mult is a diagonal scaling, but SuiteSparse takes its one-pass scaling path only when the diagonal has a value at every position; xp.arange(n) stores no 0 (its fill), and the general product it falls back to is 17x slower at 1e7 entries (0.71 s against 0.041 at 1 thread). The branch is now stored at every position first, when it is no larger than 4x the condition's selecting entries, and the result is that one product. --- graphblas/array_api/_searching.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index dac586bae..532ebbbca 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -105,9 +105,17 @@ def _where_at_condition(cond, a, b, shape, rdt): pick = select.valueeq if bool(cond._fill) else select.valuene sel = Tc.select(pick, False if Tc.dtype == BOOL else 0).new() second = binary.second[gdt] + To = other._T() if isinstance(other, Array) and other.ndim else None + if To is not None and other.shape != shape and other.size <= 4 * sel.nvals: + # (stored at every position, the broadcast operand makes the product a + # diagonal scaling, one pass; with gaps in the diagonal SuiteSparse multiplies + # in general, 17x slower at 1e7 entries) + if To.nvals < other.size: + To = _utils.with_fill(other, 0 if other._fill else 1)._T() + return sel.ewise_mult(To, second).new(dtype=gdt) R = sel.apply(second, right=_utils.gb_scalar(_value(other), rdt)).new(dtype=gdt) - if isinstance(other, Array) and other.ndim and other._T().nvals: - R(accum=second) << sel.ewise_mult(other._T(), second) + if To is not None and To.nvals: + R(accum=second) << sel.ewise_mult(To, second) return R coords, c = Tc.to_coo() c = np.asarray(c, dtype=bool) From d475b663b14c2bd9529af7ccf28f239912c8dff0 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 16:07:08 -0600 Subject: [PATCH 76/77] array_api: where over a base with entries selects and unions in GraphBLAS where(frontier, k, level), multi-source BFS's level update, has a base (level) that stores entries. It exported the condition's coordinates, gathered both operands there in numpy, rebuilt them as a Tensor and took the union with the base: on soc-LiveJournal1 with 8 sources the kernel's memory grew 3.9 GB over its inputs, against 0.85 GB for SAPS's own python-graphblas framework (which assigns under a mask). Now the condition's selecting entries are computed as for a base that holds only its fill (a select, then the other operand's values) and written over the base by one ewise_add: 0.83 GB, and the kernel takes 2.0 s instead of 3.2. A complex condition still goes through numpy. --- graphblas/array_api/_searching.py | 68 +++++++++++++-------- graphblas/tests/test_array_api_searching.py | 12 +++- 2 files changed, 51 insertions(+), 29 deletions(-) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py index 532ebbbca..33631c95d 100644 --- a/graphblas/array_api/_searching.py +++ b/graphblas/array_api/_searching.py @@ -82,8 +82,7 @@ def _where_at_condition(cond, a, b, shape, rdt): ``cond``'s stored positions written over (see ``where``), or None when that operand would have to be broadcast over stored entries, or ``cond`` itself broadcasts. """ - from .. import binary, select - from ..dtypes import BOOL + from .. import binary Tc = cond._T() if Tc.shape != shape: @@ -92,31 +91,19 @@ def _where_at_condition(cond, a, b, shape, rdt): if isinstance(base, Array) and base.ndim and base.shape != shape and base._T().nvals: return None gdt = gb_dtype(rdt) - if not (isinstance(base, Array) and base.ndim and base._T().nvals): - # The base holds only the result's fill: the condition's entries that select it - # (those equal to its fill) would store the fill, so only the others are kept, and - # the other operand is read there only (where(x > t, x, 0) stores x's entries > t) - other = b if bool(cond._fill) else a - if cond.dtype.kind != "c": - # All in GraphBLAS, in parallel: the entries that select ``other`` (the truthy - # ones, or the falsy ones under a truthy fill; NaN is truthy), holding its fill, - # then its stored values written over. A broadcast ``other`` is read through a - # diagonal scaling, without replicating it (xp.where(graph, labels, big)). - pick = select.valueeq if bool(cond._fill) else select.valuene - sel = Tc.select(pick, False if Tc.dtype == BOOL else 0).new() - second = binary.second[gdt] - To = other._T() if isinstance(other, Array) and other.ndim else None - if To is not None and other.shape != shape and other.size <= 4 * sel.nvals: - # (stored at every position, the broadcast operand makes the product a - # diagonal scaling, one pass; with gaps in the diagonal SuiteSparse multiplies - # in general, 17x slower at 1e7 entries) - if To.nvals < other.size: - To = _utils.with_fill(other, 0 if other._fill else 1)._T() - return sel.ewise_mult(To, second).new(dtype=gdt) - R = sel.apply(second, right=_utils.gb_scalar(_value(other), rdt)).new(dtype=gdt) - if To is not None and To.nvals: - R(accum=second) << sel.ewise_mult(To, second) - return R + has_base = isinstance(base, Array) and base.ndim and base._T().nvals + other = b if bool(cond._fill) else a + if cond.dtype.kind != "c": + # All in GraphBLAS, in parallel. The condition's entries that select the base (those + # equal to its fill) hold the base's values, which the base supplies; only the others + # are computed, and the other operand is read there only (where(x > t, x, 0) + # stores x's entries > t; where(frontier, k, level) is a select and a union). + R = _selecting(Tc, cond, other, shape, rdt) + if has_base: + R = base._T().ewise_add(R, binary.second[gdt]).new(dtype=gdt) + return R + if not has_base: + # (a complex condition: through numpy) coords, c = Tc.to_coo() c = np.asarray(c, dtype=bool) keep = c != bool(cond._fill) @@ -142,6 +129,33 @@ def _where_at_condition(cond, a, b, shape, rdt): return base._T().ewise_add(R, binary.second[gdt]).new(dtype=gdt) +def _selecting(Tc, cond, other, shape, rdt): + """The entries of ``Tc`` (cond's Tensor, not complex) that select ``other`` (the truthy + ones, or the falsy ones under a truthy fill; NaN is truthy), holding ``other``'s values: + its fill, then its stored values written over. A broadcast ``other`` is read through a + diagonal scaling, without replicating it (xp.where(graph, labels, big)). + """ + from .. import binary, select + from ..dtypes import BOOL + + gdt = gb_dtype(rdt) + pick = select.valueeq if bool(cond._fill) else select.valuene + sel = Tc.select(pick, False if Tc.dtype == BOOL else 0).new() + second = binary.second[gdt] + To = other._T() if isinstance(other, Array) and other.ndim else None + if To is not None and other.shape != shape and other.size <= 4 * sel.nvals: + # (stored at every position, the broadcast operand makes the product a diagonal + # scaling, one pass; with gaps in the diagonal SuiteSparse multiplies in general, + # 17x slower at 1e7 entries) + if To.nvals < other.size: + To = _utils.with_fill(other, 0 if other._fill else 1)._T() + return sel.ewise_mult(To, second).new(dtype=gdt) + R = sel.apply(second, right=_utils.gb_scalar(_value(other), rdt)).new(dtype=gdt) + if To is not None and To.nvals: + R(accum=second) << sel.ewise_mult(To, second) + return R + + def _values_at(coords, x, shape): """The values of ``x`` (an Array that broadcasts to ``shape``) at ``coords`` of ``shape``: its stored value, or its fill. O(len(coords)), read from x's own entries. diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py index e0d86efdb..5b01a3286 100644 --- a/graphblas/tests/test_array_api_searching.py +++ b/graphblas/tests/test_array_api_searching.py @@ -81,11 +81,19 @@ def test_where_stores_the_condition_and_the_selected_operand(): ]: assert_array_equal(xp.where(cond, x1, x2), e) # A condition that is not bool selects by truth value (NaN is true), with a broadcast - # operand read at its entries; a complex condition goes through numpy + # operand read at its entries, or over a base with entries (where(frontier, k, level)); + # a complex condition goes through numpy f = np.where(c, np.where(rng.random((n, n)) < 0.3, np.nan, 2.5), 0.0) f[0, 0] = -0.0 + level = make(a, 0.0) for cond, fill in [(f, 0.0), (np.where(c, 0, 3).astype(np.int8), 3), (f + 1j * f, 0j)]: - for x1, x2 in [(xp.arange(n), big), (big, xp.arange(n)[:, None]), (labels[0], -1)]: + for x1, x2 in [ + (xp.arange(n), big), + (big, xp.arange(n)[:, None]), + (labels[0], -1), + (7, level), + (xp.arange(n), level), + ]: e = np.where(cond, np.asarray(x1), np.asarray(x2)) assert_array_equal(xp.where(make(cond, fill), x1, x2), e) # The selected operand broadcasts over stored entries, or the condition broadcasts: the From f87c013bdf70aeec60d628807d574040e165afeb Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 11 Nov 2026 18:01:02 -0600 Subject: [PATCH 77/77] Tensor: wide column MAX and MIN reductions keep every thread The 2-thread cap on a reduce_columnwise into 1e5 or more outputs (603ca795f) came from sums: on FROSTT Uber (4392 x 1.96e6) 4.0 / 4.0 / 17 / 15 ms at 1 / 2 / 4 / 8 threads, and the same holds now (1e4 x 1e6 with 1e7 entries: 18 / 18 / 11 / 39 ms). MAX and MIN scale instead, at 4 and 8 threads, with no loss on any shape measured: 5e6 x 5e6 with 1e7 entries 0.122 -> 0.050 s at 8 threads, 1e7 x 1e7 with 1e8 entries 0.66 -> 0.28 s, 1e3 x 1e7 0.100 -> 0.049, Uber 4.1 -> 2.7 ms (JIT off and on alike, SuiteSparse 10.5.1, arm64). So the cap now applies to the other monoids only. --- graphblas/core/tensor_reduce.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/graphblas/core/tensor_reduce.py b/graphblas/core/tensor_reduce.py index 1129868f1..83351f2de 100644 --- a/graphblas/core/tensor_reduce.py +++ b/graphblas/core/tensor_reduce.py @@ -276,13 +276,18 @@ def _reduce_matrix(M, shape, layout, A, op, opts): M = M.ss.reshape(tl.prod(shape, seq[:q]), tl.prod(shape, seq[q:])) layout = Layout(seq[:q], seq[q:]) continue + # Into many outputs, a sum on 2 threads is as fast as on 1 or faster; on 4 and 8 it + # can be 2-4x slower (FROSTT Uber, 4392 x 1.96e6: 4.0 / 4.0 / 17 / 15 ms at 1 / 2 / + # 4 / 8 threads; 1e4 x 1e6 with 1e7 entries: 18 / 18 / 11 / 39 ms). MAX and MIN + # scale instead (5e6 x 5e6 with 1e7 entries: 122 / 120 / 51 / 50 ms; 1e7 x 1e7 + # with 1e8: 0.66 / 0.66 / 0.32 / 0.28 s; no loss on the others), so they keep + # every thread. (SuiteSparse 10.5.1, arm64.) + capped = wide and getattr(op, "name", None) not in {"max", "min"} if census.on: - census.note("reduce", "row group: reduce_columnwise" + (", 2 threads" if wide + census.note("reduce", "row group: reduce_columnwise" + (", 2 threads" if capped and nthreads > 2 else ""), nvals=nvals) # fmt: skip v = Vector(dtype, M._ncols) - # Into many outputs, 2 threads scatter as fast as 1 or faster; 4 and 8 can be - # 3-4x slower (Uber 4.0 / 4.9 / 17 / 15 ms at 2 / 1 / 4 / 8 threads) - v(**(dict(opts, nthreads=min(2, nthreads)) if wide else opts)) << M.reduce_columnwise( + v(**(dict(opts, nthreads=min(2, nthreads)) if capped else opts)) << M.reduce_columnwise( op ) M = _vector_to_row(v)