MAINT: refactor library to split up large files

meson.build

@@ -10,25 +10,45 @@ py = import('python').find_installation()
 sources = {
   'array_api_extra': files(
     'src/array_api_extra/__init__.py',
-    'src/array_api_extra/_delegation.py',
+    'src/array_api_extra/_at.py',
+    'src/array_api_extra/_creation.py',
+    'src/array_api_extra/_elementwise.py',
+    'src/array_api_extra/_indexing.py',
+    'src/array_api_extra/_lazy.py',
+    'src/array_api_extra/_linalg.py',
+    'src/array_api_extra/_manipulation.py',
+    'src/array_api_extra/_searching.py',
+    'src/array_api_extra/_set.py',
+    'src/array_api_extra/_sorting.py',
+    'src/array_api_extra/_statistical.py',
     'src/array_api_extra/py.typed',
-    'src/array_api_extra/testing.py',
+  ),
+  'array_api_extra/_agnostic': files(
+    'src/array_api_extra/_agnostic/__init__.py',
+    'src/array_api_extra/_agnostic/_creation.py',
+    'src/array_api_extra/_agnostic/_elementwise.py',
+    'src/array_api_extra/_agnostic/_indexing.py',
+    'src/array_api_extra/_agnostic/_inspection.py',
+    'src/array_api_extra/_agnostic/_linalg.py',
+    'src/array_api_extra/_agnostic/_manipulation.py',
+    'src/array_api_extra/_agnostic/_searching.py',
+    'src/array_api_extra/_agnostic/_set.py',
+    'src/array_api_extra/_agnostic/_sorting.py',
+    'src/array_api_extra/_agnostic/_statistical.py',
+  ),
+  'array_api_extra/testing': files(
+    'src/array_api_extra/testing/__init__.py',
+    'src/array_api_extra/testing/_testing.py',
   ),
   'array_api_extra/_lib': files(
     'src/array_api_extra/_lib/__init__.py',
-    'src/array_api_extra/_lib/_at.py',
     'src/array_api_extra/_lib/_backends.py',
-    'src/array_api_extra/_lib/_funcs.py',
-    'src/array_api_extra/_lib/_lazy.py',
+    'src/array_api_extra/_lib/_compat.py',
+    'src/array_api_extra/_lib/_compat.pyi',
+    'src/array_api_extra/_lib/_helpers.py',
     'src/array_api_extra/_lib/_testing.py',
-  ),
-  'array_api_extra/_lib/_utils': files(
-    'src/array_api_extra/_lib/_utils/__init__.py',
-    'src/array_api_extra/_lib/_utils/_compat.py',
-    'src/array_api_extra/_lib/_utils/_compat.pyi',
-    'src/array_api_extra/_lib/_utils/_helpers.py',
-    'src/array_api_extra/_lib/_utils/_typing.py',
-    'src/array_api_extra/_lib/_utils/_typing.pyi',
+    'src/array_api_extra/_lib/_typing.py',
+    'src/array_api_extra/_lib/_typing.pyi',
   ),
 }
 

src/array_api_extra/__init__.py

@@ -1,46 +1,22 @@
 """Extra array functions built on top of the array API standard."""
 
 from . import testing
-from ._delegation import (
-    argpartition,
-    atleast_nd,
-    broadcast_shapes,
-    cov,
-    create_diagonal,
-    deg2rad,
-    diag_indices,
-    expand_dims,
-    isclose,
-    isin,
-    kron,
-    nan_to_num,
-    nanmax,
-    nanmin,
-    nansum,
-    nunique,
-    one_hot,
-    pad,
-    partition,
-    rad2deg,
-    searchsorted,
-    setdiff1d,
-    sinc,
-    tril_indices,
-    triu_indices,
-    union1d,
-    unravel_index,
-)
-from ._lib._at import at
-from ._lib._funcs import (
-    angle,
-    apply_where,
-    default_dtype,
-)
-from ._lib._lazy import lazy_apply
+from ._agnostic._elementwise import angle, apply_where
+from ._agnostic._inspection import default_dtype
+from ._at import at
+from ._creation import create_diagonal, one_hot
+from ._elementwise import deg2rad, isclose, nan_to_num, rad2deg, sinc
+from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index
+from ._lazy import lazy_apply
+from ._linalg import kron
+from ._manipulation import atleast_nd, broadcast_shapes, expand_dims, pad
+from ._searching import searchsorted
+from ._set import isin, nunique, setdiff1d, union1d
+from ._sorting import argpartition, partition
+from ._statistical import cov, nanmax, nanmin, nansum
 
 __version__ = "0.11.2.dev0"
 
-# pylint: disable=duplicate-code
 __all__ = [
     "__version__",
     "angle",

src/array_api_extra/_agnostic/__init__.py

@@ -0,0 +1,27 @@
+"""Array-agnostic function implementations."""
+
+from . import (
+    _creation,
+    _elementwise,
+    _indexing,
+    _inspection,
+    _linalg,
+    _manipulation,
+    _searching,
+    _set,
+    _sorting,
+    _statistical,
+)
+
+__all__ = [
+    "_creation",
+    "_elementwise",
+    "_indexing",
+    "_inspection",
+    "_linalg",
+    "_manipulation",
+    "_searching",
+    "_set",
+    "_sorting",
+    "_statistical",
+]

src/array_api_extra/_agnostic/_creation.py

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