Function test_vectorize_dask_new_output_dims

xarray/tests/test_computation.py:1868–1908

()

@requires_dask
def test_vectorize_dask_new_output_dims() -> None:
    # regression test for GH3574
    # run vectorization in dask.array.gufunc by using `dask='parallelized'`
    data_array = xr.DataArray([[0, 1, 2], [1, 2, 3]], dims=("x", "y"))
    func = lambda x: x[np.newaxis, ...]
    expected = data_array.expand_dims("z")
    actual = apply_ufunc(
        func,
        data_array.chunk({"x": 1}),
        output_core_dims=[["z"]],
        vectorize=True,
        dask="parallelized",
        output_dtypes=[float],
        dask_gufunc_kwargs=dict(output_sizes={"z": 1}),
    ).transpose(*expected.dims)
    assert_identical(expected, actual)

    with pytest.raises(
        ValueError, match=r"dimension 'z1' in 'output_sizes' must correspond"
    ):
        apply_ufunc(
            func,
            data_array.chunk({"x": 1}),
            output_core_dims=[["z"]],
            vectorize=True,
            dask="parallelized",
            output_dtypes=[float],
            dask_gufunc_kwargs=dict(output_sizes={"z1": 1}),
        )

    with pytest.raises(
        ValueError, match=r"dimension 'z' in 'output_core_dims' needs corresponding"
    ):
        apply_ufunc(
            func,
            data_array.chunk({"x": 1}),
            output_core_dims=[["z"]],
            vectorize=True,
            dask="parallelized",
            output_dtypes=[float],
        )


def test_output_wrong_number() -> None: