STY/DOC: reduce from imports, fix internal docstrings

docs/api-creation.md

@@ -1,4 +1,4 @@
-# Creation functions
+# Creation Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/api-elementwise.md

@@ -1,4 +1,4 @@
-# Element-wise functions
+# Element-wise Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/api-indexing.md

@@ -1,4 +1,4 @@
-# Indexing functions
+# Indexing Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/api-linalg.md

@@ -1,4 +1,4 @@
-# Linear algebra functions
+# Linear algebra Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/api-manipulation.md

@@ -1,4 +1,4 @@
-# Manipulation functions
+# Manipulation Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/api-reference.md

@@ -1,4 +1,4 @@
-# API reference
+# API Reference
 
 ```{toctree}
 :maxdepth: 2

docs/api-searching.md

@@ -1,4 +1,4 @@
-# Searching functions
+# Searching Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/api-set.md

@@ -1,4 +1,4 @@
-# Set functions
+# Set Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/api-sorting.md

@@ -1,4 +1,4 @@
-# Sorting functions
+# Sorting Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/api-statistical.md

@@ -1,4 +1,4 @@
-# Statistical functions
+# Statistical Functions
 
 ```{eval-rst}
 .. currentmodule:: array_api_extra

docs/contributing.md

@@ -51,19 +51,19 @@ pixi shell --environment=dev
   propose the new function. You may want to wait for initial feedback on the
   issue before diving into an implementation.
   Skip this step if there is already an open issue for the function.
