Skip to content

Commit

Permalink
Merge branch 'main' into upstream_deps
Browse files Browse the repository at this point in the history
  • Loading branch information
phofl committed Feb 7, 2024
2 parents ccab90b + 7e04ff7 commit 2d32579
Show file tree
Hide file tree
Showing 19 changed files with 615 additions and 470 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/update-gpuci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ jobs:
regex: false

- name: Create Pull Request
uses: peter-evans/create-pull-request@v5
uses: peter-evans/create-pull-request@v6
if: ${{ env.UCX_PY_VER != env.NEW_UCX_PY_VER }} # make sure new ucx-py nightlies are available
with:
token: ${{ secrets.GITHUB_TOKEN }}
Expand Down
2 changes: 1 addition & 1 deletion continuous_integration/environment-3.10.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ dependencies:
- pytest-xdist
- moto<5
# Optional dependencies
- mimesis<13.1.0 # https://github.com/dask/dask/issues/10858
- mimesis
- numpy=1.23
- pandas=1.5
- numba<0.59.0 # Removed numba.generated_jit which pandas<2.1 uses
Expand Down
2 changes: 1 addition & 1 deletion continuous_integration/environment-3.11.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ dependencies:
- pytest-xdist
- moto<5
# Optional dependencies
- mimesis<13.1.0 # https://github.com/dask/dask/issues/10858
- mimesis
- numpy
- pandas
- flask
Expand Down
2 changes: 1 addition & 1 deletion continuous_integration/environment-3.12.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ dependencies:
- pytest-xdist
- moto<5
# Optional dependencies
- mimesis<13.1.0 # https://github.com/dask/dask/issues/10858
- mimesis
- numpy
- pandas
- flask
Expand Down
25 changes: 10 additions & 15 deletions dask/array/routines.py
Original file line number Diff line number Diff line change
Expand Up @@ -2056,24 +2056,19 @@ def _isnonzero_vec(v):
_isnonzero_vec = np.vectorize(_isnonzero_vec, otypes=[bool])


def _isnonzero(a):
# Output of np.vectorize can't be pickled
return _isnonzero_vec(a)


def isnonzero(a):
if a.dtype.kind in {"U", "S"}:
# NumPy treats all-whitespace strings as falsy (like in `np.nonzero`).
# but not in `.astype(bool)`. To match the behavior of numpy at least until
# 1.19, we use `_isnonzero_vec`. When NumPy changes behavior, we should just
# use the try block below.
# https://github.com/numpy/numpy/issues/9875
return a.map_blocks(_isnonzero_vec, dtype=bool)
"""Handle special cases where conversion to bool does not work correctly.
xref: https://github.com/numpy/numpy/issues/9479
"""
try:
np.zeros(tuple(), dtype=a.dtype).astype(bool)
np.zeros([], dtype=a.dtype).astype(bool)
except ValueError:
######################################################
# Handle special cases where conversion to bool does #
# not work correctly. #
# #
# xref: https://github.com/numpy/numpy/issues/9479 #
######################################################
return a.map_blocks(_isnonzero_vec, dtype=bool)
return a.map_blocks(_isnonzero, dtype=bool)
else:
return a.astype(bool)

Expand Down
16 changes: 16 additions & 0 deletions dask/array/tests/test_routines.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import contextlib
import itertools
import pickle
import sys
import warnings
from numbers import Number
Expand Down Expand Up @@ -2708,3 +2709,18 @@ def test_tril_triu_indices(n, k, m, chunks):
)
else:
assert_eq(actual, expected)


def test_pickle_vectorized_routines():
"""Test that graphs that internally use np.vectorize can be pickled"""
a = da.from_array(["foo", "bar", ""])

b = da.count_nonzero(a)
assert_eq(b, 2, check_dtype=False)
b2 = pickle.loads(pickle.dumps(b))
assert_eq(b2, 2, check_dtype=False)

c = da.argwhere(a)
assert_eq(c, [[0], [1]], check_dtype=False)
c2 = pickle.loads(pickle.dumps(c))
assert_eq(c2, [[0], [1]], check_dtype=False)
4 changes: 2 additions & 2 deletions dask/array/ufunc.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from dask import core
from dask.array.core import Array, apply_infer_dtype, asarray, blockwise, elemwise
from dask.base import is_dask_collection, normalize_function
from dask.base import is_dask_collection, normalize_token
from dask.highlevelgraph import HighLevelGraph
from dask.utils import derived_from, funcname

Expand Down Expand Up @@ -42,7 +42,7 @@ def __repr__(self):
return "da.frompyfunc<%s, %d, %d>" % (self._name, self.nin, self.nout)

def __dask_tokenize__(self):
return (normalize_function(self._func), self.nin, self.nout)
return (normalize_token(self._func), self.nin, self.nout)

def __reduce__(self):
return (da_frompyfunc, (self._func, self.nin, self.nout))
Expand Down

0 comments on commit 2d32579

Please sign in to comment.