Repository: dask/dask

Issue:
Bitwise operators mask (where parameter)
It seems that dask does not implement the `where` parameter in the bitwise operations but numpy does. It would be great to implement those to increase compatibility of the current function.

Here is the documentation in number for invert (bitwise_not)
https://numpy.org/doc/stable/reference/generated/numpy.invert.html#numpy.invert

An example in dask below would produce an error. Substituting numpy with dask would work.

```
import dask.array as da
x = da.zeros((10)).astype(np.uint8)
mask = da.ones((10)).astype(bool)
x = da.bitwise_not(x, where=mask)
```

Discussion:
GPUtester: Can one of the admins verify this patch?
quasiben: add to allowlist
ncclementi: @ek234 I see the CI hasn't run all the tests because there are conflicts. I believe that these conflicts are causing some checks not passing, notice that one says "AppVeyor was unable to build non-mergeable pull request" and the gpuCI is probably failing for the same reason. 
You might want to merge main into your branch, fix the conflicts, and push again. That way the full CI suite will run and if we get any fails we will be able to assist you better. 
ek234: Thanks for the review
I did as you said and then went through the readthedocs errors

Let me know if you would like me to make any changes.
ek234: bruh it's all green now:P
ncclementi: @ek234 thank you for your patience, I'm pinging another maintainer to see if they can review the code. 

@jrbourbeau Would you mind reviewing this PR? We are at a point that we need a maintainer to take a look. I've suggested some changes, but I think we need a deeper look.
jcrist: Thanks for all your work here @ek234, and apologies for the delayed review. This issue turned out to be a bit more complicated than a good-first-issue (apologies for the mislabeling). Since `where` interacts with `out` (which mutates the passed array), and dask arrays are effectively immutable, this required some nuanced changes to how `elemwise` works internally. I've pushed up a commit to your branch that makes the required changes on top of your initial PR. Thanks for your work here, hope to get this merged soon.
douglasdavis: Looks like this PR may take care of what https://github.com/dask/dask/pull/8322 is trying to accomplish
jcrist: Ah, that it does. I've added it to the top-level comment to auto-close.
ian-r-rose: @jcrist Is there anything left to be done here? Looks like CI is now passing, though there are now merge conflicts to be resolved.
jcrist: Thanks for the ping, I'd forgotten about this. Merge conflicts fixed, should be good to merge once tests pass.

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index aeba5aa62..18a49af87 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -4374,16 +4374,35 @@ def broadcast_shapes(*shapes):
     return tuple(reversed(out))
 
 
-def elemwise(op, *args, out=None, **kwargs):
-    """Apply elementwise function across arguments
+def elemwise(op, *args, out=None, where=True, dtype=None, name=None, **kwargs):
+    """Apply an elementwise ufunc-like function blockwise across arguments.
 
-    Respects broadcasting rules
+    Like numpy ufuncs, broadcasting rules are respected.
 
     Parameters
     ----------
-    out : dask array or None
-        If out is a dask.array then this overwrites the contents of that array with
-        the result
+    op : callable
+        The function to apply. Should be numpy ufunc-like in the parameters
+        that it accepts.
+    *args : Any
+        Arguments to pass to `op`. Non-dask array-like objects are first
+        converted to dask arrays, then all arrays are broadcast together before
+        applying the function blockwise across all arguments. Any scalar
+        arguments are passed as-is following normal numpy ufunc behavior.
+    out : dask array, optional
+        If out is a dask.array then this overwrites the contents of that array
+        with the result.
+    where : array_like, optional
+        An optional boolean mask marking locations where the ufunc should be
+        applied. Can be a scalar, dask array, or any other array-like object.
+        Mirrors the ``where`` argument to numpy ufuncs, see e.g. ``numpy.add``
+        for more information.
+    dtype : dtype, optional
+        If provided, overrides the output array dtype.
+    name : str, optional
+        A unique key name to use when building the backing dask graph. If not
+        provided, one will be automatically generated based on the input
+        arguments.
 
     Examples
     --------
