Function _wrap_func

dask/array/random.py:898–1038

Wrap numpy random function to produce dask.array random function extra_chunks should be a chunks tuple to append to the end of chunks

(
    rng, funcname, *args, size=None, chunks="auto", extra_chunks=(), **kwargs
)

def _wrap_func(
    rng, funcname, *args, size=None, chunks="auto", extra_chunks=(), **kwargs
):
    """Wrap numpy random function to produce dask.array random function
    extra_chunks should be a chunks tuple to append to the end of chunks
    """
    if size is not None and not isinstance(size, (tuple, list)):
        size = (size,)

    shapes = list(
        {
            ar.shape
            for ar in chain(args, kwargs.values())
            if isinstance(ar, (Array, np.ndarray))
        }
    )
    if size is not None:
        shapes.append(size)
    # broadcast to the final size(shape)
    size = broadcast_shapes(*shapes)
    chunks = normalize_chunks(
        chunks,
        size,  # ideally would use dtype here
        dtype=kwargs.get("dtype", np.float64),
    )
    slices = slices_from_chunks(chunks)

    def _broadcast_any(ar, shape, chunks):
        if isinstance(ar, Array):
            return broadcast_to(ar, shape).rechunk(chunks)
        elif isinstance(ar, np.ndarray):
            return np.ascontiguousarray(np.broadcast_to(ar, shape))
        else:
            raise TypeError("Unknown object type for broadcast")

    # Broadcast all arguments, get tiny versions as well
    # Start adding the relevant bits to the graph
    dsk = {}
    lookup = {}
    small_args = []
    dependencies = []
    for i, ar in enumerate(args):
        if isinstance(ar, (np.ndarray, Array)):
            res = _broadcast_any(ar, size, chunks)
            if isinstance(res, Array):
                dependencies.append(res)
                lookup[i] = res.name
            elif isinstance(res, np.ndarray):
                name = f"array-{tokenize(res)}"
                lookup[i] = name
                dsk[name] = res
            small_args.append(ar[tuple(0 for _ in ar.shape)])
        else:
            small_args.append(ar)

    small_kwargs = {}
    for key, ar in kwargs.items():
        if isinstance(ar, (np.ndarray, Array)):