| 782 | |
| 783 | |
| 784 | def arg_chunk(func, argfunc, x, axis, offset_info): |
| 785 | arg_axis = None if len(axis) == x.ndim or x.ndim == 1 else axis[0] |
| 786 | vals = func(x, axis=arg_axis, keepdims=True) |
| 787 | arg = argfunc(x, axis=arg_axis, keepdims=True) |
| 788 | if x.ndim > 0: |
| 789 | if arg_axis is None: |
| 790 | offset, total_shape = offset_info |
| 791 | ind = np.unravel_index(arg.ravel()[0], x.shape) |
| 792 | total_ind = tuple(o + i for (o, i) in zip(offset, ind)) |
| 793 | arg[:] = np.ravel_multi_index(total_ind, total_shape) |
| 794 | else: |
| 795 | arg += offset_info |
| 796 | |
| 797 | if isinstance(vals, np.ma.masked_array): |
| 798 | if "min" in argfunc.__name__: |
| 799 | fill_value = np.ma.minimum_fill_value(vals) |
| 800 | else: |
| 801 | fill_value = np.ma.maximum_fill_value(vals) |
| 802 | vals = np.ma.filled(vals, fill_value) |
| 803 | |
| 804 | try: |
| 805 | result = np.empty_like( |
| 806 | vals, shape=vals.shape, dtype=[("vals", vals.dtype), ("arg", arg.dtype)] |
| 807 | ) |
| 808 | except TypeError: |
| 809 | # Array type doesn't support structured arrays (e.g., CuPy) |
| 810 | result = dict() |
| 811 | |
| 812 | result["vals"] = vals |
| 813 | result["arg"] = arg |
| 814 | return result |
| 815 | |
| 816 | |
| 817 | def arg_combine(argfunc, data, axis=None, **kwargs): |