Function elemwise

dask/array/core.py:5089–5211

Apply an elementwise ufunc-like function blockwise across arguments. Like numpy ufuncs, broadcasting rules are respected. Parameters ---------- op : callable The function to apply. Should be numpy ufunc-like in the parameters that it accepts. *args : Any

(op, *args, out=None, where=True, dtype=None, name=None, **kwargs)

def elemwise(op, *args, out=None, where=True, dtype=None, name=None, **kwargs):
    """Apply an elementwise ufunc-like function blockwise across arguments.

    Like numpy ufuncs, broadcasting rules are respected.

    Parameters
    ----------
    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
    --------
    >>> elemwise(add, x, y)  # doctest: +SKIP
    >>> elemwise(sin, x)  # doctest: +SKIP
    >>> elemwise(sin, x, out=dask_array)  # doctest: +SKIP

    See Also
    --------
    blockwise
    """
    if kwargs:
        raise TypeError(
            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 = []
    for arg in args:
        shape = getattr(arg, "shape", ())
        if any(is_dask_collection(x) for x in shape):
            # Want to exclude Delayed shapes and dd.Scalar
            shape = ()
        shapes.append(shape)
    if isinstance(where, Array):