argsort

Contents

argsort#

class brainunit.math.argsort(a, axis=-1, *, kind=None, order=None, stable=True, descending=False, **kwargs)#

Return the indices that would sort an array or Quantity.

Units are stripped before sorting.

Parameters:

  • a (Array | ndarray | bool | number | bool | int | float | complex | saiunit.Quantity) – Array or Quantity to be sorted.

  • axis (int | None) – Axis along which to sort. Default is -1 (last axis). If None, the array is flattened first.

  • kind (None) – Sorting algorithm. Unused in JAX.

  • order (None) – Unused in JAX.

  • stable (bool) – Whether to use a stable sort. Default is True.

  • descending (bool) – Whether to sort in descending order. Default is False.

Returns:

indices – Array of indices that would sort the input.

Return type:

Array

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> u.math.argsort(jnp.array([3.0, 1.0, 2.0]))
Array([1, 2, 0], dtype=int32)
>>> q = jnp.array([3.0, 1.0, 2.0]) * u.meter
>>> u.math.argsort(q)
Array([1, 2, 0], dtype=int32)
Contents