add dask function

Author: ferrine

pymc_experimental/tests/test_histogram_trick.py

@@ -8,10 +8,14 @@
 def test_histogram_init(use_dask):
     data = np.random.randn(10000)
     if use_dask:
+        dask = pytest.importorskip("dask")
         dask_df = pytest.importorskip("dask.dataframe")
         data = dask_df.from_array(data)
     histogram = pmx.utils.quantile_histogram(data, n_quantiles=100)
+    if use_dask:
+        (histogram,) = dask.compute(histogram)
     assert isinstance(histogram, dict)
+    assert isinstance(histogram["mid"], np.ndarray)
     assert histogram["mid"].shape == (99,)
     assert histogram["low"].shape == (99,)
     assert histogram["upper"].shape == (99,)

pymc_experimental/utils/histogram_trick.py

@@ -37,10 +37,10 @@ def _(data: ArrayLike, n_quantiles=1000) -> Dict[str, ArrayLike]:
 
     @quantile_histogram.register(dask.dataframe.Series)
     def _(data: dask.dataframe.Series, n_quantiles=1000) -> Dict[str, ArrayLike]:
-        quantiles = data.quantile(np.linspace(0, 1, n_quantiles)).persist()
-        count, _ = dask.array.histogram(data, quantiles.compute())
-        low = quantiles.to_dask_array(lengths=True)[:-1]
-        upper = quantiles.to_dask_array(lengths=True)[1:]
+        quantiles = data.quantile(np.linspace(0, 1, n_quantiles)).to_dask_array(lengths=True)
+        count, _ = dask.array.histogram(data, quantiles)
+        low = quantiles[:-1]
+        upper = quantiles[1:]
         result = dict(
             low=low,
             upper=upper,