array-api-extra 0.10.0__py3-none-any.whl → 0.10.3__py3-none-any.whl

array_api_extra/__init__.py

@@ -1,13 +1,16 @@
 """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,
     expand_dims,
     isclose,
     isin,
+    kron,
     nan_to_num,
     one_hot,
     pad,
@@ -19,19 +22,19 @@ from ._delegation import (
 )
 from ._lib._at import at
 from ._lib._funcs import (
+    angle,
     apply_where,
-    broadcast_shapes,
     default_dtype,
-    kron,
     nunique,
 )
 from ._lib._lazy import lazy_apply
 
-__version__ = "0.10.0"
+__version__ = "0.10.3"
 
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
+    "angle",
     "apply_where",
     "argpartition",
     "at",
@@ -53,5 +56,6 @@ __all__ = [
     "searchsorted",
     "setdiff1d",
     "sinc",
+    "testing",
     "union1d",
 ]

array_api_extra/_delegation.py

@@ -20,10 +20,12 @@ from ._lib._utils._typing import Array, DType
 
 __all__ = [
     "atleast_nd",
+    "broadcast_shapes",
     "cov",
     "create_diagonal",
     "expand_dims",
     "isclose",
+    "kron",
     "nan_to_num",
     "one_hot",
     "pad",
@@ -81,6 +83,66 @@ def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array
     return _funcs.atleast_nd(x, ndim=ndim, xp=xp)
 
 
+def broadcast_shapes(
+    *shapes: tuple[float | None, ...], xp: ModuleType | None = None
+) -> tuple[int | None, ...]:
+    """
+    Compute the shape of the broadcasted arrays.
+
+    Duplicates :func:`numpy.broadcast_shapes`, with additional support for
+    None and NaN sizes.
+
+    This is equivalent to ``xp.broadcast_arrays(arr1, arr2, ...)[0].shape``
+    without needing to worry about the backend potentially deep copying
+    the arrays.
+
+    Parameters
+    ----------
+    *shapes : tuple[int | None, ...]
+        Shapes of the arrays to broadcast.
+    xp : array_namespace, optional
+        The standard-compatible namespace to use for native delegation.
+        Default: use the array-agnostic implementation.
+
+    Returns
+    -------
+    tuple[int | None, ...]
+        The shape of the broadcasted arrays.
+
+    See Also
+    --------
+    numpy.broadcast_shapes : Equivalent NumPy function.
+    array_api.broadcast_arrays : Function to broadcast actual arrays.
+
+    Notes
+    -----
+    This function accepts the Array API's ``None`` for unknown sizes,
+    as well as Dask's non-standard ``math.nan``.
+    Regardless of input, the output always contains ``None`` for unknown sizes.
+
+    Examples
+    --------
+    >>> import array_api_extra as xpx
+    >>> xpx.broadcast_shapes((2, 3), (2, 1))
+    (2, 3)
+    >>> xpx.broadcast_shapes((4, 2, 3), (2, 1), (1, 3))
+    (4, 2, 3)
+    """
+    if (
+        xp is not None
+        and all(isinstance(size, int) for shape in shapes for size in shape)
+        and (
+            is_numpy_namespace(xp)
+            or is_cupy_namespace(xp)
+            or is_jax_namespace(xp)
+            or is_torch_namespace(xp)
+        )
+    ):
+        return xp.broadcast_shapes(*shapes)
+
+    return _funcs.broadcast_shapes(*shapes)
+
+
 def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
     """
     Estimate a covariance matrix (or a stack of covariance matrices).
@@ -418,6 +480,101 @@ def isclose(
     return _funcs.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, xp=xp)
 
 
