TypeError: no implementation found for 'numpy.asarray' on types that implement __array_function__: [<class 'pint.Quantity'>]

Issue with flox and pint arrays. Using temporary workaround and disabling flox.

Code raises:

TypeError: no implementation found for 'numpy.asarray' on types that implement __array_function__: [<class 'pint.Quantity'>]

Workaround:

Disable flox, aggregation and compute works as expected.

xr.set_options(use_flox=False)

Package versions:

flox: 0.9.6
dask: 2024.3.1
pint: 0.23
xarray: 2024.3.0
pint-xarray: 0.3

Sample code block:

print(data)
<xarray.DataArray 'mint2m' (latitude: 721, longitude: 1440, valid_time: 120)> Size: 498MB
dask.array<xarray-<this-array>, shape=(721, 1440, 120), dtype=float32, chunksize=(362, 722, 120), chunktype=pint.Quantity>
Coordinates:
  * latitude    (latitude) float64 6kB 90.0 89.75 89.5 ... -89.5 -89.75 -90.0
  * longitude   (longitude) float64 12kB 0.0 0.25 0.5 0.75 ... 359.2 359.5 359.8
  * valid_time  (valid_time) datetime64[ns] 960B 2024-03-01 ... 2024-03-05T23...


da = data.groupby("valid_time.date").mean(dim="valid_time")
da.compute()

Traceback:

File /opt/conda/lib/python3.11/site-packages/flox/core.py:1006, in chunk_reduce(array, by, func, expected_groups, axis, fill_value, dtype, reindex, engine, kwargs, sort, user_dtype)
   1004     result = reduction(group_idx, array, **kw_func)
   1005 else:
-> 1006     result = generic_aggregate(
   1007         group_idx, array, axis=-1, engine=engine, func=reduction, **kw_func
   1008     ).astype(dt, copy=False)
   1009 if engine == "numbagg":
   1010     result = _postprocess_numbagg(
   1011         result,
   1012         func=reduction,
   (...)
   1017         seen_groups=seen_groups,
   1018     )

File /opt/conda/lib/python3.11/site-packages/flox/aggregations.py:99, in generic_aggregate(group_idx, array, engine, func, axis, size, fill_value, dtype, **kwargs)
     94 else:
     95     raise ValueError(
     96         f"Expected engine to be one of ['flox', 'numpy', 'numba', 'numbagg']. Received {engine} instead."
     97     )
---> 99 group_idx = np.asarray(group_idx, like=array)
    101 with warnings.catch_warnings():
    102     warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")

Start with the reported paths flox/core.py:1006 and flox/aggregations.py:99, then reproduce the groupby("valid_time.date").mean(dim="valid_time") example and call compute(). Trace the np.asarray(group_idx, like=array) path for the pint.Quantity chunks; done means the aggregation computes successfully without disabling flox.