Misc dtype fixes (#771)

* Add missing dtypes in cumulative_sum and scan

* Use nxp.bool not Python bool in `all` and `any` functions

* Use default array index data type rather than hardcoding.

* Simplify linspace start/size/step types, since they are just Python types

* Namespace fixes in the optimization tests

Author: tomwhite

cubed/array_api/creation_functions.py

@@ -228,16 +228,12 @@ def linspace(
 def _linspace(
     x, size, start, step, endpoint, linspace_dtype, device=None, block_id=None
 ):
-    dtypes = __array_namespace_info__().default_dtypes(device=device)
-
     bs = x.shape[0]
     i = block_id[0]
     adjusted_bs = bs - 1 if endpoint else bs
 
-    # float_ is a type casting function.
-    float_ = dtypes["real floating"].type
-    blockstart = float_(start + (i * size * step))
-    blockstop = float_(blockstart + float_(adjusted_bs * step))
+    blockstart = start + (i * size * step)
+    blockstop = blockstart + adjusted_bs * step
     return nxp.linspace(
         blockstart, blockstop, bs, endpoint=endpoint, dtype=linspace_dtype
     )

cubed/array_api/searching_functions.py

@@ -49,14 +49,15 @@ def searchsorted(x1, x2, /, *, side="left", sorter=None):
         )
 
     # call nxp.searchsorted for each pair of blocks in x1 and v
+    dtype = nxp.__array_namespace_info__().default_dtypes(device=x1.device)["indexing"]
     out = blockwise(
         _searchsorted,
         list(range(x2.ndim + 1)),
         x1,
         [0],
         x2,
         list(range(1, x2.ndim + 1)),
-        dtype=nxp.int64,  # TODO: index dtype
+        dtype=dtype,
         adjust_chunks={0: 1},  # one row for each block in x1
         side=side,
     )

cubed/array_api/statistical_functions.py

@@ -23,7 +23,14 @@ def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False, device
         fname="cumulative_sum",
         device=device,
     )
-    return scan(x, preop=nxp.sum, func=_cumulative_sum_func, binop=nxp.add, axis=axis)
+    return scan(
+        x,
+        preop=nxp.sum,
+        func=_cumulative_sum_func,
+        binop=nxp.add,
+        axis=axis,
+        dtype=dtype,
+    )
 
 
 def _cumulative_sum_func(a, /, *, axis=None, dtype=None, include_initial=False):

cubed/array_api/utility_functions.py

@@ -10,7 +10,7 @@ def all(x, /, *, axis=None, keepdims=False, split_every=None):
         x,
         nxp.all,
         axis=axis,
-        dtype=bool,
+        dtype=nxp.bool,
         keepdims=keepdims,
         split_every=split_every,
     )
@@ -23,7 +23,7 @@ def any(x, /, *, axis=None, keepdims=False, split_every=None):
         x,
         nxp.any,
         axis=axis,
-        dtype=bool,
+        dtype=nxp.bool,
         keepdims=keepdims,
         split_every=split_every,
     )

cubed/core/ops.py

@@ -1339,7 +1339,7 @@ def _partial_reduce(arrays, reduce_func=None, initial_func=None, axis=None):
 
 def arg_reduction(x, /, arg_func, axis=None, *, keepdims=False, split_every=None):
     """A reduction that returns the array indexes, not the values."""
-    dtype = nxp.int64  # index data type
+    dtype = nxp.__array_namespace_info__().default_dtypes(device=x.device)["indexing"]
     intermediate_dtype = [("i", dtype), ("v", x.dtype)]
 
     # initial map does arg reduction on each block, and uses block id to find the absolute index within whole array
@@ -1549,7 +1549,9 @@ def scan(
 
     # Blelloch (1990) out-of-core algorithm.
     # 1. First, scan blockwise
-    scanned = map_blocks(func, array, axis=axis, include_initial=include_initial)
+    scanned = map_blocks(
+        func, array, dtype=dtype, axis=axis, include_initial=include_initial
+    )
     # If there is only a single chunk, we can be done
     if array.numblocks[axis] == 1:
         return scanned
@@ -1567,6 +1569,7 @@ def identity_func(a, **kwargs):
         initial_func=partial(preop, axis=axis, keepdims=True),
         func=identity_func,
         split_every={axis: split_size},
+        dtype=dtype,
         combine_sizes={axis: split_size},
     )
 
@@ -1581,6 +1584,7 @@ def identity_func(a, **kwargs):
         preop=preop,
         binop=binop,
         axis=axis,
+        dtype=dtype,
         include_initial=True,
     )
 

cubed/tests/test_optimization.py

@@ -36,7 +36,7 @@ def spec(tmp_path):
 def test_fusion(spec, opt_fn):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.negative(a)
-    c = xp.astype(b, np.float32)
+    c = xp.astype(b, xp.float32)
     d = xp.negative(c)
 
     num_arrays = 4  # a, b, c, d
@@ -69,7 +69,7 @@ def test_fusion(spec, opt_fn):
 def test_fusion_compute_multiple(spec, opt_fn):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.negative(a)
-    c = xp.astype(b, np.float32)
+    c = xp.astype(b, xp.float32)
     d = xp.negative(c)
 
     # if we compute c and d then both have to be materialized
@@ -97,7 +97,7 @@ def test_fusion_compute_multiple(spec, opt_fn):
 def test_fusion_transpose(spec, opt_fn):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.negative(a)
-    c = xp.astype(b, np.float32)
+    c = xp.astype(b, xp.float32)
     d = c.T
 
     num_created_arrays = 3  # b, c, d
@@ -191,7 +191,7 @@ def test_no_fusion_multiple_edges(spec):
 def test_custom_optimize_function(spec):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.negative(a)
-    c = xp.astype(b, np.float32)
+    c = xp.astype(b, xp.float32)
     d = xp.negative(c)
 
     num_tasks_with_no_optimization = d.plan._finalize(optimize_graph=False).num_tasks()
@@ -992,7 +992,7 @@ def test_fuse_merge_chunks_binary(spec):
 def test_fuse_partial_reduce_unary(spec):
     a = xp.ones((3, 2), chunks=(1, 2), spec=spec)
     b = xp.negative(a)
-    c = partial_reduce(b, np.sum, split_every={0: 3})
+    c = partial_reduce(b, nxp.sum, split_every={0: 3}, dtype=xp.float64)
 
     opt_fn = fuse_multiple_levels()
 
@@ -1017,7 +1017,7 @@ def test_fuse_partial_reduce_binary(spec):
     a = xp.ones((3, 2), chunks=(1, 2), spec=spec)
     b = xp.ones((3, 2), chunks=(1, 2), spec=spec)
     c = xp.add(a, b)
-    d = partial_reduce(c, np.sum, split_every={0: 3})
+    d = partial_reduce(c, nxp.sum, split_every={0: 3}, dtype=xp.float64)
 
     opt_fn = fuse_multiple_levels()