_at meta_xp

Author: crusaderky

src/array_api_extra/_lib/_at.py

@@ -15,6 +15,7 @@
     is_jax_array,
     is_writeable_array,
 )
+from ._utils._helpers import meta_namespace
 from ._utils._typing import Array, Index
 
 
@@ -419,9 +420,16 @@ def min(
         xp: ModuleType | None = None,
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] = minimum(x[idx], y)`` and return the updated array."""
+        # On Dask, this function runs on the chunks, so we need to determine the
+        # namespace that Dask is wrapping.
+        # Note that da.minimum _incidentally_ works on numpy, cupy, and sparse
+        # thanks to all these meta-namespaces implementing the __array_ufunc__
+        # interface, but there's no guarantee that it will work for other
+        # wrapped libraries in the future.
         xp = array_namespace(self._x) if xp is None else xp
+        mxp = meta_namespace(self._x, xp=xp)
         y = xp.asarray(y)
-        return self._op(_AtOp.MIN, xp.minimum, xp.minimum, y, copy=copy, xp=xp)
+        return self._op(_AtOp.MIN, mxp.minimum, mxp.minimum, y, copy=copy, xp=xp)
 
     def max(
         self,
@@ -431,6 +439,8 @@ def max(
         xp: ModuleType | None = None,
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] = maximum(x[idx], y)`` and return the updated array."""
+        # See note on min()
         xp = array_namespace(self._x) if xp is None else xp
+        mxp = meta_namespace(self._x, xp=xp)
         y = xp.asarray(y)
-        return self._op(_AtOp.MAX, xp.maximum, xp.maximum, y, copy=copy, xp=xp)
+        return self._op(_AtOp.MAX, mxp.maximum, mxp.maximum, y, copy=copy, xp=xp)