(arr_shape, indices, axis)
| 31 | |
| 32 | |
| 33 | def _make_along_axis_idx(arr_shape, indices, axis): |
| 34 | # compute dimensions to iterate over |
| 35 | if not _nx.issubdtype(indices.dtype, _nx.integer): |
| 36 | raise IndexError('`indices` must be an integer array') |
| 37 | if len(arr_shape) != indices.ndim: |
| 38 | raise ValueError( |
| 39 | "`indices` and `arr` must have the same number of dimensions") |
| 40 | shape_ones = (1,) * indices.ndim |
| 41 | dest_dims = list(range(axis)) + [None] + list(range(axis + 1, indices.ndim)) |
| 42 | |
| 43 | # build a fancy index, consisting of orthogonal aranges, with the |
| 44 | # requested index inserted at the right location |
| 45 | fancy_index = [] |
| 46 | for dim, n in zip(dest_dims, arr_shape): |
| 47 | if dim is None: |
| 48 | fancy_index.append(indices) |
| 49 | else: |
| 50 | ind_shape = shape_ones[:dim] + (-1,) + shape_ones[dim + 1:] |
| 51 | fancy_index.append(_nx.arange(n).reshape(ind_shape)) |
| 52 | |
| 53 | return tuple(fancy_index) |
| 54 | |
| 55 | |
| 56 | def _take_along_axis_dispatcher(arr, indices, axis=None): |
no test coverage detected
searching dependent graphs…