from __future__ import annotations

import contextlib

import math

import operator

import os

import pickle

import re

import sys

import traceback

import uuid

import warnings

from bisect import bisect

from collections.abc import (

Collection,

Iterable,

Iterator,

Mapping,

MutableMapping,

Sequence,

)

from functools import partial, reduce, wraps

from itertools import product, zip_longest

from numbers import Integral, Number

from operator import add, mul

from threading import Lock

from typing import Any, TypeVar, Union, cast

import numpy as np

from numpy.typing import ArrayLike

from tlz import accumulate, concat, first, frequencies, groupby, partition

from tlz.curried import pluck

from dask import compute, config, core

from dask.array import chunk

from dask.array.chunk import getitem

from dask.array.chunk_types import is_valid_array_chunk, is_valid_chunk_type

# Keep einsum_lookup and tensordot_lookup here for backwards compatibility

from dask.array.dispatch import ( # noqa: F401

concatenate_lookup,

einsum_lookup,

tensordot_lookup,

)

from dask.array.numpy_compat import _Recurser

from dask.array.slicing import replace_ellipsis, setitem_array, slice_array

from dask.base import (

DaskMethodsMixin,

compute_as_if_collection,

dont_optimize,

is_dask_collection,

named_schedulers,

persist,

tokenize,

)

from dask.blockwise import blockwise as core_blockwise

from dask.blockwise import broadcast_dimensions

from dask.context import globalmethod

from dask.core import quote

from dask.delayed import Delayed, delayed

from dask.highlevelgraph import HighLevelGraph, MaterializedLayer

from dask.layers import ArraySliceDep, reshapelist

from dask.sizeof import sizeof

from dask.typing import Graph, Key, NestedKeys

from dask.utils import (

IndexCallable,

SerializableLock,

cached_cumsum,

cached_property,

concrete,

derived_from,

format_bytes,

funcname,

has_keyword,

is_arraylike,

is_dataframe_like,

is_index_like,

is_integer,

is_series_like,

maybe_pluralize,

ndeepmap,

ndimlist,

parse_bytes,

typename,

)

from dask.widgets import get_template

T_IntOrNaN = Union[int, float] # Should be Union[int, Literal[np.nan]]

DEFAULT_GET = named_schedulers.get("threads", named_schedulers["sync"])

unknown_chunk_message = (

"\n\n"

"A possible solution: "

"https://docs.dask.org/en/latest/array-chunks.html#unknown-chunks\n"

"Summary: to compute chunks sizes, use\n\n"

" x.compute_chunk_sizes() # for Dask Array `x`\n"

" ddf.to_dask_array(lengths=True) # for Dask DataFrame `ddf`"

)

class PerformanceWarning(Warning):

"""A warning given when bad chunking may cause poor performance"""

def getter(a, b, asarray=True, lock=None):

if isinstance(b, tuple) and any(x is None for x in b):

b2 = tuple(x for x in b if x is not None)

b3 = tuple(

None if x is None else slice(None, None)

for x in b

if not isinstance(x, Integral)

)

return getter(a, b2, asarray=asarray, lock=lock)[b3]

if lock:

lock.acquire()

try:

c = a[b]

# Below we special-case `np.matrix` to force a conversion to

# `np.ndarray` and preserve original Dask behavior for `getter`,

# as for all purposes `np.matrix` is array-like and thus

# `is_arraylike` evaluates to `True` in that case.

if asarray and (not is_arraylike(c) or isinstance(c, np.matrix)):

c = np.asarray(c)

finally:

if lock:

lock.release()

return c

def getter_nofancy(a, b, asarray=True, lock=None):

"""A simple wrapper around ``getter``.

Used to indicate to the optimization passes that the backend doesn't

support fancy indexing.

"""

return getter(a, b, asarray=asarray, lock=lock)

def getter_inline(a, b, asarray=True, lock=None):

