fix jax top_k

Author: JuliaPoo

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

@@ -7,5 +7,7 @@ jobs:
     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

array_api_compat/jax/__init__.py

@@ -10,15 +10,6 @@ def top_k(
     *,
     largest=True,
 ):
-
-    # `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
-
     # The largest keyword can't be implemented with `jax.lax.top_k`
     # efficiently so am using `jax.numpy` for now
     if k <= 0: