Merge branch 'main' into scalar-where

kgryte · web-flow · commit db2f1019d437 · 2025-01-09T02:36:37.000-08:00
diff --git a/spec/draft/API_specification/searching_functions.rst b/spec/draft/API_specification/searching_functions.rst
@@ -22,6 +22,7 @@ Objects in API
 
    argmax
    argmin
+   count_nonzero
    nonzero
    searchsorted
    where
diff --git a/spec/draft/API_specification/statistical_functions.rst b/spec/draft/API_specification/statistical_functions.rst
@@ -18,6 +18,7 @@ Objects in API
    :toctree: generated
    :template: method.rst
 
+   cumulative_prod
    cumulative_sum
    max
    mean
diff --git a/spec/draft/design_topics/copies_views_and_mutation.rst b/spec/draft/design_topics/copies_views_and_mutation.rst
@@ -6,7 +6,7 @@ Copy-view behaviour and mutability
 .. admonition:: Mutating views
    :class: important
 
-   Array API consumers are *strongly* advised to avoid *any* mutating operations when an array object may be either a "view" (i.e., an array whose data refers to memory that belongs to another array) or own memory of which one or more other array objects may be views. This admonition may become more strict in the future (e.g., this specification may require that view mutation be prohibited and trigger an exception). Accordingly, only perform mutation operations (e.g., in-place assignment) when absolutely confident that array data belongs to one, and only one, array object.
+   Array API consumers are *strongly* advised to avoid *any* mutating operations when an array object may either be a "view" (i.e., an array whose data refers to memory that belongs to another array) or own memory of which one or more other array objects may be views. This admonition may become more strict in the future (e.g., this specification may require that view mutation be prohibited and trigger an exception). Accordingly, only perform mutation operations (e.g., in-place assignment) when absolutely confident that array data belongs to one, and only one, array object.
 
 Strided array implementations (e.g. NumPy, PyTorch, CuPy, MXNet) typically
 have the concept of a "view", meaning an array containing data in memory that
diff --git a/spec/draft/purpose_and_scope.md b/spec/draft/purpose_and_scope.md
@@ -410,8 +410,8 @@ all the functions, arguments, data types, syntax, and semantics described in
 this specification.
 
 A conforming implementation of the array API standard may provide additional
-values, objects, properties, data types, and functions beyond those described
-in this specification.
+features (e.g., values, objects, properties, data types, functions, and function
+arguments) beyond those described in this specification.
 
 Libraries which aim to provide a conforming implementation but haven't yet
 completed such an implementation may, and are encouraged to, provide details on
diff --git a/src/array_api_stubs/_2021_12/linalg.py b/src/array_api_stubs/_2021_12/linalg.py
@@ -474,7 +474,7 @@ def trace(x: array, /, *, offset: int = 0) -> array:
     """
 
 
-def vecdot(x1: array, x2: array, /, *, axis: int = None) -> array:
+def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array:
     """
     Alias for :func:`~array_api.vecdot`.
     """
diff --git a/src/array_api_stubs/_2021_12/searching_functions.py b/src/array_api_stubs/_2021_12/searching_functions.py
@@ -69,7 +69,7 @@ def where(condition: array, x1: array, x2: array, /) -> array:
     Parameters
     ----------
     condition: array
-        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
+        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Should have a boolean data type. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
     x1: array
         first input array. Must be compatible with ``condition`` and ``x2`` (see :ref:`broadcasting`).
     x2: array
diff --git a/src/array_api_stubs/_2022_12/linalg.py b/src/array_api_stubs/_2022_12/linalg.py
@@ -755,7 +755,7 @@ def trace(x: array, /, *, offset: int = 0, dtype: Optional[dtype] = None) -> arr
     """
 
 
-def vecdot(x1: array, x2: array, /, *, axis: int = None) -> array:
+def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array:
     """Alias for :func:`~array_api.vecdot`."""
 
 
diff --git a/src/array_api_stubs/_2022_12/searching_functions.py b/src/array_api_stubs/_2022_12/searching_functions.py
@@ -91,7 +91,7 @@ def where(condition: array, x1: array, x2: array, /) -> array:
     Parameters
     ----------
     condition: array
-        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
+        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Should have a boolean data type. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
     x1: array
         first input array. Must be compatible with ``condition`` and ``x2`` (see :ref:`broadcasting`).
     x2: array