"""A getter function that optimizations feel comfortable inlining

Slicing operations with this function may be inlined into a graph, such as

in the following rewrite

**Before**

>>> a = x[:10] # doctest: +SKIP

>>> b = a + 1 # doctest: +SKIP

>>> c = a * 2 # doctest: +SKIP

**After**

>>> b = x[:10] + 1 # doctest: +SKIP

>>> c = x[:10] * 2 # doctest: +SKIP

This inlining can be relevant to operations when running off of disk.

"""

return getter(a, b, asarray=asarray, lock=lock)

from dask.array.optimization import fuse_slice, optimize

# __array_function__ dict for mapping aliases and mismatching names

_HANDLED_FUNCTIONS = {}

def implements(*numpy_functions):

"""Register an __array_function__ implementation for dask.array.Array

Register that a function implements the API of a NumPy function (or several

NumPy functions in case of aliases) which is handled with

``__array_function__``.

Parameters

----------

\\*numpy_functions : callables

One or more NumPy functions that are handled by ``__array_function__``

and will be mapped by `implements` to a `dask.array` function.

"""

def decorator(dask_func):

for numpy_function in numpy_functions:

_HANDLED_FUNCTIONS[numpy_function] = dask_func

return dask_func

return decorator

def _should_delegate(self, other) -> bool:

"""Check whether Dask should delegate to the other.

This implementation follows NEP-13:

https://numpy.org/neps/nep-0013-ufunc-overrides.html#behavior-in-combination-with-python-s-binary-operations

"""

if hasattr(other, "__array_ufunc__") and other.__array_ufunc__ is None:

return True

elif (

hasattr(other, "__array_ufunc__")

and not is_valid_array_chunk(other)

# don't delegate to our own parent classes

and not isinstance(self, type(other))

and type(self) is not type(other)

):

return True

return False

def check_if_handled_given_other(f):

"""Check if method is handled by Dask given type of other

Ensures proper deferral to upcast types in dunder operations without

assuming unknown types are automatically downcast types.

"""

@wraps(f)

def wrapper(self, other):

if _should_delegate(self, other):

return NotImplemented

else:

return f(self, other)

return wrapper

def slices_from_chunks(chunks):

"""Translate chunks tuple to a set of slices in product order

>>> slices_from_chunks(((2, 2), (3, 3, 3))) # doctest: +NORMALIZE_WHITESPACE

[(slice(0, 2, None), slice(0, 3, None)),

(slice(0, 2, None), slice(3, 6, None)),

(slice(0, 2, None), slice(6, 9, None)),

(slice(2, 4, None), slice(0, 3, None)),

(slice(2, 4, None), slice(3, 6, None)),

(slice(2, 4, None), slice(6, 9, None))]

"""

cumdims = [cached_cumsum(bds, initial_zero=True) for bds in chunks]

slices = [

[slice(s, s + dim) for s, dim in zip(starts, shapes)]

for starts, shapes in zip(cumdims, chunks)

]

return list(product(*slices))

def graph_from_arraylike(

arr, # Any array-like which supports slicing

chunks,

shape,

name,

getitem=getter,

lock=False,

asarray=True,

dtype=None,

inline_array=False,

) -> HighLevelGraph:

