Dask graph rewriting/ expression optimisation #11542

Activity

1. charles-turner-1

   added

   enhancement

   on Aug 24, 2026
2. dcherian commented on Aug 24, 2026

   

   dcherian

   on Aug 24, 2026

   Contributor

   More actions

   Have you seen https://github.com/mrocklin/dask-array?

   The main optimization available in ds['temperature'] * 1.8 + 32 is going to be in-place rewriting.

   res = ds["temperature"] 
   res *=1.8
   res += 32
   

   which dask could totally do.
3. charles-turner-1 commented on Aug 25, 2026

   

   charles-turner-1

   on Aug 25, 2026

   Last edited by charles-turner-1

   ContributorAuthor

   More actions

   Can't believe I'd somehow missed dask_array.

   I had a look at profiling your example against the other things I've tried:

   

   I think the in-place rewriting is doing roughly what the map-reduction approach is, right? It basically pushes f(x) = 1.8 * x + 32 into the dask blocks? It looks from the benchmark above that they give roughly the same speedup. Chunk local is the 'flox-y' mapreduce idea I was playing with above & inplace/block kernel is your snippet in the graph & the table below.

   I found out yesterday from @Kai-Striega that numpy should be able to elide doing operations in place - if I've understood correctly. So I think on a plain numpy array, this should work in place (for a sufficiently large array):

   naive_maybe_numpy_elided = (da_x * np.float64(1.8) + np.int64(32)).mean()
   naive_maybe_numpy_elided.compute()
   

   I had a look at benchmarking that too but it gives the exact same results as the naive version, so I guess maybe the numpy arrays in each individual chunk are small enough that numpy doesn't bother?

   (BTW - algebraic here is just doing the C -> F conversion on the single output number, ie. (ds['temp'].mean()) *1.8 + 32)

   case	best_s	mean_s	result	speedup_vs_naive
   naive	0.0728594595566392	0.0803545081987977	32.90003708316947	1
   naive (numpy elide)	0.07098925020545721	0.07344069983810186	32.90003708316947	1.026344965551393
   chunk-local	0.030837541446089745	0.031601191684603694	32.90003708316947	2.362687041183625
   inplace (block kernel)	0.036481207236647606	0.03783359993249178	32.90003708316947	1.9971778643182456
   algebraic (control)	0.03926900029182434	0.040669508278369904	32.90003708316948	1.855393797020299
4. Kai-Striega commented on Aug 25, 2026

   

   Kai-Striega

   on Aug 25, 2026

   More actions

   I haven't had the time to look at the entire PR, but I'm going to respond to the specific section that I was pinged for. I'll take a look at the remaining PR later on this week.

   I found out yesterday from @Kai-Striega that numpy should be able to elide doing operations in place - if I've understood correctly. So I think on a plain numpy array, this should work in place (for a sufficiently large array):

   a * 1.8 + 32.0 does elide, even without casting the Floats to NumPy types.

   You can see this using the tracemalloc.get_traced_memory function. For example this can be seen by:

   import tracemalloc
   import numpy as np
   
   MiB = 1 << 20
   N_ELEMS = 32 * MiB / 8 # 32MiB total size, 8 bytes per float64
   a = np.arange(N_ELEMS).astype(np.float64)
   
   tracemalloc.start()  # start recording
   _ = a * 1.8 + 32.0
   current, peak = tracemalloc.get_traced_memory()
   tracemalloc.stop()
   
   print(current) # prints `34437563` which is ~ 32 MiB

   This means that we allocated a total of ~32 MiB during the entire operation. If we allocated twice this would be the size of two arrays, not one (roughly 64 MiB).
