unravel_index

array_api_extra.unravel_index(indices, shape, /, *, xp=None)

Convert a flat index or array of flat indices into a tuple of coordinate arrays.

Parameters:

  • indices (object) – An integer array whose elements are indices into the flattened version of an array of dimensions shape.

  • shape (tuple[int, ...]) – The shape to use for unraveling indices.

  • xp (ModuleType | None) – The standard-compatible namespace for indices. Default: infer.

Returns:

A tuple of unraveled indices. Each array in the tuple has the same shape as the indices array.

Return type:

tuple[object, ...]

Examples

>>> import array_api_extra as xpx
>>> import array_api_strict as xp
>>> xs, ys = xpx.unravel_index(xp.asarray([1, 2, 4, 5, 6, 8]), (4, 3))
>>> xs, ys
(
    Array([0, 0, 1, 1, 2, 2], dtype=array_api_strict.int64),
    Array([1, 2, 1, 2, 0, 2], dtype=array_api_strict.int64),
)
>>> [(int(x), int(y)) for x, y in zip(xs, ys)]
[(0, 1), (0, 2), (1, 1), (1, 2), (2, 0), (2, 2)]
>>> xs, ys = xpx.unravel_index(xp.arange(6), (2, 2))
>>> [(int(x), int(y)) for x, y in zip(xs, ys)]
[(0, 0), (0, 1), (1, 0), (1, 1), (0, 0), (0, 1)]