Merge branch 'main' into searchsorted

Author: mdhaber

docs/api-reference.md

@@ -25,4 +25,5 @@
     partition
     setdiff1d
     sinc
+    union1d
 ```

pixi.lock

@@ -55,22 +55,22 @@ array-api-extra = { path = ".", editable = true }
 
 [tool.pixi.feature.lint.dependencies]
 typing-extensions = ">=4.15.0"
-pylint = ">=4.0.2"
+pylint = ">=4.0.3"
 mypy = ">=1.18.2"
-basedpyright = ">=1.33.0"
+basedpyright = ">=1.34.0"
 numpydoc = ">=1.9.0,<2"
 # import dependencies for mypy:
 array-api-strict = ">=2.4.1,<2.5"
 numpy = ">=2.1.3"
-hypothesis = ">=6.142.4"
+hypothesis = ">=6.148.2"
 dask-core = ">=2025.11.0" # No distributed, tornado, etc.
 dprint = ">=0.50.0,<0.51"
-lefthook = ">=2.0.2,<3"
-ruff = ">=0.14.4,<0.15"
-typos = ">=1.39.0,<2"
-actionlint = ">=1.7.8,<2"
+lefthook = ">=2.0.4,<3"
+ruff = ">=0.14.6,<0.15"
+typos = ">=1.40.0,<2"
+actionlint = ">=1.7.9,<2"
 blacken-docs = ">=1.20.0,<2"
-pytest = ">=8.4.2,<9"
+pytest = ">=9.0.1,<10"
 validate-pyproject = ">=0.24.1,<0.25"
 # NOTE: don't add cupy, jax, pytorch, or sparse here,
 # as they slow down mypy and are not portable across target OSs
@@ -93,9 +93,9 @@ numpydoc = { cmd = "numpydoc lint", description = "Validate docstrings with nump
 lint = { cmd = "lefthook run pre-commit --all-files --force", description = "Run all linters" }
 
 [tool.pixi.feature.tests.dependencies]
-pytest = ">=8.4.2"
+pytest = ">=9.0.1"
 pytest-cov = ">=7.0.0"
-hypothesis = ">=6.142.4"
+hypothesis = ">=6.148.2"
 array-api-strict = ">=2.4.1,<2.5"
 numpy = ">=1.22.0"
 
@@ -121,7 +121,7 @@ sphinx-copybutton = ">=0.5.2"
 sphinx-autodoc-typehints = ">=1.25.3"
 # Needed to import parsed modules with autodoc
 dask-core = ">=2025.11.0" # No distributed, tornado, etc.
-pytest = ">=8.4.2"
+pytest = ">=9.0.1"
 typing-extensions = ">=4.15.0"
 numpy = ">=2.1.3"
 
@@ -200,8 +200,8 @@ pytorch = { version = ">=2.8.0", build = "cuda12*" }
 
 [tool.pixi.feature.nogil.dependencies]
 python-freethreading = "~=3.13.0"
-pytest-run-parallel = ">=0.7.1"
-numpy = ">=2.3.4"
+pytest-run-parallel = ">=0.8.0"
+numpy = ">=2.3.5"
 # pytorch = "*"  # Not available on Python 3.13t yet
 dask-core = ">=2025.11.0" # No distributed, tornado, etc.
 # sparse = "*"  # numba not available on Python 3.13t yet

pyproject.toml

@@ -14,6 +14,7 @@
     partition,
     setdiff1d,
     sinc,
+    union1d,
 )
 from ._lib._at import at
 from ._lib._funcs import (
@@ -52,4 +53,5 @@
     "searchsorted",
     "setdiff1d",
     "sinc",
+    "union1d",
 ]

src/array_api_extra/__init__.py

@@ -226,13 +226,14 @@ def create_diagonal(
     if is_torch_namespace(xp):
         return xp.diag_embed(x, offset=offset, dim1=-2, dim2=-1)
 
-    if (is_dask_namespace(xp) or is_cupy_namespace(xp)) and x.ndim < 2:
+    if (
+        is_dask_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_numpy_namespace(xp)
+        or is_jax_namespace(xp)
+    ) and (x.ndim < 2):
         return xp.diag(x, k=offset)
 
-    if (is_jax_namespace(xp) or is_numpy_namespace(xp)) and x.ndim < 3:
-        batch_dim, n = eager_shape(x)[:-1], eager_shape(x, -1)[0] + abs(offset)
-        return xp.reshape(xp.diag(x, k=offset), (*batch_dim, n, n))
-
     return _funcs.create_diagonal(x, offset=offset, xp=xp)
 
 
@@ -1026,3 +1027,37 @@ def isin(
         return xp.isin(a, b, assume_unique=assume_unique, invert=invert)
 
     return _funcs.isin(a, b, assume_unique=assume_unique, invert=invert, xp=xp)
+
+
+def union1d(a: Array, b: Array, /, *, xp: ModuleType | None = None) -> Array:
+    """
+    Find the union of two arrays.
+
+    Return the unique, sorted array of values that are in either of the two
+    input arrays.
+
+    Parameters
+    ----------
+    a, b : Array
+        Input arrays. They are flattened internally if they are not already 1D.
+
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a` and `b`. Default: infer.
+
+    Returns
+    -------
+    Array
+        Unique, sorted union of the input arrays.
+    """
+    if xp is None:
+        xp = array_namespace(a, b)
+
+    if (
+        is_numpy_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_dask_namespace(xp)
+        or is_jax_namespace(xp)
+    ):
+        return xp.union1d(a, b)
+
+    return _funcs.union1d(a, b, xp=xp)

src/array_api_extra/_delegation.py

@@ -837,3 +837,12 @@ def isin(  # numpydoc ignore=PR01,RT01
         _helpers.in1d(a, b, assume_unique=assume_unique, invert=invert, xp=xp),
         original_a_shape,
     )
+
+
+def union1d(a: Array, b: Array, /, *, xp: ModuleType) -> Array:
+    # numpydoc ignore=PR01,RT01
+    """See docstring in `array_api_extra._delegation.py`."""
+    a = xp.reshape(a, (-1,))
+    b = xp.reshape(b, (-1,))
+    # XXX: `sparse` returns NumPy arrays from `unique_values`
+    return xp.asarray(xp.unique_values(xp.concat([a, b])))

src/array_api_extra/_lib/_funcs.py

@@ -32,6 +32,7 @@
     searchsorted,
     setdiff1d,
     sinc,
+    union1d,
 )
 from array_api_extra._lib._backends import NUMPY_VERSION, Backend
 from array_api_extra._lib._testing import xfail, xp_assert_close, xp_assert_equal
@@ -60,8 +61,6 @@
 lazy_xp_function(searchsorted)
 lazy_xp_function(sinc)
 
-NestedFloatList = list[float] | list["NestedFloatList"]
-
 
 class TestApplyWhere:
     @staticmethod
@@ -716,6 +715,7 @@ def test_0d_raises(self, xp: ModuleType):
             (0, 1),
             (1, 0),
             (0, 0),
+            (2, 3),
             (4, 2, 1),
             (1, 1, 7),
             (0, 0, 1),
@@ -1815,3 +1815,36 @@ def test_nd(
         x, y = xp.asarray(x.copy()), xp.asarray(y.copy())
         res = searchsorted(x, y, side=side, xp=xp)
         xp_assert_equal(res, ref)
+
+
+@pytest.mark.skip_xp_backend(
+    Backend.ARRAY_API_STRICTEST,
+    reason="data_dependent_shapes flag for unique_values is disabled",
+)
+class TestUnion1d:
+    def test_simple(self, xp: ModuleType):
+        a = xp.asarray([-1, 1, 0])
+        b = xp.asarray([2, -2, 0])
+        expected = xp.asarray([-2, -1, 0, 1, 2])
+        res = union1d(a, b)
+        xp_assert_equal(res, expected)
+
+    def test_2d(self, xp: ModuleType):
+        a = xp.asarray([[-1, 1, 0], [1, 2, 0]])
+        b = xp.asarray([[1, 0, 1], [-2, -1, 0]])
+        expected = xp.asarray([-2, -1, 0, 1, 2])
+        res = union1d(a, b)
+        xp_assert_equal(res, expected)
+
+    def test_3d(self, xp: ModuleType):
+        a = xp.asarray([[[-1, 0], [1, 2]], [[-1, 0], [1, 2]]])
+        b = xp.asarray([[[0, 1], [-1, 2]], [[1, -2], [0, 2]]])
+        expected = xp.asarray([-2, -1, 0, 1, 2])
+        res = union1d(a, b)
+        xp_assert_equal(res, expected)
+
+    @pytest.mark.skip_xp_backend(Backend.TORCH, reason="materialize 'meta' device")
+    def test_device(self, xp: ModuleType, device: Device):
+        a = xp.asarray([-1, 1, 0], device=device)
+        b = xp.asarray([2, -2, 0], device=device)
+        assert get_device(union1d(a, b)) == device