Array API 2024 additions (#779)

* Add 'max dimensions' to capabilities

* Add `nextafter`

* Add `reciprocal`

* Add `count_nonzero`

Author: tomwhite

api_status.md

@@ -69,7 +69,7 @@ This table shows which parts of the the [Array API](https://data-apis.org/array-
 |                          | `unstack`           | :white_check_mark: | 2023.12    |                              |
 | Searching Functions      | `argmax`            | :white_check_mark: |            |                              |
 |                          | `argmin`            | :white_check_mark: |            |                              |
-|                          | `count_nonzero`     | :x:                | 2024.12    |                              |
+|                          | `count_nonzero`     | :white_check_mark: | 2024.12    |                              |
 |                          | `nonzero`           | :x:                |            | Shape is data dependent      |
 |                          | `searchsorted`      | :white_check_mark: | 2023.12    |                              |
 |                          | `where`             | :white_check_mark: |            |                              |

cubed/__init__.py

@@ -193,10 +193,12 @@
     minimum,
     multiply,
     negative,
+    nextafter,
     not_equal,
     positive,
     pow,
     real,
+    reciprocal,
     remainder,
     round,
     sign,
@@ -260,11 +262,13 @@
     "maximum",
     "minimum",
     "multiply",
+    "nextafter",
     "negative",
     "not_equal",
     "positive",
     "pow",
     "real",
+    "reciprocal",
     "remainder",
     "round",
     "sign",
@@ -326,9 +330,15 @@
     "unstack",
 ]
 
-from .array_api.searching_functions import argmax, argmin, searchsorted, where
+from .array_api.searching_functions import (
+    argmax,
+    argmin,
+    count_nonzero,
+    searchsorted,
+    where,
+)
 
-__all__ += ["argmax", "argmin", "searchsorted", "where"]
+__all__ += ["argmax", "argmin", "count_nonzero", "searchsorted", "where"]
 
 from .array_api.statistical_functions import (
     cumulative_sum,

cubed/array_api/__init__.py

@@ -136,10 +136,12 @@
     minimum,
     multiply,
     negative,
+    nextafter,
     not_equal,
     positive,
     pow,
     real,
+    reciprocal,
     remainder,
     round,
     sign,
@@ -203,11 +205,13 @@
     "maximum",
     "minimum",
     "multiply",
+    "nextafter",
     "negative",
     "not_equal",
     "positive",
     "pow",
     "real",
+    "reciprocal",
     "remainder",
     "round",
     "sign",
@@ -264,9 +268,9 @@
     "unstack",
 ]
 
-from .searching_functions import argmax, argmin, searchsorted, where
+from .searching_functions import argmax, argmin, count_nonzero, searchsorted, where
 
-__all__ += ["argmax", "argmin", "searchsorted", "where"]
+__all__ += ["argmax", "argmin", "count_nonzero", "searchsorted", "where"]
 
 from .statistical_functions import cumulative_sum, max, mean, min, prod, std, sum, var
 

cubed/array_api/elementwise_functions.py

@@ -372,6 +372,13 @@ def multiply(x1, x2, /):
     return elemwise(nxp.multiply, x1, x2, dtype=result_type(x1, x2))
 
 
+def nextafter(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "nextafter")
+    if x1.dtype not in _real_floating_dtypes or x2.dtype not in _real_floating_dtypes:
+        raise TypeError("Only real floating-point dtypes are allowed in nextafter")
+    return elemwise(nxp.nextafter, x1, x2, dtype=x1.dtype)
+
+
 def negative(x, /):
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in negative")
@@ -406,6 +413,12 @@ def real(x, /):
     return elemwise(nxp.real, x, dtype=dtype)
 
 
+def reciprocal(x, /):
+    if x.dtype not in _floating_dtypes:
+        raise TypeError("Only floating-point dtypes are allowed in reciprocal")
+    return elemwise(nxp.reciprocal, x, dtype=x.dtype)
+
+
 def remainder(x1, x2, /):
     x1, x2 = _promote_scalars(x1, x2, "remainder")
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:

cubed/array_api/inspection.py

@@ -2,11 +2,13 @@
 
 
 class __array_namespace_info__:
-    # capabilities are determined by Cubed, not the backend array API
     def capabilities(self):
         return {
-            "boolean indexing": False,
-            "data-dependent shapes": False,
+            "boolean indexing": False,  # not supported in Cubed (#73)
+            "data-dependent shapes": False,  # not supported in Cubed
+            "max dimensions": nxp.__array_namespace_info__().capabilities()[
+                "max dimensions"
+            ],
         }
 
     # devices and dtypes are determined by the backend array API

cubed/array_api/searching_functions.py

@@ -1,8 +1,8 @@
 from cubed.array_api.creation_functions import asarray, zeros_like
-from cubed.array_api.data_type_functions import result_type
+from cubed.array_api.data_type_functions import astype, result_type
 from cubed.array_api.dtypes import _promote_scalars, _real_numeric_dtypes
 from cubed.array_api.manipulation_functions import reshape
-from cubed.array_api.statistical_functions import max
+from cubed.array_api.statistical_functions import max, sum
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.ops import arg_reduction, blockwise, elemwise
 
@@ -39,6 +39,17 @@ def argmin(x, /, *, axis=None, keepdims=False, split_every=None):
     )
 
 
+def count_nonzero(x, /, *, axis=None, keepdims=False, split_every=None):
+    dtype = nxp.__array_namespace_info__().default_dtypes(device=x.device)["indexing"]
+    return sum(
+        astype(x, nxp.bool),
+        axis=axis,
+        dtype=dtype,
+        keepdims=keepdims,
+        split_every=split_every,
+    )
+
+
 def searchsorted(x1, x2, /, *, side="left", sorter=None):
     if x1.ndim != 1:
         raise ValueError("Input array x1 must be one dimensional")