Generic function for argreduction. Parameters ---------- x : Array chunk : callable Partialed ``arg_chunk``. combine : callable Partialed ``arg_combine``. agg : callable Partialed ``arg_agg``. axis : int, optional split_every : int or dict, op
(
x, chunk, combine, agg, axis=None, keepdims=False, split_every=None, out=None
)
| 840 | |
| 841 | |
| 842 | def arg_reduction( |
| 843 | x, chunk, combine, agg, axis=None, keepdims=False, split_every=None, out=None |
| 844 | ): |
| 845 | """Generic function for argreduction. |
| 846 | |
| 847 | Parameters |
| 848 | ---------- |
| 849 | x : Array |
| 850 | chunk : callable |
| 851 | Partialed ``arg_chunk``. |
| 852 | combine : callable |
| 853 | Partialed ``arg_combine``. |
| 854 | agg : callable |
| 855 | Partialed ``arg_agg``. |
| 856 | axis : int, optional |
| 857 | split_every : int or dict, optional |
| 858 | """ |
| 859 | if axis is None: |
| 860 | axis = tuple(range(x.ndim)) |
| 861 | ravel = True |
| 862 | elif isinstance(axis, Integral): |
| 863 | axis = validate_axis(axis, x.ndim) |
| 864 | axis = (axis,) |
| 865 | ravel = x.ndim == 1 |
| 866 | else: |
| 867 | raise TypeError(f"axis must be either `None` or int, got '{axis}'") |
| 868 | |
| 869 | for ax in axis: |
| 870 | chunks = x.chunks[ax] |
| 871 | if len(chunks) > 1 and np.isnan(chunks).any(): |
| 872 | raise ValueError( |
| 873 | "Arg-reductions do not work with arrays that have " |
| 874 | "unknown chunksizes. At some point in your computation " |
| 875 | "this array lost chunking information.\n\n" |
| 876 | "A possible solution is with \n" |
| 877 | " x.compute_chunk_sizes()" |
| 878 | ) |
| 879 | |
| 880 | # Map chunk across all blocks |
| 881 | name = f"arg-reduce-{tokenize(axis, x, chunk, combine, split_every)}" |
| 882 | old = x.name |
| 883 | keys = list(product(*map(range, x.numblocks))) |
| 884 | offsets = list(product(*(accumulate(operator.add, bd[:-1], 0) for bd in x.chunks))) |
| 885 | if ravel: |
| 886 | offset_info = zip(offsets, repeat(x.shape)) |
| 887 | else: |
| 888 | offset_info = pluck(axis[0], offsets) |
| 889 | |
| 890 | chunks = tuple((1,) * len(c) if i in axis else c for (i, c) in enumerate(x.chunks)) |
| 891 | dsk = { |
| 892 | (name,) + k: (chunk, (old,) + k, axis, off) |
| 893 | for (k, off) in zip(keys, offset_info) |
| 894 | } |
| 895 | |
| 896 | dtype = np.argmin(asarray_safe([1], like=meta_from_array(x))) |
| 897 | meta = None |
| 898 | if is_arraylike(dtype): |
| 899 | # This case occurs on non-NumPy types (e.g., CuPy), where the returned |
no test coverage detected