idk what im doing

Author: JuliaPoo

array_api_compat/jax/__init__.py

@@ -1,19 +1,13 @@
 import jax
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from typing import Optional, Tuple
-
-    from ..common._typing import Array
-
 
 def top_k(
-    x: Array,
-    k: int,
+    x,
+    k,
     /,
-    axis: Optional[int] = None,
+    axis=None,
     *,
-    largest: bool = True,
-) -> Tuple[Array, Array]:
+    largest=True,
+):
 
     # `swapaxes` is used to implement
     # the `axis` kwarg

array_api_compat/numpy/_aliases.py

@@ -62,7 +62,7 @@
 tensordot = get_xp(np)(_aliases.tensordot)
 
 
-def top_k(a, k, /, *, axis=-1, largest=True):
+def top_k(a, k, /, axis=-1, *, largest=True):
     """
     Returns the ``k`` largest/smallest elements and corresponding
     indices along the given ``axis``.