@@ -4395,12 +4414,14 @@ def elemwise(op, *args, out=None, **kwargs):
     --------
     blockwise
     """
-    if not {"name", "dtype"}.issuperset(kwargs):
-        msg = "%s does not take the following keyword arguments %s"
+    if kwargs:
         raise TypeError(
-            msg % (op.__name__, str(sorted(set(kwargs) - {"name", "dtype"})))
+            f"{op.__name__} does not take the following keyword arguments "
+            f"{sorted(kwargs)}"
         )
 
+    out = _elemwise_normalize_out(out)
+    where = _elemwise_normalize_where(where)
     args = [np.asarray(a) if isinstance(a, (list, tuple)) else a for a in args]
 
     shapes = []
@@ -4410,6 +4431,10 @@ def elemwise(op, *args, out=None, **kwargs):
             # Want to exclude Delayed shapes and dd.Scalar
             shape = ()
         shapes.append(shape)
+    if isinstance(where, Array):
+        shapes.append(where.shape)
+    if isinstance(out, Array):
+        shapes.append(out.shape)
 
     shapes = [s if isinstance(s, Iterable) else () for s in shapes]
     out_ndim = len(
@@ -4417,10 +4442,8 @@ def elemwise(op, *args, out=None, **kwargs):
     )  # Raises ValueError if dimensions mismatch
     expr_inds = tuple(range(out_ndim))[::-1]
 
-    need_enforce_dtype = False
-    if "dtype" in kwargs:
+    if dtype is not None:
         need_enforce_dtype = True
-        dt = kwargs["dtype"]
     else:
         # We follow NumPy's rules for dtype promotion, which special cases
         # scalars and 0d ndarrays (which it considers equivalent) by using
@@ -4437,20 +4460,28 @@ def elemwise(op, *args, out=None, **kwargs):
             for a in args
         ]
         try:
-            dt = apply_infer_dtype(op, vals, {}, "elemwise", suggest_dtype=False)
+            dtype = apply_infer_dtype(op, vals, {}, "elemwise", suggest_dtype=False)
         except Exception:
             return NotImplemented
         need_enforce_dtype = any(
             not is_scalar_for_elemwise(a) and a.ndim == 0 for a in args
         )
 
-    name = kwargs.get("name", None) or f"{funcname(op)}-{tokenize(op, dt, *args)}"
+    if not name:
+        name = f"{funcname(op)}-{tokenize(op, dtype, *args, where)}"
+
+    blockwise_kwargs = dict(dtype=dtype, name=name, token=funcname(op).strip("_"))
+
+    if where is not True:
+        blockwise_kwargs["elemwise_where_function"] = op
+        op = _elemwise_handle_where
+        args.extend([where, out])
 
-    blockwise_kwargs = dict(dtype=dt, name=name, token=funcname(op).strip("_"))
     if need_enforce_dtype:
-        blockwise_kwargs["enforce_dtype"] = dt
+        blockwise_kwargs["enforce_dtype"] = dtype
         blockwise_kwargs["enforce_dtype_function"] = op
         op = _enforce_dtype
+
     result = blockwise(
         op,
         expr_inds,
@@ -4464,12 +4495,23 @@ def elemwise(op, *args, out=None, **kwargs):
     return handle_out(out, result)
 
 
-def handle_out(out, result):
-    """Handle out parameters
+def _elemwise_normalize_where(where):
+    if where is True:
+        return True
+    elif where is False or where is None:
+        return False
+    return asarray(where)
 
