Function safe_sqrt

dask/array/reductions.py:699–710

A version of sqrt that properly handles scalar masked arrays. To mimic ``np.ma`` reductions, we need to convert scalar masked arrays that have an active mask to the ``np.ma.masked`` singleton. This is properly handled automatically for reduction code, but not for ufuncs. We implement

(a)

def safe_sqrt(a):
    """A version of sqrt that properly handles scalar masked arrays.

    To mimic ``np.ma`` reductions, we need to convert scalar masked arrays that
    have an active mask to the ``np.ma.masked`` singleton. This is properly
    handled automatically for reduction code, but not for ufuncs. We implement
    a simple version here, since calling `np.ma.sqrt` everywhere is
    significantly more expensive.
    """
    if hasattr(a, "_elemwise"):
        return a._elemwise(_sqrt, a)
    return _sqrt(a)


@derived_from(np)