Dask support

Like xarray, xarray-einstats also aims to support using their functions on NumPy or Dask backed xarray objects.

Functions in xarray-einstats should support Dask backed xarray objects. But Dask support is still not extensively tested. If you have issues using Dask backed xarray objects please open an issue

Dask support within xarray-einstats requires explicit activation by users, exactly like it happens with xarray.apply_ufunc. In fact, what is happening is xarray-einstats forwards any unrecognized keyword argument (e.g. dask or dask_gufunc_kwargs) to apply_ufunc.

import dask.array as da
import xarray as xr
from xarray_einstats import linalg

xr.set_options(display_expand_data=False);

matrix_batch = xr.DataArray(da.random.normal(size=(1000, 77, 77), chunks="auto"), dims=["batch", "dim1", "dim2"])

Trying to use xarray-einstats functions directly on Dask arrays triggers an error:

linalg.norm(matrix_batch, dims=["dim1", "dim2"])

Hide code cell output

 
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[3], line 1
----> 1 linalg.norm(matrix_batch, dims=["dim1", "dim2"])

File ~/Documents/repos_oss/xarray-einstats/src/xarray_einstats/linalg.py:801, in norm(da, dims, ord, **kwargs)
    799     in_dims = dims
    800     norm_kwargs["axis"] = (-2, -1)
--> 801 return xr.apply_ufunc(
    802     np.linalg.norm, da, input_core_dims=[in_dims], kwargs=norm_kwargs, **kwargs
    803 )

File ~/bin/miniforge3/envs/general/lib/python3.14/site-packages/xarray/computation/apply_ufunc.py:1267, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, on_missing_core_dim, *args)
   1265 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
   1266 elif any(isinstance(a, DataArray) for a in args):
-> 1267     return apply_dataarray_vfunc(
   1268         variables_vfunc,
   1269         *args,
   1270         signature=signature,
   1271         join=join,
   1272         exclude_dims=exclude_dims,
   1273         keep_attrs=keep_attrs,
   1274     )
   1275 # feed Variables directly through apply_variable_ufunc
   1276 elif any(isinstance(a, Variable) for a in args):

File ~/bin/miniforge3/envs/general/lib/python3.14/site-packages/xarray/computation/apply_ufunc.py:312, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
    307 result_coords, result_indexes = build_output_coords_and_indexes(
    308     args, signature, exclude_dims, combine_attrs=keep_attrs
    309 )
    311 data_vars = [getattr(a, "variable", a) for a in args]
--> 312 result_var = func(*data_vars)
    314 out: tuple[DataArray, ...] | DataArray
    315 if signature.num_outputs > 1:

File ~/bin/miniforge3/envs/general/lib/python3.14/site-packages/xarray/computation/apply_ufunc.py:741, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
    739 if any(is_chunked_array(array) for array in input_data):
    740     if dask == "forbidden":
--> 741         raise ValueError(
    742             "apply_ufunc encountered a chunked array on an "
    743             "argument, but handling for chunked arrays has not "
    744             "been enabled. Either set the ``dask`` argument "
    745             "or load your data into memory first with "
    746             "``.load()`` or ``.compute()``"
    747         )
    748     elif dask == "parallelized":
    749         chunkmanager = get_chunked_array_type(*input_data)

ValueError: apply_ufunc encountered a chunked array on an argument, but handling for chunked arrays has not been enabled. Either set the ``dask`` argument or load your data into memory first with ``.load()`` or ``.compute()``

We need to allow dask inputs through dask="allowed" or dask="parallelized". In general, allowed is preferred, but it only works if there is a Dask implementation of the requested function available, parallelized works in most cases but is usually slower and still not a silver bullet. Consequently, xarray-einstats does not set the dask argument automatically when the input is a dask array, and finding which of the two will work and be better might require a bit of trial an error an also be dependent on the Dask version available.

Here is one example with linalg.norm for which both alternatives work:

norms = linalg.norm(matrix_batch, dims=["dim1", "dim2"], dask="allowed")
norms.data
Array Chunk
Bytes 7.81 kiB 7.81 kiB
Shape (1000,) (1000,)
Dask graph 1 chunks in 6 graph layers
Data type float64 numpy.ndarray
1000 1 
norms.compute()

<xarray.DataArray 'normal-758a57bb75ed3127c94212c3e83a5ed8' (batch: 1000)> Size: 8kB
77.07 77.61 76.86 77.53 77.99 77.22 76.9 ... 77.92 78.03 77.96 76.6 77.31 76.26
Dimensions without coordinates: batch
linalg.norm(matrix_batch, dims=["dim1", "dim2"], dask="parallelized").compute()

<xarray.DataArray 'normal-758a57bb75ed3127c94212c3e83a5ed8' (batch: 1000)> Size: 8kB
77.07 77.61 76.86 77.53 77.99 77.22 76.9 ... 77.92 78.03 77.96 76.6 77.31 76.26
Dimensions without coordinates: batch

See also

Xarray documentation

Dask examples

%load_ext watermark
%watermark -n -u -v -iv -w

Last updated: Mon, 20 Jul 2026

Python implementation: CPython
Python version       : 3.14.6
IPython version      : 9.15.0

dask           : 2026.7.1
xarray         : 2026.7.0
xarray_einstats: 0.11.0

Watermark: 2.6.0