+def kron(
+    a: Array | complex,
+    b: Array | complex,
+    /,
+    *,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Kronecker product of two arrays.
+
+    Computes the Kronecker product, a composite array made of blocks of the
+    second array scaled by the first.
+
+    Equivalent to ``numpy.kron`` for NumPy arrays.
+
+    Parameters
+    ----------
+    a, b : Array | int | float | complex
+        Input arrays or scalars. At least one must be an array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a` and `b`. Default: infer.
+
+    Returns
+    -------
+    array
+        The Kronecker product of `a` and `b`.
+
+    Notes
+    -----
+    The function assumes that the number of dimensions of `a` and `b`
+    are the same, if necessary prepending the smallest with ones.
+    If ``a.shape = (r0,r1,..,rN)`` and ``b.shape = (s0,s1,...,sN)``,
+    the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``.
+    The elements are products of elements from `a` and `b`, organized
+    explicitly by::
+
+        kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN]
+
+    where::
+
+        kt = it * st + jt,  t = 0,...,N
+
+    In the common 2-D case (N=1), the block structure can be visualized::
+
+        [[ a[0,0]*b,   a[0,1]*b,  ... , a[0,-1]*b  ],
+         [  ...                              ...   ],
+         [ a[-1,0]*b,  a[-1,1]*b, ... , a[-1,-1]*b ]]
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> xpx.kron(xp.asarray([1, 10, 100]), xp.asarray([5, 6, 7]), xp=xp)
+    Array([  5,   6,   7,  50,  60,  70, 500,
+           600, 700], dtype=array_api_strict.int64)
+
+    >>> xpx.kron(xp.asarray([5, 6, 7]), xp.asarray([1, 10, 100]), xp=xp)
+    Array([  5,  50, 500,   6,  60, 600,   7,
+            70, 700], dtype=array_api_strict.int64)
+
+    >>> xpx.kron(xp.eye(2), xp.ones((2, 2)), xp=xp)
+    Array([[1., 1., 0., 0.],
+           [1., 1., 0., 0.],
+           [0., 0., 1., 1.],
+           [0., 0., 1., 1.]], dtype=array_api_strict.float64)
+
+    >>> a = xp.reshape(xp.arange(100), (2, 5, 2, 5))
+    >>> b = xp.reshape(xp.arange(24), (2, 3, 4))
+    >>> c = xpx.kron(a, b, xp=xp)
+    >>> c.shape
+    (2, 10, 6, 20)
+    >>> I = (1, 3, 0, 2)
+    >>> J = (0, 2, 1)
+    >>> J1 = (0,) + J             # extend to ndim=4
+    >>> S1 = (1,) + b.shape
+    >>> K = tuple(xp.asarray(I) * xp.asarray(S1) + xp.asarray(J1))
+    >>> c[K] == a[I]*b[J]
+    Array(True, dtype=array_api_strict.bool)
+    """
+    if xp is None:
+        xp = array_namespace(a, b)
+
+    a, b = asarrays(a, b, xp=xp)
+
+    if (
+        is_cupy_namespace(xp)
+        or is_jax_namespace(xp)
+        or is_numpy_namespace(xp)
+        or is_torch_namespace(xp)
+    ):
+        return xp.kron(a, b)
+
+    return _funcs.kron(a, b, xp=xp)
+
+
 def nan_to_num(
     x: Array | float | complex,
     /,

array_api_extra/_lib/_at.py

@@ -37,7 +37,7 @@ class _AtOp(Enum):
     MAX = "max"
 
     # @override from Python 3.12
-    def __str__(self) -> str:  # pyright: ignore[reportImplicitOverride]
+    def __str__(self) -> str:  # pyright: ignore[reportImplicitOverride] # pyrefly: ignore[missing-override-decorator]
         """
         Return string representation (useful for pytest logs).
 
@@ -381,7 +381,7 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
         # Note for this and all other methods based on _iop:
         # operator.iadd and operator.add subtly differ in behaviour, as
         # only iadd will trigger exceptions when y has an incompatible dtype.
-        return self._op(_AtOp.ADD, operator.iadd, operator.add, y, copy=copy, xp=xp)
+        return self._op(_AtOp.ADD, operator.iadd, operator.add, y, copy=copy, xp=xp)  # pyright: ignore[reportUnknownArgumentType]
 
     def subtract(
         self,
@@ -392,7 +392,12 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] -= y`` and return the updated array."""
         return self._op(
-            _AtOp.SUBTRACT, operator.isub, operator.sub, y, copy=copy, xp=xp
+            _AtOp.SUBTRACT,
+            operator.isub,  # pyright: ignore[reportUnknownArgumentType]
+            operator.sub,
+            y,
+            copy=copy,
+            xp=xp,
         )
 
     def multiply(
@@ -404,7 +409,12 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] *= y`` and return the updated array."""
         return self._op(
-            _AtOp.MULTIPLY, operator.imul, operator.mul, y, copy=copy, xp=xp
+            _AtOp.MULTIPLY,
+            operator.imul,  # pyright: ignore[reportUnknownArgumentType]
+            operator.mul,
+            y,
+            copy=copy,
+            xp=xp,
         )
 
     def divide(
@@ -416,7 +426,12 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] /= y`` and return the updated array."""
         return self._op(
-            _AtOp.DIVIDE, operator.itruediv, operator.truediv, y, copy=copy, xp=xp
+            _AtOp.DIVIDE,
+            operator.itruediv,  # pyright: ignore[reportUnknownArgumentType]
+            operator.truediv,  # pyright: ignore[reportUnknownArgumentType]
+            y,
+            copy=copy,
+            xp=xp,
         )
 
     def power(
@@ -427,7 +442,7 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
         xp: ModuleType | None = None,
     ) -> Array:  # numpydoc ignore=PR01,RT01
         """Apply ``x[idx] **= y`` and return the updated array."""
-        return self._op(_AtOp.POWER, operator.ipow, operator.pow, y, copy=copy, xp=xp)
+        return self._op(_AtOp.POWER, operator.ipow, operator.pow, y, copy=copy, xp=xp)  # pyright: ignore[reportUnknownArgumentType]
 
     def min(
         self,

array_api_extra/_lib/_backends.py

@@ -47,11 +47,12 @@ class Backend(Enum):  # numpydoc ignore=PR02
 
     def pytest_param(self) -> Any:
         """
-        Backend as a pytest parameter
+        Backend as a pytest parameter.
 
         Returns
         -------
         pytest.mark.ParameterSet
+            The backend as a pytest parameter.
         """
         id_ = (
             self.name.lower().replace("_gpu", ":gpu").replace("_readonly", ":readonly")

array_api_extra/_lib/_funcs.py

@@ -23,6 +23,7 @@ from ._utils._helpers import (
 from ._utils._typing import Array, Device, DType
 
 __all__ = [
+    "angle",
     "apply_where",
     "atleast_nd",
     "broadcast_shapes",
@@ -220,46 +221,10 @@ def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType) -> Array:
 
 # `float` in signature to accept `math.nan` for Dask.
 # `int`s are still accepted as `float` is a superclass of `int` in typing
-def broadcast_shapes(*shapes: tuple[float | None, ...]) -> tuple[int | None, ...]:
-    """
-    Compute the shape of the broadcasted arrays.
-
-    Duplicates :func:`numpy.broadcast_shapes`, with additional support for
-    None and NaN sizes.
-
-    This is equivalent to ``xp.broadcast_arrays(arr1, arr2, ...)[0].shape``
-    without needing to worry about the backend potentially deep copying
-    the arrays.
-
-    Parameters
-    ----------
-    *shapes : tuple[int | None, ...]
-        Shapes of the arrays to broadcast.
-
-    Returns
-    -------
-    tuple[int | None, ...]
-        The shape of the broadcasted arrays.
-
-    See Also
-    --------
-    numpy.broadcast_shapes : Equivalent NumPy function.
-    array_api.broadcast_arrays : Function to broadcast actual arrays.
-
-    Notes
-    -----
-    This function accepts the Array API's ``None`` for unknown sizes,
-    as well as Dask's non-standard ``math.nan``.
-    Regardless of input, the output always contains ``None`` for unknown sizes.
-
-    Examples
-    --------
-    >>> import array_api_extra as xpx
-    >>> xpx.broadcast_shapes((2, 3), (2, 1))
-    (2, 3)
-    >>> xpx.broadcast_shapes((4, 2, 3), (2, 1), (1, 3))
-    (4, 2, 3)
-    """
+def broadcast_shapes(  # numpydoc ignore=PR01,RT01
+    *shapes: tuple[float | None, ...],
+) -> tuple[int | None, ...]:
+    """See docstring in array_api_extra._delegation."""
     if not shapes:
         return ()  # Match NumPy output
 
@@ -291,7 +256,7 @@ def cov(m: Array, /, *, xp: ModuleType) -> Array:  # numpydoc ignore=PR01,RT01
     m = atleast_nd(m, ndim=2, xp=xp)
     m = xp.astype(m, dtype)
 
-    avg = _helpers.mean(m, axis=-1, keepdims=True, xp=xp)
+    avg = xp.mean(m, axis=-1, keepdims=True)
 
     m_shape = eager_shape(m)
     fact = m_shape[-1] - 1
@@ -440,87 +405,13 @@ def isclose(
 
 
 def kron(
-    a: Array | complex,
-    b: Array | complex,
+    a: Array,
+    b: Array,
     /,
     *,
-    xp: ModuleType | None = None,
-) -> Array:
-    """
-    Kronecker product of two arrays.
-
-    Computes the Kronecker product, a composite array made of blocks of the
-    second array scaled by the first.
-
-    Equivalent to ``numpy.kron`` for NumPy arrays.
-
-    Parameters
-    ----------
-    a, b : Array | int | float | complex
-        Input arrays or scalars. At least one must be an array.
-    xp : array_namespace, optional
-        The standard-compatible namespace for `a` and `b`. Default: infer.
-
-    Returns
-    -------
-    array
-        The Kronecker product of `a` and `b`.
-
-    Notes
-    -----
-    The function assumes that the number of dimensions of `a` and `b`
-    are the same, if necessary prepending the smallest with ones.
-    If ``a.shape = (r0,r1,..,rN)`` and ``b.shape = (s0,s1,...,sN)``,
-    the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``.
-    The elements are products of elements from `a` and `b`, organized
-    explicitly by::
-
-        kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN]
-
-    where::
-
-        kt = it * st + jt,  t = 0,...,N
-
-    In the common 2-D case (N=1), the block structure can be visualized::
-
-        [[ a[0,0]*b,   a[0,1]*b,  ... , a[0,-1]*b  ],
-         [  ...                              ...   ],
-         [ a[-1,0]*b,  a[-1,1]*b, ... , a[-1,-1]*b ]]
-
-    Examples
-    --------
-    >>> import array_api_strict as xp
-    >>> import array_api_extra as xpx
-    >>> xpx.kron(xp.asarray([1, 10, 100]), xp.asarray([5, 6, 7]), xp=xp)
-    Array([  5,   6,   7,  50,  60,  70, 500,
-           600, 700], dtype=array_api_strict.int64)
-
-    >>> xpx.kron(xp.asarray([5, 6, 7]), xp.asarray([1, 10, 100]), xp=xp)
-    Array([  5,  50, 500,   6,  60, 600,   7,
-            70, 700], dtype=array_api_strict.int64)
-
-    >>> xpx.kron(xp.eye(2), xp.ones((2, 2)), xp=xp)
-    Array([[1., 1., 0., 0.],
-           [1., 1., 0., 0.],
-           [0., 0., 1., 1.],
-           [0., 0., 1., 1.]], dtype=array_api_strict.float64)
-
-    >>> a = xp.reshape(xp.arange(100), (2, 5, 2, 5))
-    >>> b = xp.reshape(xp.arange(24), (2, 3, 4))
-    >>> c = xpx.kron(a, b, xp=xp)
-    >>> c.shape
-    (2, 10, 6, 20)
-    >>> I = (1, 3, 0, 2)
-    >>> J = (0, 2, 1)
-    >>> J1 = (0,) + J             # extend to ndim=4
-    >>> S1 = (1,) + b.shape
-    >>> K = tuple(xp.asarray(I) * xp.asarray(S1) + xp.asarray(J1))
-    >>> c[K] == a[I]*b[J]
-    Array(True, dtype=array_api_strict.bool)
-    """
-    if xp is None:
-        xp = array_namespace(a, b)
-    a, b = asarrays(a, b, xp=xp)
+    xp: ModuleType,
+) -> Array:  # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
 
     singletons = (1,) * (b.ndim - a.ndim)
     a = cast(Array, xp.broadcast_to(a, singletons + a.shape))
@@ -818,3 +709,51 @@ def union1d(a: Array, b: Array, /, *, xp: ModuleType) -> Array:
     b = xp.reshape(b, (-1,))
     # XXX: `sparse` returns NumPy arrays from `unique_values`
     return xp.asarray(xp.unique_values(xp.concat([a, b])))
+
+
+def angle(z: Array, /, *, deg: bool = False, xp: ModuleType | None = None) -> Array:
+    """
+    Return the angle of the complex argument.
+
+    Parameters
+    ----------
+    z : Array
+        Input array.
+    deg : bool, optional
+        Return angle in degrees if True, radians if False (default).
+    xp : array_namespace, optional
+        The standard-compatible namespace for `z`. Default: infer.
+
+    Returns
+    -------
+    array
+        The counterclockwise angle from the positive real axis on the complex
+        plane in the range ``(-pi, pi]``.
+
+    Notes
+    -----
+    Real input ``x`` is interpreted as ``x + 0j``.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), xp=xp)
+    Array([0.        , 1.57079633, 0.78539816], dtype=array_api_strict.float64)
+    >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), deg=True, xp=xp)
+    Array([ 0., 90., 45.], dtype=array_api_strict.float64)
+    """
+    if xp is None:
+        xp = 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)))
+        zimag = xp.zeros_like(z)
+        zreal = z
+    a = xp.atan2(zimag, zreal)
+    if deg:
+        a = a * 180 / xp.pi
+    return a

array_api_extra/_lib/_lazy.py

@@ -30,7 +30,7 @@ else:
 P = ParamSpec("P")
 
 
-@overload
+@overload  # pyrefly: ignore[invalid-param-spec]
 def lazy_apply(  # type: ignore[valid-type]
     func: Callable[P, Array | ArrayLike],
     *args: Array | complex | None,
@@ -42,7 +42,7 @@ def lazy_apply(  # type: ignore[valid-type]
 ) -> Array: ...  # numpydoc ignore=GL08
 
 
-@overload
+@overload  # pyrefly: ignore[invalid-param-spec]
 def lazy_apply(  # type: ignore[valid-type]
     func: Callable[P, Sequence[Array | ArrayLike]],
     *args: Array | complex | None,
@@ -54,7 +54,7 @@ def lazy_apply(  # type: ignore[valid-type]
 ) -> tuple[Array, ...]: ...  # numpydoc ignore=GL08
 
 
-def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
+def lazy_apply(  # type: ignore[valid-type]  # pyrefly: ignore[invalid-param-spec]  # numpydoc ignore=GL07,SA04
     func: Callable[P, Array | ArrayLike | Sequence[Array | ArrayLike]],
     *args: Array | complex | None,
     shape: tuple[int | None, ...] | Sequence[tuple[int | None, ...]] | None = None,
@@ -240,7 +240,7 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
     if is_dask_namespace(xp):
         import dask
 
-        metas: list[Array] = [arg._meta for arg in array_args]  # pylint: disable=protected-access    # pyright: ignore[reportAttributeAccessIssue]
+        metas: list[Array] = [arg._meta for arg in array_args]  # pylint: disable=protected-access # pyright: ignore[reportAttributeAccessIssue] # pyrefly: ignore[missing-attribute]
         meta_xp = array_namespace(*metas)
 
         wrapped = dask.delayed(  # type: ignore[attr-defined]  # pyright: ignore[reportPrivateImportUsage]