| 1231 | |
| 1232 | @requires_dask |
| 1233 | def test_apply_dask_parallelized_two_args() -> None: |
| 1234 | import dask.array as da |
| 1235 | |
| 1236 | array = da.ones((2, 2), chunks=(1, 1), dtype=np.int64) |
| 1237 | data_array = xr.DataArray(array, dims=("x", "y")) |
| 1238 | data_array.name = None |
| 1239 | |
| 1240 | def parallel_add(x, y): |
| 1241 | return apply_ufunc( |
| 1242 | operator.add, x, y, dask="parallelized", output_dtypes=[np.int64] |
| 1243 | ) |
| 1244 | |
| 1245 | def check(x, y): |
| 1246 | actual = parallel_add(x, y) |
| 1247 | assert isinstance(actual.data, da.Array) |
| 1248 | assert actual.data.chunks == array.chunks |
| 1249 | assert_identical(data_array, actual) |
| 1250 | |
| 1251 | check(data_array, 0) |
| 1252 | check(0, data_array) |
| 1253 | check(data_array, xr.DataArray(0)) |
| 1254 | check(data_array, 0 * data_array) |
| 1255 | check(data_array, 0 * data_array[0]) |
| 1256 | check(data_array[:, 0], 0 * data_array[0]) |
| 1257 | check(data_array, 0 * data_array.compute()) |
| 1258 | |
| 1259 | |
| 1260 | @requires_dask |