Merge branch 'main' of https://github.com/data-apis/array-api into feat/broadcast_shapes

kgryte · kgryte · commit 56064e7ae540 · 2026-02-18T22:21:55.000-08:00
diff --git a/spec/draft/API_specification/set_functions.rst b/spec/draft/API_specification/set_functions.rst
@@ -18,6 +18,7 @@ Objects in API
    :toctree: generated
    :template: method.rst
 
+   isin
    unique_all
    unique_counts
    unique_inverse
diff --git a/src/array_api_stubs/_2021_12/manipulation_functions.py b/src/array_api_stubs/_2021_12/manipulation_functions.py
@@ -24,7 +24,7 @@ def concat(
     """
 
 
-def expand_dims(x: array, /, *, axis: int = 0) -> array:
+def expand_dims(x: array, /, axis: int) -> array:
     """
     Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by ``axis``.
 
diff --git a/src/array_api_stubs/_2022_12/elementwise_functions.py b/src/array_api_stubs/_2022_12/elementwise_functions.py
@@ -2290,8 +2290,8 @@ def sqrt(x: array, /) -> array:
     - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, the result is ``+0 + 0j``.
     - If ``a`` is any value (including ``NaN``) and ``b`` is ``+infinity``, the result is ``+infinity + infinity j``.
     - If ``a`` is a finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``.
-    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``.
-    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+infinity + 0 j``.
     - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, the result is ``NaN + infinity j`` (sign of the imaginary component is unspecified).
     - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, the result is ``+infinity + NaN j``.
     - If ``a`` is ``NaN`` and ``b`` is any value, the result is ``NaN + NaN j``.
diff --git a/src/array_api_stubs/_2022_12/manipulation_functions.py b/src/array_api_stubs/_2022_12/manipulation_functions.py
@@ -58,7 +58,7 @@ def concat(
     """
 
 
-def expand_dims(x: array, /, *, axis: int = 0) -> array:
+def expand_dims(x: array, /, axis: int) -> array:
     """
     Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by ``axis``.
 
diff --git a/src/array_api_stubs/_2023_12/elementwise_functions.py b/src/array_api_stubs/_2023_12/elementwise_functions.py
@@ -2598,8 +2598,8 @@ def sqrt(x: array, /) -> array:
     - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, the result is ``+0 + 0j``.
     - If ``a`` is any value (including ``NaN``) and ``b`` is ``+infinity``, the result is ``+infinity + infinity j``.
     - If ``a`` is a finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``.
-    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``.
-    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+infinity + 0 j``.
     - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, the result is ``NaN + infinity j`` (sign of the imaginary component is unspecified).
     - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, the result is ``+infinity + NaN j``.
     - If ``a`` is ``NaN`` and ``b`` is any value, the result is ``NaN + NaN j``.
diff --git a/src/array_api_stubs/_2023_12/manipulation_functions.py b/src/array_api_stubs/_2023_12/manipulation_functions.py
@@ -76,7 +76,7 @@ def concat(
     """
 
 
-def expand_dims(x: array, /, *, axis: int = 0) -> array:
+def expand_dims(x: array, /, axis: int) -> array:
     """
     Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by ``axis``.
 
diff --git a/src/array_api_stubs/_2024_12/elementwise_functions.py b/src/array_api_stubs/_2024_12/elementwise_functions.py
@@ -2872,8 +2872,8 @@ def sqrt(x: array, /) -> array:
     - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, the result is ``+0 + 0j``.
     - If ``a`` is any value (including ``NaN``) and ``b`` is ``+infinity``, the result is ``+infinity + infinity j``.
     - If ``a`` is a finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``.
-    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``.
-    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+infinity + 0 j``.
     - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, the result is ``NaN + infinity j`` (sign of the imaginary component is unspecified).
     - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, the result is ``+infinity + NaN j``.
     - If ``a`` is ``NaN`` and ``b`` is any value, the result is ``NaN + NaN j``.
diff --git a/src/array_api_stubs/_2024_12/manipulation_functions.py b/src/array_api_stubs/_2024_12/manipulation_functions.py
@@ -79,7 +79,7 @@ def concat(
     """
 
 
-def expand_dims(x: array, /, *, axis: int = 0) -> array:
+def expand_dims(x: array, /, axis: int) -> array:
     """
     Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by ``axis``.
 
diff --git a/src/array_api_stubs/_draft/creation_functions.py b/src/array_api_stubs/_draft/creation_functions.py
@@ -451,7 +451,7 @@ def linspace(
     """
 
 
-def meshgrid(*arrays: array, indexing: Literal["xy", "ij"] = "xy") -> List[array]:
+def meshgrid(*arrays: array, indexing: Literal["xy", "ij"] = "xy") -> Tuple[array, ...]:
     """
     Returns coordinate matrices from coordinate vectors.
 
@@ -464,8 +464,8 @@ def meshgrid(*arrays: array, indexing: Literal["xy", "ij"] = "xy") -> List[array
 
     Returns
     -------
-    out: List[array]
-        list of N arrays, where ``N`` is the number of provided one-dimensional input arrays. Each returned array must have rank ``N``. For ``N`` one-dimensional arrays having lengths ``Ni = len(xi)``,
+    out: Tuple[array, ...]
+        tuple of N arrays, where ``N`` is the number of provided one-dimensional input arrays. Each returned array must have rank ``N``. For ``N`` one-dimensional arrays having lengths ``Ni = len(xi)``,
 
         - if matrix indexing ``ij``, then each returned array must have the shape ``(N1, N2, N3, ..., Nn)``.
         - if Cartesian indexing ``xy``, then each returned array must have shape ``(N2, N1, N3, ..., Nn)``.
diff --git a/src/array_api_stubs/_draft/elementwise_functions.py b/src/array_api_stubs/_draft/elementwise_functions.py
@@ -2872,8 +2872,8 @@ def sqrt(x: array, /) -> array:
     - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, the result is ``+0 + 0j``.
     - If ``a`` is any value (including ``NaN``) and ``b`` is ``+infinity``, the result is ``+infinity + infinity j``.
     - If ``a`` is a finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``.
-    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``.
-    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+infinity + 0 j``.
     - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, the result is ``NaN + infinity j`` (sign of the imaginary component is unspecified).
     - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, the result is ``+infinity + NaN j``.
     - If ``a`` is ``NaN`` and ``b`` is any value, the result is ``NaN + NaN j``.
diff --git a/src/array_api_stubs/_draft/info.py b/src/array_api_stubs/_draft/info.py
@@ -174,19 +174,19 @@ def dtypes(
     """
 
 
-def devices() -> List[device]:
+def devices() -> Tuple[device, ...]:
     """
     Returns a list of supported devices which are available at runtime.
 
     Returns
     -------
-    out: List[device]
-        a list of supported devices.
+    out: Tuple[device, ...]
+        tuple containing a list of supported devices.
 
     Notes
     -----
 
-    -   Each device object (see :ref:`device-support`) in the list of returned devices **must** be an object which can be provided as a valid keyword-argument to array creation functions.
+    -   Each device object (see :ref:`device-support`) in the tuple of returned devices **must** be an object which can be provided as a valid keyword-argument to array creation functions.
 
     Notes
     -----
diff --git a/src/array_api_stubs/_draft/linalg.py b/src/array_api_stubs/_draft/linalg.py
@@ -187,9 +187,6 @@ def eigh(x: array, /) -> Tuple[array, array]:
     .. note::
        Whether an array library explicitly checks whether an input array is Hermitian or a symmetric matrix (or a stack of matrices) is implementation-defined.
 
-    .. note::
-       The function ``eig`` will be added in a future version of the specification.
-
     Parameters
     ----------
     x: array
@@ -201,7 +198,7 @@ def eigh(x: array, /) -> Tuple[array, array]:
         a namedtuple (``eigenvalues``, ``eigenvectors``) whose
 
         -   first element must have the field name ``eigenvalues`` (corresponding to :math:`\operatorname{diag}\Lambda` above) and must be an array consisting of computed eigenvalues. The array containing the eigenvalues must have shape ``(..., M)`` and must have a real-valued floating-point data type whose precision matches the precision of ``x`` (e.g., if ``x`` is ``complex128``, then ``eigenvalues`` must be ``float64``).
-        -   second element have have the field name ``eigenvectors`` (corresponding to :math:`Q` above) and must be an array where the columns of the inner most matrices contain the computed eigenvectors. These matrices must be orthogonal. The array containing the eigenvectors must have shape ``(..., M, M)`` and must have the same data type as ``x``.
+        -   second element must have the field name ``eigenvectors`` (corresponding to :math:`Q` above) and must be an array where the columns of the inner most matrices contain the computed eigenvectors. These matrices must be orthogonal. The array containing the eigenvectors must have shape ``(..., M, M)`` and must have the same data type as ``x``.
 
     Notes
     -----
@@ -227,15 +224,12 @@ def eigvalsh(x: array, /) -> array:
 
     where :math:`\lambda \in \mathbb{R}` and where :math:`I_n` is the *n*-dimensional identity matrix.
 
-    .. note:;
+    .. note::
        The eigenvalues of a complex Hermitian or real symmetric matrix are always real.
 
     .. note::
        Whether an array library explicitly checks whether an input array is Hermitian or a symmetric matrix (or a stack of matrices) is implementation-defined.
 
-    .. note::
-       The function ``eigvals`` will be added in a future version of the specification.
-
     Parameters
     ----------
     x: array
@@ -257,6 +251,114 @@ def eigvalsh(x: array, /) -> array:
     """
 
 
+def eig(x: array, /) -> Tuple[array, array]:
+    r"""
+    Returns eigenvalues and eigenvectors of a real or complex matrix (or stack of matrices) ``x``.
+
+    Let :math:`\mathbb{K}` be the union of the set of real numbers :math:`\mathbb{R}`
+    and the set of complex numbers, :math:`\mathbb{C}`.
+
+    A real or complex value :math:`lambda \in \mathbb{K}` is an **eigenvalue** of a real
+    or complex matrix :math:`x \in \mathbb{K}^{n \times n}` if there exists a real or complex vector
+    :math:`v \in \mathbb{K}^{n}`, such that
+
+    .. math::
+
+            x v = \lambda v
+
+    Then, :math:`v` is referred to as an **eigenvector** of :math:`x`, corresponding to
+    the eigenvalue :math:`\lambda`.
+
+    A general matrix :math:`x \in \mathbb{K}^{n \times n}`
+
+    - has :math:`n` eigenvalues, which are defined as the roots (counted with multiplicity) of the
+      polynomial :math:`p` of degree :math:`n` given by
+
+       .. math::
+
+          p(\lambda) = \operatorname{det}(x - \lambda I_n)
+
+    - does not in general have :math:`n` linearly independent eigenvectors if it has
+      eigenvalues with multiplicity larger than one.
+
+    .. note::
+       The eigenvalues of a non-symmetric real matrix are in general complex: for
+       :math:x \in \mathbb{R}^{n \times n}`, the eigenvalues, :math:`\lambda \in \mathbb{C}`,
+       may or may not reside on the real axis of the complex plane.
+
+    .. warning::
+       The eigenvectors of a general matrix are not unique and are not continuous with respect to ``x``. Because eigenvectors are not unique, different hardware and software may compute different eigenvectors.
+
+       For eigenvalues of multiplity :math:`s=1`, the non-uniqueness stems from the fact that multiplying an eigenvector by :math:`-1` when ``x`` is real-valued and by :math:`e^{\phi j}` (:math:`\phi \in \mathbb{R}`) when ``x`` is complex produces another set of valid eigenvectors.
+
+       For eigenvalues of multiplity :math:`s > 1`, the :math:`s` computed eigenvectors may be repeated or
+       have entries differ by an order of machine epsilon for the data type of :math:`x`.
+
+    Parameters
+    ----------
+    x: array
+        input array having shape ``(..., M, M)`` and whose innermost two dimensions form square matrices. Should have a floating-point data type.
+
+    Returns
+    -------
+    out: Tuple[array, array]
+        a namedtuple (``eigenvalues``, ``eigenvectors``) whose
+
+        -   first element must have the field name ``eigenvalues`` (corresponding to :math:`\lambda` above) and must be an array consisting of computed eigenvalues. The array containing the eigenvalues must have shape ``(..., M)`` and must have a complex floating-point array data type having the same precision as that of ``x`` (e.g., if ``x`` has a ``float32`` data type, ``eigenvalues`` must have the ``complex64`` data type; if ``x`` has a ``float64`` data type, ``eigenvalues`` have the ``complex128`` data type).
+
+        -   second element must have the field name ``eigenvectors`` (corresponding to :math:`v` above) and must be an array where the columns of the inner most matrices contain the computed eigenvectors. These matrices must be orthogonal. The array containing the eigenvectors must have shape ``(..., M, M)`` and must have the same data type as ``eigenvalues``.
+
+    Notes
+    -----
+
+    -   Eigenvalue sort order is left unspecified and is thus implementation-dependent.
+
+    .. note::
+        For real symmetric or complex Hermitian matrices, prefer using the ``eigh`` routine.
+
+    .. versionadded:: 2025.12
+    """
+
+
+def eigvals(x: array, /) -> array:
+    r"""
+    Returns the eigenvalues of a real or complex matrix (or a stack of matrices) ``x``.
+
+    If ``x`` is real-valued, let :math:`\mathbb{K}` be the union of the set of real numbers :math:`\mathbb{R}`
+    and the set of complex numbers, :math:`\mathbb{C}`;  if ``x`` is complex-valued, let :math:`\mathbb{K}` be the set of complex numbers :math:`\mathbb{C}`.
+
+    The **eigenvalues** of a real or complex matrix :math:`x \in\ \mathbb{K}^{n \times n}` are defined as the roots (counted with multiplicity) of the polynomial :math:`p` of degree :math:`n` given by
+
+    .. math::
+       p(\lambda) = \operatorname{det}(x - \lambda I_n)
+
+    where :math:`\lambda \in \mathbb{K}` and where :math:`I_n` is the *n*-dimensional identity matrix.
+
+    .. note::
+       The eigenvalues of a non-symmetric real matrix are in general complex: for
+       :math:x \in \mathbb{R}^{n \times n}`, the eigenvalues :math:`\lambda \in \mathbb{C}`,
+       may or may not reside on the real axis of the complex plane.
+
+
+    Parameters
+    ----------
+    x: array
+        input array having shape ``(..., M, M)`` and whose innermost two dimensions form square matrices. Should have a floating-point data type.
+
+    Returns
+    -------
+    out: array
+        an array containing the computed eigenvalues. The returned array must have shape ``(..., M)`` and have a complex floating-point array data type having the same precision as that of ``x`` (e.g., if ``x`` has a ``float32`` data type, the array must have the ``complex64`` data type; if ``x`` has a ``float64`` data type, the array has the ``complex128`` data type).
+
+    Notes
+    -----
+
+    -   Eigenvalue sort order is left unspecified and is thus implementation-dependent.
+
+    .. versionadded:: 2025.12
+    """
+
+
 def inv(x: array, /) -> array:
     r"""
     Returns the multiplicative inverse of a square matrix (or a stack of square matrices) ``x``.
diff --git a/src/array_api_stubs/_draft/manipulation_functions.py b/src/array_api_stubs/_draft/manipulation_functions.py
@@ -20,7 +20,7 @@
 from ._types import List, Optional, Tuple, Union, array
 
 
-def broadcast_arrays(*arrays: array) -> List[array]:
+def broadcast_arrays(*arrays: array) -> Tuple[array, ...]:
     """
     Broadcasts one or more arrays against one another.
 
@@ -31,8 +31,8 @@ def broadcast_arrays(*arrays: array) -> List[array]:
 
     Returns
     -------
-    out: List[array]
-        a list of broadcasted arrays. Each array **must** have the same shape. Each array **must** have the same dtype as its corresponding input array.
+    out: Tuple[array, ...]
+        tuple of broadcasted arrays. Each array **must** have the same shape. Each array **must** have the same dtype as its corresponding input array.
     """
 
 
@@ -107,26 +107,42 @@ def concat(
     """
 
 
-def expand_dims(x: array, /, *, axis: int = 0) -> array:
+def expand_dims(x: array, /, axis: Union[int, Tuple[int, ...]]) -> array:
     """
-    Expands the shape of an array by inserting a new axis of size one at the position specified by ``axis``.
+    Expands the shape of an array by inserting a new axis of size one at the position (or positions) specified by ``axis``.
 
     Parameters
     ----------
     x: array
         input array.
-    axis: int
-        axis position (zero-based). A valid ``axis`` **must** reside on the closed-interval ``[-N-1, N]``, where ``N`` is the number of axes in ``x``. If an axis is specified as a negative integer, the axis position at which to insert a singleton dimension **must** be computed as ``N + axis + 1``. Hence, if provided ``-1``, the resolved axis position **must** be ``N`` (i.e., a singleton dimension **must** be appended to the input array ``x``). If provided ``-N-1``, the resolved axis position **must** be ``0`` (i.e., a singleton dimension **must** be prepended to the input array ``x``). If provided an invalid axis, the function **must** raise an exception. Default: ``0``.
+    axis: Union[int, Tuple[int, ...]]
+        axis position(s) (zero-based). If ``axis`` is an integer, ``axis`` **must** be equivalent to the tuple ``(axis,)``. If ``axis`` is a tuple,
+
+        -   a valid axis position **must** reside on the half-open interval ``[-M, M)``, where ``M = N + len(axis)`` and ``N`` is the number of dimensions in ``x``.
+        -   if the i-th entry is a negative integer, the axis position of the inserted singleton dimension in the output array **must** be computed as ``M + axis[i]``.
+        -   each entry of ``axis`` must resolve to a unique positive axis position.
+        -   for each entry of ``axis``, the corresponding dimension in the expanded output array **must** be a singleton dimension.
+        -   for the remaining dimensions of the expanded output array, the output array dimensions **must** correspond to the dimensions of ``x`` in order.
+        -   if provided an invalid axis position, the function **must** raise an exception.
 
     Returns
     -------
     out: array
-        an expanded output array. **Must** have the same data type as ``x``.
+        an expanded output array. **Must** have the same data type as ``x``. If ``axis`` is an integer, the output array must have ``N + 1`` dimensions. If ``axis`` is a tuple, the output array must have ``N + len(axis)`` dimensions.
 
     Raises
     ------
     IndexError
         If provided an invalid ``axis``, an ``IndexError`` **should** be raised.
+
+    Notes
+    -----
+
+    -   Calling this function with a tuple of axis positions **must** be semantically equivalent to calling this function repeatedly with a single axis position only when the following three conditions are met:
+
+        -   each entry of the tuple is normalized to positive axis positions according to the number of dimensions in the expanded output array.
+        -   the normalized positive axis positions are sorted in ascending order.
+        -   the normalized positive axis positions are unique.
     """
 
 
diff --git a/src/array_api_stubs/_draft/set_functions.py b/src/array_api_stubs/_draft/set_functions.py
