Repository: dask/dask
File: dask/array/tests/test_reductions.py (Python)

Code before:
    x = da.reduction(dx, w_sum, np.sum, dtype=dx.dtype)
    assert_eq(x, np.sum(a), check_shape=True)

    # Weighted sum
    x = da.reduction(dx, w_sum, np.sum, dtype="f8", weights=w)
    assert_eq(x, np.sum(a * w), check_shape=True)

    # Non-broadcastable weights (short axis)
    with pytest.raises(ValueError):
        da.reduction(dx, w_sum, np.sum, weights=[1, 2, 3])

    # Non-broadcastable weights (too many dims)
    with pytest.raises(ValueError):
        da.reduction(dx, w_sum, np.sum, weights=[[[2]]])


@pytest.mark.parametrize("axis", [3, 0, [1, 3]])
@pytest.mark.parametrize("q", [0.75, [0.75], [0.75, 0.4]])
@pytest.mark.parametrize("rechunk", [True, False])
def test_nanquantile(rechunk, q, axis):
    shape = 10, 15, 10, 15
    arr = np.random.randn(*shape)
    indexer = np.random.randint(0, 10, size=shape)
    arr[indexer >= 8] = np.nan

    darr = da.from_array(arr, chunks=(2, 3, 4, (-1 if not rechunk else 5)))
    assert_eq(da.nanquantile(darr, q, axis=axis), np.nanquantile(arr, q, axis=axis))
    assert_eq(
        da.nanquantile(darr, q, axis=axis, keepdims=True),
        np.nanquantile(arr, q, axis=axis, keepdims=True),
    )


@pytest.mark.parametrize("axis", [3, [1, 3]])
@pytest.mark.parametrize("q", [0.75, [0.75]])
@pytest.mark.parametrize("rechunk", [True, False])
def test_quantile(rechunk, q, axis):
    shape = 10, 15, 20, 15
    arr = np.random.randn(*shape)
    indexer = np.random.randint(0, 10, size=shape)
    arr[indexer >= 8] = np.nan

    darr = da.from_array(arr, chunks=(2, 3, 4, (-1 if not rechunk else 5)))
    assert_eq(da.quantile(darr, q, axis=axis), np.quantile(arr, q, axis=axis))
    assert_eq(
        da.quantile(darr, q, axis=axis, keepdims=True),
        np.quantile(arr, q, axis=axis, keepdims=True),
    )


def test_nanquantile_all_nan():

Review comment (suggestion):
nit:
```suggestion
    darr = da.from_array(arr, chunks=(2, 3, 4, (5 if rechunk else -1)))
```

Code after:
    x = da.reduction(dx, w_sum, np.sum, dtype=dx.dtype)
    assert_eq(x, np.sum(a), check_shape=True)

    # Weighted sum
    x = da.reduction(dx, w_sum, np.sum, dtype="f8", weights=w)
    assert_eq(x, np.sum(a * w), check_shape=True)

    # Non-broadcastable weights (short axis)
    with pytest.raises(ValueError):
        da.reduction(dx, w_sum, np.sum, weights=[1, 2, 3])

    # Non-broadcastable weights (too many dims)
    with pytest.raises(ValueError):
        da.reduction(dx, w_sum, np.sum, weights=[[[2]]])


@pytest.mark.parametrize("axis", [3, 0, [1, 3]])
@pytest.mark.parametrize("q", [0.75, [0.75], [0.75, 0.4]])
@pytest.mark.parametrize("rechunk", [True, False])
def test_nanquantile(rechunk, q, axis):
    shape = 10, 15, 10, 15
    arr = np.random.randn(*shape)
    indexer = np.random.randint(0, 10, size=shape)
    arr[indexer >= 8] = np.nan

    darr = da.from_array(arr, chunks=(2, 3, 4, (5 if rechunk else -1)))
    assert_eq(da.nanquantile(darr, q, axis=axis), np.nanquantile(arr, q, axis=axis))
    assert_eq(
        da.nanquantile(darr, q, axis=axis, keepdims=True),
        np.nanquantile(arr, q, axis=axis, keepdims=True),
    )


@pytest.mark.parametrize("axis", [3, [1, 3]])
@pytest.mark.parametrize("q", [0.75, [0.75]])
@pytest.mark.parametrize("rechunk", [True, False])
def test_quantile(rechunk, q, axis):
    shape = 10, 15, 20, 15
    arr = np.random.randn(*shape)
    indexer = np.random.randint(0, 10, size=shape)
    arr[indexer >= 8] = np.nan

    darr = da.from_array(arr, chunks=(2, 3, 4, (5 if rechunk else -1)))
    assert_eq(da.quantile(darr, q, axis=axis), np.quantile(arr, q, axis=axis))
    assert_eq(
        da.quantile(darr, q, axis=axis, keepdims=True),
        np.quantile(arr, q, axis=axis, keepdims=True),
    )


def test_nanquantile_all_nan():