"""

HighLevelGraph for slicing chunks from an array-like according to a chunk pattern.

If ``inline_array`` is True, this make a Blockwise layer of slicing tasks where the

array-like is embedded into every task.,

If ``inline_array`` is False, this inserts the array-like as a standalone value in

a MaterializedLayer, then generates a Blockwise layer of slicing tasks that refer

to it.

>>> dict(graph_from_arraylike(arr, chunks=(2, 3), shape=(4, 6), name="X", inline_array=True)) # doctest: +SKIP

{(arr, 0, 0): (getter, arr, (slice(0, 2), slice(0, 3))),

(arr, 1, 0): (getter, arr, (slice(2, 4), slice(0, 3))),

(arr, 1, 1): (getter, arr, (slice(2, 4), slice(3, 6))),

(arr, 0, 1): (getter, arr, (slice(0, 2), slice(3, 6)))}

>>> dict( # doctest: +SKIP

graph_from_arraylike(arr, chunks=((2, 2), (3, 3)), shape=(4,6), name="X", inline_array=False)

)

{"original-X": arr,

('X', 0, 0): (getter, 'original-X', (slice(0, 2), slice(0, 3))),

('X', 1, 0): (getter, 'original-X', (slice(2, 4), slice(0, 3))),

('X', 1, 1): (getter, 'original-X', (slice(2, 4), slice(3, 6))),

('X', 0, 1): (getter, 'original-X', (slice(0, 2), slice(3, 6)))}

"""

chunks = normalize_chunks(chunks, shape, dtype=dtype)

out_ind = tuple(range(len(shape)))

if (

has_keyword(getitem, "asarray")

and has_keyword(getitem, "lock")

and (not asarray or lock)

):

kwargs = {"asarray": asarray, "lock": lock}

else:

# Common case, drop extra parameters

kwargs = {}

if inline_array:

layer = core_blockwise(

getitem,

name,

out_ind,

arr,

None,

ArraySliceDep(chunks),

out_ind,

numblocks={},

**kwargs,

)

return HighLevelGraph.from_collections(name, layer)

else:

original_name = "original-" + name

layers = {}

layers[original_name] = MaterializedLayer({original_name: arr})

layers[name] = core_blockwise(

getitem,

name,

out_ind,

original_name,

None,

ArraySliceDep(chunks),

out_ind,

numblocks={},

**kwargs,

)

deps = {

original_name: set(),

name: {original_name},

}

return HighLevelGraph(layers, deps)

def dotmany(A, B, leftfunc=None, rightfunc=None, **kwargs):

"""Dot product of many aligned chunks

>>> x = np.array([[1, 2], [1, 2]])

>>> y = np.array([[10, 20], [10, 20]])

>>> dotmany([x, x, x], [y, y, y])

array([[ 90, 180],

[ 90, 180]])

Optionally pass in functions to apply to the left and right chunks

>>> dotmany([x, x, x], [y, y, y], rightfunc=np.transpose)

array([[150, 150],

[150, 150]])

"""

if leftfunc:

A = map(leftfunc, A)

if rightfunc:

B = map(rightfunc, B)

return sum(map(partial(np.dot, **kwargs), A, B))

def _concatenate2(arrays, axes=None):

"""Recursively concatenate nested lists of arrays along axes

Each entry in axes corresponds to each level of the nested list. The

length of axes should correspond to the level of nesting of arrays.

If axes is an empty list or tuple, return arrays, or arrays[0] if

arrays is a list.

>>> x = np.array([[1, 2], [3, 4]])

>>> _concatenate2([x, x], axes=[0])

array([[1, 2],

[3, 4],

[1, 2],

[3, 4]])

>>> _concatenate2([x, x], axes=[1])

array([[1, 2, 1, 2],

[3, 4, 3, 4]])

>>> _concatenate2([[x, x], [x, x]], axes=[0, 1])

array([[1, 2, 1, 2],

[3, 4, 3, 4],

[1, 2, 1, 2],

[3, 4, 3, 4]])

Supports Iterators

>>> _concatenate2(iter([x, x]), axes=[1])

array([[1, 2, 1, 2],

[3, 4, 3, 4]])

Special Case

>>> _concatenate2([x, x], axes=())

array([[1, 2],

[3, 4]])

"""

if axes is None:

axes = []

if axes == ():

if isinstance(arrays, list):

return arrays[0]

else:

return arrays

if isinstance(arrays, Iterator):

arrays = list(arrays)

if not isinstance(arrays, (list, tuple)):

return arrays

if len(axes) > 1:

arrays = [_concatenate2(a, axes=axes[1:]) for a in arrays]

