raise on incompatible cast

Author: crusaderky

src/array_api_extra/_funcs.py

@@ -848,7 +848,11 @@ def add(
         xp: ModuleType | None = None,
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] += y`` and return the updated array."""
-        return self._iop("add", operator.add, y, copy=copy, xp=xp)
+
+        # Note for this and all other methods based on _iop:
+        # operator.iadd and operator.add subtly differ in behaviour, as
+        # only iadd will trigger exceptions when y has an incompatible dtype.
+        return self._iop("add", operator.iadd, y, copy=copy, xp=xp)
 
     def subtract(
         self,
@@ -858,7 +862,7 @@ def subtract(
         xp: ModuleType | None = None,
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] -= y`` and return the updated array."""
-        return self._iop("subtract", operator.sub, y, copy=copy, xp=xp)
+        return self._iop("subtract", operator.isub, y, copy=copy, xp=xp)
 
     def multiply(
         self,
@@ -868,7 +872,7 @@ def multiply(
         xp: ModuleType | None = None,
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] *= y`` and return the updated array."""
-        return self._iop("multiply", operator.mul, y, copy=copy, xp=xp)
+        return self._iop("multiply", operator.imul, y, copy=copy, xp=xp)
 
     def divide(
         self,
@@ -878,7 +882,7 @@ def divide(
         xp: ModuleType | None = None,
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] /= y`` and return the updated array."""
-        return self._iop("divide", operator.truediv, y, copy=copy, xp=xp)
+        return self._iop("divide", operator.itruediv, y, copy=copy, xp=xp)
 
     def power(
         self,
@@ -888,7 +892,7 @@ def power(
         xp: ModuleType | None = None,
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] **= y`` and return the updated array."""
-        return self._iop("power", operator.pow, y, copy=copy, xp=xp)
+        return self._iop("power", operator.ipow, y, copy=copy, xp=xp)
 
     def min(
         self,

tests/test_at.py

@@ -136,3 +136,24 @@ def test_alternate_index_syntax():
         at(a).set(4)
     with pytest.raises(ValueError, match="Index"):
         at(a, 0)[0].set(4)
+
+
+@pytest.mark.parametrize("copy", [True, False])
+@pytest.mark.parametrize("op", ["add", "subtract", "multiply", "divide", "power"])
+def test_iops_incompatible_dtype(op, copy):
+    """Test that at() replicates the backend's behaviour for
+    in-place operations with incompatible dtypes.
+
+    Note:
+    >>> a = np.asarray([1, 2, 3])
+    >>> a / 1.5
+    a / 1.5
+    Out[11]: array([0.        , 0.66666667, 1.33333333])
+    >>> a /= 1.5
+    UFuncTypeError: Cannot cast ufunc 'divide' output from dtype('float64')
+    to dtype('int64') with casting rule 'same_kind'
+    """
+    a = np.asarray([2, 4])
+    func = getattr(at(a)[:], op)
+    with pytest.raises(TypeError, match="Cannot cast ufunc"):
+        func(1.1, copy=copy)