| 507 | @pytest.mark.parametrize("right_is_da", [False, True]) |
| 508 | @pytest.mark.parametrize("where_kind", [True, False, "numpy", "dask"]) |
| 509 | def test_ufunc_where(dtype, left_is_da, right_is_da, where_kind): |
| 510 | left = np.arange(12).reshape((3, 4)) |
| 511 | right = np.arange(4) |
| 512 | out = np.zeros_like(left, dtype=dtype) |
| 513 | d_out = da.zeros_like(left, dtype=dtype) |
| 514 | |
| 515 | if where_kind in (True, False): |
| 516 | d_where = where = where_kind |
| 517 | else: |
| 518 | d_where = where = np.array([False, True, True, False]) |
| 519 | if where_kind == "dask": |
| 520 | d_where = da.from_array(where, chunks=2) |
| 521 | |
| 522 | d_left = da.from_array(left, chunks=2) if left_is_da else left |
| 523 | d_right = da.from_array(right, chunks=2) if right_is_da else right |
| 524 | |
| 525 | expected = np.add(left, right, where=where, out=out, dtype=dtype) |
| 526 | result = da.add(d_left, d_right, where=d_where, out=d_out, dtype=dtype) |
| 527 | assert result is d_out |
| 528 | assert_eq(expected, result) |
| 529 | |
| 530 | |
| 531 | @pytest.mark.parametrize("left_is_da", [False, True]) |