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