-    If out is a dask.array then this overwrites the contents of that array with
-    the result
-    """
+
+def _elemwise_handle_where(*args, **kwargs):
+    function = kwargs.pop("elemwise_where_function")
+    *args, where, out = args
+    if hasattr(out, "copy"):
+        out = out.copy()
+    return function(*args, where=where, out=out, **kwargs)
+
+
+def _elemwise_normalize_out(out):
     if isinstance(out, tuple):
         if len(out) == 1:
             out = out[0]
@@ -4477,6 +4519,21 @@ def handle_out(out, result):
             raise NotImplementedError("The out parameter is not fully supported")
         else:
             out = None
+    if not (out is None or isinstance(out, Array)):
+        raise NotImplementedError(
+            f"The out parameter is not fully supported."
+            f" Received type {type(out).__name__}, expected Dask Array"
+        )
+    return out
+
+
+def handle_out(out, result):
+    """Handle out parameters
+
+    If out is a dask.array then this overwrites the contents of that array with
+    the result
+    """
+    out = _elemwise_normalize_out(out)
     if isinstance(out, Array):
         if out.shape != result.shape:
             raise ValueError(
@@ -4487,12 +4544,7 @@ def handle_out(out, result):
         out.dask = result.dask
         out._meta = result._meta
         out._name = result.name
-    elif out is not None:
-        msg = (
-            "The out parameter is not fully supported."
-            " Received type %s, expected Dask Array" % type(out).__name__
-        )
-        raise NotImplementedError(msg)
+        return out
     else:
         return result
 


Test patch:
diff --git a/dask/array/tests/test_ufunc.py b/dask/array/tests/test_ufunc.py
index 1c398d11e..2e52321e4 100644
--- a/dask/array/tests/test_ufunc.py
+++ b/dask/array/tests/test_ufunc.py
@@ -490,3 +490,90 @@ def test_dtype_kwarg(dt):
 
     result = da.add(darr1, darr2, dtype=dt)
     assert_eq(expected, result)
+
+
+@pytest.mark.parametrize("dtype", [None, "f8"])
+@pytest.mark.parametrize("left_is_da", [False, True])
+@pytest.mark.parametrize("right_is_da", [False, True])
+@pytest.mark.parametrize("where_kind", [True, False, "numpy", "dask"])
+def test_ufunc_where(dtype, left_is_da, right_is_da, where_kind):
+    left = np.arange(12).reshape((3, 4))
+    right = np.arange(4)
+    out = np.zeros_like(left, dtype=dtype)
+    d_out = da.zeros_like(left, dtype=dtype)
+
+    if where_kind in (True, False):
+        d_where = where = where_kind
+    else:
+        d_where = where = np.array([False, True, True, False])
+        if where_kind == "dask":
+            d_where = da.from_array(where, chunks=2)
+
+    d_left = da.from_array(left, chunks=2) if left_is_da else left
+    d_right = da.from_array(right, chunks=2) if right_is_da else right
+
+    expected = np.add(left, right, where=where, out=out, dtype=dtype)
+    result = da.add(d_left, d_right, where=d_where, out=d_out, dtype=dtype)
+    assert result is d_out
+    assert_eq(expected, result)
+
+
+@pytest.mark.parametrize("left_is_da", [False, True])
+@pytest.mark.parametrize("right_is_da", [False, True])
+@pytest.mark.parametrize("where_is_da", [False, True])
+def test_ufunc_where_broadcasts(left_is_da, right_is_da, where_is_da):
+    left = np.arange(4)
+    right = np.arange(4, 8)
+    where = np.array([[0, 1, 1, 0], [1, 0, 0, 1], [0, 1, 0, 1]]).astype("bool")
+    out = np.zeros(where.shape, dtype=left.dtype)
+
+    d_out = da.zeros(where.shape, dtype=left.dtype)
+    d_where = da.from_array(where, chunks=2) if where_is_da else where
+    d_left = da.from_array(left, chunks=2) if left_is_da else left
+    d_right = da.from_array(right, chunks=2) if right_is_da else right
+
+    expected = np.add(left, right, where=where, out=out)
+    result = da.add(d_left, d_right, where=d_where, out=d_out)
+    assert result is d_out
+    assert_eq(expected, result)
+
+
+def test_ufunc_where_no_out():
+    left = np.arange(4)
+    right = np.arange(4, 8)
+    where = np.array([[0, 1, 1, 0], [1, 0, 0, 1], [0, 1, 0, 1]]).astype("bool")
+
+    d_where = da.from_array(where, chunks=2)
+    d_left = da.from_array(left, chunks=2)
+    d_right = da.from_array(right, chunks=2)
+
+    expected = np.add(left, right, where=where)
+    result = da.add(d_left, d_right, where=d_where)
+
+    # If no `out` is provided, numpy leaves elements that don't match `where`
+    # uninitialized, so they effectively may be any random value.  We test that
+    # the set values match, and that the unset values aren't equal to if
+    # `where` wasn't provided (to test that `where` was actually passed).
+
+    expected_masked = np.where(where, expected, 0)
+    result_masked = np.where(where, expected, 0)
+    assert_eq(expected_masked, result_masked)
+
+    expected_no_where = np.add(left, right)
+    assert not np.equal(result.compute(), expected_no_where).all()
+
+
+def test_ufunc_where_doesnt_mutate_out():
+    """Dask array's are immutable, ensure that the backing numpy array for
+    `out` isn't actually mutated"""
+    left = da.from_array(np.arange(4, dtype="i8"), chunks=2)
+    right = da.from_array(np.arange(4, 8, dtype="i8"), chunks=2)
+    where = da.from_array(np.array([1, 0, 0, 1], dtype="bool"), chunks=2)
+    out_np = np.zeros(4, dtype="i8")
+    out = da.from_array(out_np, chunks=2)
+    result = da.add(left, right, where=where, out=out)
+    assert out is result
+    assert_eq(out, np.array([4, 0, 0, 10], dtype="i8"))
+
+    # Check that original `out` array isn't mutated
+    assert np.equal(out_np, 0).all()