-- Add the implementation of your function to
-  `src/array_api_extra/_lib/_funcs.py`.
+- Add the implementation of your array-agnostic function to an appropriate file in
+  `src/array_api_extra/_agnostic/`.
   - Ensure that your function includes type annotations and a
     [numpydoc-style docstring](https://numpydoc.readthedocs.io/en/latest/format.html).
   - Add your function to `__all__` at the top of the file.
-- Import your function to `src/array_api_extra/__init__.py` and add it to
-  `__all__` there.
-- Add a test class for your function in `tests/test_funcs.py`.
+- Import your function to `src/array_api_extra/__init__.py`
+  and add it to`__all__` there.
+- Add a test class for your function in the corresponding file under `tests/`.
   - Ensure that `lazy_xp_function` is called on the function if lazy backends
     are supposed to be tested.
-- Add your function to a suitable page under `docs/`.
+- Add your function to the corresponding API page under `docs/`.
 - Don't worry if you are not sure how to do some of the above steps or think you
-  might have done something wrong -
+  might have done something wrong —
   [make a PR!](https://github.com/data-apis/array-api-extra/pulls)
 
 ### Delegation
@@ -72,17 +72,18 @@ Many new functions should also have 'delegation' to existing implementations in
 known array libraries added. This can happen in the same PR which adds the function,
 or in a follow-up PR.
 
-- Create a function in `src/array_api_extra/_delegation.py` with a signature
-  matching the function in `src/array_api_extra/_lib/_funcs.py`, and move the
-  docstring to the new function. Leave a one-line docstring in `_funcs.py`,
-  pointing to `_delegation.py` to see the full docstring.
-- Also move the initial `array_namespace` call and any input validation over to
-  the new function.
-- Add delegation to backends using the `if_*_namespace` functions. See
-  `src/array_api_extra/_lib/_backends.py` for the full list of backends we have
-  worked with so far.
-- After all delegation layers, return the result from the implementation in
-  `_funcs`.
-- Simplify the signature in `_funcs.py` to remove impossible arguments now that
-  it is only called internally via `_delegation`. For example, the `xp`
-  parameter can be changed from type `ArrayNamespace | None` to `ArrayNamespace`.
+- Create a function in the corresponding file under `src/array_api_extra/`
+  with a signature matching the function under `src/array_api_extra/_agnostic/`,
+  and move the docstring to the new function. Leave a one-line docstring in the
+  array-agnostic function, pointing to the new function.
+- Also move the initial `_compat.array_namespace` call
+  and any input validation over to the new function.
+- Add delegation to backends using the `_compat.if_*_namespace` functions.
+  See `src/array_api_extra/_lib/_backends.py` for the full list of backends
+  we have worked with so far.
+- After all delegation layers,
+  return the result from the implementation under `_agnostic/`
+- Simplify the signature of the function under `_agnostic/`
+  to remove impossible arguments now that it is only called internally.
+  For example, the `xp` parameter can be changed from type `ArrayNamespace | None`
+  to `ArrayNamespace`.

src/array_api_extra/_agnostic/_creation.py

@@ -1,8 +1,7 @@
 """Array-agnostic implementations for creation functions."""
 
 from .._at import at
-from .._lib import _compat
-from .._lib._helpers import eager_shape, ndindex
+from .._lib import _compat, _helpers
 from .._lib._typing import Array, ArrayNamespace
 
 __all__ = ["create_diagonal", "one_hot"]
@@ -11,8 +10,8 @@
 def create_diagonal(
     x: Array, /, *, offset: int = 0, xp: ArrayNamespace
 ) -> Array:  # numpydoc ignore=PR01,RT01
-    """See docstring in array_api_extra._delegation."""
-    x_shape = eager_shape(x)
+    """See docstring in `array_api_extra._creation`."""
+    x_shape = _helpers.eager_shape(x)
     batch_dims = x_shape[:-1]
     n = x_shape[-1] + abs(offset)
     diag = xp.zeros((*batch_dims, n**2), dtype=x.dtype, device=_compat.device(x))
@@ -22,7 +21,7 @@ def create_diagonal(
         min(n * (n - offset), diag.shape[-1]),
         n + 1,
     )
-    for index in ndindex(*batch_dims):
+    for index in _helpers.ndindex(*batch_dims):
         diag = at(diag)[(*index, target_slice)].set(x[(*index, slice(None))])
     return xp.reshape(diag, (*batch_dims, n, n))
 
@@ -34,7 +33,7 @@ def one_hot(
     *,
     xp: ArrayNamespace,
 ) -> Array:  # numpydoc ignore=PR01,RT01
-    """See docstring in `array_api_extra._delegation.py`."""
+    """See docstring in `array_api_extra._creation`."""
     # TODO: Benchmark whether this is faster on the NumPy backend:
     # if is_numpy_array(x):
     #     out = xp.zeros((x.size, num_classes), dtype=dtype)

src/array_api_extra/_agnostic/_elementwise.py

@@ -1,15 +1,13 @@
 """Array-agnostic implementations for element-wise functions."""
 
+import typing
 from collections.abc import Callable
 from types import NoneType
-from typing import overload
 
 from .._at import at
-from .._lib import _compat
-from .._lib._compat import array_namespace, is_dask_namespace
-from .._lib._helpers import asarrays, capabilities, meta_namespace
+from .._lib import _compat, _helpers
 from .._lib._typing import Array, ArrayNamespace
-from ._inspection import default_dtype
+from . import _inspection
 
 __all__ = [
     "angle",
@@ -22,7 +20,7 @@
 ]
 
 
-@overload
+@typing.overload
 def apply_where(  # numpydoc ignore=GL08
     cond: Array,
     args: Array | tuple[Array, ...],
@@ -35,7 +33,7 @@ def apply_where(  # numpydoc ignore=GL08
 ) -> Array: ...
 
 
-@overload
+@typing.overload
 def apply_where(  # numpydoc ignore=GL08
     cond: Array,
     args: Array | tuple[Array, ...],
@@ -129,15 +127,15 @@ def apply_where(  # numpydoc ignore=PR01,PR02
     args_ = [*args_, *kwargs_.values()]
     del kwargs
 
-    xp = array_namespace(cond, fill_value, *args_) if xp is None else xp
+    xp = _compat.array_namespace(cond, fill_value, *args_) if xp is None else xp
 
     if isinstance(fill_value, int | float | complex | NoneType):
         cond, *args_ = xp.broadcast_arrays(cond, *args_)
     else:
         cond, fill_value, *args_ = xp.broadcast_arrays(cond, fill_value, *args_)
 
-    if is_dask_namespace(xp):
-        meta_xp = meta_namespace(cond, fill_value, *args_, xp=xp)
+    if _compat.is_dask_namespace(xp):
+        meta_xp = _helpers.meta_namespace(cond, fill_value, *args_, xp=xp)
         # map_blocks doesn't descend into tuples of Arrays
         return xp.map_blocks(
             _apply_where, cond, f1, f2, fill_value, *args_, kwkeys=kwkeys, xp=meta_xp
@@ -161,7 +159,7 @@ def _apply_where(  # numpydoc ignore=PR01,RT01
     kwargs = dict(zip(kwkeys, args[nargs:], strict=True))
     args = args[:nargs]
 
-    if not capabilities(xp, device=_compat.device(cond))["boolean indexing"]:
+    if not _helpers.capabilities(xp, cond)["boolean indexing"]:
         # jax.jit does not support assignment by boolean mask
         return xp.where(
             cond,
@@ -201,14 +199,14 @@ def isclose(
     equal_nan: bool = False,
     xp: ArrayNamespace,
 ) -> Array:  # numpydoc ignore=PR01,RT01
-    """See docstring in array_api_extra._delegation."""
-    a, b = asarrays(a, b, xp=xp)
+    """See docstring in `array_api_extra._elementwise`."""
+    a, b = _helpers.asarrays(a, b, xp=xp)
 
     a_inexact = xp.isdtype(a.dtype, ("real floating", "complex floating"))
     b_inexact = xp.isdtype(b.dtype, ("real floating", "complex floating"))
     if a_inexact or b_inexact:
         # prevent warnings on NumPy and Dask on inf - inf
-        mxp = meta_namespace(a, b, xp=xp)
+        mxp = _helpers.meta_namespace(a, b, xp=xp)
         out = apply_where(
             xp.isinf(a) | xp.isinf(b),
             (a, b),
@@ -246,7 +244,7 @@ def nan_to_num(  # numpydoc ignore=PR01,RT01
     *,
     xp: ArrayNamespace,
 ) -> Array:
-    """See docstring in `array_api_extra._delegation.py`."""
+    """See docstring in `array_api_extra._elementwise`."""
 
     def perform_replacements(  # numpydoc ignore=PR01,RT01
         x: Array,
@@ -282,7 +280,7 @@ def perform_replacements(  # numpydoc ignore=PR01,RT01
 
 def sinc(x: Array, /, *, xp: ArrayNamespace) -> Array:
     # numpydoc ignore=PR01,RT01
-    """See docstring in `array_api_extra._delegation.py`."""
+    """See docstring in `array_api_extra._elementwise`."""
 
     # no scalars in `where` - array-api#807
     y = xp.pi * xp.where(
@@ -326,13 +324,13 @@ def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -
     Array([ 0., 90., 45.], dtype=array_api_strict.float64)
     """
     if xp is None:
-        xp = array_namespace(z)
+        xp = _compat.array_namespace(z)
     if xp.isdtype(z.dtype, "complex floating"):
         zimag = xp.imag(z)
         zreal = xp.real(z)
     else:
         if not xp.isdtype(z.dtype, "real floating"):
-            z = xp.astype(z, default_dtype(xp, device=_compat.device(z)))
+            z = xp.astype(z, _inspection.default_dtype(xp, device=_compat.device(z)))
         zimag = xp.zeros_like(z)
         zreal = z
     a = xp.atan2(zimag, zreal)
@@ -343,11 +341,11 @@ def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -
 
 def deg2rad(x: Array, /, *, xp: ArrayNamespace) -> Array:
     # numpydoc ignore=PR01,RT01
-    """See docstring in `array_api_extra._delegation.py`."""
+    """See docstring in `array_api_extra._elementwise`."""
     return x * xp.pi / 180
 
 
 def rad2deg(x: Array, /, *, xp: ArrayNamespace) -> Array:
     # numpydoc ignore=PR01,RT01
-    """See docstring in `array_api_extra._delegation.py`."""
+    """See docstring in `array_api_extra._elementwise`."""
     return x * 180 / xp.pi

src/array_api_extra/_agnostic/_indexing.py

@@ -8,7 +8,7 @@
 def diag_indices(
     n: int, /, *, ndim: int, device: Device | None, xp: ArrayNamespace
 ) -> tuple[Array, ...]:  # numpydoc ignore=PR01,RT01
-    """See docstring in array_api_extra._delegation."""
+    """See docstring in `array_api_extra._indexing`."""
     idx = xp.arange(n, device=device)
     return (idx,) * ndim
 
@@ -41,7 +41,7 @@ def tril_indices(
     device: Device | None,
     xp: ArrayNamespace,
 ) -> tuple[Array, Array]:  # numpydoc ignore=PR01,RT01
-    """See docstring in array_api_extra._delegation."""
+    """See docstring in `array_api_extra._indexing`."""
     return _tri_indices(n, offset=offset, m=m, upper=False, device=device, xp=xp)
 
 
@@ -54,13 +54,13 @@ def triu_indices(
     device: Device | None,
     xp: ArrayNamespace,
 ) -> tuple[Array, Array]:  # numpydoc ignore=PR01,RT01
-    """See docstring in array_api_extra._delegation."""
+    """See docstring in `array_api_extra._indexing`."""
     return _tri_indices(n, offset=offset, m=m, upper=True, device=device, xp=xp)
 
 
 def unravel_index(indices: Array, shape: tuple[int, ...], /) -> tuple[Array, ...]:
     # numpydoc ignore=PR01,RT01
-    """See docstring in `array_api_extra._delegation.py`."""
+    """See docstring in `array_api_extra._indexing`."""
     coords: list[Array] = []
     for dim in reversed(shape):
         coords.append(indices % dim)

src/array_api_extra/_agnostic/_linalg.py

@@ -1,10 +1,10 @@
 """Array-agnostic implementations for linear algebra functions."""
 
-from typing import cast
+import typing
 
-from .._lib._helpers import eager_shape
+from .._lib import _helpers
 from .._lib._typing import Array, ArrayNamespace
-from ._manipulation import expand_dims
+from . import _manipulation
 
 __all__ = ["kron"]
 
@@ -16,30 +16,30 @@ def kron(
     *,
     xp: ArrayNamespace,
 ) -> Array:  # numpydoc ignore=PR01,RT01
-    """See docstring in array_api_extra._delegation."""
+    """See docstring in `array_api_extra._linalg`."""
 
     singletons = (1,) * (b.ndim - a.ndim)
-    a = cast(Array, xp.broadcast_to(a, singletons + a.shape))
+    a = typing.cast(Array, xp.broadcast_to(a, singletons + a.shape))
 
     nd_b, nd_a = b.ndim, a.ndim
     nd_max = max(nd_b, nd_a)
     if nd_a == 0 or nd_b == 0:
         return xp.multiply(a, b)
 
-    a_shape = eager_shape(a)
-    b_shape = eager_shape(b)
+    a_shape = _helpers.eager_shape(a)
+    b_shape = _helpers.eager_shape(b)
 
     # Equalise the shapes by prepending smaller one with 1s
     a_shape = (1,) * max(0, nd_b - nd_a) + a_shape
     b_shape = (1,) * max(0, nd_a - nd_b) + b_shape
 
     # Insert empty dimensions
-    a_arr = expand_dims(a, axis=tuple(range(nd_b - nd_a)), xp=xp)
-    b_arr = expand_dims(b, axis=tuple(range(nd_a - nd_b)), xp=xp)
+    a_arr = _manipulation.expand_dims(a, axis=tuple(range(nd_b - nd_a)), xp=xp)
+    b_arr = _manipulation.expand_dims(b, axis=tuple(range(nd_a - nd_b)), xp=xp)
 
     # Compute the product
-    a_arr = expand_dims(a_arr, axis=tuple(range(1, nd_max * 2, 2)), xp=xp)
-    b_arr = expand_dims(b_arr, axis=tuple(range(0, nd_max * 2, 2)), xp=xp)
+    a_arr = _manipulation.expand_dims(a_arr, axis=tuple(range(1, nd_max * 2, 2)), xp=xp)
+    b_arr = _manipulation.expand_dims(b_arr, axis=tuple(range(0, nd_max * 2, 2)), xp=xp)
     result = xp.multiply(a_arr, b_arr)
 
     # Reshape back and return