Move args to second position

Author: crusaderky

src/array_api_extra/_lib/_at.py

@@ -299,7 +299,7 @@ def _op(
                 if out_of_place_op:
                     # suppress inf warnings on Dask
                     out = apply_where(
-                        idx, out_of_place_op, (x, y_xp), fill_value=x, xp=xp
+                        idx, (x, y_xp), out_of_place_op, fill_value=x, xp=xp
                     )
                     # Undo int->float promotion on JAX after _AtOp.DIVIDE
                     out = xp.astype(out, x.dtype, copy=False)

src/array_api_extra/_lib/_funcs.py

@@ -38,9 +38,9 @@
 @overload
 def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ignore=GL08
     cond: Array,
+    args: Array | tuple[Array, ...],
     f1: Callable[..., Array],
     f2: Callable[..., Array],
-    args: Array | tuple[Array, ...],
     /,
     *,
     xp: ModuleType | None = None,
@@ -50,8 +50,8 @@ def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ig
 @overload
 def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ignore=GL08
     cond: Array,
-    f1: Callable[..., Array],
     args: Array | tuple[Array, ...],
+    f1: Callable[..., Array],
     /,
     *,
     fill_value: Array | int | float | complex | bool,
@@ -61,9 +61,9 @@ def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ig
 
 def apply_where(  # type: ignore[no-any-explicit] # numpydoc ignore=PR01,PR02
     cond: Array,
+    args: Array | tuple[Array, ...],
     f1: Callable[..., Array],
-    f2: Callable[..., Array] | Array | tuple[Array],  # optional positional argument
-    args: Array | tuple[Array, ...] | None = None,
+    f2: Callable[..., Array] | None = None,
     /,
     *,
     fill_value: Array | int | float | complex | bool | None = None,
@@ -79,15 +79,15 @@ def apply_where(  # type: ignore[no-any-explicit] # numpydoc ignore=PR01,PR02
     ----------
     cond : array
         The condition, expressed as a boolean array.
+    args : Array or tuple of Arrays
+        Argument(s) to `f1` (and `f2`). Must be broadcastable with `cond`.
     f1 : callable
         Elementwise function of `args`, returning a single array.
         Where `cond` is True, output will be ``f1(arg0[cond], arg1[cond], ...)``.
     f2 : callable, optional
         Elementwise function of `args`, returning a single array.
         Where `cond` is False, output will be ``f2(arg0[cond], arg1[cond], ...)``.
         Mutually exclusive with `fill_value`.
-    args : Array or tuple of Arrays
-        Argument(s) to `f1` (and `f2`). Must be broadcastable with `cond`.
     fill_value : Array or scalar, optional
         If provided, value with which to fill output array where `cond` is False.
         It does not need to be scalar; it needs however to be broadcastable with
@@ -119,25 +119,15 @@ def apply_where(  # type: ignore[no-any-explicit] # numpydoc ignore=PR01,PR02
     >>> b = xp.asarray([0, 2, 2])
     >>> def f(a, b):
     ...     return a // b
-    >>> apply_where(b != 0, f, (a, b), fill_value=xp.nan)
+    >>> apply_where(b != 0, (a, b), f, fill_value=xp.nan)
     array([ nan,  2., 1.])
     """
     # Parse and normalize arguments
-    mutually_exc_msg = "Exactly one of `fill_value` or `f2` must be given."
-    if args is None:
-        f2, args = None, f2
-        if fill_value is None:
-            raise TypeError(mutually_exc_msg)
-    else:
-        if not callable(f2):
-            msg = "Third parameter must be a callable, Array, or tuple of Arrays."
-            raise TypeError(msg)
-        if fill_value is not None:
-            raise TypeError(mutually_exc_msg)
-
+    if (f2 is None) == (fill_value is None):
+        msg = "Exactly one of `fill_value` or `f2` must be given."
+        raise TypeError(msg)
     if not isinstance(args, tuple):
         args = (args,)
-    f2 = cast(Callable[..., Array] | None, f2)  # type: ignore[no-any-explicit]
     args = cast(tuple[Array, ...], args)
 
     xp = array_namespace(cond, *args) if xp is None else xp
@@ -547,10 +537,10 @@ def isclose(
         mxp = meta_namespace(a, b, xp=xp)
         out = apply_where(
             xp.isinf(a) | xp.isinf(b),
+            (a, b),
             lambda a, b: mxp.isinf(a) & mxp.isinf(b) & (mxp.sign(a) == mxp.sign(b)),  # pyright: ignore[reportUnknownArgumentType]
             # Note: inf <= inf is True!
             lambda a, b: mxp.abs(a - b) <= (atol + rtol * mxp.abs(b)),  # pyright: ignore[reportUnknownArgumentType]
-            (a, b),
             xp=xp,
         )
         if equal_nan:

tests/test_funcs.py

@@ -1,7 +1,6 @@
 import contextlib
 import math
 import warnings
-from collections.abc import Callable
 from types import ModuleType
 
 import hypothesis
@@ -36,6 +35,7 @@
 # some xp backends are untyped
 # mypy: disable-error-code=no-untyped-def
 
+lazy_xp_function(apply_where, static_argnums=(2, 3), static_argnames="xp")
 lazy_xp_function(atleast_nd, static_argnames=("ndim", "xp"))
 lazy_xp_function(cov, static_argnames="xp")
 # FIXME .device attribute https://github.com/data-apis/array-api-compat/pull/238
@@ -50,29 +50,6 @@
 lazy_xp_function(sinc, jax_jit=False, static_argnames="xp")
 
 
-def apply_where_jit(  # type: ignore[no-any-explicit]
-    cond: Array,
-    f1: Callable[..., Array],
-    f2: Callable[..., Array] | None,
-    args: Array | tuple[Array, ...],
-    fill_value: Array | int | float | complex | bool | None = None,
-    xp: ModuleType | None = None,
-) -> Array:
-    """
-    Work around jax.jit's inability to handle variadic positional arguments.
-
-    This is a lazy_xp_function artefact for when jax.jit is applied directly
-    to apply_where, which would not happen in real life.
-    """
-    if f2 is None:
-        return apply_where(cond, f1, args, fill_value=fill_value, xp=xp)
-    assert fill_value is None
-    return apply_where(cond, f1, f2, args, xp=xp)
-
-
-lazy_xp_function(apply_where_jit, static_argnames=("f1", "f2", "xp"))
-
-
 class TestApplyWhere:
     @staticmethod
     def f1(x: Array, y: Array | int = 10) -> Array:
@@ -86,27 +63,27 @@ def f2(x: Array, y: Array | int = 10) -> Array:
     def test_f1_f2(self, xp: ModuleType):
         x = xp.asarray([1, 2, 3, 4])
         cond = x % 2 == 0
-        actual = apply_where_jit(cond, self.f1, self.f2, x)
+        actual = apply_where(cond, x, self.f1, self.f2)
         expect = xp.where(cond, self.f1(x), self.f2(x))
         xp_assert_equal(actual, expect)
 
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
     def test_fill_value(self, xp: ModuleType):
         x = xp.asarray([1, 2, 3, 4])
         cond = x % 2 == 0
-        actual = apply_where_jit(x % 2 == 0, self.f1, None, x, fill_value=0)
+        actual = apply_where(x % 2 == 0, x, self.f1, fill_value=0)
         expect = xp.where(cond, self.f1(x), xp.asarray(0))
         xp_assert_equal(actual, expect)
 
-        actual = apply_where_jit(x % 2 == 0, self.f1, None, x, fill_value=xp.asarray(0))
+        actual = apply_where(x % 2 == 0, x, self.f1, fill_value=xp.asarray(0))
         xp_assert_equal(actual, expect)
 
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
     def test_args_tuple(self, xp: ModuleType):
         x = xp.asarray([1, 2, 3, 4])
         y = xp.asarray([10, 20, 30, 40])
         cond = x % 2 == 0
-        actual = apply_where_jit(cond, self.f1, self.f2, (x, y))
+        actual = apply_where(cond, (x, y), self.f1, self.f2)
         expect = xp.where(cond, self.f1(x, y), self.f2(x, y))
         xp_assert_equal(actual, expect)
 
@@ -116,21 +93,21 @@ def test_broadcast(self, xp: ModuleType):
         y = xp.asarray([[10], [20], [30]])
         cond = xp.broadcast_to(xp.asarray(True), (4, 1, 1))
 
-        actual = apply_where_jit(cond, self.f1, self.f2, (x, y))
+        actual = apply_where(cond, (x, y), self.f1, self.f2)
         expect = xp.where(cond, self.f1(x, y), self.f2(x, y))
         xp_assert_equal(actual, expect)
 
-        actual = apply_where_jit(
+        actual = apply_where(
             cond,
+            (x, y),
             lambda x, _: x,  # pyright: ignore[reportUnknownArgumentType]
             lambda _, y: y,  # pyright: ignore[reportUnknownArgumentType]
-            (x, y),
         )
         expect = xp.where(cond, x, y)
         xp_assert_equal(actual, expect)
 
         # Shaped fill_value
-        actual = apply_where_jit(cond, self.f1, None, x, fill_value=y)
+        actual = apply_where(cond, x, self.f1, fill_value=y)
         expect = xp.where(cond, self.f1(x), y)
         xp_assert_equal(actual, expect)
 
@@ -141,15 +118,15 @@ def test_dtype_propagation(self, xp: ModuleType, library: Backend):
         cond = x % 2 == 0
 
         mxp = np if library is Backend.DASK else xp
-        actual = apply_where_jit(
+        actual = apply_where(
             cond,
+            (x, y),
             self.f1,
             lambda x, y: mxp.astype(x - y, xp.int64),  # pyright: ignore[reportUnknownArgumentType]
-            (x, y),
         )
         assert actual.dtype == xp.int64
 
-        actual = apply_where_jit(cond, self.f1, None, y, fill_value=5)
+        actual = apply_where(cond, y, self.f1, fill_value=5)
         assert actual.dtype == xp.int16
 
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
@@ -168,14 +145,14 @@ def test_dtype_propagation_fill_value(
         cond = x % 2 == 0
         fill_value = xp.asarray(fill_value_raw, dtype=getattr(xp, fill_value_dtype))
 
-        actual = apply_where_jit(cond, self.f1, None, x, fill_value=fill_value)
+        actual = apply_where(cond, x, self.f1, fill_value=fill_value)
         assert actual.dtype == getattr(xp, expect_dtype)
 
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
     def test_dont_overwrite_fill_value(self, xp: ModuleType):
         x = xp.asarray([1, 2])
         fill_value = xp.asarray([100, 200])
-        actual = apply_where_jit(x % 2 == 0, self.f1, None, x, fill_value=fill_value)
+        actual = apply_where(x % 2 == 0, x, self.f1, fill_value=fill_value)
         xp_assert_equal(actual, xp.asarray([100, 12]))
         xp_assert_equal(fill_value, xp.asarray([100, 200]))
 
@@ -184,11 +161,11 @@ def test_dont_run_on_false(self, xp: ModuleType):
         x = xp.asarray([1.0, 2.0, 0.0])
         y = xp.asarray([0.0, 3.0, 4.0])
         # On NumPy, division by zero will trigger warnings
-        actual = apply_where_jit(
+        actual = apply_where(
             x == 0,
+            (x, y),
             lambda x, y: x / y,  # pyright: ignore[reportUnknownArgumentType]
             lambda x, y: y / x,  # pyright: ignore[reportUnknownArgumentType]
-            (x, y),
         )
         xp_assert_equal(actual, xp.asarray([0.0, 1.5, 0.0]))
 
@@ -197,29 +174,28 @@ def test_bad_args(self, xp: ModuleType):
         cond = x % 2 == 0
         # Neither f2 nor fill_value
         with pytest.raises(TypeError, match="Exactly one of"):
-            apply_where(cond, self.f1, x)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
+            apply_where(cond, x, self.f1)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
         # Both f2 and fill_value
         with pytest.raises(TypeError, match="Exactly one of"):
-            apply_where(cond, self.f1, self.f2, x, fill_value=0)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
-        # Multiple args; forgot to wrap them in a tuple
-        with pytest.raises(TypeError, match="takes from 3 to 4 positional arguments"):
-            apply_where(cond, self.f1, self.f2, x, x)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
-        with pytest.raises(TypeError, match="callable"):
-            apply_where(cond, self.f1, x, x, fill_value=0)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
+            apply_where(cond, x, self.f1, self.f2, fill_value=0)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
 
     @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp")
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
     def test_xp(self, xp: ModuleType):
         x = xp.asarray([1, 2, 3, 4])
         cond = x % 2 == 0
-        actual = apply_where_jit(cond, self.f1, self.f2, x, xp=xp)
+        actual = apply_where(cond, x, self.f1, self.f2, xp=xp)
         expect = xp.where(cond, self.f1(x), self.f2(x))
         xp_assert_equal(actual, expect)
 
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
     def test_device(self, xp: ModuleType, device: Device):
         x = xp.asarray([1, 2, 3, 4], device=device)
-        y = apply_where_jit(x % 2 == 0, self.f1, self.f2, x)
+        y = apply_where(x % 2 == 0, x, self.f1, self.f2)
+        assert get_device(y) == device
+        y = apply_where(x % 2 == 0, x, self.f1, fill_value=0)
+        assert get_device(y) == device
+        y = apply_where(x % 2 == 0, x, self.f1, fill_value=x)
         assert get_device(y) == device
 
     # skip instead of xfail in order not to waste time
@@ -273,10 +249,9 @@ def f2(*args: Array) -> Array:
         rng = np.random.default_rng(rng_seed)
         cond = xp.asarray(rng.random(size=cond_shape) > p)
 
-        # Use apply_where instead of apply_where_jit to speed the test up
-        res1 = apply_where(cond, f1, arrays, fill_value=fill_value)
-        res2 = apply_where(cond, f1, f2, arrays)
-        res3 = apply_where(cond, f1, arrays, fill_value=float_fill_value)
+        res1 = apply_where(cond, arrays, f1, fill_value=fill_value)
+        res2 = apply_where(cond, arrays, f1, f2)
+        res3 = apply_where(cond, arrays, f1, fill_value=float_fill_value)
 
         ref1 = xp.where(cond, f1(*arrays), fill_value)
         ref2 = xp.where(cond, f1(*arrays), f2(*arrays))