5. charles-turner-1 commented on Aug 27, 2026

   

   charles-turner-1

   on Aug 27, 2026

   Last edited by charles-turner-1

   ContributorAuthor

   More actions

   Okay, I've been quite a long way down the rabbit hole with this today, including trying these approaches with regular dask.array versus dask-array:

   It turns out all of this is wrong, see comment below. Still interesting though (at least to me). The TLDR; is that dask is recording individual ops in the tasks, eg. [mul, add, mean], which adds a reference & stops numpy elision.

   In hindsight, using a benchmark which was trying to discriminate between eliding two buffers to one or cutting the amount of data moving between dask chunks by ~50% was a bad choice.

   Details

   # /// script
   # requires-python = ">=3.12"
   # dependencies = [
   #   "xarray[complete]@git+https://github.com/pydata/xarray.git@main",
   #   "dask-array",
   #   "dask",
   #   "tabulate",
   #
   # ]
   # ///
   import time
   from collections.abc import Callable
   from typing import Any, Literal
   
   import dask.array as da
   import numpy as np
   import pandas as pd
   import xarray as xr
   
   BenchmarkFn = Callable[[xr.DataArray], Any]
   
   type One = Literal[1]
   type Two = Literal[2]
   type TwoElementArray = np.ndarray[tuple[Two, *tuple[One, ...]], np.float64]
   
   
   def naive(data: xr.DataArray) -> Any:
       """Do conversion before aggregation. Should be slowest."""
       return ((data * 1.8) + 32).mean().compute()
   
   
   def algebraic_control(data: xr.DataArray) -> Any:
       """Do conversion after aggregation. Since we only have a single element array
       coming out of"""
       return (data.mean() * 1.8 + 32).compute()
   
   
   def inplace_simple(data: xr.DataArray) -> Any:
       """
       Deepak's suggestion, straight up - no block kernel or anything, just do the
       operations in place.
   
       In theory, I think this should force everything to be done in place, but it
       still has the overhead of having to spit a full chunk out of each block,
       rather than just a 2 element array.
       """
       data *= 1.8
       data += 32
       return data.mean().compute()
   
   
   def block_local_mapreduce(data: xr.DataArray) -> Any:
       stats = map_blocks_for(data)(
           _elide_kernel,
           data.data,
           dtype=np.float64,
           chunks=mean_and_count_chunks(data),
           new_axis=0,
       )
       return (stats[0].sum() / stats[1].sum()).compute()
   
   
   def inplace_block_kernel(data: xr.DataArray) -> Any:
       stats = map_blocks_for(data)(
           _inplace_kernel,
           data.data,
           dtype=np.float64,
           chunks=mean_and_count_chunks(data),
           new_axis=0,
       )
       return (stats[0].sum() / stats[1].sum()).compute()
   
   
   # For both of the _..._kernel functions, the only data coming out of each block
   # is a 2 element array, rather than a full block. This should reduce the amount
   # of data moving between blocks & speed things up a lot.
   def _elide_kernel(block: np.ndarray) -> TwoElementArray:
       """Numpy should elide the intermediates here.
   
       Question - numpy should stop elision when chunks are are larger than 256KiB"""
       f = block * 1.8 + 32
       return np.array([f.sum(dtype=np.float64), f.size], dtype=np.float64).reshape(
           (2,) + (1,) * block.ndim
       )
       # (2,) is the new axis we specify in map_blocks
   
   
   def _inplace_kernel(block: np.ndarray) -> TwoElementArray:
       """Deepak's suggestion, applied as a block kernel. Manually forces the operations to be in place."""
       f = block.copy()
       f *= 1.8
       f += 32
       return np.array([f.sum(dtype=np.float64), f.size], dtype=np.float64).reshape(
           (2,) + (1,) * block.ndim
       )
   
   
   BENCHMARKS: tuple[tuple[str, BenchmarkFn], ...] = (
       ("naive", naive),
       ("block-local (Charles Version)", block_local_mapreduce),
       ("inplace (Deepak Version)", inplace_simple),
       ("inplace (Deepak Version, block local)", inplace_block_kernel),
       ("algebraic (control)", algebraic_control),
   )
   
   
   def map_blocks_for(data: xr.DataArray) -> Callable[..., Any]:
       """If the data is a dask-array, use the dask-array version of map_blocks, otherwise use the stock dask.array version."""
       import dask_array as new_dask_array
   
       if isinstance(data.data, new_dask_array.Array):
           return new_dask_array.map_blocks
       return da.map_blocks
   
   
   def mean_and_count_chunks(data: xr.DataArray) -> tuple[tuple[int, ...], ...]:
       """Insert the extra dim as dim zero, with length two. Gets filled with sum, count. This is then each chunk in our
       post reduction array."""
       return ((2,),) + tuple(tuple(1 for _ in axis) for axis in data.data.chunks)
   
   
   def build_stock_data(n: int, chunk_edge: int, seed: int) -> xr.DataArray:
       shape = (n, n, n)
       chunks = (chunk_edge, chunk_edge, chunk_edge)
       random = da.random.RandomState(seed)
       data = xr.DataArray(
           random.random_sample(shape, chunks=chunks),
           dims=("time", "y", "x"),
           name="temperature_c",
       ).persist()
       data.compute()  # Pregenerate the random array but keep data lazy for the benchmarks.
       return data
   
   
   def register_dask_array() -> None:
       import dask_array as dexpr
   
       dexpr.xarray.register()
   
   
   def build_dask_array_data(stock_data: xr.DataArray) -> xr.DataArray:
       import dask_array as dexpr
   
       chunks = tuple(axis[0] for axis in stock_data.data.chunks)
       data = xr.DataArray(
           dexpr.from_array(stock_data.data.compute(), chunks=chunks),
           dims=stock_data.dims,
           name=stock_data.name,
       ).persist()
       data.compute()
       return data
   
   
   def benchmark_phase(phase: str, data: xr.DataArray, repeats: int) -> pd.DataFrame:
       rows = []
   
       for name, run_benchmark in BENCHMARKS:
           mutates_wrapper = run_benchmark is inplace_simple
           if mutates_wrapper:
               # The expression-level inplace case mutates the xarray wrapper. Prebuild
               # fresh wrappers around the same lazy array so allocation is not mixed into
               # the timing loop and later benchmarks still see the original Celsius data.
               warmup_data = data.copy(deep=False)
               run_inputs = [data.copy(deep=False) for _ in range(repeats)]
           else:
               warmup_data = data
               run_inputs = [data] * repeats
   
           run_benchmark(warmup_data)
   
           result, times = None, []
           for run_data in run_inputs:
               start = time.perf_counter()
               result = float(np.asarray(run_benchmark(run_data)))
               times.append(time.perf_counter() - start)
   
           rows.append(
               (
                   phase,
                   name,
                   min(times),
                   float(np.mean(times)),
                   float(np.median(times)),
                   result,
               )
           )
   
       table = pd.DataFrame(
           rows,
           columns=["phase", "case", "best_s", "mean_s", "median_s", "result"],
       )
       if not np.allclose(table["result"], table["result"].iloc[0], rtol=1e-6):
           raise AssertionError(f"{phase} benchmark results disagree:\n{table}")
       return table
   
   
   def markdown_table(table: pd.DataFrame) -> str:
       table = table.drop(columns=["sort_speedup"])
       return table.to_markdown(
           index=False,
           floatfmt=(".6f", ".6f", ".6f", ".12g"),
       )
   
   
   def main() -> None:
       n = 400
       chunk_edge = 200
       repeats = 5
   
       nbytes = n**3 * 8
       chunk_nbytes = chunk_edge**3 * 8
   
       print(
           f"Array: {n}^3 = {nbytes / 1e9:.2f} GB; "
           f"chunks: {chunk_edge}^3 = {chunk_nbytes / 1e6:.0f} MB; "
           f"repeats: {repeats}"
       )
       print("Note: stock means dask.array, not dask-array.\n\n")
   
       stock_data = build_stock_data(n, chunk_edge, seed=42)
       before = benchmark_phase("before dask-array", stock_data, repeats)
       baseline_median_s = float(before["median_s"].iloc[0])
   
       register_dask_array()
       stock_after_register = benchmark_phase(
           "stock after dask-array register", stock_data, repeats
       )
   
       dask_array_data = build_dask_array_data(stock_data)
       after = benchmark_phase("after dask-array", dask_array_data, repeats)
   
       table = pd.concat([before, stock_after_register, after], ignore_index=True)
       table["sort_speedup"] = baseline_median_s / table["median_s"]
       table["speedup_vs_pre_dask_naive"] = table["sort_speedup"].map(
           lambda speedup: f"{speedup:.2f}x"
       )
   
       if not np.allclose(table["result"], table["result"].iloc[0], rtol=1e-6):
           raise AssertionError(f"benchmark results disagree:\n{table}")
   
       table = table.sort_values("sort_speedup", ascending=False)
       print(markdown_table(table))
   
   
   if __name__ == "__main__":
       main()

   Details

   phase	case	best_s	mean_s	median_s	result	speedup_vs_pre_dask_naive
   stock after dask-array register	inplace (Deepak Version, block local)	0.038680	0.0393476335332	0.0393129	32.9	2.21x
   stock after dask-array register	block-local (Charles Version)	0.041497	0.0423958832398	0.0420139	32.9	2.07x
   before dask-array	block-local (Charles Version)	0.042717	0.0434249166399	0.043249	32.9	2.01x
   before dask-array	inplace (Deepak Version, block local)	0.037783	0.0432256001979	0.0436069	32.9	1.99x
   before dask-array	algebraic (control)	0.043410	0.0454272253439	0.0454796	32.9	1.91x
   after dask-array	block-local (Charles Version)	0.046414	0.0480979338288	0.048228	32.9	1.80x
   after dask-array	inplace (Deepak Version, block local)	0.049398	0.0506826827303	0.0505513	32.9	1.72x
   after dask-array	algebraic (control)	0.057243	0.0632531996816	0.065048	32.9	1.34x
   before dask-array	inplace (Deepak Version)	0.075677	0.0791818581522	0.0792139	32.9	1.10x
   after dask-array	inplace (Deepak Version)	0.076630	0.0813624249771	0.0826737	32.9	1.05x
   after dask-array	naive	0.077564	0.0830438083038	0.0835446	32.9	1.04x
   before dask-array	naive	0.077132	0.0886118249968	0.0869357	32.9	1.00x
   stock after dask-array register	algebraic (control)	0.227121	0.231938382983	0.23167	32.9	0.38x
   stock after dask-array register	inplace (Deepak Version)	0.258865	0.261013432965	0.26078	32.9	0.33x
   stock after dask-array register	naive	0.260266	0.277592425048	0.26953	32.9	0.32x

   There's a heap going on there, but my interpretation is:
   1. Block local approaches are the fastest at ~2x. I figure the explanation is that the majority of the run time is just bussing the data from one dask block to another, and the block local approaches are ~halving the amount of data bussed out. It doesn't seem to matter matter if that's written as arr * 1.8 + 32 or arr += 1.8; arr += 32, so long as the computation gets pushed into the chunk~~
   2. Using a regular dask.array after registering dask-array is a bit of a footgun and slows things down here.
   3. dask-array isn't pushing the reduction into the blocks, or this line: | after dask-array | inplace (Deepak Version)| 0.076630 | 0.0813624249771 | 0.0826737 | 32.9 | 1.05x | would be sitting around the 2* speedup too. It seems like trying to make dask-array do that would be a good idea to me, I'll dig into that.
   4~~. I think the numpy elision might have been a bit of a red herring here (Sorry Kai! ). I suspect what's happening is probably that dask is holding a bunch of references to the numpy arrays in each block - so what looks like it should elide in a purely numpy array gets blocked by dask. I did a very cursory check (used codex to write the script below):

   This is all wrong

   Details

   # /// script
   # requires-python = "==3.12"
   # dependencies = [
   #   "dask",
   #   "numpy",
   #   "pandas",
   #   "tabulate",
   # ]
   # ///
   """Check whether dask prevents NumPy temporary elision inside a single chunk."""
   
   from __future__ import annotations
   
   # ruff: noqa: D101,D103
   import argparse
   import gc
   import multiprocessing as mp
   import platform
   import time
   import tracemalloc
   from collections.abc import Callable
   from dataclasses import dataclass
   from typing import Any
   
   import dask.array as da
   import numpy as np
   import pandas as pd
   
   CaseFn = Callable[[np.ndarray], tuple[str, Callable[[], Any]]]
   
   
   @dataclass(frozen=True)
   class CaseResult:
       case: str
       seconds: float
       peak_mb: float
       result: float
   
   
   def pure_numpy(x: np.ndarray) -> tuple[str, Callable[[], Any]]:
       return "numpy: (x * 1.8 + 32).mean()", lambda: (x * 1.8 + 32).mean()
   
   
   def pure_numpy_inplace(x: np.ndarray) -> tuple[str, Callable[[], Any]]:
       def run() -> Any:
           out = x.copy()
           out *= 1.8
           out += 32
           return out.mean()
   
       return "numpy: copy; *=; +=; mean", run
   
   
   def dask_one_chunk_optimized(x: np.ndarray) -> tuple[str, Callable[[], Any]]:
       dx = da.from_array(x, chunks=x.shape)
       expr = ((dx * 1.8) + 32).mean()
       return (
           "dask one chunk: optimized",
           lambda: expr.compute(scheduler="single-threaded", optimize_graph=True),
       )
   
   
   def dask_one_chunk_unoptimized(x: np.ndarray) -> tuple[str, Callable[[], Any]]:
       dx = da.from_array(x, chunks=x.shape)
       expr = ((dx * 1.8) + 32).mean()
       return (
           "dask one chunk: optimize_graph=False",
           lambda: expr.compute(scheduler="single-threaded", optimize_graph=False),
       )
   
   
   def dask_one_chunk_augmented_assignment_optimized(
       x: np.ndarray,
   ) -> tuple[str, Callable[[], Any]]:
       dx = da.from_array(x, chunks=x.shape)
       expr = dx.copy()
       expr *= 1.8
       expr += 32
       mean = expr.mean()
       return (
           "dask one chunk: copy; *=; +=; optimized",
           lambda: mean.compute(scheduler="single-threaded", optimize_graph=True),
       )
   
   
   def dask_one_chunk_augmented_assignment_unoptimized(
       x: np.ndarray,
   ) -> tuple[str, Callable[[], Any]]:
       dx = da.from_array(x, chunks=x.shape)
       expr = dx.copy()
       expr *= 1.8
       expr += 32
       mean = expr.mean()
       return (
           "dask one chunk: copy; *=; +=; optimize_graph=False",
           lambda: mean.compute(scheduler="single-threaded", optimize_graph=False),
       )
   
   
   def dask_one_chunk_block_local(x: np.ndarray) -> tuple[str, Callable[[], Any]]:
       dx = da.from_array(x, chunks=x.shape)
   
       def stats(block: np.ndarray) -> np.ndarray:
           f = block * 1.8 + 32
           return np.array([f.sum(dtype=np.float64), f.size], dtype=np.float64)
   
       expr = da.map_blocks(
           stats,
           dx,
           dtype=np.float64,
           chunks=((2,),),
           drop_axis=(0, 1, 2),
           new_axis=0,
       )
       mean = expr[0] / expr[1]
       return (
           "dask one chunk: block-local sum/count",
           lambda: mean.compute(scheduler="single-threaded", optimize_graph=True),
       )
   
   
   CASES: tuple[CaseFn, ...] = (
       pure_numpy,
       pure_numpy_inplace,
       dask_one_chunk_optimized,
       dask_one_chunk_unoptimized,
       dask_one_chunk_augmented_assignment_optimized,
       dask_one_chunk_augmented_assignment_unoptimized,
       dask_one_chunk_block_local,
   )
   
   
   def run_case(case_name: str, n: int, seed: int, queue: mp.Queue[CaseResult]) -> None:
       case_by_name = {case.__name__: case for case in CASES}
       case = case_by_name[case_name]
   
       rng = np.random.default_rng(seed)
       x = rng.random((n, n, n), dtype=np.float64)
       label, run = case(x)
   
       gc.collect()
       tracemalloc.start()
       start = time.perf_counter()
       result = float(np.asarray(run()))
       seconds = time.perf_counter() - start
       _, peak = tracemalloc.get_traced_memory()
       tracemalloc.stop()
   
       queue.put(
           CaseResult(
               case=label,
               seconds=seconds,
               peak_mb=peak / 1e6,
               result=result,
           )
       )
   
   
   def measure(case: CaseFn, n: int, seed: int) -> CaseResult:
       context = mp.get_context("spawn")
       queue: mp.Queue[CaseResult] = context.Queue()
       process = context.Process(target=run_case, args=(case.__name__, n, seed, queue))
       process.start()
       process.join()
   
       if process.exitcode != 0:
           raise RuntimeError(f"{case.__name__} failed with exit code {process.exitcode}")
       return queue.get()
   
   
   def result_table(results: list[CaseResult], input_mb: float) -> pd.DataFrame:
       table = pd.DataFrame([result.__dict__ for result in results])
       baseline_peak = float(
           table.loc[table["case"].str.startswith("numpy:"), "peak_mb"].iloc[0]
       )
       table["peak_vs_numpy"] = table["peak_mb"] / baseline_peak
       table["peak_vs_input"] = table["peak_mb"] / input_mb
       return table.sort_values("peak_mb")
   
   
   def parse_args() -> argparse.Namespace:
       parser = argparse.ArgumentParser(
           description="Compare NumPy temporary elision with one-chunk dask execution."
       )
       parser.add_argument("--n", type=int, default=400, help="cube edge length")
       parser.add_argument("--seed", type=int, default=42, help="random seed")
       return parser.parse_args()
   
   
   def main() -> None:
       args = parse_args()
       input_mb = args.n**3 * 8 / 1e6
       print(f"Array: {args.n}^3 float64 = {input_mb:.1f} MB")
       print(f"Python: {platform.python_version()}, NumPy: {np.__version__}")
       print(
           "tracemalloc starts after the input array and dask graph are built; each case runs in a fresh process.\n"
       )
   
       results = [measure(case, args.n, args.seed) for case in CASES]
       table = result_table(results, input_mb)
       print(
           table.to_markdown(
               index=False,
               floatfmt=(".6f", ".1f", ".12g", ".2f", ".2f"),
           )
       )
   
   
   if __name__ == "__main__":
       main()

   Array: 400^3 float64 = 512.0 MB
   tracemalloc starts after the input array and dask graph are built; each case runs in a fresh process.

   case	seconds	peak_mb	result	peak_vs_numpy	peak_vs_input
   numpy: (x * 1.8 + 32).mean()	0.1	512.001001	32.90	1.00	1
   numpy: copy; *=; +=; mean	0.1	512.001001	32.90	1.00	1
   dask one chunk: copy; *=; +=; optimize_graph=False	0.2	1024.054835	32.90	2.00	2.00011
   dask one chunk: optimize_graph=False	0.2	1024.054948	32.90	2.00	2.00011
   dask one chunk: block-local sum/count	0.2	1024.068484	32.90	2.00	2.00013
   dask one chunk: optimized	0.2	1024.074	32.90	2.00	2.00014
   dask one chunk: copy; *=; +=; optimized	0.2	1024.074062	32.90	2.00	2.00014

   I haven't audited it a great deal of depth, but it aligns with what I would expect, so I'm satisfied. Also, pinning to python 3.14 stops numpy eliding in pure numpy and also pushes the block local dask version up to 3x memory usage - presumably there's some meaningful link there (threading or something) that I might look into when I get the chance.

   All of that's a lot of work to say that I think it'd definitely be worth trying to figure out how to push the reductions into chunks in every possible instance. I think a generic roughly 50% runtime decrease across any kind of map-reduction problem is something that would be worthwhile pursuing?

   ---

   Everything above this is wrong

   I wasn't super convinced by the answer I came to so I did some more testing. This time, taking a 3d array and then reducing over either 1,2 or 3 dimensions after the conversion - the logic being that if this were about the full block being written out before being reduced then doing partial reductions should show that up, since it'll vary the number of bytes crossing the block boundary from one buffer to another (no idea if this is the right way to state this but this is how I've been visualising it).

   # /// script
   # requires-python = ">=3.11,<3.13"
   # dependencies = [
   #   "numpy<2.4",
   #   "dask>=2026.4.0",
   #   "pandas",
   # ]
   # ///
   import time
   import tracemalloc
   
   import dask.array as da
   import numpy as np
   import pandas as pd
   
   base_np_arr = np.random.random((400, 400, 400))
   chunk_len = 100
   
   
   def trace_dask(build):
       _ = build().compute()  # warm up
       expr = build()
       tracemalloc.start()
       t0 = time.perf_counter()
       out = np.asarray(expr.compute())
       dt = time.perf_counter() - t0
       peak = tracemalloc.get_traced_memory()[1] / 1014 / 1024  # MB
       tracemalloc.stop()
       return peak, dt, float(out.sum())
   
   
   def make_dask_arr():
       return da.from_array(base_np_arr, chunks=(chunk_len, chunk_len, chunk_len))
   
   
   def naive(axis):
       return ((make_dask_arr() * 1.8) + 32).sum(axis=tuple(axis))
   
   
   def block_local(axis):
       x = make_dask_arr()
       axis = tuple(axis)
   
       def kern(b):
           f = b * 1.8 + 32  # one buffer
           return f.sum(axis=axis, keepdims=True)
   
       oc = tuple(
           (1,) * len(x.chunks[i]) if i in axis else x.chunks[i] for i in range(x.ndim)
       )
       return da.map_blocks(kern, x, dtype=float, chunks=oc).sum(axis=axis)
   
   
   def best_time(expr, reps=5):
       # Build the graph ONCE at the call site and pass the expression in, so the
       # timed region is only .compute(). Timing build() re-tokenizes the 512 MB
       # source array in da.from_array every call (~0.23 s), swamping the ~0.05 s
       # compute and washing the ratio out toward 1.
       expr.compute()  # warm up
       _ts = []
       for _ in range(reps):
           _t0 = time.perf_counter()
           expr.compute()
           _ts.append(time.perf_counter() - _t0)
       return min(_ts)  # least noisy run
   
   
   res_dims = []
   for axis, label in [
       ((2,), "1 axis  (x)"),
       ((1, 2), "2 axes  (x,y)"),
       ((0, 1, 2), "3 axes  (x,y,z)"),
   ]:
       print(f"reducing over {label}...")
       # memory: tracemalloc run (overhead slows execution, so time it separately)
       mem_naive, _, output_naive = trace_dask(lambda a=axis: naive(a))
       mem_blockwise, _, output_blockwise = trace_dask(lambda a=axis: block_local(a))
       assert np.isclose(
           output_naive, output_blockwise, rtol=1e-6
       )  # same totals -> same result
       # time: clean best-of-3, no tracemalloc; build the graphs once, outside timing
       time_naive = best_time(naive(axis))
       time_blockwise = best_time(block_local(axis))
       res_dims.append(
           {
               "reduction over": label,
               "naive MB": round(mem_naive),
               "block-local MB": round(mem_blockwise),
               "mem ratio": round(mem_naive / mem_blockwise, 2),
               "naive s": round(time_naive, 3),
               "block-local s": round(time_blockwise, 3),
               "time ratio": round(time_naive / time_blockwise, 2),
           }
       )
   
   print(pd.DataFrame(res_dims))

   Which turns out

   reducing over 1 axis  (x)...
   reducing over 2 axes  (x,y)...
   reducing over 3 axes  (x,y,z)...
       reduction over  naive MB  block-local MB  mem ratio  naive s  block-local s  time ratio
   0      1 axis  (x)       124              70       1.76    0.069          0.055        1.26
   1    2 axes  (x,y)        93              78       1.20    0.058          0.046        1.27
   2  3 axes  (x,y,z)       101              78       1.30    0.058          0.045        1.28
   

   If this was about the number of bytes coming out of each block computation and being written into an intermediate buffer, I would expect the 3 axis reduction to be the fastest, but there's no clear difference.

   Anyway, then I realised it was a bit weird these all sat around ~1.3, not two. It looks like that's because DataArray.mean() is doing a nanmean - so extra work. (N.B the benchmarks I'm doing in this bit are just reducing as sums, not means - just a bit simpler to do).

   It looks like what I thought might have been a bit of low hanging fruit with this optimisation might not be quite so worthwhile. Maybe trying to combine the elementwise operations would bump performance by up to ~30% in this sort of case. With that said, I think it would also potentially halve the peak memory footprint, which is potentially well worthwhile...