concatenate = concatenate_lookup.dispatch(

type(max(arrays, key=lambda x: getattr(x, "__array_priority__", 0)))

)

if isinstance(arrays[0], dict):

# Handle concatenation of `dict`s, used as a replacement for structured

# arrays when that's not supported by the array library (e.g., CuPy).

keys = list(arrays[0].keys())

assert all(list(a.keys()) == keys for a in arrays)

ret = dict()

for k in keys:

ret[k] = concatenate(list(a[k] for a in arrays), axis=axes[0])

return ret

else:

return concatenate(arrays, axis=axes[0])

def apply_infer_dtype(func, args, kwargs, funcname, suggest_dtype="dtype", nout=None):

"""

Tries to infer output dtype of ``func`` for a small set of input arguments.

Parameters

----------

func: Callable

Function for which output dtype is to be determined

args: List of array like

Arguments to the function, which would usually be used. Only attributes

``ndim`` and ``dtype`` are used.

kwargs: dict

Additional ``kwargs`` to the ``func``

funcname: String

Name of calling function to improve potential error messages

suggest_dtype: None/False or String

If not ``None`` adds suggestion to potential error message to specify a dtype

via the specified kwarg. Defaults to ``'dtype'``.

nout: None or Int

``None`` if function returns single output, integer if many.

Deafults to ``None``.

Returns

-------

: dtype or List of dtype

One or many dtypes (depending on ``nout``)

"""

from dask.array.utils import meta_from_array

# make sure that every arg is an evaluated array

args = [

np.ones_like(meta_from_array(x), shape=((1,) * x.ndim), dtype=x.dtype)

if is_arraylike(x)

else x

for x in args

]

try:

with np.errstate(all="ignore"):

o = func(*args, **kwargs)

except Exception as e:

exc_type, exc_value, exc_traceback = sys.exc_info()

tb = "".join(traceback.format_tb(exc_traceback))

suggest = (

(

"Please specify the dtype explicitly using the "

"`{dtype}` kwarg.\n\n".format(dtype=suggest_dtype)

)

if suggest_dtype

else ""

)

msg = (

f"`dtype` inference failed in `{funcname}`.\n\n"

f"{suggest}"

"Original error is below:\n"

"------------------------\n"

f"{e!r}\n\n"

"Traceback:\n"

"---------\n"

f"{tb}"

)

else:

msg = None

if msg is not None:

raise ValueError(msg)

return getattr(o, "dtype", type(o)) if nout is None else tuple(e.dtype for e in o)

def normalize_arg(x):

"""Normalize user provided arguments to blockwise or map_blocks

We do a few things:

1. If they are string literals that might collide with blockwise_token then we

quote them

2. IF they are large (as defined by sizeof) then we put them into the

graph on their own by using dask.delayed

"""

if is_dask_collection(x):

return x

elif isinstance(x, str) and re.match(r"_\d+", x):

return delayed(x)

elif isinstance(x, list) and len(x) >= 10:

return delayed(x)

elif sizeof(x) > 1e6:

return delayed(x)

else:

return x

def _pass_extra_kwargs(func, keys, *args, **kwargs):

"""Helper for :func:`dask.array.map_blocks` to pass `block_info` or `block_id`.

For each element of `keys`, a corresponding element of args is changed

to a keyword argument with that key, before all arguments re passed on

to `func`.

"""

kwargs.update(zip(keys, args))

return func(*args[len(keys) :], **kwargs)

def map_blocks(

func,

*args,

name=None,

token=None,

dtype=None,

chunks=None,

drop_axis=None,

new_axis=None,

enforce_ndim=False,

meta=None,

**kwargs,

):