diff --git a/src/array_api_stubs/_2023_12/linalg.py b/src/array_api_stubs/_2023_12/linalg.py
@@ -781,7 +781,7 @@ def trace(x: array, /, *, offset: int = 0, dtype: Optional[dtype] = None) -> arr
     """
 
 
-def vecdot(x1: array, x2: array, /, *, axis: int = None) -> array:
+def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array:
     """Alias for :func:`~array_api.vecdot`."""
 
 
diff --git a/src/array_api_stubs/_2023_12/searching_functions.py b/src/array_api_stubs/_2023_12/searching_functions.py
@@ -146,7 +146,7 @@ def where(condition: array, x1: array, x2: array, /) -> array:
     Parameters
     ----------
     condition: array
-        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
+        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Should have a boolean data type. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
     x1: array
         first input array. Must be compatible with ``condition`` and ``x2`` (see :ref:`broadcasting`).
     x2: array
diff --git a/src/array_api_stubs/_draft/data_type_functions.py b/src/array_api_stubs/_draft/data_type_functions.py
@@ -209,20 +209,25 @@ def isdtype(
     """
 
 
-def result_type(*arrays_and_dtypes: Union[array, dtype]) -> dtype:
+def result_type(
+    *arrays_and_dtypes: Union[array, int, float, complex, bool, dtype]
+) -> dtype:
     """
-    Returns the dtype that results from applying the type promotion rules (see :ref:`type-promotion`) to the arguments.
-
-    .. note::
-       If provided mixed dtypes (e.g., integer and floating-point), the returned dtype will be implementation-specific.
+    Returns the dtype that results from applying type promotion rules (see :ref:`type-promotion`) to the arguments.
 
     Parameters
     ----------
-    arrays_and_dtypes: Union[array, dtype]
-        an arbitrary number of input arrays and/or dtypes.
+    arrays_and_dtypes: Union[array, int, float, complex, bool, dtype]
+        an arbitrary number of input arrays, scalars, and/or dtypes.
 
     Returns
     -------
     out: dtype
-        the dtype resulting from an operation involving the input arrays and dtypes.
+        the dtype resulting from an operation involving the input arrays, scalars, and/or dtypes.
+
+    Notes
+    -----
+
+    -   At least one argument must be an array or a dtype.
+    -   If provided array and/or dtype arguments having mixed data type kinds (e.g., integer and floating-point), the returned dtype is unspecified and is implementation-dependent.
     """
diff --git a/src/array_api_stubs/_draft/indexing_functions.py b/src/array_api_stubs/_draft/indexing_functions.py
@@ -13,7 +13,7 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None) -> array:
     Parameters
     ----------
     x: array
-        input array.
+        input array. Should have one or more dimensions (axes).
     indices: array
         array indices. The array must be one-dimensional and have an integer data type.
 
@@ -33,6 +33,8 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None) -> array:
     Notes
     -----
 
+    -   When ``x`` is a zero-dimensional array, behavior is unspecified and thus implementation-defined.
+
     .. versionadded:: 2022.12
 
     .. versionchanged:: 2023.12
diff --git a/src/array_api_stubs/_draft/linalg.py b/src/array_api_stubs/_draft/linalg.py
@@ -781,7 +781,7 @@ def trace(x: array, /, *, offset: int = 0, dtype: Optional[dtype] = None) -> arr
     """
 
 
-def vecdot(x1: array, x2: array, /, *, axis: int = None) -> array:
+def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array:
     """Alias for :func:`~array_api.vecdot`."""
 
 
diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py
@@ -1,4 +1,4 @@
-__all__ = ["argmax", "argmin", "nonzero", "searchsorted", "where"]
+__all__ = ["argmax", "argmin", "count_nonzero", "nonzero", "searchsorted", "where"]
 
 
 from ._types import Optional, Tuple, Literal, Union, array
@@ -54,15 +54,41 @@ def argmin(x: array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
     """
 
 
