ENH: size() to return None on dask instead of nan

Author: crusaderky

array_api_compat/common/_helpers.py

@@ -788,19 +788,24 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
     return x.to_device(device, stream=stream)
 
 
-def size(x):
+def size(x: Array) -> int | None:
     """
     Return the total number of elements of x.
 
     This is equivalent to `x.size` according to the `standard
     <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.size.html>`__.
+
     This helper is included because PyTorch defines `size` in an
     :external+torch:meth:`incompatible way <torch.Tensor.size>`.
-
+    It also fixes dask.array's behaviour which returns nan for unknown sizes, whereas
+    the standard requires None.
     """
-    if None in x.shape:
+    if None in x.shape:  # this happens e.g. in ndonnx
+        return None
+    out = math.prod(x.shape)
+    if math.isnan(out):  # this happens e.g. in dask
         return None
-    return math.prod(x.shape)
+    return out
 
 
 def is_writeable_array(x) -> bool:

tests/_helpers.py

@@ -11,7 +11,7 @@
     all_libraries.append('sparse')
 
 def import_(library, wrapper=False):
-    if library == 'cupy':
+    if library in ('cupy', 'ndonnx'):
         pytest.importorskip(library)
     if wrapper:
         if 'jax' in library:

tests/test_common.py

@@ -5,8 +5,9 @@
     is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace,
 )
 
-from array_api_compat import device, is_array_api_obj, is_writeable_array, to_device
-
+from array_api_compat import (
+    device, is_array_api_obj, is_writeable_array, size, to_device
+)
 from ._helpers import import_, wrapped_libraries, all_libraries
 
 import pytest
@@ -92,6 +93,27 @@ def test_is_writeable_array_numpy():
     assert not is_writeable_array(x)
 
 
+@pytest.mark.parametrize("library", all_libraries)
+def test_size(library):
+    xp = import_(library)
+    x = xp.asarray([1, 2, 3])
+    assert size(x) == 3
+
+
+def test_size_nan():
+    xp = import_("dask.array")
+    x = xp.arange(10)
+    x = x[x < 5]
+    assert size(x) is None  # NaNs in the shape have been special-cased
+
+
+def test_size_none():
+    xp = import_("ndonnx")
+    x = xp.arange(10)
+    x = x[x < 5]
+    assert size(x) is None
+
+
 @pytest.mark.parametrize("library", all_libraries)
 def test_device(library):
     xp = import_(library, wrapper=True)