"""Map a function across all blocks of a dask array.

Note that ``map_blocks`` will attempt to automatically determine the output

array type by calling ``func`` on 0-d versions of the inputs. Please refer to

the ``meta`` keyword argument below if you expect that the function will not

succeed when operating on 0-d arrays.

Parameters

----------

func : callable

Function to apply to every block in the array.

If ``func`` accepts ``block_info=`` or ``block_id=``

as keyword arguments, these will be passed dictionaries

containing information about input and output chunks/arrays

during computation. See examples for details.

args : dask arrays or other objects

dtype : np.dtype, optional

The ``dtype`` of the output array. It is recommended to provide this.

If not provided, will be inferred by applying the function to a small

set of fake data.

chunks : tuple, optional

Chunk shape of resulting blocks if the function does not preserve

shape. If not provided, the resulting array is assumed to have the same

block structure as the first input array.

drop_axis : number or iterable, optional

Dimensions lost by the function.

new_axis : number or iterable, optional

New dimensions created by the function. Note that these are applied

after ``drop_axis`` (if present).

enforce_ndim : bool, default False

Whether to enforce at runtime that the dimensionality of the array

produced by ``func`` actually matches that of the array returned by

``map_blocks``.

If True, this will raise an error when there is a mismatch.

token : string, optional

The key prefix to use for the output array. If not provided, will be

determined from the function name.

name : string, optional

The key name to use for the output array. Note that this fully

specifies the output key name, and must be unique. If not provided,

will be determined by a hash of the arguments.

meta : array-like, optional

The ``meta`` of the output array, when specified is expected to be an

array of the same type and dtype of that returned when calling ``.compute()``

on the array returned by this function. When not provided, ``meta`` will be

inferred by applying the function to a small set of fake data, usually a

0-d array. It's important to ensure that ``func`` can successfully complete

computation without raising exceptions when 0-d is passed to it, providing

``meta`` will be required otherwise. If the output type is known beforehand

(e.g., ``np.ndarray``, ``cupy.ndarray``), an empty array of such type dtype

can be passed, for example: ``meta=np.array((), dtype=np.int32)``.

**kwargs :

Other keyword arguments to pass to function. Values must be constants

(not dask.arrays)

See Also

--------

dask.array.map_overlap : Generalized operation with overlap between neighbors.

dask.array.blockwise : Generalized operation with control over block alignment.

Examples

--------

>>> import dask.array as da

>>> x = da.arange(6, chunks=3)

>>> x.map_blocks(lambda x: x * 2).compute()

array([ 0, 2, 4, 6, 8, 10])

The ``da.map_blocks`` function can also accept multiple arrays.

>>> d = da.arange(5, chunks=2)

>>> e = da.arange(5, chunks=2)

>>> f = da.map_blocks(lambda a, b: a + b**2, d, e)

>>> f.compute()

array([ 0, 2, 6, 12, 20])

If the function changes shape of the blocks then you must provide chunks

explicitly.

>>> y = x.map_blocks(lambda x: x[::2], chunks=((2, 2),))

You have a bit of freedom in specifying chunks. If all of the output chunk

sizes are the same, you can provide just that chunk size as a single tuple.

>>> a = da.arange(18, chunks=(6,))

>>> b = a.map_blocks(lambda x: x[:3], chunks=(3,))

If the function changes the dimension of the blocks you must specify the

created or destroyed dimensions.

>>> b = a.map_blocks(lambda x: x[None, :, None], chunks=(1, 6, 1),

... new_axis=[0, 2])

If ``chunks`` is specified but ``new_axis`` is not, then it is inferred to

add the necessary number of axes on the left.

Note that ``map_blocks()`` will concatenate chunks along axes specified by

the keyword parameter ``drop_axis`` prior to applying the function.

This is illustrated in the figure below:

.. image:: /images/map_blocks_drop_axis.png

Due to memory-size-constraints, it is often not advisable to use ``drop_axis``

on an axis that is chunked. In that case, it is better not to use

``map_blocks`` but rather

``dask.array.reduction(..., axis=dropped_axes, concatenate=False)`` which

maintains a leaner memory footprint while it drops any axis.

Map_blocks aligns blocks by block positions without regard to shape. In the

following example we have two arrays with the same number of blocks but

with different shape and chunk sizes.

>>> x = da.arange(1000, chunks=(100,))

>>> y = da.arange(100, chunks=(10,))

The relevant attribute to match is numblocks.

>>> x.numblocks

(10,)

>>> y.numblocks

(10,)

If these match (up to broadcasting rules) then we can map arbitrary

functions across blocks

>>> def func(a, b):

... return np.array([a.max(), b.max()])

>>> da.map_blocks(func, x, y, chunks=(2,), dtype='i8')

dask.array<func, shape=(20,), dtype=int64, chunksize=(2,), chunktype=numpy.ndarray>

>>> _.compute()

array([ 99, 9, 199, 19, 299, 29, 399, 39, 499, 49, 599, 59, 699,

69, 799, 79, 899, 89, 999, 99])

Your block function can get information about where it is in the array by

accepting a special ``block_info`` or ``block_id`` keyword argument.

During computation, they will contain information about each of the input

and output chunks (and dask arrays) relevant to each call of ``func``.

>>> def func(block_info=None):

... pass

This will receive the following information:

>>> block_info # doctest: +SKIP

{0: {'shape': (1000,),

'num-chunks': (10,),

'chunk-location': (4,),

'array-location': [(400, 500)]},

None: {'shape': (1000,),

'num-chunks': (10,),

'chunk-location': (4,),

'array-location': [(400, 500)],

'chunk-shape': (100,),

'dtype': dtype('float64')}}

The keys to the ``block_info`` dictionary indicate which is the input and

output Dask array:

- **Input Dask array(s):** ``block_info[0]`` refers to the first input Dask array.

The dictionary key is ``0`` because that is the argument index corresponding

to the first input Dask array.

In cases where multiple Dask arrays have been passed as input to the function,

you can access them with the number corresponding to the input argument,

eg: ``block_info[1]``, ``block_info[2]``, etc.

(Note that if you pass multiple Dask arrays as input to map_blocks,

the arrays must match each other by having matching numbers of chunks,

along corresponding dimensions up to broadcasting rules.)

- **Output Dask array:** ``block_info[None]`` refers to the output Dask array,

and contains information about the output chunks.

The output chunk shape and dtype may may be different than the input chunks.

For each dask array, ``block_info`` describes:

- ``shape``: the shape of the full Dask array,

- ``num-chunks``: the number of chunks of the full array in each dimension,

- ``chunk-location``: the chunk location (for example the fourth chunk over

in the first dimension), and

- ``array-location``: the array location within the full Dask array

(for example the slice corresponding to ``40:50``).

In addition to these, there are two extra parameters described by

``block_info`` for the output array (in ``block_info[None]``):

- ``chunk-shape``: the output chunk shape, and

- ``dtype``: the output dtype.

These features can be combined to synthesize an array from scratch, for

example:

>>> def func(block_info=None):

... loc = block_info[None]['array-location'][0]

... return np.arange(loc[0], loc[1])

>>> da.map_blocks(func, chunks=((4, 4),), dtype=np.float64)

dask.array<func, shape=(8,), dtype=float64, chunksize=(4,), chunktype=numpy.ndarray>

>>> _.compute()

array([0, 1, 2, 3, 4, 5, 6, 7])

``block_id`` is similar to ``block_info`` but contains only the ``chunk_location``:

>>> def func(block_id=None):

... pass

This will receive the following information:

>>> block_id # doctest: +SKIP

(4, 3)

You may specify the key name prefix of the resulting task in the graph with

the optional ``token`` keyword argument.

>>> x.map_blocks(lambda x: x + 1, name='increment')

dask.array<increment, shape=(1000,), dtype=int64, chunksize=(100,), chunktype=numpy.ndarray>

For functions that may not handle 0-d arrays, it's also possible to specify

``meta`` with an empty array matching the type of the expected result. In

the example below, ``func`` will result in an ``IndexError`` when computing

``meta``:

>>> rng = da.random.default_rng()

>>> da.map_blocks(lambda x: x[2], rng.random(5), meta=np.array(()))

dask.array<lambda, shape=(5,), dtype=float64, chunksize=(5,), chunktype=numpy.ndarray>

Similarly, it's possible to specify a non-NumPy array to ``meta``, and provide

a ``dtype``:

>>> import cupy # doctest: +SKIP

>>> rng = da.random.default_rng(cupy.random.default_rng()) # doctest: +SKIP

>>> dt = np.float32

>>> da.map_blocks(lambda x: x[2], rng.random(5, dtype=dt), meta=cupy.array((), dtype=dt)) # doctest: +SKIP

dask.array<lambda, shape=(5,), dtype=float32, chunksize=(5,), chunktype=cupy.ndarray>

"""

