Skip to content

Commit

Permalink
implement ravel_multi_index & define test struct
Browse files Browse the repository at this point in the history
  • Loading branch information
amoungui committed Aug 14, 2023
1 parent 1b1c66d commit 853e764
Show file tree
Hide file tree
Showing 3 changed files with 86 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -98,3 +98,8 @@ def compress(condition, a, axis=None, out=None):
)
arr = arr[: condition_arr.shape[0]]
return ivy.moveaxis(arr[condition_arr], 0, axis)

@to_ivy_arrays_and_back
def ravel_multi_index(arr, dims, mode='raise', order='C'):
"""ravel_multi_index"""
return ivy.ravel_multi_index(arr, dims, mode=mode, order=order)
74 changes: 74 additions & 0 deletions ivy/functional/ivy/experimental/manipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -2550,3 +2550,77 @@ def soft_thresholding(
if ivy.exists(out):
return ivy.inplace_update(out, res)
return res


@handle_exceptions
@handle_backend_invalid
@handle_nestable
@handle_array_like_without_promotion
@handle_out_argument
@to_native_arrays_and_back
@handle_array_function
@handle_device_shifting
def ravel_multi_index(
arr: Union[
Tuple[Union[ivy.Array, ivy.NativeArray]],
Union[ivy.Array, ivy.NativeArray],
],
dims: Tuple[int],
/,
*,
mode: str | None = None,
order: str | None = None,
out: Optional[ivy.Array] = None,
) -> ivy.Array:
"""
Converts a tuple of index arrays into an array of flat indices,
applying boundary modes to the multi-index.
Parameters
----------
arr
tuple of array_like: A tuple of integer arrays, one array for each dimension.
dims
tuple of ints: The shape of array into which the indices from multi_index apply.
out
The output array.
mode
{'raise', 'wrap', 'clip'}, optional
Specifies how out-of-bounds indices are handled. Can specify either one mode or
a tuple of modes, one mode per index.
'raise' - raise an error (default)
'wrap' - wrap around
'clip' - clip to the range
In 'clip' mode, a negative index which would normally
wrap will clip to 0 instead.
order{'C', 'F'}, optional
Determines whether the multi-index should be viewed
as indexing in row-major (C-style)
or column-major (Fortran-style) order.
Returns
-------
ret
ndarray An array of indices into the flattened version
of an array of dimensions dims
Examples
--------
>>> arr = np.array([[3,6,6],[4,5,1]])
>>> ivy.ravel_multi_index(arr, (7,6))
array([22, 41, 37])
>>> ivy.ravel_multi_index(arr, (7,6), order='F')
array([31, 41, 13])
>>> ivy.ravel_multi_index(arr, (4,6), mode='clip')
array([22, 23, 19])
>>> ivy.ravel_multi_index(arr, (4,4), mode=('clip','wrap'))
array([12, 13, 13])
"""
return current_backend(arr).ravel_multi_index(arr, dims, mode=mode, order=order)
Original file line number Diff line number Diff line change
Expand Up @@ -188,3 +188,10 @@ def test_numpy_compress(
a=arr[0],
axis=ax,
)

@handle_frontend_test(
fn_tree="numpy.ravel_multi_index",
)
def test_numpy_ravel_multi_index():
"""test_numpy_ravel_multi_index"""
pass

0 comments on commit 853e764

Please sign in to comment.