Repository: dask/dask
File: dask/array/core.py (Python)

Code before:

            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_arr[block_id] = info

        cs = tuple((1,) * len(c) for c in out.chunks)
        dsk = graph_from_arraylike(
            block_info_arr,
            cs,
            block_info_arr.shape,
            block_info_name,
            getitem=_getter_item,
            dtype=block_info_arr.dtype,
            inline_array=False,
        )
        block_info = Array(dsk, block_info_name, chunks=cs, 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,
            align_arrays=False,
            adjust_chunks=dict(zip(out_ind, out.chunks)),
            meta=meta,
            **kwargs,

Review comment (refactor):
`graph_from_arraylike` isn't used anywhere (it's imported in dask-expr, we might want to nuke the array modules there to simplify work)

Code after:
            name=out.name,
            dtype=out.dtype,
            concatenate=True,
            align_arrays=False,
            adjust_chunks=dict(zip(out_ind, out.chunks)),
            meta=meta,
            **kwargs,
        )

    return out


def apply_and_enforce(*args, **kwargs):
    """Apply a function, and enforce the output.ndim to match expected_ndim

    Ensures the output has the expected dimensionality."""
    func = kwargs.pop("_func")
    expected_ndim = kwargs.pop("expected_ndim")
    out = func(*args, **kwargs)
    if getattr(out, "ndim", 0) != expected_ndim:
        out_ndim = getattr(out, "ndim", 0)
        raise ValueError(
            f"Dimension mismatch: expected output of {func} "
            f"to have dims = {expected_ndim}.  Got {out_ndim} instead."
        )
    return out


def broadcast_chunks(*chunkss):
    """Construct a chunks tuple that broadcasts many chunks tuples

    >>> a = ((5, 5),)
    >>> b = ((5, 5),)
    >>> broadcast_chunks(a, b)
    ((5, 5),)

    >>> a = ((10, 10, 10), (5, 5),)
    >>> b = ((5, 5),)
    >>> broadcast_chunks(a, b)
    ((10, 10, 10), (5, 5))

    >>> a = ((10, 10, 10), (5, 5),)
    >>> b = ((1,), (5, 5),)
    >>> broadcast_chunks(a, b)
    ((10, 10, 10), (5, 5))

    >>> a = ((10, 10, 10), (5, 5),)
    >>> b = ((3, 3,), (5, 5),)
    >>> broadcast_chunks(a, b)
    Traceback (most recent call last):
        ...