re-enable tests

Author: JuliaPoo

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

@@ -40,9 +40,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.9', 
-          # '3.10', '3.11', '3.12'
-        ]
+        python-version: ['3.9', '3.10', '3.11', '3.12']
 
     steps:
     - name: Checkout array-api-compat

array_api_compat/jax/__init__.py

@@ -13,11 +13,39 @@ def top_k(
 
     # `swapaxes` is used to implement
     # the `axis` kwarg
-    x = numpy.swapaxes(x, axis, -1)
-    vals, args = lax.top_k(x, k)
-    vals = numpy.swapaxes(vals, axis, -1)
-    args = numpy.swapaxes(args, axis, -1)
-    return vals, args
+    # x = numpy.swapaxes(x, axis, -1)
+    # vals, args = lax.top_k(x, k)
+    # vals = numpy.swapaxes(vals, axis, -1)
+    # args = numpy.swapaxes(args, axis, -1)
+    # return vals, args
+
+    # 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 = 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 = (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_values = take_along_axis(arr, topk_indices, axis=axis)
+
+    return (topk_values, topk_indices)
 
 
 __all__ = ['top_k']

array_api_compat/numpy/_aliases.py

@@ -63,78 +63,6 @@
 
 
 def top_k(a, k, /, axis=-1, *, largest=True):
-    """
-    Returns the ``k`` largest/smallest elements and corresponding
-    indices along the given ``axis``.
-
-    When ``axis`` is None, a flattened array is used.
-
-    If ``largest`` is false, then the ``k`` smallest elements are returned.
-
-    A tuple of ``(values, indices)`` is returned, where ``values`` and
-    ``indices`` of the largest/smallest elements of each row of the input
-    array in the given ``axis``.
-
-    Parameters
-    ----------
-    a: array_like
-        The source array
-    k: int
-        The number of largest/smallest elements to return. ``k`` must
-        be a positive integer and within indexable range specified by
-        ``axis``.
-    axis: int, optional
-        Axis along which to find the largest/smallest elements.
-        The default is -1 (the last axis).
-        If None, a flattened array is used.
-    largest: bool, optional
-        If True, largest elements are returned. Otherwise the smallest
-        are returned.
-
-    Returns
-    -------
-    tuple_of_array: tuple
-        The output tuple of ``(topk_values, topk_indices)``, where
-        ``topk_values`` are returned elements from the source array
-        (not necessarily in sorted order), and ``topk_indices`` are
-        the corresponding indices.
-
-    See Also
-    --------
-    argpartition : Indirect partition.
-    sort : Full sorting.
-
-    Notes
-    -----
-    The returned indices are not guaranteed to be sorted according to
-    the values. Furthermore, the returned indices are not guaranteed
-    to be the earliest/latest occurrence of the element. E.g.,
-    ``np.top_k([3,3,3], 1)`` can return ``(array([3]), array([1]))``
-    rather than ``(array([3]), array([0]))`` or
-    ``(array([3]), array([2]))``.
-
-    Warning: The treatment of ``np.nan`` in the input array is undefined.
-
-    Examples
-    --------
-    >>> a = np.array([[1,2,3,4,5], [5,4,3,2,1], [3,4,5,1,2]])
-    >>> np.top_k(a, 2)
-    (array([[4, 5],
-            [4, 5],
-            [4, 5]]),
-     array([[3, 4],
-            [1, 0],
-            [1, 2]]))
-    >>> np.top_k(a, 2, axis=0)
-    (array([[3, 4, 3, 2, 2],
-            [5, 4, 5, 4, 5]]),
-     array([[2, 1, 1, 1, 2],
-            [1, 2, 2, 0, 0]]))
-    >>> a.flatten()
-    array([1, 2, 3, 4, 5, 5, 4, 3, 2, 1, 3, 4, 5, 1, 2])
-    >>> np.top_k(a, 2, axis=None)
-    (array([5, 5]), array([ 5, 12]))
-    """
     if k <= 0:
         raise ValueError(f'k(={k}) provided must be positive.')