if drop_axis is None:

drop_axis = []

if not callable(func):

msg = (

"First argument must be callable function, not %s\n"

"Usage: da.map_blocks(function, x)\n"

" or: da.map_blocks(function, x, y, z)"

)

raise TypeError(msg % type(func).__name__)

if token:

warnings.warn(

"The `token=` keyword to `map_blocks` has been moved to `name=`. "

"Please use `name=` instead as the `token=` keyword will be removed "

"in a future release.",

category=FutureWarning,

)

name = token

name = f"{name or funcname(func)}-{tokenize(func, dtype, chunks, drop_axis, new_axis, *args, **kwargs)}"

new_axes = {}

if isinstance(drop_axis, Number):

drop_axis = [drop_axis]

if isinstance(new_axis, Number):

new_axis = [new_axis] # TODO: handle new_axis

arrs = [a for a in args if isinstance(a, Array)]

argpairs = [

(a, tuple(range(a.ndim))[::-1]) if isinstance(a, Array) else (a, None)

for a in args

]

if arrs:

out_ind = tuple(range(max(a.ndim for a in arrs)))[::-1]

else:

out_ind = ()

original_kwargs = kwargs

if dtype is None and meta is None:

try:

meta = compute_meta(func, dtype, *args, **kwargs)

except Exception:

pass

dtype = apply_infer_dtype(func, args, original_kwargs, "map_blocks")

if drop_axis:

ndim_out = len(out_ind)

if any(i < -ndim_out or i >= ndim_out for i in drop_axis):

raise ValueError(

f"drop_axis out of range (drop_axis={drop_axis}, "

f"but output is {ndim_out}d)."

)

drop_axis = [i % ndim_out for i in drop_axis]

out_ind = tuple(x for i, x in enumerate(out_ind) if i not in drop_axis)

if new_axis is None and chunks is not None and len(out_ind) < len(chunks):

new_axis = range(len(chunks) - len(out_ind))

if new_axis:

# new_axis = [x + len(drop_axis) for x in new_axis]

out_ind = list(out_ind)

for ax in sorted(new_axis):

n = len(out_ind) + len(drop_axis)

out_ind.insert(ax, n)

if chunks is not None:

new_axes[n] = chunks[ax]

else:

new_axes[n] = 1

out_ind = tuple(out_ind)

if max(new_axis) > max(out_ind):

raise ValueError("New_axis values do not fill in all dimensions")

if chunks is not None:

if len(chunks) != len(out_ind):

raise ValueError(

f"Provided chunks have {len(chunks)} dims; expected {len(out_ind)} dims"

)

adjust_chunks = dict(zip(out_ind, chunks))

else:

adjust_chunks = None

if enforce_ndim:

out = blockwise(

apply_and_enforce,

out_ind,

*concat(argpairs),

expected_ndim=len(out_ind),

_func=func,

name=name,

new_axes=new_axes,

dtype=dtype,

concatenate=True,

align_arrays=False,

adjust_chunks=adjust_chunks,

meta=meta,

**kwargs,

)

else:

out = blockwise(

func,

out_ind,

*concat(argpairs),

name=name,

new_axes=new_axes,

dtype=dtype,

concatenate=True,

align_arrays=False,

adjust_chunks=adjust_chunks,

meta=meta,

**kwargs,

)

extra_argpairs = []

extra_names = []

# If func has block_id as an argument, construct an array of block IDs and

# prepare to inject it.

if has_keyword(func, "block_id"):

block_id_name = "block-id-" + out.name

block_id_dsk = {

(block_id_name,) + block_id: block_id

for block_id in product(*(range(len(c)) for c in out.chunks))

}

block_id_array = Array(

block_id_dsk,

block_id_name,

chunks=tuple((1,) * len(c) for c in out.chunks),

dtype=np.object_,

)

extra_argpairs.append((block_id_array, out_ind))

extra_names.append("block_id")

# If func has block_info as an argument, construct an array of block info

# objects and prepare to inject it.

if has_keyword(func, "block_info"):

starts = {}

num_chunks = {}

shapes = {}

for i, (arg, in_ind) in enumerate(argpairs):

if in_ind is not None:

shapes[i] = arg.shape

if drop_axis:

# We concatenate along dropped axes, so we need to treat them

# as if there is only a single chunk.

starts[i] = [

(

cached_cumsum(arg.chunks[j], initial_zero=True)

if ind in out_ind

else [0, arg.shape[j]]

)

for j, ind in enumerate(in_ind)

]

num_chunks[i] = tuple(len(s) - 1 for s in starts[i])

else:

starts[i] = [

cached_cumsum(c, initial_zero=True) for c in arg.chunks

]

num_chunks[i] = arg.numblocks

out_starts = [cached_cumsum(c, initial_zero=True) for c in out.chunks]

block_info_name = "block-info-" + out.name

block_info_dsk = {}

for block_id in product(*(range(len(c)) for c in out.chunks)):

# Get position of chunk, indexed by axis labels

location = {out_ind[i]: loc for i, loc in enumerate(block_id)}

info = {}

for i, shape in shapes.items():

# Compute chunk key in the array, taking broadcasting into

# account. We don't directly know which dimensions are

# broadcast, but any dimension with only one chunk can be

# treated as broadcast.

arr_k = tuple(

location.get(ind, 0) if num_chunks[i][j] > 1 else 0

for j, ind in enumerate(argpairs[i][1])

)

info[i] = {

"shape": shape,

"num-chunks": num_chunks[i],

"array-location": [

(starts[i][ij][j], starts[i][ij][j + 1])

for ij, j in enumerate(arr_k)

],

"chunk-location": arr_k,

}

info[None] = {

"shape": out.shape,

"num-chunks": out.numblocks,

"array-location": [

(out_starts[ij][j], out_starts[ij][j + 1])

for ij, j in enumerate(block_id)

],

"chunk-location": block_id,

"chunk-shape": tuple(

out.chunks[ij][j] for ij, j in enumerate(block_id)

),

"dtype": dtype,

}

block_info_dsk[(block_info_name,) + block_id] = info

block_info = Array(

block_info_dsk,

block_info_name,

chunks=tuple((1,) * len(c) for c in out.chunks),

dtype=np.object_,

)

extra_argpairs.append((block_info, out_ind))

extra_names.append("block_info")

if extra_argpairs:

# Rewrite the Blockwise layer. It would be nice to find a way to

# avoid doing it twice, but it's currently needed to determine

# out.chunks from the first pass. Since it constructs a Blockwise

# rather than an expanded graph, it shouldn't be too expensive.

out = blockwise(

_pass_extra_kwargs,

out_ind,

func,

None,

tuple(extra_names),

None,

*concat(extra_argpairs),

*concat(argpairs),

name=out.name,

dtype=out.dtype,

concatenate=True,