-def nonzero(x: array, /) -> Tuple[array, ...]:
+def count_nonzero(
+    x: array,
+    /,
+    *,
+    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    keepdims: bool = False,
+) -> array:
     """
-    Returns the indices of the array elements which are non-zero.
+    Counts the number of array elements which are non-zero.
 
-    .. note::
-       If ``x`` has a complex floating-point data type, non-zero elements are those elements having at least one component (real or imaginary) which is non-zero.
+    Parameters
+    ----------
+    x: array
+        input array.
+    axis: Optional[Union[int, Tuple[int, ...]]]
+        axis or axes along which to count non-zero values. By default, the number of non-zero values must be computed over the entire array. If a tuple of integers, the number of non-zero values must be computed over multiple axes. Default: ``None``.
+    keepdims: bool
+        if ``True``, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes (dimensions) must not be included in the result. Default: ``False``.
 
-    .. note::
-       If ``x`` has a boolean data type, non-zero elements are those elements which are equal to ``True``.
+    Returns
+    -------
+    out: array
+        if the number of non-zeros values was computed over the entire array, a zero-dimensional array containing the total number of non-zero values; otherwise, a non-zero-dimensional array containing the counts along the specified axes. The returned array must have the default array index data type.
+
+    Notes
+    -----
+
+    -   If ``x`` has a complex floating-point data type, non-zero elements are those elements having at least one component (real or imaginary) which is non-zero.
+    -   If ``x`` has a boolean data type, non-zero elements are those elements which are equal to ``True``.
+    """
+
+
+def nonzero(x: array, /) -> Tuple[array, ...]:
+    """
+    Returns the indices of the array elements which are non-zero.
 
     .. admonition:: Data-dependent output shape
        :class: admonition important
@@ -76,12 +102,15 @@ def nonzero(x: array, /) -> Tuple[array, ...]:
 
     Returns
     -------
-    out: Typle[array, ...]
+    out: Tuple[array, ...]
         a tuple of ``k`` arrays, one for each dimension of ``x`` and each of size ``n`` (where ``n`` is the total number of non-zero elements), containing the indices of the non-zero elements in that dimension. The indices must be returned in row-major, C-style order. The returned array must have the default array index data type.
 
     Notes
     -----
 
+    -   If ``x`` has a complex floating-point data type, non-zero elements are those elements having at least one component (real or imaginary) which is non-zero.
+    -   If ``x`` has a boolean data type, non-zero elements are those elements which are equal to ``True``.
+
     .. versionchanged:: 2022.12
        Added complex data type support.
     """
@@ -151,7 +180,7 @@ def where(
     Parameters
     ----------
     condition: array
-        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
+        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Should have a boolean data type. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
     x1: Union[array, int, float, complex, bool]
         first input array. Must be compatible with ``condition`` and ``x2`` (see :ref:`broadcasting`).
     x2: Union[array, int, float, complex, bool]
diff --git a/src/array_api_stubs/_draft/statistical_functions.py b/src/array_api_stubs/_draft/statistical_functions.py
@@ -1,9 +1,71 @@
-__all__ = ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]
+__all__ = [
+    "cumulative_sum",
+    "cumulative_prod",
+    "max",
+    "mean",
+    "min",
+    "prod",
+    "std",
+    "sum",
+    "var",
+]
 
 
 from ._types import Optional, Tuple, Union, array, dtype
 
 
