merge

Author: crusaderky

pixi.lock

@@ -15,9 +15,8 @@
     is_jax_array,
     is_writeable_array,
 )
-from ._utils._typing import Array, SetIndex
 from ._utils._helpers import meta_namespace
-from ._utils._typing import Array, Index
+from ._utils._typing import Array, SetIndex
 
 
 class _AtOp(Enum):

src/array_api_extra/_lib/_at.py

@@ -7,19 +7,17 @@
 import warnings
 from collections.abc import Callable, Sequence
 from types import ModuleType
-from typing import cast
+from typing import cast, overload
 
 from ._at import at
 from ._utils import _compat, _helpers
-from ._utils._compat import array_namespace, is_jax_array
-from ._utils._helpers import asarrays, eager_shape, ndindex
 from ._utils._compat import (
     array_namespace,
     is_dask_namespace,
     is_jax_array,
     is_jax_namespace,
 )
-from ._utils._helpers import asarrays, meta_namespace, ndindex
+from ._utils._helpers import asarrays, eager_shape, meta_namespace, ndindex
 from ._utils._typing import Array
 
 __all__ = [

src/array_api_extra/_lib/_funcs.py

@@ -28,9 +28,7 @@
 from array_api_extra._lib import Backend
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
-from array_api_extra._lib._utils._helpers import eager_shape, ndindex
-from array_api_extra._lib._utils._typing import Array, Device
-from array_api_extra._lib._utils._helpers import asarrays, ndindex
+from array_api_extra._lib._utils._helpers import asarrays, eager_shape, ndindex
 from array_api_extra._lib._utils._typing import Array, Device, DType
 from array_api_extra.testing import lazy_xp_function
 

tests/test_funcs.py

@@ -8,10 +8,14 @@
 from array_api_extra._lib._testing import xp_assert_equal
 from array_api_extra._lib._utils._compat import array_namespace
 from array_api_extra._lib._utils._compat import device as get_device
-from array_api_extra._lib._utils._helpers import asarrays, eager_shape, in1d, ndindex
+from array_api_extra._lib._utils._helpers import (
+    asarrays,
+    eager_shape,
+    in1d,
+    meta_namespace,
+    ndindex,
+)
 from array_api_extra._lib._utils._typing import Array, Device, DType
-from array_api_extra._lib._utils._helpers import asarrays, in1d, meta_namespace, ndindex
-from array_api_extra._lib._utils._typing import Device
 from array_api_extra.testing import lazy_xp_function
 
 # mypy: disable-error-code=no-untyped-usage