heat.array_api._searching_functions

Module Contents

argmax(x: heat.array_api._array_object.Array, /, *, axis: int | None = None, keepdims: bool = False) → heat.array_api._array_object.Array[source]

Returns the indices of the maximum values along a specified axis. When the maximum value occurs multiple times, only the indices corresponding to the first occurrence are returned.

Parameters:

  • x (Array) – Input array. Must have a numeric data type.

  • axis (Optional[int]) – Axis along which to search. If None, the function returns the index of the maximum value of the flattened array. Default: None.

  • keepdims (bool) – If True, the reduced axes (dimensions) are included in the result as singleton dimensions. Otherwise, if False, the reduced axes (dimensions) are not be included in the result. Default: False.

argmin(x: heat.array_api._array_object.Array, /, *, axis: int | None = None, keepdims: bool = False) → heat.array_api._array_object.Array[source]

Returns the indices of the minimum values along a specified axis. When the minimum value occurs multiple times, only the indices corresponding to the first occurrence are returned.

Parameters:

  • x (Array) – Input array. Must have a numeric data type.

  • axis (Optional[int]) – Axis along which to search. If None, the function returns the index of the minimum value of the flattened array. Default: None.

  • keepdims (bool) – If True, the reduced axes (dimensions) are included in the result as singleton dimensions. Otherwise, if False, the reduced axes (dimensions) are not be included in the result. Default: False.

nonzero(x: heat.array_api._array_object.Array, /) → Tuple[heat.array_api._array_object.Array, Ellipsis][source]

Returns the indices of the array elements which are non-zero.

Parameters:

x (Array) – Input array. Must have a positive rank.

where(condition: heat.array_api._array_object.Array, x1: heat.array_api._array_object.Array, x2: heat.array_api._array_object.Array, /) → heat.array_api._array_object.Array[source]

Returns elements chosen from x1 or x2 depending on condition.

Parameters:

  • condition (Array) – When True, yield x1_i; otherwise, yield x2_i. Must be compatible with x1 and x2.

  • x1 (Array) – First input array. Must be compatible with condition and x2.

  • x2 (Array) – Second input array. Must be compatible with condition and x1.