+def cumulative_prod(
+    x: array,
+    /,
+    *,
+    axis: Optional[int] = None,
+    dtype: Optional[dtype] = None,
+    include_initial: bool = False,
+) -> array:
+    """
+    Calculates the cumulative product of elements in the input array ``x``.
+
+    Parameters
+    ----------
+    x: array
+        input array. Should have one or more dimensions (axes). Should have a numeric data type.
+    axis: Optional[int]
+        axis along which a cumulative product must be computed. If ``axis`` is negative, the function must determine the axis along which to compute a cumulative product by counting from the last dimension.
+
+        If ``x`` is a one-dimensional array, providing an ``axis`` is optional; however, if ``x`` has more than one dimension, providing an ``axis`` is required.
+
+    dtype: Optional[dtype]
+        data type of the returned array. If ``None``, the returned array must have the same data type as ``x``, unless ``x`` has an integer data type supporting a smaller range of values than the default integer data type (e.g., ``x`` has an ``int16`` or ``uint32`` data type and the default integer data type is ``int64``). In those latter cases:
+
+        -   if ``x`` has a signed integer data type (e.g., ``int16``), the returned array must have the default integer data type.
+        -   if ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned array must have an unsigned integer data type having the same number of bits as the default integer data type (e.g., if the default integer data type is ``int32``, the returned array must have a ``uint32`` data type).
+
+        If the data type (either specified or resolved) differs from the data type of ``x``, the input array should be cast to the specified data type before computing the product (rationale: the ``dtype`` keyword argument is intended to help prevent overflows). Default: ``None``.
+
+    include_initial: bool
+        boolean indicating whether to include the initial value as the first value in the output. By convention, the initial value must be the multiplicative identity (i.e., one). Default: ``False``.
+
+    Returns
+    -------
+    out: array
+        an array containing the cumulative products. The returned array must have a data type as described by the ``dtype`` parameter above.
+
+        Let ``N`` be the size of the axis along which to compute the cumulative product. The returned array must have a shape determined according to the following rules:
+
+        -   if ``include_initial`` is ``True``, the returned array must have the same shape as ``x``, except the size of the axis along which to compute the cumulative product must be ``N+1``.
+        -   if ``include_initial`` is ``False``, the returned array must have the same shape as ``x``.
+
+    Notes
+    -----
+
+    -   When ``x`` is a zero-dimensional array, behavior is unspecified and thus implementation-defined.
+
+    **Special Cases**
+
+    For both real-valued and complex floating-point operands, special cases must be handled as if the operation is implemented by successive application of :func:`~array_api.multiply`.
+    """
+
+
 def cumulative_sum(
     x: array,
     /,
@@ -18,7 +80,7 @@ def cumulative_sum(
     Parameters
     ----------
     x: array
-        input array. Should have a numeric data type.
+        input array. Should have one or more dimensions (axes). Should have a numeric data type.
     axis: Optional[int]
         axis along which a cumulative sum must be computed. If ``axis`` is negative, the function must determine the axis along which to compute a cumulative sum by counting from the last dimension.
 
@@ -48,6 +110,8 @@ def cumulative_sum(
     Notes
     -----
 
+    -   When ``x`` is a zero-dimensional array, behavior is unspecified and thus implementation-defined.
+
     **Special Cases**
 
     For both real-valued and complex floating-point operands, special cases must be handled as if the operation is implemented by successive application of :func:`~array_api.add`.
@@ -113,7 +177,7 @@ def mean(
     Parameters
     ----------
     x: array
-        input array. Should have a real-valued floating-point data type.
+        input array. Should have a floating-point data type.
     axis: Optional[Union[int, Tuple[int, ...]]]
         axis or axes along which arithmetic means must be computed. By default, the mean must be computed over the entire array. If a tuple of integers, arithmetic means must be computed over multiple axes. Default: ``None``.
     keepdims: bool
@@ -125,17 +189,26 @@ def mean(
         if the arithmetic mean was computed over the entire array, a zero-dimensional array containing the arithmetic mean; otherwise, a non-zero-dimensional array containing the arithmetic means. The returned array must have the same data type as ``x``.
 
         .. note::
-           While this specification recommends that this function only accept input arrays having a real-valued floating-point data type, specification-compliant array libraries may choose to accept input arrays having an integer data type. While mixed data type promotion is implementation-defined, if the input array ``x`` has an integer data type, the returned array must have the default real-valued floating-point data type.
+           While this specification recommends that this function only accept input arrays having a floating-point data type, specification-compliant array libraries may choose to accept input arrays having an integer data type. While mixed data type promotion is implementation-defined, if the input array ``x`` has an integer data type, the returned array must have the default real-valued floating-point data type.
 
     Notes
     -----
 
     **Special Cases**
 
-    Let ``N`` equal the number of elements over which to compute the arithmetic mean.
+    Let ``N`` equal the number of elements over which to compute the arithmetic mean. For real-valued operands,
 
     -   If ``N`` is ``0``, the arithmetic mean is ``NaN``.
     -   If ``x_i`` is ``NaN``, the arithmetic mean is ``NaN`` (i.e., ``NaN`` values propagate).
+
+    For complex floating-point operands, real-valued floating-point special cases should independently apply to the real and imaginary component operations involving real numbers. For example, let ``a = real(x_i)`` and ``b = imag(x_i)``, and
+
+    -   If ``N`` is ``0``, the arithmetic mean is ``NaN + NaN j``.
+    -   If ``a`` is ``NaN``, the real component of the result is ``NaN``.
+    -   Similarly, if ``b`` is ``NaN``, the imaginary component of the result is ``NaN``.
+
+    .. note::
+       Array libraries, such as NumPy, PyTorch, and JAX, currently deviate from this specification in their handling of components which are ``NaN`` when computing the arithmetic mean. In general, consumers of array libraries implementing this specification should use :func:`~array_api.isnan` to test whether the result of computing the arithmetic mean over an array have a complex floating-point data type is ``NaN``, rather than relying on ``NaN`` propagation of individual components.
     """
 
 
