WIP: Add top_k compatibility

This references the PR data-apis/array-api-tests#274.

Author: JuliaPoo

.github/workflows/array-api-tests-jax.yml

@@ -0,0 +1,13 @@
+name: Array API Tests (JAX)
+
+on: [push, pull_request]
+
+jobs:
+  array-api-tests-jax:
+    uses: ./.github/workflows/array-api-tests.yml
+    with:
+      package-name: jax
+      # See https://github.com/google/jax/issues/22137 for reason behind skipped dtypes
+      extra-env-vars: |
+        JAX_ENABLE_X64=1
+        ARRAY_API_TESTS_SKIP_DTYPES=uint8,uint16,uint32,uint64

.github/workflows/array-api-tests.yml

@@ -33,7 +33,7 @@ on:
         description: "Multiline string of environment variables to set for the test run."
 
 env:
-  PYTEST_ARGS: "--max-examples 200 -v -rxXfE --ci ${{ inputs.pytest-extra-args }} --hypothesis-disable-deadline"
+  PYTEST_ARGS: "--max-examples 200 -v -rxXfE --ci ${{ inputs.pytest-extra-args }} -k top_k --hypothesis-disable-deadline"
 
 jobs:
   tests:
@@ -50,9 +50,10 @@ jobs:
     - name: Checkout array-api-tests
       uses: actions/checkout@v4
       with:
-        repository: data-apis/array-api-tests
+        repository: JuliaPoo/array-api-tests
         submodules: 'true'
         path: array-api-tests
+        ref: ci-wip-topk-tests
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
@@ -77,6 +78,7 @@ jobs:
         # This enables the NEP 50 type promotion behavior (without it a lot of
         # tests fail on bad scalar type promotion behavior)
         NPY_PROMOTION_STATE: weak
+        ARRAY_API_TESTS_VERSION: draft
       run: |
         export PYTHONPATH="${GITHUB_WORKSPACE}/array-api-compat"
         cd ${GITHUB_WORKSPACE}/array-api-tests

array_api_compat/dask/array/_aliases.py

