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
)
| 900 | |
| 901 | |
| 902 | def _wrap_func( |
| 903 | rng, funcname, *args, size=None, chunks="auto", extra_chunks=(), **kwargs |
| 904 | ): |
| 905 | """Wrap numpy random function to produce dask.array random function |
| 906 | extra_chunks should be a chunks tuple to append to the end of chunks |
| 907 | """ |
| 908 | if size is not None and not isinstance(size, (tuple, list)): |
| 909 | size = (size,) |
| 910 | |
| 911 | shapes = list( |
| 912 | { |
| 913 | ar.shape |
| 914 | for ar in chain(args, kwargs.values()) |
| 915 | if isinstance(ar, (Array, np.ndarray)) |
| 916 | } |
| 917 | ) |
| 918 | if size is not None: |
| 919 | shapes.append(size) |
| 920 | # broadcast to the final size(shape) |
| 921 | size = broadcast_shapes(*shapes) |
| 922 | chunks = normalize_chunks( |
| 923 | chunks, |
| 924 | size, # ideally would use dtype here |
| 925 | dtype=kwargs.get("dtype", np.float64), |
| 926 | ) |
| 927 | slices = slices_from_chunks(chunks) |
| 928 | |
| 929 | def _broadcast_any(ar, shape, chunks): |
| 930 | if isinstance(ar, Array): |
| 931 | return broadcast_to(ar, shape).rechunk(chunks) |
| 932 | elif isinstance(ar, np.ndarray): |
| 933 | return np.ascontiguousarray(np.broadcast_to(ar, shape)) |
| 934 | else: |
| 935 | raise TypeError("Unknown object type for broadcast") |
| 936 | |
| 937 | # Broadcast all arguments, get tiny versions as well |
| 938 | # Start adding the relevant bits to the graph |
| 939 | dsk = {} |
| 940 | lookup = {} |
| 941 | small_args = [] |
| 942 | dependencies = [] |
| 943 | for i, ar in enumerate(args): |
| 944 | if isinstance(ar, (np.ndarray, Array)): |
| 945 | res = _broadcast_any(ar, size, chunks) |
| 946 | if isinstance(res, Array): |
| 947 | dependencies.append(res) |
| 948 | lookup[i] = res.name |
| 949 | elif isinstance(res, np.ndarray): |
| 950 | name = f"array-{tokenize(res)}" |
| 951 | lookup[i] = name |
| 952 | dsk[name] = res |
| 953 | small_args.append(ar[tuple(0 for _ in ar.shape)]) |
| 954 | else: |
| 955 | small_args.append(ar) |
| 956 | |
| 957 | small_kwargs = {} |
| 958 | for key, ar in kwargs.items(): |
| 959 | if isinstance(ar, (np.ndarray, Array)): |
no test coverage detected