documentation nit

Author: crusaderky

Diff for: src/array_api_extra/_lib/_funcs.py

@@ -713,13 +713,14 @@ def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array:
 
     # There are 3 general use cases:
     # 1. backend has unique_counts and it returns an array with known shape
-    # 2. backend has unique_counts and it returns a None-shaped array
-    # 3. backend does not have unique_counts
+    # 2. backend has unique_counts and it returns a None-sized array;
+    #    e.g. Dask, ndonnx
+    # 3. backend does not have unique_counts; e.g. wrapped JAX
     if capabilities(xp)["data-dependent shapes"]:
         # xp has unique_counts; O(n) complexity
         _, counts = xp.unique_counts(x)
         n = _compat.size(counts)
-        if n is None:  # e.g. Dask, ndonnx
+        if n is None:
             return xp.sum(xp.ones_like(counts))
         return xp.asarray(n, device=_compat.device(x))