@@ -150,6 +150,28 @@ def asarray(
 
     return da.asarray(obj, dtype=dtype, **kwargs)
 
+
+def top_k(
+    x: Array,
+    k: int,
+    /,
+    axis: Optional[int] = None,
+    *,
+    largest: bool = True,
+) -> tuple[Array, Array]:
+
+    if not largest:
+        k = -k
+
+    # For now, perform the computation twice,
+    # since an equivalent to numpy's `take_along_axis`
+    # does not exist.
+    # See https://github.com/dask/dask/issues/3663.
+    args = da.argtopk(x, k, axis=axis).compute()
+    vals = da.topk(x, k, axis=axis).compute()
+    return vals, args
+
+
 from dask.array import (
     # Element wise aliases
     arccos as acos,
@@ -178,6 +200,7 @@ def asarray(
                             'bitwise_right_shift', 'concat', 'pow',
                             'e', 'inf', 'nan', 'pi', 'newaxis', 'float32', 'float64', 'int8',
                             'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'uint64',
-                            'complex64', 'complex128', 'iinfo', 'finfo', 'can_cast', 'result_type']
+                            'complex64', 'complex128', 'iinfo', 'finfo', 'can_cast', 'result_type',
+                            'top_k']
 
 _all_ignore = ['get_xp', 'da', 'partial', 'common_aliases', 'np']

array_api_compat/jax/__init__.py

@@ -0,0 +1,85 @@
+from jax.numpy import (
+    # Constants
+    e,
+    inf,
+    nan,
+    pi,
+    newaxis,
+    # Dtypes
+    bool,
+    float32,
+    float64,
+    int8,
+    int16,
+    int32,
+    int64,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+    complex64,
+    complex128,
+    iinfo,
+    finfo,
+    can_cast,
+    result_type,
+    # functions
+    zeros,
+    all,
+    any,
+    isnan,
+    isfinite,
+    reshape
+)
+from jax.numpy import (
+    asarray,
+    s_,
+    int_,
+    argpartition,
+    take_along_axis
+)
+
+
+def top_k(
+    x,
+    k,
+    /,
+    axis=None,
+    *,
+    largest=True,
+):
+    # The largest keyword can't be implemented with `jax.lax.top_k`
+    # efficiently so am using `jax.numpy` for now
+    if k <= 0:
+        raise ValueError(f'k(={k}) provided must be positive.')
+
+    positive_axis: int
+    _arr = asarray(x)
+    if axis is None:
+        arr = _arr.ravel()
+        positive_axis = 0
+    else:
+        arr = _arr
+        positive_axis = axis if axis > 0 else axis % arr.ndim
+
+    slice_start = (s_[:],) * positive_axis
+    if largest:
+        indices_array = argpartition(arr, -k, axis=axis)
+        slice = slice_start + (s_[-k:],)
+        topk_indices = indices_array[slice]
+    else:
+        indices_array = argpartition(arr, k-1, axis=axis)
+        slice = slice_start + (s_[:k],)
+        topk_indices = indices_array[slice]
+
+    topk_indices = topk_indices.astype(int_)
+    topk_values = take_along_axis(arr, topk_indices, axis=axis)
+    return (topk_values, topk_indices)
+
+
+__all__ = ['top_k', 'e', 'inf', 'nan', 'pi', 'newaxis', 'bool',
+           'float32', 'float64', 'int8', 'int16', 'int32',
+           'int64', 'uint8', 'uint16', 'uint32', 'uint64',
+           'complex64', 'complex128', 'iinfo', 'finfo',
+           'can_cast', 'result_type', 'zeros', 'all', 'isnan',
+           'isfinite', 'reshape', 'any']

array_api_compat/numpy/_aliases.py

@@ -61,6 +61,35 @@
 matrix_transpose = get_xp(np)(_aliases.matrix_transpose)
 tensordot = get_xp(np)(_aliases.tensordot)
 
+
+def top_k(a, k, /, axis=-1, *, largest=True):
+    if k <= 0:
+        raise ValueError(f'k(={k}) provided must be positive.')
+
+    positive_axis: int
+    _arr = np.asanyarray(a)
+    if axis is None:
+        arr = _arr.ravel()
+        positive_axis = 0
+    else:
+        arr = _arr
+        positive_axis = axis if axis > 0 else axis % arr.ndim
+
+    slice_start = (np.s_[:],) * positive_axis
+    if largest:
+        indices_array = np.argpartition(arr, -k, axis=axis)
+        slice = slice_start + (np.s_[-k:],)
+        topk_indices = indices_array[slice]
+    else:
+        indices_array = np.argpartition(arr, k-1, axis=axis)
+        slice = slice_start + (np.s_[:k],)
+        topk_indices = indices_array[slice]
+
+    topk_values = np.take_along_axis(arr, topk_indices, axis=axis)
+
+    return (topk_values, topk_indices)
+
+
 def _supports_buffer_protocol(obj):
     try:
         memoryview(obj)
@@ -126,6 +155,6 @@ def asarray(
 __all__ = _aliases.__all__ + ['asarray', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
                               'atanh', 'bitwise_left_shift', 'bitwise_invert',
-                              'bitwise_right_shift', 'concat', 'pow']
+                              'bitwise_right_shift', 'concat', 'pow', 'top_k']
 
 _all_ignore = ['np', 'get_xp']

array_api_compat/torch/_aliases.py

@@ -700,6 +700,8 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
         axis = 0
     return torch.index_select(x, axis, indices, **kwargs)
 
+top_k = torch.topk
+
 __all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert',
            'newaxis', 'add', 'atan2', 'bitwise_and', 'bitwise_left_shift',
            'bitwise_or', 'bitwise_right_shift', 'bitwise_xor', 'divide',
@@ -713,6 +715,6 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'matmul', 'matrix_transpose', 'vecdot', 'tensordot', 'isdtype',
-           'take']
+           'take', 'top_k']
 
 _all_ignore = ['torch', 'get_xp']