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